graph test

app.py

@@ -64,6 +64,7 @@ Last sentence should of your response shoule either:
 - start **EXACTLY** with `FINAL ANSWER: ` followed by the answer.
 - OR start **EXACTLY** with `TOOL: ` followed by the tool name and the request to the tool, example `TOOL: DuckDuckGoSearchAgent(request='search query')`
 
+Before answering reread the question and make sure you are answering the exact question asked.
 """
 
 
@@ -75,6 +76,7 @@ class AnswerState(TypedDict):
     is_final_answer: bool
     search_request: Optional[str]
     wiki_request: Optional[str]
+    python_request: Optional[str]
 
 
 class BasicAgent:
@@ -85,6 +87,7 @@ class BasicAgent:
         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("python_exec", self.python_exec)
         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)
@@ -93,6 +96,7 @@ class BasicAgent:
         self.graph.add_edge("log_question", "invoke_model")
         self.graph.add_edge("web_search", "invoke_model")
         self.graph.add_edge("wiki_search", "invoke_model")
+        self.graph.add_edge("python_exec", "invoke_model")
 
         # Add conditional edges
         self.graph.add_conditional_edges(
@@ -103,6 +107,7 @@ class BasicAgent:
                 "EXCEEDED_ATTEMPTS": "exceeded_attempts",
                 "WEB_SEARCH": "web_search",
                 "WIKI_SEARCH": "wiki_search",
+                "PYTHON_EXEX": "python_exec",
             },
         )
 
@@ -122,6 +127,7 @@ class BasicAgent:
                 "is_final_answer": False,
                 "search_request": None,
                 "wiki_request": None,
+                "python_request": None,
             }
         )
         return res.get("answer", "N/A")
@@ -175,6 +181,31 @@ class BasicAgent:
             )
         return {"messages": state["messages"]}
 
+    def python_exec(self, state: AnswerState) -> Dict[str, Any]:
+        print("[python_exec] Executing: " + str(state["python_request"]))
+        import io
+        import contextlib
+
+
+
+        # Redirect stdout to capture output
+        output = io.StringIO()
+        sandbox_globals = {}
+        sandbox_locals = {}
+
+        with contextlib.redirect_stdout(output):
+            exec(str(state["python_request"]), sandbox_globals, sandbox_locals)
+
+        res = output.getvalue()
+        print(f"Python results: {res}")
+        state["messages"].append(
+            ChatMessage(
+                role="user",
+                content="CodeExecutorAgent tool results are: " + res,
+            )
+        )
+        return {"messages": state["messages"]}
+
     def web_search(self, state: AnswerState) -> Dict[str, Any]:
         print("[web_search] Searching for: " + str(state["search_request"]))
         search_tool = GoogleSearchAPIWrapper()
@@ -223,12 +254,16 @@ class BasicAgent:
             return "WEB_SEARCH"
         elif state["wiki_request"]:
             return "WIKI_SEARCH"
+        elif state["python_request"]:
+            return "PYTHON_EXEC"
         return "FINAL_ANSWER"
 
+
     def invoke_model(self, state: AnswerState) -> Dict[str, Any]:
         print("[invoke_model] Agent trying to answer")
         state["search_request"] = None
         state["wiki_request"] = None
+        state["python_request"] = None
 
         response = self.model.invoke(state["messages"])
         print(f"Agent response: {response.content}")
@@ -259,6 +294,13 @@ class BasicAgent:
             request = request_full.split("request=")[-1].strip("'")
             print(f"Wiki request: {request}")
             state["wiki_request"] = request
+
+        elif "TOOL: CodeExecutorAgent" 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"Python request: {request}")
+            state["python_request"] = request
         return {
             "messages": state["messages"],
             "attempt": state.get("attempt", 0) + 1,