Refactor BasicAgent to support async processing and rate limiting; implement concurrent question handling.

app.py

@@ -1,11 +1,35 @@
 import os
 import gradio as gr
 import requests
-import inspect
 import pandas as pd
+from langchain_openai import ChatOpenAI
+from os import getenv
+from dotenv import load_dotenv
+from typing import Annotated
+from pydantic import SecretStr
+
+from typing_extensions import TypedDict
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+import asyncio  # Added for async processing
+import time  # Added for rate limiting
+
+from langchain_community.tools import WikipediaQueryRun
+from langchain_community.utilities.wikipedia import WikipediaAPIWrapper
+from langgraph.prebuilt import tools_condition
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import create_react_agent
+
 # Phoenix imports
 from phoenix.otel import register
 import logging
+
+from agent.agent import get_agent
+from agent.config import create_agent_config
+
+load_dotenv()
+
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 APP_NAME = "OracleBot"
@@ -25,16 +49,85 @@ start_phoenix()
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
+
+# class State(TypedDict):
+#     # Messages have the type "list". The `add_messages` function
+#     # in the annotation defines how this state key should be updated
+#     # (in this case, it appends messages to the list, rather than overwriting them)
+#     messages: Annotated[list, add_messages]
+
 class BasicAgent:
     def __init__(self):
-        print("BasicAgent initialized.")
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        fixed_answer = "This is a default answer."
-        print(f"Agent returning fixed answer: {fixed_answer}")
-        return fixed_answer
-
-def run_and_submit_all( profile: gr.OAuthProfile | None):
+        self.agent = get_agent()
+        self._last_request_time = 0
+        self._request_lock = asyncio.Lock()
+
+    async def __call__(self, question: str) -> str:
+        print(f"Agent received question: {question}")
+        
+        # Rate limiting: ensure at least 1 second between requests
+        async with self._request_lock:
+            current_time = time.time()
+            time_since_last_request = current_time - self._last_request_time
+            if time_since_last_request < 1.0:
+                sleep_time = 1.0 - time_since_last_request
+                print(f"Rate limiting: sleeping for {sleep_time:.2f} seconds")
+                await asyncio.sleep(sleep_time)
+            
+            self._last_request_time = time.time()
+        
+        # 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": [{"role": "user", "content": question}]},
+            config=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):
+    print(f"Running agent on {len(questions_data)} questions concurrently (simple fan-out)...")
+
+    async def handle(item: dict):
+        task_id = item.get("task_id")
+        question_text = item.get("question")
+        if not task_id or question_text is None:
+            print(f"Skipping item with missing task_id or question: {item}")
+            return None
+        try:
+            submitted_answer = await agent(question_text)
+            return {
+                "log": {"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer},
+                "payload": {"task_id": task_id, "submitted_answer": submitted_answer},
+            }
+        except Exception as e:
+            print(f"Error running agent on task {task_id}: {e}")
+            return {
+                "log": {"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"},
+                "payload": {"task_id": task_id, "submitted_answer": f"AGENT ERROR: {e}"},
+            }
+
+    results = await asyncio.gather(*(handle(item) for item in questions_data))
+    results_log = [r["log"] for r in results if r]
+    answers_payload = [r["payload"] for r in results if r]
+    return results_log, answers_payload
+
+async def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
     Fetches all questions, runs the BasicAgent on them, submits all answers,
     and displays the results.
@@ -73,34 +166,24 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
              print("Fetched questions list is empty.")
              return "Fetched questions list is empty or invalid format.", None
         print(f"Fetched {len(questions_data)} questions.")
-    except requests.exceptions.RequestException as e:
-        print(f"Error fetching questions: {e}")
-        return f"Error fetching questions: {e}", None
     except requests.exceptions.JSONDecodeError as e:
          print(f"Error decoding JSON response from questions endpoint: {e}")
          print(f"Response text: {response.text[:500]}")
          return f"Error decoding server response for questions: {e}", None
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching questions: {e}")
+        return f"Error fetching questions: {e}", None
     except Exception as e:
         print(f"An unexpected error occurred fetching questions: {e}")
         return f"An unexpected error occurred fetching questions: {e}", None
 
-    # 3. Run your Agent
-    results_log = []
-    answers_payload = []
-    print(f"Running agent on {len(questions_data)} questions...")
-    for item in questions_data:
-        task_id = item.get("task_id")
-        question_text = item.get("question")
-        if not task_id or question_text is None:
-            print(f"Skipping item with missing task_id or question: {item}")
-            continue
-        try:
-            submitted_answer = agent(question_text)
-            answers_payload.append({"task_id": task_id, "submitted_answer": submitted_answer})
-            results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": submitted_answer})
-        except Exception as e:
-             print(f"Error running agent on task {task_id}: {e}")
-             results_log.append({"Task ID": task_id, "Question": question_text, "Submitted Answer": f"AGENT ERROR: {e}"})
+    # 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.")
@@ -208,4 +291,4 @@ if __name__ == "__main__":
     print("-"*(60 + len(" App Starting ")) + "\n")
 
     print("Launching Gradio Interface for Basic Agent Evaluation...")
-    demo.launch(debug=True, share=False)
+    demo.launch(debug=True, share=False)