added final graph function

agent.py

@@ -141,39 +141,38 @@ retriever_tool = create_retriever_tool(
     description="Retrieve similar questions from a vector store."
 )
 
-## ----- LLM WITH TOOLS ----- ##
-
-llm = ChatOpenAI(model="gpt-4o", temperature=0)
-llm_with_tools = llm.bind_tools(tools)
-
 ## ----- GRAPH PIPELINE ----- ##
 
-def assistant(state: MessagesState):
-    """Assistant node to generate answers."""
-    return {"messages": [llm_with_tools.invoke(state["messages"])]}
-
-# Use a retriever node to inject a similar example
-def retriever(state: MessagesState):
-    """Retriever node to provide example context."""
-    similar = vector_store.similarity_search(state["messages"][0].content)
-    if not similar:
-        return {"messages": [sys_msg] + state["messages"]}
-    example = HumanMessage(content=f"Similar Q&A for context:\n\n{similar[0].page_content}")
-    return {"messages": [sys_msg] + state["messages"] + [example]}
-
-# 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", tools_condition)
-builder.add_edge("tools", "assistant")
-
-graph = builder.compile()
-return graph
+def graph():
+    llm = ChatOpenAI(model="gpt-4o", temperature=0)
+    llm_with_tools = llm.bind_tools(tools)
+
+    def assistant(state: MessagesState):
+        """Assistant node to generate answers."""
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+    
+    # Use a retriever node to inject a similar example
+    def retriever(state: MessagesState):
+        """Retriever node to provide example context."""
+        similar = vector_store.similarity_search(state["messages"][0].content)
+        if not similar:
+            return {"messages": [sys_msg] + state["messages"]}
+        example = HumanMessage(content=f"Similar Q&A for context:\n\n{similar[0].page_content}")
+        return {"messages": [sys_msg] + state["messages"] + [example]}
+    
+    # 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", tools_condition)
+    builder.add_edge("tools", "assistant")
+    
+    return builder.compile()
+
 
 ## ----- TESTING (Optional) ----- ##