Proper way to communicate with agent via websockets?

[JasonS09]

It's me again. For my next step on my app development I need to communicate with my langchain chatbot instance with a mobile interface via websockets.

I have an implemented it the following way:

import asyncio
from agent import ChatBot
from websockets.server import serve

async def reply(websocket):
    async for message in websocket:
        reply = chatbot.chat(message)
        await websocket.send(reply)

async def main():
    async with serve(reply, "192.168.0.109", 8765):
        await asyncio.Future()  # run forever

chatbot = ChatBot()
asyncio.run(main())

The connection works and you can communicate with the agent back and forth successfully. However, there is something odd: the agent is unable to call actions. There is no error, the actions simply don't kick and the agent gives a straight reply. I know there is no flaw in the way I implemented the tools because if I run the agent directly (via Python script), agent is able to run the actions as expected. Agent doesn't call the actions only when I communicate with it via websockets, which makes me suspect it has something to do with the way websockets run the chatbot asynchronously.

For reference, here is my current chatbot implementation:

class ChatBot():
    def __init__(self):
        #Filled with images from Stable Diffusion tool output
        self.images = []
        llm = KoboldApiLLM(endpoint="http://localhost:5001", max_length=512)

        wrapper = DuckDuckGoSearchAPIWrapper()
        search = DuckDuckGoSearchResults(wrapper=wrapper)
        sd = StableDiffusionGenerator()

        def run_stable_diffusion(prompt):
            try:
                msg, images = sd.get_images(prompt)
                self.images = images
                return msg
            except:
                return "Error when generating images."  

        tools = [
            Tool(
                func=search.run,
                name="Search",
                description="Useful for when you need to answer questions about current events or concepts.",
                handle_tool_error=True
                # coroutine= ... <- you can specify an async method if desired as well
            ),
            Tool(
                func=run_stable_diffusion,
                name="Generate Images",
                description="Useful for when you need to generate images or visually illustrate a concept.",
                handle_tool_error=True
                # coroutine= ... <- you can specify an async method if desired as well
            )
        ]

        # Set up the base template
        template_with_history = """Complete the objective as best you can. You have access to the following tools:

        {tools}

        Use the following format:

        Question: the input question you must answer
        Thought: Do you need to take an action? If not, skip to Final Answer.
        Action: the name of the action to take, should be one of the following: {tool_names}
        Action Input: the input to the action
        Observation: the result of the action
        ... (this Thought/Action/Action Input/Observation can repeat N times)
        Thought: I now know the final answer
        Final Answer: the final answer to the original input question. Make sure to give long, detailed final answers if needed.
        Begin!

        Previous conversation history:
        {history}

        Question: {input}
        {agent_scratchpad}"""

        prompt_with_history = CustomPromptTemplate(
            template=template_with_history,
            tools=tools,
            # This omits the `agent_scratchpad`, `tools`, and `tool_names` variables because those are generated dynamically
            # This includes the `intermediate_steps` variable because that is needed
            input_variables=["input", "intermediate_steps", "history"]
        )

        output_parser = CustomOutputParser()
        llm_chain = LLMChain(llm=llm, prompt=prompt_with_history)

        tool_names = [tool.name for tool in tools]
        agent = LLMSingleActionAgent(
            llm_chain=llm_chain,
            output_parser=output_parser,
            stop=["Begin!", "Question:"],
            allowed_tools=tool_names
        )

        memory=ConversationBufferWindowMemory(k=2)

        self.agent_executor = AgentExecutor(
            agent=agent, tools=tools, verbose=True, 
            memory=memory, handle_parsing_errors=True)
    
    def chat(self, input):
        output = self.agent_executor.run(input)
        if len(self.images) > 0:
            return f"{{'output': '{output}', 'images': {self.images}}}"
        return output

[dosubot[bot]]

🤖

Hello again,

From your description, it seems like the issue might be related to the asynchronous nature of websockets and how LangChain handles tool execution in such an environment.

In LangChain, the execution of tools in an asynchronous environment like websockets is handled through the use of Python's async and await keywords. This is evident in the _acall method of the Chain class. In the _acall method, the _call method is executed asynchronously using the await keyword. This means that the execution of the _acall method will pause and wait for the _call method to complete before moving on to the next line of code. This allows for non-blocking execution of code, which is crucial in an asynchronous environment like websockets.

However, in your websocket server implementation, you are not using the async version of the chat method. This might be the reason why the actions are not being called as expected. The chat method of your ChatBot class is a synchronous method, and when it's called inside the reply coroutine, it might be blocking the event loop, preventing the actions from being executed properly.

To solve this issue, you could try making the chat method of your ChatBot class an async method and use the await keyword when calling it inside the reply coroutine. Here's how you could modify your code:

import asyncio
from agent import ChatBot
from websockets.server import serve

async def reply(websocket):
    async for message in websocket:
        reply = await chatbot.chat(message)
        await websocket.send(reply)

async def main():
    async with serve(reply, "192.168.0.109", 8765):
        await asyncio.Future()  # run forever

chatbot = ChatBot()
asyncio.run(main())

And in your ChatBot class:

class ChatBot():
    # ... rest of your code ...

    async def chat(self, input):
        output = await self.agent_executor.run(input)
        if len(self.images) > 0:
            return f"{{'output': '{output}', 'images': {self.images}}}"
        return output

Please note that you might need to adjust other parts of your code to make them compatible with async operations.

I hope this helps! Let me know if you have any other questions.

Sources

• libs/langchain/langchain/chains/api/base.py
• libs/langchain/langchain/chat_models/jinachat.py
• libs/langchain/langchain/chains/base.py

---

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[JasonS09]

This solution worked a couple of times, then mysteriously went back to the previous state, no action running. This is my current relevant code:

import asyncio
from agent import ChatBot
from websockets.server import serve

async def reply(websocket):
    async for message in websocket:
        reply = await chatbot.chat(message)
        await websocket.send(reply)

async def main():
    async with serve(reply, "192.168.0.109", 8765):
        await asyncio.Future()  # run forever

chatbot = ChatBot()
asyncio.run(main())

    async def chat(self, input):
        output = self.agent_executor.run(input)
        if len(self.images) > 0:
            return f"{{'output': {output}, 'images': {self.images}}}"
        return output

[pr0mila]

LangChain WebSocket streaming often lags or breaks under real-time load.
GroqStreamChain fixes that with a fully async FastAPI backend and smooth token-by-token streaming from Groq.

No delays. No dropped connections. Just fast, real-time AI chats.
Check it out: https://github.com/pr0mila/GroqStreamChain