updated gemini agent

app.py

@@ -4,109 +4,32 @@ import requests
 import inspect
 import pandas as pd
 
-from smolagents import CodeAgent, DuckDuckGoSearchTool, OpenAIServerModel, LiteLLMModel
-from langchain_openai import AzureChatOpenAI
 from typing import Optional, Dict
+from gemini_agent import GeminiAgent
 
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 
-class AzureOpenAIServerModel(OpenAIServerModel):
-    """This model connects to an Azure OpenAI deployment.
-
-    Parameters:
-        model_id (`str`):
-            The model identifier to use on the server (e.g. "gpt-3.5-turbo").
-        azure_endpoint (`str`, *optional*):
-            The Azure endpoint, including the resource, e.g. `https://example-resource.azure.openai.com/`
-        api_key (`str`, *optional*):
-            The API key to use for authentication.
-        custom_role_conversions (`Dict{str, str]`, *optional*):
-            Custom role conversion mapping to convert message roles in others.
-            Useful for specific models that do not support specific message roles like "system".
-        **kwargs:
-            Additional keyword arguments to pass to the Azure OpenAI API.
-    """
-
-    def __init__(
-        self,
-        model_id: str,
-        azure_endpoint: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_version: Optional[str] = None,
-        custom_role_conversions: Optional[Dict[str, str]] = None,
-        **kwargs,
-    ):
-        super().__init__(model_id=model_id, api_key=api_key, custom_role_conversions=custom_role_conversions, **kwargs)
-        # if we've reached this point, it means the openai package is available (baseclass check) so go ahead and import it
-        import openai
-
-        self.client = openai.AzureOpenAI(
-            api_key=api_key,
-            api_version=api_version,
-            azure_endpoint=azure_endpoint
-        )
-
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
 class BasicAgent:
     def __init__(self):
         print("Initializing the BasicAgent")
-        # model = OpenAIServerModel(model_id="gpt-4o")
-        # model = OpenAIServerModel(model_id="gpt-4o-mini")
-        # model = OpenAIServerModel(model_id="o1-mini")
-        # google_api_key = os.getenv('GOOGLE_API_KEY')
-        # if google_api_key:
-        #     print(f"GOOGLE_API_KEY found. Last 10 characters: ...{google_api_key[-10:]}")
-        # else:
-        #     print("GOOGLE_API_KEY environment variable not set.")
+        
+        # Load environment variables
+        load_dotenv()
+        
+        # Get Gemini API key
+        api_key = os.getenv('GOOGLE_API_KEY')
+        if not api_key:
+            raise ValueError("GOOGLE_API_KEY environment variable not set.")
             
-        # try:
-        #     gemini_model_id = "google/gemini-2.0-flash"
-        #     model = LiteLLMModel(model_id=gemini_model_id)
-        #     # print("GPT-4o LLM is instantiated")
-        #     print(f"{gemini_model_id} LLM instantiated using LiteLLM")
-        # except Exception as e:
-        #     print(f"Error instantiating LLM: {e}")
-        #     print("Please ensure you have installed smolagents[litellm] and google-generativeai,")
-        #     print("and that your GOOGLE_API_KEY environment variable is set.")
-        #     raise e # Re-raise the exception to fail initialization if model fails
-
-        print(f"AZURE_API_ENDPOINT: {os.getenv('AZURE_API_ENDPOINT')}")
-        print(f"AZURE_DEPLOYMENT_NME: {os.getenv('AZURE_DEPLOYMENT_NME')}")
-        print(f"AZURE_API_VERSION: {os.getenv('AZURE_API_VERSION')}")
-        # Be cautious printing the key, maybe just check if it exists
-        print(f"AZURE_OPENAI_KEY set: {'AZURE_OPENAI_KEY' in os.environ}")
-
-        # model = AzureChatOpenAI(
-        # azure_endpoint=os.getenv("AZURE_API_ENDPOINT"),
-        # azure_deployment=os.getenv("AZURE_DEPLOYMENT_NME"),  # or your deployment
-        # api_key=os.getenv("AZURE_OPENAI_KEY"),  
-        # api_version=os.getenv("AZURE_API_VERSION"),
-        # temperature=0.5,
-        # max_tokens=None,
-        # timeout=None,
-        # max_retries=2,
-        # )
-        model = AzureOpenAIServerModel(
-            model_id = os.environ.get("AZURE_DEPLOYMENT_NME"),
-            api_key=os.environ.get("AZURE_OPENAI_KEY"),
-            api_version=os.environ.get("AZURE_API_VERSION"),
-            azure_endpoint=os.environ.get("AZURE_API_ENDPOINT")
-        )
-        # Initialize the search tool
-        search_tool = DuckDuckGoSearchTool()
-        print("DuckDuckSearchTool instantiated")
+        # Initialize GeminiAgent
+        self.agent = GeminiAgent(api_key=api_key)
+        print("GeminiAgent initialized successfully")
 
-        print("Agent is now created")
-        # Initializing Agent
-        self.agent = CodeAgent(
-            model = model,
-            tools=[search_tool]
-        )
-        
     def __call__(self, question: str) -> str:
         print(f"Agent received question (first 50 chars): {question[:50]}...")
         final_answer = self.agent.run(question)