LiteLM Update

app.py

@@ -8,6 +8,7 @@ from src.agent import CustomAgent
 # (Keep Constants as is)
 # --- Constants ---
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+SYSTEM_PROMPT = "You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.\n"
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """
@@ -65,7 +66,7 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     print(f"Running agent on {len(questions_data)} questions...")
     for item in questions_data:
         task_id = item.get("task_id")
-        question_text = item.get("question")
+        question_text = SYSTEM_PROMPT + item.get("question")
         if not task_id or question_text is None:
             print(f"Skipping item with missing task_id or question: {item}")
             continue

requirements.txt

@@ -1,3 +1,4 @@
 gradio
 requests
-smolagents
+smolagents
+smolagents[litellm]

src/agent.py

@@ -1,13 +1,30 @@
-from smolagents import CodeAgent, InferenceClientModel, DuckDuckGoSearchTool
-# import asyncio
+from smolagents import CodeAgent, LiteLLMModel, DuckDuckGoSearchTool
 import datetime
+import os
+# import asyncio
+from dotenv import load_dotenv
+load_dotenv()
 
 
+PROMPT = """
+You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
+"""
 class CustomAgent:
     def __init__(
         self, 
-        model_name="Qwen/Qwen2.5-Omni-7B", 
-        authorized_imports=["wiki", "json"], 
+        model_id="gemini/gemini-2.0-flash-exp", 
+        authorized_imports=[
+            "wiki",
+            "json",
+            "BeautifulSoup",
+            "bs4",
+            "requests",
+            "os",
+            "re",
+            "pandas",
+            "numpy",
+            "datetime",
+        ], 
         max_steps=5,
         logging=False
     ):
@@ -16,9 +33,17 @@ class CustomAgent:
         If no model is provided, a default one is used.
         """
         self.logging = logging
-        model = InferenceClientModel(model_name)
+        
+        # Initialize the model
+        model = LiteLLMModel(
+            model_id=model_id,
+            api_key=os.getenv("GEMINI_API_KEY")
+        )
+
+        # Initialize tools
         search_tool = DuckDuckGoSearchTool()
         
+        # Initialize the agent
         self.agent = CodeAgent(
             model = model,
             tools=[search_tool],
@@ -33,9 +58,12 @@ class CustomAgent:
             fixed_answer = self.agent.run(question)
             
             if self.logging:
+                # print current directory
+                if not os.path.exists("logging"):
+                    os.makedirs("logging")
                 timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                 with open("logging/log.txt", "a") as log_file:
-                    log_file.write(f"{timestamp}: {question} -> {fixed_answer}\n")
+                    log_file.write(f"{timestamp}: {question}\nAnswer: {fixed_answer}\n")
 
             return question, fixed_answer