Update agent.py

agent.py

@@ -1,4 +1,4 @@
-from smolagents import CodeAgent, LiteLLMModel, Tool
+from smolagents import CodeAgent, LiteLLMModel, Tool, DuckDuckGoSearchTool, WikipediaSearchTool
 from token_bucket import Limiter, MemoryStorage
 from tenacity import retry, stop_after_attempt, wait_exponential
 from langchain_community.document_loaders import ArxivLoader
@@ -165,6 +165,141 @@ class UniversalLoader(Tool):
     def _fallback(self, source: str, context: str) -> str:
         return CrossVerifiedSearch()(f"{source} {context}")
 
+
+
+# --------------------------
+# Validation Pipeline
+# --------------------------
+class ValidationPipeline:
+    VALIDATORS = {
+        'numeric': {
+            'check': lambda x: pd.api.types.is_numeric_dtype(x),
+            'error': "Non-numeric value found in numeric field"
+        },
+        'temporal': {
+            'check': lambda x: pd.api.types.is_datetime64_any_dtype(x),
+            'error': "Invalid date format detected"
+        },
+        'categorical': {
+            'check': lambda x: x.isin(x.dropna().unique()),
+            'error': "Invalid category value detected"
+        }
+    }
+
+    def validate(self, data, schema: dict):
+        errors = []
+        for field, config in schema.items():
+            validator = self.VALIDATORS.get(config['type'])
+            if not validator['check'](data[field]):
+                errors.append(f"{field}: {validator['error']}")
+        return {
+            'valid': len(errors) == 0,
+            'errors': errors,
+            'confidence': 1.0 - (len(errors) / len(schema))
+        }
+
+# --------------------------
+# Tool Router
+# --------------------------
+class ToolRouter:
+    def __init__(self):
+        self.encoder = SentenceTransformer('all-MiniLM-L6-v2')
+        self.domain_embeddings = {
+            'music': self.encoder.encode("music album release artist track"),
+            'sports': self.encoder.encode("athlete team score tournament"),
+            'science': self.encoder.encode("chemistry biology physics research")
+        }
+        self.ddg = DuckDuckGoSearchTool()
+        self.wiki = WikipediaSearchTool()
+        self.arxiv = ArxivSearchTool()
+    
+    def forward(self, query: str, domain: str = None) -> str:
+        """Smart search with domain prioritization"""
+        if domain == "academic":
+            return self.arxiv(query)
+        elif domain == "general":
+            return self.ddg(query)
+        elif domain == "encyclopedic":
+            return self.wiki(query)
+        
+        # Fallback: Search all sources
+        results = {
+            "web": self.ddg(query),
+            "wikipedia": self.wiki(query),
+            "arxiv": self.arxiv(query)
+        }
+        return json.dumps(results)
+        
+    def route(self, question: str):
+        query_embed = self.encoder.encode(question)
+        scores = {
+            domain: np.dot(query_embed, domain_embed)
+            for domain, domain_embed in self.domain_embeddings.items()
+        }
+        return max(scores, key=scores.get)
+
+# --------------------------
+# Temporal Search
+# --------------------------
+class HistoricalSearch:
+    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=2, max=10))
+    def get_historical_content(self, url: str, target_date: str):
+        return requests.get(
+            f"http://archive.org/wayback/available?url={url}&timestamp={target_date}"
+        ).json()
+
+# --------------------------
+# Enhanced Excel Reader
+# --------------------------
+class EnhancedExcelReader(Tool):
+    def forward(self, path: str):
+        df = pd.read_excel(path)
+        validation = ValidationPipeline().validate(df, self._detect_schema(df))
+        if not validation['valid']:
+            raise ValueError(f"Data validation failed: {validation['errors']}")
+        return df.to_markdown()
+
+    def _detect_schema(self, df: pd.DataFrame):
+        schema = {}
+        for col in df.columns:
+            dtype = 'categorical'
+            if pd.api.types.is_numeric_dtype(df[col]):
+                dtype = 'numeric'
+            elif pd.api.types.is_datetime64_any_dtype(df[col]):
+                dtype = 'temporal'
+            schema[col] = {'type': dtype}
+        return schema
+
+# --------------------------
+# Cross-Verified Search
+# --------------------------
+class CrossVerifiedSearch:
+    SOURCES = [
+        DuckDuckGoSearchTool(),
+        WikipediaSearchTool(),
+        ArxivSearchTool()
+    ]
+
+    def __call__(self, query: str):
+        results = []
+        for source in self.SOURCES:
+            try:
+                results.append(source(query))
+            except Exception as e:
+                continue
+        return self._consensus(results)
+
+    def _consensus(self, results):
+        # Simple majority voting implementation
+        counts = {}
+        for result in results:
+            key = str(result)[:100]  # Simple hash for demo
+            counts[key] = counts.get(key, 0) + 1
+        return max(counts, key=counts.get)
+
+
+
+
 # --------------------------
 # Main Agent Class (Integrated)
 # --------------------------
@@ -180,11 +315,12 @@ class MagAgent:
 
         self.tools = [
             UniversalLoader(),
+            EnhancedSearchTool(),  # Replaces individual search tools
             ValidatedExcelReader(),
-            ArxivSearchTool(),
             VisitWebpageTool(),
             DownloadTaskAttachmentTool(),
-            SpeechToTextTool()
+            SpeechToTextTool(),
+            CrossVerifiedSearch()
         ]
 
         with open("prompts.yaml") as f: