gemini

agent.py

@@ -128,9 +128,9 @@ class ImageAnalysisTool(Tool):
 
 
 def final_answer_formatting(answer, question):
-    client = InferenceClient(
-        provider="nebius",
-        api_key=os.getenv("HF_TOKEN"),
+    model = LiteLLMModel(
+        model_id="gemini/gemini-2.0-flash",
+        api_key=os.getenv("GOOGLE_API_KEY"),
     )
 
     prompt = f"""
@@ -140,28 +140,18 @@ def final_answer_formatting(answer, question):
     Answer: {answer}
     """
 
-    completion = client.chat.completions.create(
-        model="mistralai/Mistral-Small-3.1-24B-Instruct-2503",
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": prompt
-                    }
-                ]
-            }
-        ],
-    )
+    messages = [
+    {"role": "user", "content": [{"type": "text", "text": prompt}]}
+    ]
 
-    return completion.choices[0].message.content
+    output = model(messages).content
+    return output
 
 
 web_agent = CodeAgent(
     model=LiteLLMModel(
-        model_id="mistral/mistral-large-latest",
-        api_key=os.getenv("MISTRAL_API_KEY"),
+        model_id="gemini/gemini-2.0-flash",
+        api_key=os.getenv("GOOGLE_API_KEY"),
     ),
     tools=[
         WikipediaSearchTool(),
@@ -179,8 +169,9 @@ web_agent = CodeAgent(
 )
 
 audio_agent = CodeAgent(
-    model=HfApiModel(
-        model_id='Qwen/Qwen3-32B'
+    model=LiteLLMModel(
+        model_id="gemini/gemini-2.0-flash",
+        api_key=os.getenv("GOOGLE_API_KEY"),
     ),
     tools=[AudioToTextTool()],
     add_base_tools=False,
@@ -192,8 +183,9 @@ audio_agent = CodeAgent(
 
 manager_agent = CodeAgent(
     name="manager_agent",
-    model=HfApiModel(
-        model_id='Qwen/Qwen3-235B-A22B'
+    model=LiteLLMModel(
+        model_id="gemini/gemini-2.5-flash-preview-04-17",
+        api_key=os.getenv("GOOGLE_API_KEY"),
     ),
     tools=[getFile(), LoadXlsxFileTool(), LoadTextFileTool(), ImageAnalysisTool()],
     managed_agents=[web_agent, audio_agent],