Update core_agent.py

core_agent.py

@@ -6,7 +6,8 @@ from smolagents import (
     OpenAIServerModel,
     PythonInterpreterTool,
     tool,
-    InferenceClientModel
+    InferenceClientModel,
+    LiteLLMModel
 )
 from typing import List, Dict, Any, Optional
 import os
@@ -211,80 +212,15 @@ class GAIAAgent:
         self.system_prompt = system_prompt  # Store for potential future use
         
         # Initialize model based on configuration
-        if model_type == "HfApiModel":
-            if api_key is None:
-                api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-                if not api_key:
-                    raise ValueError("No Hugging Face token provided. Please set HUGGINGFACEHUB_API_TOKEN environment variable or pass api_key parameter.")
-            
-            if self.verbose:
-                print(f"Using Hugging Face token: {api_key[:5]}...")
-                
-            self.model = HfApiModel(
-                model_id=model_id or "meta-llama/Llama-3-70B-Instruct",
-                token=api_key,
-                temperature=temperature
-            )
-        elif model_type == "InferenceClientModel":
-            if api_key is None:
-                api_key = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-                if not api_key:
-                    raise ValueError("No Hugging Face token provided. Please set HUGGINGFACEHUB_API_TOKEN environment variable or pass api_key parameter.")
-            
-            if self.verbose:
-                print(f"Using Hugging Face token: {api_key[:5]}...")
+        
                 
-            self.model = InferenceClientModel(
-                model_id=model_id or "meta-llama/Llama-3-70B-Instruct",
-                provider=provider or "hf-inference",
-                token=api_key,
-                timeout=timeout or 120,
-                temperature=temperature
-            )
-        elif model_type == "LiteLLMModel":
-            from smolagents import LiteLLMModel
-            self.model = LiteLLMModel(
-                model_id=model_id or "gpt-4o",
-                api_key=api_key or os.getenv("OPENAI_API_KEY"),
-                temperature=temperature
-            )
-        elif model_type == "OpenAIServerModel":
-            # Check for xAI API key and base URL first
-            xai_api_key = os.getenv("XAI_API_KEY")
-            xai_api_base = os.getenv("XAI_API_BASE")
+        model_type == "LiteLLMModel":
             
-            # If xAI credentials are available, use them
-            if xai_api_key and api_key is None:
-                api_key = xai_api_key
-                if self.verbose:
-                    print(f"Using xAI API key: {api_key[:5]}...")
-            
-            # If no API key specified, fall back to OPENAI_API_KEY
-            if api_key is None:
-                api_key = os.getenv("OPENAI_API_KEY")
-                if not api_key:
-                    raise ValueError("No OpenAI API key provided. Please set OPENAI_API_KEY or XAI_API_KEY environment variable or pass api_key parameter.")
-            
-            # If xAI API base is available and no api_base is provided, use it
-            if xai_api_base and api_base is None:
-                api_base = xai_api_base
-                if self.verbose:
-                    print(f"Using xAI API base URL: {api_base}")
-            
-            # If no API base specified but environment variable available, use it
-            if api_base is None:
-                api_base = os.getenv("AGENT_API_BASE")
-                if api_base and self.verbose:
-                    print(f"Using API base from AGENT_API_BASE: {api_base}")
-            
-            self.model = OpenAIServerModel(
-                model_id=model_id or "gpt-4o",
-                api_key=api_key,
-                api_base=api_base,
-                temperature=temperature
-            )
-        else:
-            raise ValueError(f"Unknown model type: {model_type}")
+        self.model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash",
+            api_key=os.getenv("GEMINI_API_KEY"),
+            temperature=temperature
+        )
         
         if self.verbose:
             print(f"Initialized model: {model_type} - {model_id}")