Add some rate limiting, and image support.

app.py

@@ -2,14 +2,17 @@ import asyncio
 import os
 from pathlib import Path
 import gradio as gr
+import mimetypes
 import requests
 import pandas as pd
+from llama_index.core.llms import ChatMessage, TextBlock, ImageBlock, AudioBlock
 from llama_index.llms.gemini import Gemini
-from llama_index.core.agent.workflow import ReActAgent
+from llama_index.core.agent.workflow import ReActAgent, AgentOutput
 from llama_index.core.tools import FunctionTool
 from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec
 from llama_index.tools.wikipedia import WikipediaToolSpec
 from dotenv import load_dotenv
+from pydantic import ValidationError
 
 try:
     import mlflow
@@ -28,7 +31,7 @@ GOOGLE_API_KEY = os.environ['GOOGLE_API_KEY']
 
 # --- Basic Agent Definition ---
 class BasicAgent:
-    def __init__(self):
+    def __init__(self, max_calls_per_minute=15):
         search_tool = FunctionTool.from_defaults(DuckDuckGoSearchToolSpec().duckduckgo_full_search)
         wikipedia_load_tool = FunctionTool.from_defaults(WikipediaToolSpec().load_data)
         wikipedia_search_tool = FunctionTool.from_defaults(WikipediaToolSpec().search_data)
@@ -39,10 +42,17 @@ class BasicAgent:
         # Modify the react prompt.
         self._agent.update_prompts({"react_header": SYSTEM_PROMPT})
         print("BasicAgent initialized.")
-
-    async def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        agent_output = await self._agent.run(user_msg=question)
+        self._min_call_interval = 1/max_calls_per_minute
+
+    async def __call__(self, question: ChatMessage) -> str:
+        question.blocks[0].text
+        print(f"Agent received question (first 50 chars): {question.blocks[0].text[:50]}...")
+        # Here, we need to rate limit
+        handler = self._agent.run(user_msg=question)
+        async for event in handler.stream_events():
+            if isinstance(event, AgentOutput):
+                await asyncio.sleep(self._min_call_interval)
+        agent_output = await handler
         print(f"Agent returning answer: {agent_output}")
         response_parts = str(agent_output).split('FINAL ANSWER: ')
         if len(response_parts) > 1:
@@ -65,14 +75,46 @@ def fetch_questions(api_url: str = DEFAULT_API_URL):
     return questions_data
 
 
+def get_media_type(filename: str):
+    media_type = mimetypes.guess_type(filename)[0]
+    if media_type is not None:
+        return media_type.split('/')[0]
+
+
+def get_media_content(item):
+    if item.get('file_name'):
+        file_response = requests.get(f"{DEFAULT_API_URL}/files/{item.get('task_id')}")
+        if file_response:
+            media_type = get_media_type(item.get('file_name'))
+            if media_type == 'image':
+                return ImageBlock(image=file_response.content)
+            # Audio currently not supported
+            #elif media_type == 'audio':
+            #    return AudioBlock(audio=file_response.content)
+
+
+def create_question_message(item):
+    question_text = item.get("question")
+    msg_blocks = [TextBlock(text=question_text)]
+    media_block = get_media_content(item)
+    if media_block is not None:
+        msg_blocks.append(media_block)
+    question_message = ChatMessage(role="user", blocks=msg_blocks)
+    return question_message
+
+
 async def answer_question(agent, item, answers_payload, results_log):
     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}")
+    try:
+        question_message = create_question_message(item)
+    except ValidationError:
+        print(f"Skipping item for which the question could not be processed: {item}")
+        return
+    if not task_id:
+        print(f"Skipping item with missing task_id: {item}")
         return
     try:
-        submitted_answer = await agent(question_text)
+        submitted_answer = await agent.run(question_message)
 
         # Avoid hitting the Google rate limits. 
         await asyncio.sleep(60)