Update agent.py

agent.py

@@ -3,7 +3,6 @@ from dotenv import load_dotenv
 from langgraph.graph import START, END, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
@@ -190,11 +189,13 @@ def build_graph(provider: str = "openai"):
     # Define nodes
     def assistant(state: MessagesState):
         """Assistant node - generates responses"""
-        messages = state["messages"]
+        messages = llm_with_tools.invoke(state["messages"])
         # Generate response using LLM
-        response = llm_with_tools.invoke(messages)
+        # response = llm_with_tools.invoke(messages)
         # Return new state with appended message
-        return {"messages": messages + [response]}
+        return {"messages": messages}
+
+
     def retriever(state: MessagesState):
         """Retriever node - provides context from vector store"""
         messages = state["messages"]
@@ -215,26 +216,13 @@ def build_graph(provider: str = "openai"):
     builder = StateGraph(MessagesState)
     
     # Add nodes
-    builder.add_node("retriever", retriever)
+    # builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
     
     # Set up edges
-    builder.set_entry_point("retriever")
-    builder.add_edge("retriever", "assistant")
-    
-    # Conditional tool usage
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-        {
-            "continue": "tools",  # If tools needed
-            "end": END            # If no tools needed
-        }
-    )
-    
-    # After tools execute, go back to assistant
-    builder.add_edge("tools", "assistant")
+    builder.set_entry_point("assistant")
+    builder.set_finish_point("assistant")
     
     # Compile graph
     return builder.compile()