Update agent.py

agent.py

@@ -36,7 +36,7 @@ import whisper
 #        return f"Error performing Google search: {str(e)}"
 
 @tool
-def ImageAnalysisTool(question: str) -> str:
+def ImageAnalysisTool(question: str, model: LiteLLMModel) -> str:
     """Tool for analyzing images mentioned in the question.
     Args:
         question (str): The question text which may contain an image URL.
@@ -60,16 +60,10 @@ def ImageAnalysisTool(question: str) -> str:
     except Exception as e:
         return f"Error fetching image: {e}"
 
-    model = LiteLLMModel(
-        model_id="gemini/gemini-2.5-pro",
-        api_key=os.environ.get("GEMINI_KEY"),
-        max_tokens=8192
-    )
-
     agent = CodeAgent(
         tools=[],
         model=model,
-        max_steps=20,
+        max_steps=10,
         verbosity_level=2
     )
 
@@ -96,17 +90,17 @@ def SpeechToTextTool(audio_path: str) -> str:
     return result.get("text", "")
 
 
-#@tool
-#def youtube_transcript(url: str) -> str:
-#    """
-#    Get transcript of YouTube video.
-#    Args:
-#        url: YouTube video url in ""
-#    """    
-#    video_id = url.partition("https://www.youtube.com/watch?v=")[2]
-#    transcript = YouTubeTranscriptApi.get_transcript(video_id)
-#    transcript_text = " ".join([item["text"] for item in transcript])
-#    return {"youtube_transcript": transcript_text}
+@tool
+def youtube_transcript(url: str) -> str:
+    """
+    Get transcript of YouTube video.
+    Args:
+        url: YouTube video url in ""
+    """    
+    video_id = url.partition("https://www.youtube.com/watch?v=")[2]
+    transcript = YouTubeTranscriptApi.get_transcript(video_id)
+    transcript_text = " ".join([item["text"] for item in transcript])
+    return {"youtube_transcript": transcript_text}
     
 #@tool
 #class LocalFileAudioTool:
@@ -123,8 +117,9 @@ def SpeechToTextTool(audio_path: str) -> str:
 #        return f"Transcribed audio from '{file_path}' (simulated)."
 
 class MagAgent:
-    def __init__(self):
+    def __init__(self, rate_limiter: Optional[Limiter] = None):
         """Initialize the MagAgent with search tools."""
+        self.rate_limiter = rate_limiter
         print("Initializing MagAgent with search tools...")
         model = LiteLLMModel(
             model_id="gemini/gemini-2.0-flash",
@@ -139,7 +134,7 @@ class MagAgent:
         self.agent = CodeAgent(
             model= model,
             tools=[
-#                youtube_transcript,
+                youtube_transcript,
 #                GoogleSearchTool,
                 DuckDuckGoSearchTool(),
                 WikipediaSearchTool(),
@@ -153,7 +148,11 @@ class MagAgent:
     async def __call__(self, question: str) -> str:
         """Process a question asynchronously using the MagAgent."""
         print(f"MagAgent received question (first 50 chars): {question[:50]}...")
+        
         try:
+            if self.rate_limiter:
+                while not self.rate_limiter.consume(1):
+                    await asyncio.sleep(60 / RATE_LIMIT)
             # Define a task with fallback search logic
             task = (
                 f"Answer the following question accurately and concisely: {question}\n"