Update agent.py

agent.py

@@ -41,6 +41,13 @@ def build_graph(provider: str = "openai"):
         raise ValueError("Invalid provider.")
 
     llm_with_tools = llm.bind_tools(tools)
+    def assistant(state: MessagesState):
+        user_message = state["messages"][-1]
+        # You must have llm_with_tools defined earlier in build_graph
+        result = llm_with_tools.invoke([sys_msg, user_message])
+        return {"messages": [result]}
+
+    
 
     def retriever(state: MessagesState):
         query = state["messages"][-1].content
@@ -55,11 +62,6 @@ def build_graph(provider: str = "openai"):
         else:
             return {"messages": [AIMessage(content=content.strip())]}
 
-    def assistant(state: MessagesState):
-        user_message = state["messages"][-1]
-        # Make sure you send both system and user message
-        result = llm_with_tools.invoke([sys_msg, user_message])
-        return {"messages": [result]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
@@ -67,5 +69,6 @@ def build_graph(provider: str = "openai"):
     builder.add_edge(START, "retriever")
     builder.add_edge("retriever", "assistant")
     builder.set_finish_point("assistant")
+    return builder.compile()
 
     return builder.compile()