Enhance OracleBot to accept optional file path for context in answers; add utility to fetch task files from API.

agent/agent.py

@@ -1,3 +1,4 @@
+import os
 from typing import Literal
 from typing_extensions import TypedDict
 from langgraph.graph import StateGraph, START, END
@@ -17,10 +18,18 @@ class OracleBot:
         self.config = create_agent_config(self.name, self.thread_id)
         self.graph = self._build_agent(self.name)
 
-    def answer_question(self, question: str):
+    def answer_question(self, question: str, file_path: str | None = None):
         """
         Answer a question using the LangGraph agent.
+        
+        Args:
+            question: The question to answer
+            file_path: Optional path to a file associated with this question
         """
+        # Enhance question with file context if available
+        if file_path and os.path.exists(file_path):
+            question = f"{question}\n\nNote: There is an associated file at: {file_path}\nYou can use the file management tools to read and analyze this file."
+        
         messages = [HumanMessage(content=question)]
 
         for mode, chunk in self.graph.stream({"messages": messages}, config=self.config, stream_mode=["messages", "updates"]): # type: ignore
@@ -48,7 +57,8 @@ class OracleBot:
                             # Handle final answer messages (no tool calls)
                             elif hasattr(message, 'content') and message.content:
                                 cprint(f"\n{message.content}\n", color="black", on_color="on_white", attrs=["bold"])
-                
+                                return message.content # Return final answer
+
                 # Look for tool outputs in updates
                 elif isinstance(chunk, dict) and 'tools' in chunk:
                     tools_update = chunk['tools']
@@ -57,6 +67,36 @@ class OracleBot:
                             if hasattr(message, 'content') and message.content:
                                 cprint(f"\n📤 Tool output:\n{message.content}\n", color="green")
 
+    async def answer_question_async(self, question: str, file_path: str | None = None) -> str:
+        """
+        Answer a question using the LangGraph agent asynchronously.
+        
+        Args:
+            question: The question to answer
+            file_path: Optional path to a file associated with this question
+            
+        Returns the final answer as a string.
+        """
+        from langchain_core.runnables import RunnableConfig
+        from typing import cast
+        
+        # Enhance question with file context if available
+        if file_path and os.path.exists(file_path):
+            question = f"{question}\n\nNote: There is an associated file at: {file_path}\nYou can use the file management tools to read and analyze this file."
+        
+        messages = [HumanMessage(content=question)]
+
+        # Use LangGraph's built-in ainvoke method
+        result = await self.graph.ainvoke({"messages": messages}, config=cast(RunnableConfig, self.config)) # type: ignore
+        
+        # Extract the content from the last message
+        if "messages" in result and result["messages"]:
+            last_message = result["messages"][-1]
+            if hasattr(last_message, 'content'):
+                return last_message.content or ""
+        
+        return ""
+
     def _build_agent(self, name: str):
         """
         Get our LangGraph agent with the given model and tools.
@@ -77,10 +117,7 @@ class OracleBot:
         graph.add_conditional_edges("agent", tools_condition)
         graph.add_edge("tools", "agent")
 
-        # Add memory
-        memory = InMemorySaver()
-            
-        return graph.compile(checkpointer=memory)
+        return graph.compile()
 
 # test
 if __name__ == "__main__":
@@ -92,7 +129,7 @@ if __name__ == "__main__":
         from config import start_phoenix
         start_phoenix()
         bot = OracleBot()
-        bot.answer_question(question)
+        bot.answer_question(question, None)
 
     except Exception as e:
         print(f"Error running agent: {e}")

app.py

@@ -4,14 +4,12 @@ import requests
 import pandas as pd
 from os import getenv
 from dotenv import load_dotenv
-from langchain_core.messages import HumanMessage
-from langchain_core.runnables import RunnableConfig
 import asyncio
-from typing import cast
+import tempfile
 
 from agent.agent import OracleBot
-from agent.config import create_agent_config
 from config import start_phoenix, APP_NAME, DEFAULT_API_URL
+from utils import fetch_task_file, extract_task_id_from_question_data
 
 load_dotenv()
 
@@ -28,38 +26,8 @@ start_phoenix()
 #     # (in this case, it appends messages to the list, rather than overwriting them)
 #     messages: Annotated[list, add_messages]
 
-class BasicAgent:
-    def __init__(self):
-        self.agent = OracleBot()
-
-    async def __call__(self, question: str) -> str:
-        print(f"Agent received question: {question}")
-        
-        # Create configuration like in main.py
-        config = create_agent_config(app_name=APP_NAME)
-        
-        # Call the agent with the question and config (like main.py)
-        answer = await self.agent.ainvoke(
-            {"messages": [HumanMessage(content=question)]},
-            cast(RunnableConfig, config)
-        )
-
-        print(f"Agent returning answer: {answer}")
-        
-        # Extract content from the last message in the response
-        if "messages" in answer and answer["messages"]:
-            last_message = answer["messages"][-1]
-            if hasattr(last_message, 'content'):
-                content = last_message.content
-            else:
-                content = str(last_message)
-        else:
-            content = str(answer)
-            
-        return str(content) if content is not None else ""
-
 # Simplified concurrent processor: launch all tasks immediately and await them together
-async def process_questions(agent: BasicAgent, questions_data: list):
+async def process_questions(agent: OracleBot, questions_data: list, working_dir: str):
     print(f"Running agent on {len(questions_data)} questions concurrently (simple fan-out)...")
 
     async def handle(item: dict):
@@ -69,7 +37,16 @@ async def process_questions(agent: BasicAgent, questions_data: list):
             print(f"Skipping item with missing task_id or question: {item}")
             return None
         try:
-            submitted_answer = await agent(question_text)
+            # Fetch associated file if it exists
+            file_path = fetch_task_file(task_id, working_dir)
+            if file_path:
+                print(f"Found file for task {task_id}: {file_path}")
+            
+            # Pass file_path to agent
+            submitted_answer = await agent.answer_question_async(question_text, file_path)
+            # Extract everything after "FINAL ANSWER: "
+            if "FINAL ANSWER: " in submitted_answer:
+                submitted_answer = submitted_answer.split("FINAL ANSWER: ", 1)[-1].strip()
             return {
                 "log": {"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer},
                 "payload": {"task_id": task_id, "submitted_answer": submitted_answer},
@@ -107,7 +84,7 @@ async def run_and_submit_all( profile: gr.OAuthProfile | None):
 
     # 1. Instantiate Agent ( modify this part to create your agent)
     try:
-        agent = BasicAgent()
+        agent = OracleBot()
     except Exception as e:
         print(f"Error instantiating agent: {e}")
         return f"Error initializing agent: {e}", None
@@ -137,64 +114,66 @@ async def run_and_submit_all( profile: gr.OAuthProfile | None):
         return f"An unexpected error occurred fetching questions: {e}", None
 
     # 3. Run your Agent concurrently (simple gather)
-    results_log, answers_payload = await process_questions(agent, questions_data)
-
-    # Remove everything before "FINAL ANSWER: " in submitted answers
-    for answer in answers_payload:
-        if "submitted_answer" in answer:
-            answer["submitted_answer"] = answer["submitted_answer"].split("FINAL ANSWER: ", 1)[-1].strip()
-
-    if not answers_payload:
-        print("Agent did not produce any answers to submit.")
-        return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
-
-    # 4. Prepare Submission 
-    submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
-    status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
-    print(status_update)
-
-    # 5. Submit
-    print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
-    try:
-        response = requests.post(submit_url, json=submission_data, timeout=60)
-        response.raise_for_status()
-        result_data = response.json()
-        final_status = (
-            f"Submission Successful!\n"
-            f"User: {result_data.get('username')}\n"
-            f"Overall Score: {result_data.get('score', 'N/A')}% "
-            f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
-            f"Message: {result_data.get('message', 'No message received.')}"
-        )
-        print("Submission successful.")
-        results_df = pd.DataFrame(results_log)
-        return final_status, results_df
-    except requests.exceptions.HTTPError as e:
-        error_detail = f"Server responded with status {e.response.status_code}."
+    # Create a temporary working directory for this session
+    with tempfile.TemporaryDirectory() as working_dir:
+        results_log, answers_payload = await process_questions(agent, questions_data, working_dir)
+
+        # Remove everything before "FINAL ANSWER: " in submitted answers
+        for answer in answers_payload:
+            if "submitted_answer" in answer:
+                answer["submitted_answer"] = answer["submitted_answer"].split("FINAL ANSWER: ", 1)[-1].strip()
+
+        if not answers_payload:
+            print("Agent did not produce any answers to submit.")
+            return "Agent did not produce any answers to submit.", pd.DataFrame(results_log)
+
+        # 4. Prepare Submission 
+        submission_data = {"username": username.strip(), "agent_code": agent_code, "answers": answers_payload}
+        status_update = f"Agent finished. Submitting {len(answers_payload)} answers for user '{username}'..."
+        print(status_update)
+
+        # 5. Submit
+        print(f"Submitting {len(answers_payload)} answers to: {submit_url}")
         try:
-            error_json = e.response.json()
-            error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
-        except requests.exceptions.JSONDecodeError:
-            error_detail += f" Response: {e.response.text[:500]}"
-        status_message = f"Submission Failed: {error_detail}"
-        print(status_message)
-        results_df = pd.DataFrame(results_log)
-        return status_message, results_df
-    except requests.exceptions.Timeout:
-        status_message = "Submission Failed: The request timed out."
-        print(status_message)
-        results_df = pd.DataFrame(results_log)
-        return status_message, results_df
-    except requests.exceptions.RequestException as e:
-        status_message = f"Submission Failed: Network error - {e}"
-        print(status_message)
-        results_df = pd.DataFrame(results_log)
-        return status_message, results_df
-    except Exception as e:
-        status_message = f"An unexpected error occurred during submission: {e}"
-        print(status_message)
-        results_df = pd.DataFrame(results_log)
-        return status_message, results_df
+            response = requests.post(submit_url, json=submission_data, timeout=60)
+            response.raise_for_status()
+            result_data = response.json()
+            final_status = (
+                f"Submission Successful!\n"
+                f"User: {result_data.get('username')}\n"
+                f"Overall Score: {result_data.get('score', 'N/A')}% "
+                f"({result_data.get('correct_count', '?')}/{result_data.get('total_attempted', '?')} correct)\n"
+                f"Message: {result_data.get('message', 'No message received.')}"
+            )
+            print("Submission successful.")
+            results_df = pd.DataFrame(results_log)
+            return final_status, results_df
+        except requests.exceptions.HTTPError as e:
+            error_detail = f"Server responded with status {e.response.status_code}."
+            try:
+                error_json = e.response.json()
+                error_detail += f" Detail: {error_json.get('detail', e.response.text)}"
+            except requests.exceptions.JSONDecodeError:
+                error_detail += f" Response: {e.response.text[:500]}"
+            status_message = f"Submission Failed: {error_detail}"
+            print(status_message)
+            results_df = pd.DataFrame(results_log)
+            return status_message, results_df
+        except requests.exceptions.Timeout:
+            status_message = "Submission Failed: The request timed out."
+            print(status_message)
+            results_df = pd.DataFrame(results_log)
+            return status_message, results_df
+        except requests.exceptions.RequestException as e:
+            status_message = f"Submission Failed: Network error - {e}"
+            print(status_message)
+            results_df = pd.DataFrame(results_log)
+            return status_message, results_df
+        except Exception as e:
+            status_message = f"An unexpected error occurred during submission: {e}"
+            print(status_message)
+            results_df = pd.DataFrame(results_log)
+            return status_message, results_df
 
 
 # --- Build Gradio Interface using Blocks ---

utils.py

@@ -0,0 +1,68 @@
+import os
+import requests
+import tempfile
+from pathlib import Path
+from config import DEFAULT_API_URL
+
+
+def fetch_task_file(task_id: str, working_dir: str) -> str | None:
+    """
+    Fetch the file associated with a task_id from the API and save it to the working directory.
+    
+    Args:
+        task_id: The task ID to fetch the file for
+        working_dir: The working directory to save the file to
+        
+    Returns:
+        The path to the downloaded file, or None if no file exists or error occurred
+    """
+    try:
+        files_url = f"{DEFAULT_API_URL}/files/{task_id}"
+        response = requests.get(files_url, timeout=30)
+        
+        if response.status_code == 404:
+            # No file associated with this task
+            return None
+        elif response.status_code == 200:
+            # Try to determine filename from content-disposition header
+            filename = f"task_{task_id}_file"
+            if 'content-disposition' in response.headers:
+                content_disp = response.headers['content-disposition']
+                if 'filename=' in content_disp:
+                    filename = content_disp.split('filename=')[1].strip('"')
+            
+            # If content type suggests a specific extension
+            content_type = response.headers.get('content-type', '')
+            if 'json' in content_type and not filename.endswith('.json'):
+                filename += '.json'
+            elif 'text' in content_type and not filename.endswith('.txt'):
+                filename += '.txt'
+            elif 'csv' in content_type and not filename.endswith('.csv'):
+                filename += '.csv'
+            
+            # Save file to working directory
+            file_path = os.path.join(working_dir, filename)
+            with open(file_path, 'wb') as f:
+                f.write(response.content)
+            
+            print(f"Downloaded file for task {task_id}: {file_path}")
+            return file_path
+        else:
+            response.raise_for_status()
+            
+    except Exception as e:
+        print(f"Error fetching file for task {task_id}: {e}")
+        return None
+
+
+def extract_task_id_from_question_data(question_data: dict) -> str | None:
+    """
+    Extract task_id from question data dictionary.
+    
+    Args:
+        question_data: Dictionary containing question information
+        
+    Returns:
+        The task_id if found, None otherwise
+    """
+    return question_data.get("task_id")