Update agent.py

agent.py

@@ -1,7 +1,7 @@
-from smolagents import CodeAgent,  LiteLLMModel, tool, load_tool # DuckDuckGoSearchTool, WikipediaSearchTool, HfApiModel, OpenAIServerModel
-from langchain_community.tools.tavily_search import TavilySearchResults
-from langchain_community.document_loaders import WikipediaLoader
-from youtube_transcript_api import YouTubeTranscriptApi
+from smolagents import CodeAgent,  LiteLLMModel, tool, load_tool, DuckDuckGoSearchTool, WikipediaSearchTool #, HfApiModel, OpenAIServerModel
+#from langchain_community.tools.tavily_search import TavilySearchResults
+#from langchain_community.document_loaders import WikipediaLoader
+#from youtube_transcript_api import YouTubeTranscriptApi
 import asyncio
 import os
 import yaml
@@ -36,33 +36,33 @@ from io import BytesIO
 #    except Exception as e:
 #        return f"Error performing Google search: {str(e)}"
 
-@tool
-def wiki_search(query: str) -> str:
-    """Search Wikipedia for a query and return maximum 2 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"wiki_results": formatted_search_docs}
+#@tool
+#def wiki_search(query: str) -> str:
+#    """Search Wikipedia for a query and return maximum 2 results.
+#    
+#    Args:
+#        query: The search query."""
+#    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+#    formatted_search_docs = "\n\n---\n\n".join(
+#        [
+#            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+#            for doc in search_docs
+#        ])
+#    return {"wiki_results": formatted_search_docs}
 
-@tool
-def web_search(query: str) -> str:
-    """Search Tavily for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"web_results": formatted_search_docs}
+#@tool
+#def web_search(query: str) -> str:
+#    """Search Tavily for a query and return maximum 3 results.
+#    
+#    Args:
+#        query: The search query."""
+#    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+#    formatted_search_docs = "\n\n---\n\n".join(
+#        [
+#            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+#            for doc in search_docs
+#        ])
+#    return {"web_results": formatted_search_docs}
 
 @tool
 def ImageAnalysisTool(image_path: str) -> str:
@@ -94,7 +94,7 @@ def ImageAnalysisTool(image_path: str) -> str:
 
     response = agent.run(
         """
-        Describe what you see in the image in details.
+        Describe in details the chess position you see in the image.
         """,
         images=image
     )
@@ -102,17 +102,17 @@ def ImageAnalysisTool(image_path: str) -> str:
     return f"The image description: '{response}'"
 
 
-@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:
@@ -133,7 +133,7 @@ class MagAgent:
         """Initialize the MagAgent with search tools."""
         print("Initializing MagAgent with search tools...")
         model = LiteLLMModel(
-            model_id="gemini/gemini-2.0-flash-lite",
+            model_id="gemini/gemini-2.0-flash",
             api_key= os.environ.get("GEMINI_KEY"),
             max_tokens=8192
         )