Updates

app.py

@@ -9,6 +9,12 @@ from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 import uvicorn
 import json
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_openai import ChatOpenAI
+from langchain_core.tools import tool
+from langchain.agents import AgentExecutor, create_openai_functions_agent
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.schema import SystemMessage
 
 # LangGraph
 from langgraph.graph import END, StateGraph
@@ -60,12 +66,25 @@ def generate_image_prompt(description: str) -> str:
     return f"Generated image prompt: {enhanced_prompt}"
 
 
+# Create tools list
+tools = [search_web, calculate, generate_image_prompt]
+
 # Set up the language model
 model = ChatOpenAI(temperature=0.5)
 
-# Create tools list
-tools = [search_web, calculate, generate_image_prompt]
-tool_executor = create_agent_executor(tools)
+# Create the prompt template
+prompt = ChatPromptTemplate.from_messages([
+    SystemMessage(content="You are a helpful AI assistant with access to tools. Use them when appropriate."),
+    MessagesPlaceholder(variable_name="chat_history"),
+    ("human", "{input}"),
+    MessagesPlaceholder(variable_name="agent_scratchpad"),
+])
+
+# Create the agent
+agent = create_openai_functions_agent(model, tools, prompt)
+
+# Create the agent executor
+agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
 
 
 # Define agent nodes
@@ -123,7 +142,7 @@ def build_agent_graph():
     
     # Add the nodes
     graph.add_node("agent", create_agent_node())
-    graph.add_node("tool", ToolNode(tool_executor))
+    graph.add_node("tool", ToolNode(agent_executor))
     
     # Add the edges
     graph.add_conditional_edges(
@@ -157,32 +176,14 @@ agent_executor = build_agent_graph()
 async def websocket_endpoint(websocket: WebSocket):
     await websocket.accept()
     try:
-        # Initialize the state
-        state = {
-            "messages": [],
-            "tools": tools,
-        }
-        
         while True:
             data = await websocket.receive_text()
-            # Add the user message to the state
-            state["messages"].append(HumanMessage(content=data))
-            
-            # Run the agent
-            result = agent_executor.invoke(state)
-            
-            # Get the final AI message
-            for message in result["messages"]:
-                if isinstance(message, AIMessage):
-                    await websocket.send_json({"type": "ai_message", "content": message.content})
-                elif isinstance(message, ToolMessage):
-                    await websocket.send_json({"type": "tool_message", "content": message.content})
-            
-            # Update the state
-            state = result
-    
-    except WebSocketDisconnect:
-        print("Client disconnected")
+            # Process the message with the agent
+            response = agent_executor.invoke({"input": data, "chat_history": []})
+            await websocket.send_json({"type": "ai_message", "content": response["output"]})
+    except Exception as e:
+        print(f"Error in WebSocket: {str(e)}")
+        await websocket.close()
 
 
 # Serve the HTML frontend

requirements.txt

@@ -2,7 +2,8 @@ fastapi>=0.109.0
 uvicorn>=0.27.0
 python-dotenv>=1.0.0
 openai>=1.12.0
-langchain-core>=0.2.0
+langchain>=0.1.0
+langchain-core>=0.1.0
 langgraph>=0.1.0
 langchain-openai>=0.0.5
 websockets>=12.0