Update app.py

app.py

@@ -7,7 +7,10 @@ import re
 import time
 import base64
 import json
+import sys
+import contextlib
 #from tavily import TavilyClient
+from io import StringIO
 from langchain_tavily import TavilySearch
 from langgraph.prebuilt import create_react_agent
 from langgraph.graph.message import add_messages
@@ -53,6 +56,7 @@ prompt_recomendado = """You are a general AI assistant. I will ask you a questio
                         To assist in your task, you can supervise other agents who perform specific tasks that could not be handled by tools, since they require the processing of another LLM. Below, I will inform you about your assistants:
                         - web_research_agent. Assign web research related tasks to this agent, prioritizing the use of Wikipedia sources
                         - chess_position_review_agent. Assign chess position review related tasks to this agent
+                        - python_code_runner_agent. Assign python code execution related tasks to this agent
                         Assign work to one agent at a time, do not call agents in parallel.
                         Priorize the use of tools and another agents to help in reasoning.
                         When a file or URL is entered at the prompt, use it in tools or other agents, both are prepared to handle files and URLs."""
@@ -70,12 +74,48 @@ prompt_chess = """You are a chess position reviewing agent.
                    After you're done with your tasks, respond to the supervisor directly
                    Respond ONLY with the results of your work, do NOT include ANY other text."""
 
+prompt_python_execute = """You are a python code execution agent.
+                           INSTRUCTIONS:
+                           Assist ONLY with tasks related to running python code, DO NOT do any math
+                           After you're done with your tasks, respond to the supervisor directly
+                           Respond ONLY with the results of your work, do NOT include ANY other text."""
+
 #TOOLS
 web_search = TavilySearch(
         max_results=5,
         topic="general",
     )
 
+@contextlib.contextmanager
+def stdoutIO(stdout=None):
+    old = sys.stdout
+    if stdout is None:
+        stdout = StringIO()
+    sys.stdout = stdout
+    yield stdout
+    sys.stdout = old
+
+def python_code_runner_tool(task_id:str) -> str:
+    """
+        Download and run python code, capturing the output.
+        Args:
+            task_id: Task ID necessary to retrieve the python code to be run.
+        Returns:
+            String with the output of the python code.
+    """
+    print(f"python code runner invocada com os seguintes parametros:")
+    print(f"task_id: {task_id}")
+
+    python_code = download_file_as_string(task_id)
+
+    with stdoutIO() as s:
+        exec(python_code)
+
+    output = s.getvalue()
+    print(f"Captured output: {output}") 
+
+    return output
+
 def chess_image_to_fen_tool(task_id:str, current_player: Literal["black", "white"]) -> Dict[str,str]:
     """
         Convert chess image to FEN (Forsyth-Edwards Notation) notation.
@@ -214,6 +254,21 @@ def download_file_as_base64(task_id: str) -> str:
         return encoded_str
     else:
         raise Exception(f"Failed to download the file. Status code: {response.status_code}")
+
+def download_file_as_string(task_id: str) -> str:
+    # Construct the URL
+    url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
+    
+    # Send the request to download the file
+    response = requests.get(url)
+    
+    if response.status_code == 200:
+        # Encode the content to Base64
+        bytes = response.content
+        encoded_str = bytes.decode('utf-8')  # Convert bytes to string
+        return encoded_str
+    else:
+        raise Exception(f"Failed to download the file. Status code: {response.status_code}")        
         
 
 
@@ -247,6 +302,13 @@ chess_position_review_agent = create_react_agent(
     name="chess_position_review_agent"
 )
 
+python_code_runner_agent = create_react_agent(
+    model=gemini_llm,
+    tools=[python_code_runner_tool],
+    prompt=prompt_python_execute,
+    name="python_code_runner_agent"
+)
+
 supervisor = create_supervisor(
     model=gemini_llm,
     agents=[web_research_agent,chess_position_review_agent],