Update agent.py

agent.py

@@ -697,44 +697,43 @@ def get_youtube_transcript(url: str) -> str:
         except:
             pass
 
+from litellm import LiteLLMModel
+import os
+
 class BasicAgent:
     def __init__(self):
         print("BasicAgent initialized.")
-        # Initialize the Anthropic models
-        # Standard model for supervisor and validator
+        # Initialize callback handler
         self.langfuse_handler = CallbackHandler()
 
-        self.supervisor_model = ChatAnthropic(
-            model="claude-3-7-sonnet-20250219",
-            max_tokens=20000,
-            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"),
-            temperature=0.6,
-#            thinking={
-#                    "type": "enabled",
-#                    "budget_tokens": 5000
-#            }
+        # Supervisor model using Gemini
+        self.supervisor_model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash-lite",
+            api_key=os.getenv("GEMINI_API_KEY"),
+            temperature=0.5,
+            max_tokens=1024,
         )
         
-        # Standard model for validator
-        self.validator_model = ChatAnthropic(
-            model="claude-3-7-sonnet-20250219",
-            max_tokens=20000,
-            temperature=0.5,  # Lower temperature for more consistent validation
-            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY")
+        # Validator model using Gemini
+        self.validator_model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash-lite",
+            api_key=os.getenv("GEMINI_API_KEY"),
+            temperature=0.5,
+            max_tokens=1024,
         )
         
-        # Tool-enabled model for worker
-        self.worker_model_base = ChatAnthropic(
-            model="claude-3-7-sonnet-20250219",
-            max_tokens=20000,
+        # Worker base model using Gemini
+        self.worker_model_base = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash-lite",
+            api_key=os.getenv("GEMINI_API_KEY"),
             temperature=0.75,
-            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY")
+            max_tokens=20000,
         )
         
-        # Initialize tools
-        self.tools = [search_web_tavily, search_web_serper, execute_code_securely, execute_shell_command, sandbox_file_operation, extract_document_data, extract_image_data, extract_url_content, get_youtube_transcript]
-        
-        # Bind tools only to the worker model
+        # Bind tools to the worker model
+        self.tools = [search_web_tavily, search_web_serper, execute_code_securely,
+                      execute_shell_command, sandbox_file_operation, extract_document_data,
+                      extract_image_data, extract_url_content, get_youtube_transcript]
         self.worker_model = self.worker_model_base.bind_tools(self.tools)
         
         # Create the tool node for executing tools