graph test

app.py

@@ -7,13 +7,14 @@ from typing import Any, Dict, List, Optional, TypedDict
 import gradio as gr
 import pandas as pd
 import requests
-from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.tools import DuckDuckGoSearchRun, WikipediaQueryRunTool
 from langchain_core import messages
 from langchain_core.messages import (
     BaseMessage,
     ChatMessage,
     HumanMessage,
     SystemMessage,
+    ToolCall,
     ToolMessage,
 )
 from langchain_openai import ChatOpenAI
@@ -69,7 +70,7 @@ class AnswerState(TypedDict):
     messages: List[BaseMessage]
     is_final_answer: bool
     search_request: Optional[str]
-    search_response: Optional[str]
+    wiki_request: Optional[str]
 
 
 class BasicAgent:
@@ -79,6 +80,7 @@ class BasicAgent:
         self.graph.add_node("log_question", self.log_question)
         self.graph.add_node("invoke_model", self.invoke_model)
         self.graph.add_node("web_search", self.web_search)
+        self.graph.add_node("wiki_search", self.wiki_search)
         self.graph.add_node("decide", self.decide)
         self.graph.add_node("final_answer", self.final_answer)
         self.graph.add_node("exceeded_attempts", self.exceeded_attempts)
@@ -95,6 +97,7 @@ class BasicAgent:
                 "FINAL_ANSWER": "final_answer",
                 "EXCEEDED_ATTEMPTS": "exceeded_attempts",
                 "WEB_SEARCH": "web_search",
+                "WIKI_SEARCH": "wiki_search",
             },
         )
 
@@ -113,7 +116,7 @@ class BasicAgent:
                 "messages": [],
                 "is_final_answer": False,
                 "search_request": None,
-                "search_response": None,
+                "wiki_request": None,
             }
         )
         return res.get("answer", "N/A")
@@ -134,15 +137,27 @@ class BasicAgent:
             "is_final_answer": True,
         }
 
+    def wiki_search(self, state: AnswerState) -> Dict[str, Any]:
+        print("[wiki_search] Searching for: " + str(state["search_request"]))
+        wiki_tool = WikipediaQueryRunTool()
+        results = wiki_tool.invoke(str(state["search_request"]))
+        print(f"Wiki results: {results}")
+        state["messages"].append(
+            ChatMessage(
+                role="WikipediaAgent",
+                content=results,
+            )
+        )
+        return {"messages": state["messages"]}
+
     def web_search(self, state: AnswerState) -> Dict[str, Any]:
         print("[web_search] Searching for: " + str(state["search_request"]))
         search_tool = DuckDuckGoSearchRun()
         results = search_tool.invoke(str(state["search_request"]))
         print(f"Search results: {results}")
         state["messages"].append(
-            ToolMessage(
-                tool="DuckDuckGoSearchAgent",
-                tool_call_id=str(state["search_request"])[:20],
+            ChatMessage(
+                role="DuckDuckGoSearchAgent",
                 content=results,
             )
         )
@@ -173,6 +188,7 @@ class BasicAgent:
     def invoke_model(self, state: AnswerState) -> Dict[str, Any]:
         print("[invoke_model] Agent trying to answer")
         state["search_request"] = None
+        state["wiki_request"] = None
 
         response = self.model.invoke(state["messages"])
         print(f"Agent response: {response.content}")
@@ -189,8 +205,7 @@ class BasicAgent:
             print(f"Agent returning answer: {state['answer']}")
             state["answer"] = answer
             state["is_final_answer"] = True
-
-        if "TOOL: DuckDuckGoSearchAgent" in response.content:
+        elif "TOOL: DuckDuckGoSearchAgent" in response.content:
             request_full = (
                 response.text().split("TOOL: DuckDuckGoSearchAgent(")[-1].strip(")")
             )
@@ -198,13 +213,19 @@ class BasicAgent:
             request = request_full.split("request=")[-1].strip("'")
             print(f"Search request: {request}")
             state["search_request"] = request
-
+        elif "TOOL: WikipediaAgent" in response.content:
+            request_full = response.text().split("TOOL: WikipediaAgent(")[-1].strip(")")
+            print(f"Tool invocation request: {request_full}")
+            request = request_full.split("request=")[-1].strip("'")
+            print(f"Wiki request: {request}")
+            state["wiki_request"] = request
         return {
             "messages": state["messages"],
             "attempt": state.get("attempt", 0) + 1,
             "answer": state["answer"],
             "is_final_answer": state["is_final_answer"],
             "search_request": state["search_request"],
+            "wiki_request": state["wiki_request"],
         }