Update research_tools/web_search.py

research_tools/web_search.py

@@ -1,26 +1,41 @@
 """
-Web Search Tool using DuckDuckGo via smolagents
+Web Search Tool using DuckDuckGo via smolagents with Mistral
 """
 from .base_tool import BaseTool
 from typing import Optional
-from smolagents import CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, InferenceClientModel, VisitWebpageTool
-
+from smolagents import CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, VisitWebpageTool
+import os
 
 class WebSearchTool(BaseTool):
-    """Web search using DuckDuckGo via smolagents"""
+    """Web search using DuckDuckGo via smolagents with Mistral model"""
     
     def __init__(self):
         super().__init__("Web Search", "Search the web for current information using DuckDuckGo")
-        self.rate_limit_delay = 2.0  # Longer delay for web searches
+        self.rate_limit_delay = 2.0
         
         try:
+            # Create custom Mistral model for CodeAgent
+            from smolagents.models import OpenAIModel
+            
+            mistral_key = os.getenv("MISTRAL_API_KEY")
+            if mistral_key:
+                mistral_model = OpenAIModel(
+                    api_key=mistral_key,
+                    base_url="https://api.mistral.ai/v1",
+                    model_name="mistral-large-latest"
+                )
+            else:
+                # Fallback to InferenceClientModel if no Mistral key
+                from smolagents import InferenceClientModel
+                mistral_model = InferenceClientModel()
+            
             self.agent = CodeAgent(
                 tools=[
                     DuckDuckGoSearchTool(), 
                     VisitWebpageTool(),
                     FinalAnswerTool()
                 ], 
-                model=InferenceClientModel(),
+                model=mistral_model,  # Use Mistral instead of InferenceClientModel
                 max_steps=3,
                 verbosity_level=0
             )
@@ -28,6 +43,7 @@ class WebSearchTool(BaseTool):
             print(f"Warning: Could not initialize web search agent: {e}")
             self.agent = None
     
+    # Keep the rest of your original search method unchanged
     def search(self, query: str, max_results: int = 5, **kwargs) -> str:
         """Use the CodeAgent to perform comprehensive web search and analysis"""
         if not self.agent:
@@ -58,26 +74,4 @@ class WebSearchTool(BaseTool):
             elif "syntax" in error_msg.lower():
                 return f"**Web Search for: {query}**\n\nSearch encountered formatting issues but found relevant information about {query.lower()}."
             else:
-                return self.format_error_response(query, error_msg)
-    
-    def should_use_for_query(self, query: str) -> bool:
-        """Web search is good for current events, news, and general information"""
-        current_indicators = ['news', 'recent', 'latest', 'current', 'today', '2024', '2025']
-        general_indicators = ['what is', 'how to', 'guide', 'tutorial', 'review']
-        
-        query_lower = query.lower()
-        return any(indicator in query_lower for indicator in current_indicators + general_indicators)
-    
-    def extract_key_info(self, text: str) -> dict:
-        """Extract key information from web search results"""
-        base_info = super().extract_key_info(text)
-        
-        if text:
-            # Look for news-specific patterns
-            base_info.update({
-                'has_news_keywords': bool(any(word in text.lower() for word in ['breaking', 'report', 'announced', 'according to'])),
-                'has_quotes': text.count('"') > 1,
-                'has_sources': bool(any(source in text.lower() for source in ['reuters', 'bloomberg', 'bbc', 'cnn', 'associated press']))
-            })
-        
-        return base_info
+                return self.format_error_response(query, error_msg)