Update agent.py

agent.py

@@ -181,43 +181,44 @@ tools = [
     arvix_search,
     create_retriever_tool
 ]
-
-llm = init_chat_model("google_genai:gemini-2.0-flash",google_api_key=os.environ["GOOGLE_API_KEY"])
-llm_with_tools = llm.bind_tools(tools)
-
-sys_msg = SystemMessage(content=system_prompt)
-
-class MessagesState(TypedDict):
-  messages: Annotated[list[AnyMessage], add_messages]
-# Node
-def assistant(state: MessagesState):
-    """Assistant node"""
-    return {"messages": [llm_with_tools.invoke(state["messages"])]}
-def retriever(state: MessagesState):
-        """Retriever node"""
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
-        example_msg = HumanMessage(
-            content=f"Here I provide a question and answer using query for reference if it is similar to question below:  \n\n{similar_question[0].page_content}",
-        )
-        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-
-# Build graph
-builder = StateGraph(MessagesState)
-builder.add_node("retriever", retriever)
-builder.add_node("assistant", assistant)
-builder.add_node("tools", ToolNode(tools))
-builder.add_edge(START, "retriever")
-builder.add_edge("retriever", "assistant")
-builder.add_conditional_edges(
-    "assistant",
-    # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
-    # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
-    tools_condition,
-)
-builder.add_edge("tools", "assistant")
-
-# Compile graph
-graph = builder.compile()
+def build_graph()
+    """Build the graph"""
+    llm = init_chat_model("google_genai:gemini-2.0-flash",google_api_key=os.environ["GOOGLE_API_KEY"])
+    llm_with_tools = llm.bind_tools(tools)
+    
+    sys_msg = SystemMessage(content=system_prompt)
+    
+    class MessagesState(TypedDict):
+      messages: Annotated[list[AnyMessage], add_messages]
+    # Node
+    def assistant(state: MessagesState):
+        """Assistant node"""
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+    def retriever(state: MessagesState):
+            """Retriever node"""
+            similar_question = vector_store.similarity_search(state["messages"][0].content)
+            example_msg = HumanMessage(
+                content=f"Here I provide a question and answer using query for reference if it is similar to question below:  \n\n{similar_question[0].page_content}",
+            )
+            return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+    
+    # Build graph
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges(
+        "assistant",
+        # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+        # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
+        tools_condition,
+    )
+    builder.add_edge("tools", "assistant")
+    
+    # Compile graph
+    return builder.compile()
 
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"