Update agent.py

agent.py

@@ -145,7 +145,7 @@ def wiki_search(query: str) -> str:
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
             for doc in search_docs
         ])
-    return {"wiki_results": formatted_search_docs}
+    return formatted_search_docs
 
 
     
@@ -175,7 +175,7 @@ def web_search(query: str) -> str:
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
             for doc in search_docs
         ])
-    return {"web_results": formatted_search_docs}
+    return formatted_search_docs
 
 @tool
 def arvix_search(query: str) -> str:
@@ -189,7 +189,7 @@ def arvix_search(query: str) -> str:
             f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
             for doc in search_docs
         ])
-    return {"arvix_results": formatted_search_docs}
+    return formatted_search_docs
 
 
 
@@ -414,26 +414,6 @@ question_retriever_tool = create_retriever_tool(
 )
 
 
-# -------------------------------
-# Step 6: Create LangChain Tools
-# -------------------------------
-calc_tool = calculator
-file_tool = analyze_attachment
-web_tool = web_search
-wiki_tool = wiki_search
-arvix_tool = arvix_search
-youtube_tool = get_youtube_transcript
-video_tool = extract_video_id
-analyze_tool = analyze_attachment
-wikiq_tool = wikidata_query
-
-
-# -------------------------------
-# Step 7: Create the Planner-Agent Logic
-# -------------------------------
-
-# Define the tools (as you've already done)
-tools = [wiki_tool, calc_tool, file_tool, web_tool, arvix_tool, youtube_tool, video_tool, analyze_tool, wikiq_tool]
 
 # Define the LLM before using it
 #llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo") # or "gpt-3.5-turbo" "gpt-4"
@@ -693,8 +673,7 @@ model_config = {
 }
 
 def build_graph(provider, model_config):
-#    from langgraph.prebuilt.tool_node import ToolNode
-
+    # Step 1: Initialize the LLM
     def get_llm(provider: str, config: dict):
         if provider == "huggingface":
             from langchain_huggingface import HuggingFaceEndpoint
@@ -707,24 +686,53 @@ def build_graph(provider, model_config):
             )
         else:
             raise ValueError(f"Unsupported provider: {provider}")
-            
-
+    
     llm = get_llm(provider, model_config)
-    return llm
     
+    # -------------------------------
+    # Step 6: Define LangChain Tools
+    # -------------------------------
+    calc_tool = calculator  # Math operations tool
+    web_tool = web_search    # Web search tool
+    wiki_tool = wiki_search  # Wikipedia search tool
+    arvix_tool = arvix_search  # Arxiv search tool
+    youtube_tool = get_youtube_transcript  # YouTube transcript extraction
+    video_tool = extract_video_id  # Video ID extraction tool
+    analyze_tool = analyze_attachment  # File analysis tool
+    wikiq_tool = wikidata_query  # Wikidata query tool
+    
+    # -------------------------------
+    # Step 7: Create the Planner-Agent Logic
+    # -------------------------------
+    # Define tools list
+    tools = [
+        wiki_tool, 
+        calc_tool, 
+        web_tool, 
+        arvix_tool,
+        youtube_tool, 
+        video_tool, 
+        analyze_tool, 
+        wikiq_tool
+    ]
+    
+    # Step 8: Bind tools to the LLM
     llm_with_tools = llm.bind_tools(tools)
 
+    # Return the LLM with tools bound
+    return llm_with_tools
 
 
-    sys_msg = SystemMessage(content="You are a helpful assistant.")
-
-    def assistant(state: MessagesState):
-        return {"messages": [llm_with_tools.invoke(state["messages"])]}
 
+    
+    # Initialize system message
+    sys_msg = SystemMessage(content="You are a helpful assistant.")
+    
+    # Define the retriever function
     def retriever(state: MessagesState):
         user_query = state["messages"][0].content
         similar_docs = vector_store.similarity_search(user_query)
-
+    
         if not similar_docs:
             wiki_result = wiki_tool.run(user_query)
             return {
@@ -742,18 +750,30 @@ def build_graph(provider, model_config):
                     HumanMessage(content=f"Reference:\n\n{similar_docs[0].page_content}")
                 ]
             }
-
+    
+    # Define the assistant function
+    def assistant(state: MessagesState):
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+    
+    # Define condition for tools usage
     def tools_condition(state: MessagesState) -> str:
         if "use tool" in state["messages"][-1].content.lower():
             return "tools"
         else:
-            return END
-
+            return "END"
+    
+    # Initialize the StateGraph
     builder = StateGraph(MessagesState)
+    
+    # Add nodes to the graph
     builder.add_node("retriever", retriever)
     builder.add_node("assistant", assistant)
     builder.add_node("tools", ToolNode(tools))
+    
+    # Set the entry point
     builder.set_entry_point("retriever")
+    
+    # Define edges
     builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")