Update agent.py

agent.py

@@ -1,176 +1,47 @@
+from smolagents import LiteLLMModel, CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool, PythonInterpreterTool, tool
+from youtube_transcript_api import YouTubeTranscriptApi
 import os
-import re
-import requests
-import tempfile
-import pandas as pd
-from openai import OpenAI
 
-try:
-    from duckduckgo_search import DDGS
-except ImportError:
-    DDGS = None
-
-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."
-)
+@tool
+def reverse_sentence_tool(reverse_sentence: str) -> str:
+    inverted_words = reverse_sentence.split(" ")[::-1]
+    correct_words = [word[::-1] for word in inverted_words]
+    return " ".join(correct_words)
+
+@tool
+def get_youtube_transcript(video_url: str) -> str:
+    video_id = video_url.split("v=")[-1]
+    transcript = YouTubeTranscriptApi.get_transcript(video_id)
+    full_text = " ".join([entry['text'] for entry in transcript])
+    return full_text
+
+@tool
+def check_answer(answer: str) -> str:
+    if answer and answer[-1] == '.':
+        answer = answer[:-1]
+    if "St." in answer:
+        answer = answer.replace("St.", "Saint")
+    return answer
 
 class BasicAgent:
     def __init__(self):
-        self.llm = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        print("BasicAgent initialized.")
-
-    def web_search(self, query: str, max_results: int = 5) -> str:
-        if not DDGS:
-            return ""
-        try:
-            with DDGS() as ddgs:
-                results = list(ddgs.text(query, max_results=max_results))
-            if not results:
-                return ""
-            formatted_results = ""
-            for i, result in enumerate(results, 1):
-                title = result.get('title', '')
-                body = result.get('body', '')
-                href = result.get('href', '')
-                formatted_results += f"{i}. {title}\n   URL: {href}\n   Description: {body}\n\n"
-            return formatted_results
-        except Exception:
-            return ""
-
-    def excel_tool(self, file_url: str) -> str:
-        try:
-            r = requests.get(file_url, timeout=20)
-            r.raise_for_status()
-            with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as f:
-                f.write(r.content)
-                f.flush()
-                excel_path = f.name
-            df = pd.read_excel(excel_path)
-            if "Type" in df.columns and "Sales" in df.columns:
-                total = df[df["Type"].str.lower() == "food"]["Sales"].sum()
-                return f"{round(total, 2)}"
-            total = df.select_dtypes(include='number').sum().sum()
-            return f"{round(total, 2)}"
-        except Exception:
-            return ""
-
-    def fetch_file_url(self, task_id):
-        DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-        try:
-            url = f"{DEFAULT_API_URL}/files/{task_id}"
-            r = requests.head(url, timeout=5)
-            if r.status_code == 200:
-                return url
-        except Exception:
-            pass
-        return None
-
-    def solve_chess_image(self, image_url: str) -> str:
-        prompt = (
-            "You are a chess engine. Only answer with the best move for Black in algebraic notation (e.g., Qd1#). "
-            "Do not explain your reasoning, do not include any commentary, only the move."
+        api_key = os.environ.get("OPENAI_API_KEY", "")
+        # Użyj modelu openai, np. gpt-4o lub gpt-3.5-turbo
+        self.model = LiteLLMModel(model_id="gpt-4o", api_key=api_key)
+        self.agent = CodeAgent(
+            tools=[
+                DuckDuckGoSearchTool(),
+                PythonInterpreterTool(),
+                VisitWebpageTool(),
+                reverse_sentence_tool,
+                get_youtube_transcript,
+                check_answer
+            ],
+            model=self.model
         )
-        try:
-            response = self.llm.chat.completions.create(
-                model="gpt-4o",
-                messages=[
-                    {"role": "system", "content": prompt},
-                    {
-                        "role": "user",
-                        "content": [
-                            {"type": "text", "text": prompt},
-                            {"type": "image_url", "image_url": {"url": image_url}},
-                        ],
-                    }
-                ],
-                max_tokens=32,
-                temperature=0.0,
-            )
-            result = response.choices[0].message.content.strip()
-            move = re.findall(r"\b([KQRNB]?[a-h]?[1-8]?x?[a-h][1-8](?:=[QRNB])?#?)\b", result)
-            if move:
-                return move[0]
-            return result
-        except Exception:
-            return ""
-
-    def __call__(self, question: str, task_id: str = None) -> str:
-        file_url = self.fetch_file_url(task_id) if task_id else None
-        file_result = None
-
-        # --- Chess image detection (priority) ---
-        if file_url and ("chess" in question.lower() or "move" in question.lower() or "image" in question.lower() or "position" in question.lower()):
-            move = self.solve_chess_image(file_url)
-            if move and len(move) <= 6:
-                return move
-
-        # --- Excel file detection ---
-        ext = file_url.split('.')[-1].lower() if file_url else ""
-        if file_url and (ext in ["xlsx", "xls"] or "excel" in question.lower() or "spreadsheet" in question.lower()):
-            file_result = self.excel_tool(file_url)
-            if file_result and re.match(r'^\d+(\.\d+)?$', file_result):
-                return file_result
-
-        # --- Web search + LLM ---
-        search_snippet = self.web_search(question)
-        prompt = PROMPT + f"\n\nWeb search results:\n{search_snippet}\n\nQuestion: {question}"
-        response = self.llm.chat.completions.create(
-            model="gpt-4o",
-            messages=[{"role": "system", "content": prompt}],
-            temperature=0.0,
-            max_tokens=512,
-        )
-        answer = response.choices[0].message.content.strip()
-        final_line = ""
-        for line in answer.splitlines():
-            if line.strip().lower().startswith("final answer:"):
-                final_line = line.split(":", 1)[-1].strip(" .\"'")
-                break
+        print("BasicAgent initialized.")
 
-        # --- Fallback for non-answers ---
-        bads = [
-            "", "unknown", "unable to determine", "unable to provide page numbers",
-            "unable to access video content directly", "unable to analyze video content",
-            "unable to determine without code", "unable to determine without file",
-            "follow the steps to locate the paper and find the nasa award number in the acknowledgment section",
-            "i am unable to view images or access external content directly", "unable to determine without access to the file",
-            "no results found", "n/a", "[your final answer]", "i'm sorry", "i apologize"
-        ]
-        norm_final = (final_line or "").lower()
-        if norm_final in bads or norm_final.startswith("unable") or norm_final.startswith("i'm sorry") or norm_final.startswith("i apologize"):
-            numbers = re.findall(r'\b\d{2,}\b', search_snippet)
-            if numbers:
-                return numbers[0]
-            words = re.findall(r'\b[A-Z][a-z]{2,}\b', search_snippet)
-            if words:
-                return words[0]
-            if file_result:
-                file_numbers = re.findall(r'\b\d{2,}\b', str(file_result))
-                if file_numbers:
-                    return file_numbers[0]
-                file_words = re.findall(r'\b[A-Z][a-z]{2,}\b', str(file_result))
-                if file_words:
-                    return file_words[0]
-            retry_prompt = (
-                "Based ONLY on the search results and/or file content above, return a direct answer to the question. "
-                "If you do not know, make your best plausible guess. Do NOT apologize or say you cannot assist. "
-                f"File: {file_result}\n\nWeb: {search_snippet}\n\nQuestion: {question}\nFINAL ANSWER:"
-            )
-            response2 = self.llm.chat.completions.create(
-                model="gpt-4o",
-                messages=[{"role": "system", "content": retry_prompt}],
-                temperature=0.1,
-                max_tokens=128,
-            )
-            retry_answer = response2.choices[0].message.content.strip()
-            for line in retry_answer.splitlines():
-                if line.strip().lower().startswith("final answer:"):
-                    return line.split(":", 1)[-1].strip(" .\"'")
-            if retry_answer:
-                return retry_answer.strip(" .\"'")
-        return final_line or answer
+    def __call__(self, question: str) -> str:
+        print(f"Agent received question: {question[:50]}...")
+        answer = self.agent.run(question)
+        return answer