Update app.py

app.py

@@ -22,429 +22,188 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 #Load environment variables
 load_dotenv()
 
-import io
-import contextlib
-import traceback
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from smolagents import Tool, CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, HfApiModel
-
-
-class CodeLlamaTool(Tool):
-    name = "code_llama_tool"
-    description = "Solves reasoning/code questions using Meta Code Llama 7B Instruct"
-
-    inputs = {
-        "question": {
-            "type": "string",
-            "description": "The question requiring code-based or reasoning-based solution"
-        }
-    }
-    output_type = "string"
-
-    def __init__(self):
-        self.model_id = "codellama/CodeLlama-7b-Instruct-hf"
-        token = os.getenv("HF_TOKEN")
-
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, token=token)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_id,
-            device_map="auto",
-            torch_dtype="auto",
-            token=token
-        )
-        self.pipeline = pipeline(
-            "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
-            max_new_tokens=512,
-            temperature=0.2,
-            truncation=True
-        )
-
-    def forward(self, question: str) -> str:
-        prompt = f"""You are an AI that uses Python code to answer questions.
-Question: {question}
-Instructions:
-- If solving requires code, use a block like <tool>code</tool>.
-- Always end with <final>FINAL ANSWER</final> containing the final number or string.
-Example:
-Question: What is 5 * sqrt(36)?
-Answer:
-<tool>
-import math
-print(5 * math.sqrt(36))
-</tool>
-<final>30.0</final>
-Answer:"""
-
-        response = self.pipeline(prompt)[0]["generated_text"]
-        return self.parse_and_execute(response)
-
-    def parse_and_execute(self, response: str) -> str:
+from duckduckgo_search import DDGS
+import wikipedia
+import arxiv
+from transformers import pipeline
+import os
+import re
+import ast
+import subprocess
+import sys
+
+# ===== Search Tools =====
+class DuckDuckGoSearchTool:
+    def __init__(self, max_results=3):
+        self.description = "Search web using DuckDuckGo. Input: search query"
+        self.max_results = max_results
+        
+    def run(self, query: str) -> str:
         try:
-            # Extract and run code if exists
-            if "<tool>" in response and "</tool>" in response:
-                code = response.split("<tool>")[1].split("</tool>")[0].strip()
-                result = self._run_code(code)
-                return f"FINAL ANSWER (code output): {result}"
-
-            # Extract final result directly
-            elif "<final>" in response and "</final>" in response:
-                final = response.split("<final>")[1].split("</final>")[0].strip()
-                return f"FINAL ANSWER: {final}"
-
-            return f"Could not extract final answer.\n\n{response}"
-
+            with DDGS() as ddgs:
+                results = [r for r in ddgs.text(query, max_results=self.max_results)]
+                return "\n\n".join(
+                    f"Title: {res['title']}\nURL: {res['href']}\nSnippet: {res['body']}" 
+                    for res in results
+                )
         except Exception as e:
-            return f"Error in parse_and_execute: {str(e)}\n\nFull response:\n{response}"
-
-    def _run_code(self, code: str) -> str:
-        buffer = io.StringIO()
-        try:
-            with contextlib.redirect_stdout(buffer):
-                exec(code, {})
-            return buffer.getvalue().strip()
-        except Exception:
-            return f"Error executing code:\n{traceback.format_exc()}"
-
-
-
+            return f"Search error: {str(e)}"
 
-import requests
-from smolagents import Tool
-
-class ArxivSearchTool(Tool):
-    name = "arxiv_search"
-    description = "Search Arxiv for papers matching a query and return titles and links."
-    inputs = {
-        "query": {"type": "string", "description": "Search query for Arxiv papers"}
-    }
-    output_type = "string"
-
-    def forward(self, query: str) -> str:
-        url = "http://export.arxiv.org/api/query"
-        params = {
-            "search_query": query,
-            "start": 0,
-            "max_results": 3,
-            "sortBy": "relevance",
-            "sortOrder": "descending"
-        }
+class WikiSearchTool:
+    def __init__(self, sentences=3):
+        self.description = "Get Wikipedia summaries. Input: search phrase"
+        self.sentences = sentences
+        
+    def run(self, query: str) -> str:
         try:
-            response = requests.get(url, params=params, timeout=10)
-            response.raise_for_status()
-            # Simple parse titles and links (basic, for demo)
-            import xml.etree.ElementTree as ET
-            root = ET.fromstring(response.content)
-            ns = {"atom": "http://www.w3.org/2005/Atom"}
-
-            entries = root.findall("atom:entry", ns)
-            results = []
-            for entry in entries:
-                title = entry.find("atom:title", ns).text.strip().replace('\n', ' ')
-                link = entry.find("atom:id", ns).text.strip()
-                results.append(f"{title}\n{link}")
-            return "\n\n".join(results) if results else "No results found."
+            return wikipedia.summary(query, sentences=self.sentences)
+        except wikipedia.DisambiguationError as e:
+            return f"Disambiguation error. Options: {', '.join(e.options[:5])}"
+        except wikipedia.PageError:
+            return "Page not found"
         except Exception as e:
-            return f"Error during Arxiv search: {e}"
-
-
-from transformers import pipeline
-from smolagents import Tool
-from PIL import Image
-
-class HuggingFaceDocumentQATool(Tool):
-    name = "document_qa"
-    description = "Answer questions from document images (e.g., scanned invoices)."
-    inputs = {
-        "image_path": {"type": "string", "description": "Path to the image file"},
-        "question": {"type": "string", "description": "Question to ask about the document"}
-    }
-    output_type = "string"
-
-    def __init__(self):
-        self.pipeline = pipeline("document-question-answering", model="impira/layoutlm-document-qa")
-
-    def forward(self, image_path: str, question: str) -> str:
-        image = Image.open(image_path)
-        result = self.pipeline(image, question=question)
-        return result[0]['answer']
-
-
-from transformers import BlipProcessor, BlipForQuestionAnswering
-
-class HuggingFaceImageQATool(Tool):
-    name = "image_qa"
-    description = "Answer questions about an image."
-    inputs = {
-        "image_path": {"type": "string", "description": "Path to image"},
-        "question": {"type": "string", "description": "Question about the image"}
-    }
-    output_type = "string"
-
-    def __init__(self):
-        self.processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
-        self.model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
-
-    def forward(self, image_path: str, question: str) -> str:
-        image = Image.open(image_path)
-        inputs = self.processor(image, question, return_tensors="pt")
-        out = self.model.generate(**inputs)
-        return self.processor.decode(out[0], skip_special_tokens=True)
-
-
-from transformers import pipeline
-
-class HuggingFaceTranslationTool(Tool):
-    name = "translate"
-    description = "Translate text from English to another language."
-    inputs = {
-        "text": {"type": "string", "description": "Text to translate"}
-    }
-    output_type = "string"
-
-    def __init__(self):
-        self.translator = pipeline("translation", model="Helsinki-NLP/opus-mt-en-fr")
-
-    def forward(self, text: str) -> str:
-        return self.translator(text)[0]["translation_text"]
+            return f"Wikipedia error: {str(e)}"
 
-
-import io
-import contextlib
-
-class PythonCodeExecutionTool(Tool):
-    name = "run_python"
-    description = "Execute Python code and return result."
-    inputs = {
-        "code": {"type": "string", "description": "Python code to execute"}
-    }
-    output_type = "string"
-
-    def forward(self, code: str) -> str:
-        output = io.StringIO()
+class ArxivSearchTool:
+    def __init__(self, max_results=3):
+        self.description = "Search academic papers on arXiv. Input: search query"
+        self.max_results = max_results
+        
+    def run(self, query: str) -> str:
         try:
-            with contextlib.redirect_stdout(output):
-                exec(code, {})
-            return output.getvalue().strip()
+            results = arxiv.Search(
+                query=query, 
+                max_results=self.max_results,
+                sort_by=arxiv.SortCriterion.Relevance
+            ).results()
+            
+            output = []
+            for r in results:
+                output.append(
+                    f"Title: {r.title}\n"
+                    f"Authors: {', '.join(a.name for a in r.authors)}\n"
+                    f"Published: {r.published.strftime('%Y-%m-%d')}\n"
+                    f"Summary: {r.summary[:250]}...\n"
+                    f"URL: {r.entry_id}"
+                )
+            return "\n\n".join(output)
         except Exception as e:
-            return f"Error: {str(e)}"
-
-
+            return f"arXiv error: {str(e)}"
 
-
-#from smolagents import Tool
-#from langchain_community.document_loaders import WikipediaLoader
-
-class WikiSearchTool(Tool):
-    name = "wiki_search"
-    description = "Search Wikipedia for a query and return up to 2 results."
-    inputs = {
-        "query": {"type": "string", "description": "The search term for Wikipedia."}
-    }
-    output_type = "string"
-
-    def forward(self, query: str) -> str:
-        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
-
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'<Document source="{doc.metadata.get("source", "Wikipedia")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-                for doc in search_docs
-            ]
+# ===== QA Tools =====
+class HuggingFaceDocumentQATool:
+    def __init__(self):
+        self.description = "Answer questions from documents. Input: 'document_text||question'"
+        self.model = pipeline(
+            'question-answering', 
+            model='deepset/roberta-base-squad2',
+            tokenizer='deepset/roberta-base-squad2'
         )
-        return formatted_search_docs
-
     
-
-
-class StringReverseTool(Tool):
-    name = "reverse_message"
-    description = "When you received a strange text, try to reverse it and perform action described in reversed message."
-    inputs = {
-        "message": {
-            "type": "string",
-            "description": "A message, which looks like strange and can be reversed to get actions to execute."
-        }
-    }
-    output_type = "string"
-
-    def __init__(self):
-        return
-
-    def forward(self, message: str):
-        return message[::-1]
-
-class KeywordsExtractorTool(Tool):
-    """Extracts top 5 keywords from a given text based on frequency."""
-
-    name = "keywords_extractor"
-    description = "This tool returns the 5 most frequent keywords occur in provided block of text."
-
-    inputs = {
-        "text": {
-            "type": "string",
-            "description": "Text to analyze for keywords.",
-        }
-    }
-    output_type = "string"
-
-    def forward(self, text: str) -> str:
+    def run(self, input_str: str) -> str:
         try:
-            all_words = re.findall(r'\b\w+\b', text.lower())
-            conjunctions = {'a', 'and', 'of', 'is', 'in', 'to', 'the'}
-            filtered_words = []
-            for w in all_words:
-                if w not in conjunctions:
-                    filtered_words.push(w)
-            word_counts = Counter(filtered_words)
-            k = 5
-            return heapq.nlargest(k, word_counts.items(), key=lambda x: x[1])
+            if '||' not in input_str:
+                return "Invalid format. Use: 'document_text||question'"
+                
+            context, question = input_str.split('||', 1)
+            result = self.model(question=question, context=context)
+            return result['answer']
         except Exception as e:
-            return f"Error during extracting most common words: {e}"
+            return f"QA error: {str(e)}"
 
-@tool
-def parse_excel_to_json(task_id: str) -> dict:
-    """
-    For a given task_id fetch and parse an Excel file and save parsed data in structured JSON file.
-    Args:
-        task_id: An task ID to fetch.
+# ===== Code Execution =====
+class PythonCodeExecutionTool:
+    def __init__(self):
+        self.description = "Execute Python code. Input: valid Python code"
         
-    Returns:
-        {
-            "task_id": str,
-            "sheets": {
-                "SheetName1": [ {col1: val1, col2: val2, ...}, ... ],
-                ...
-            },
-            "status": "Success" | "Error"
-        }
-    """
-    url = f"https://agents-course-unit4-scoring.hf.space/files/{task_id}"
-
-    try:
-        response = requests.get(url, timeout=100)
-        if response.status_code != 200:
-            return {"task_id": task_id, "sheets": {}, "status": f"{response.status_code} - Failed"}
-
-        xls_content = pd.ExcelFile(BytesIO(response.content))
-        json_sheets = {}
-
-        for sheet in xls_content.sheet_names:
-            df = xls_content.parse(sheet)
-            df = df.dropna(how="all")  
-            rows = df.head(20).to_dict(orient="records")
-            json_sheets[sheet] = rows
-
-        return {
-            "task_id": task_id,
-            "sheets": json_sheets,
-            "status": "Success"
-        }
-
-    except Exception as e:
-        return {
-            "task_id": task_id,
-            "sheets": {},
-            "status": f"Error in parsing Excel file: {str(e)}"
-        }
-
-
-
-class VideoTranscriptionTool(Tool):
-    """Fetch transcripts from YouTube videos"""
-    name = "transcript_video"
-    description = "Fetch text transcript from YouTube movies with optional timestamps"
-    inputs = {
-        "url": {"type": "string", "description": "YouTube video URL or ID"},
-        "include_timestamps": {"type": "boolean", "description": "If timestamps should be included in output", "nullable": True}
-    }
-    output_type = "string"
-
-    def forward(self, url: str, include_timestamps: bool = False) -> str:
-
-        if "youtube.com/watch" in url:
-            video_id = url.split("v=")[1].split("&")[0]
-        elif "youtu.be/" in url:
-            video_id = url.split("youtu.be/")[1].split("?")[0]
-        elif len(url.strip()) == 11:  # Direct ID
-            video_id = url.strip()
-        else:
-            return f"YouTube URL or ID: {url} is invalid!"
-
+    def run(self, code: str) -> str:
         try:
-            transcription = YouTubeTranscriptApi.get_transcript(video_id)
+            # Isolate code in a clean environment
+            env = {}
+            exec(f"def __temp_func__():\n    {indent_code(code)}", env)
+            output = env['__temp_func__']()
+            return str(output)
+        except Exception as e:
+            return f"Execution error: {str(e)}"
 
-            if include_timestamps:
-                formatted_transcription = []
-                for part in transcription:
-                    timestamp = f"{int(part['start']//60)}:{int(part['start']%60):02d}"
-                    formatted_transcription.append(f"[{timestamp}] {part['text']}")
-                return "\n".join(formatted_transcription)
-            else:
-                return " ".join([part['text'] for part in transcription])
+def indent_code(code: str) -> str:
+    """Add proper indentation for multiline code"""
+    return '\n    '.join(code.splitlines())
 
-        except Exception as e:
-            return f"Error in extracting YouTube transcript: {str(e)}"
+# ===== Answer Formatting =====
+class FinalAnswerTool:
+    def __init__(self):
+        self.description = "Format final answer. Input: answer content"
+        
+    def run(self, answer: str) -> str:
+        return f"FINAL ANSWER: {answer}"
 
 class BasicAgent:
     def __init__(self):
         token = os.environ.get("HF_API_TOKEN")
         model = HfApiModel(
-            temperature=0.1,
+            temperature=0.0,  # Reduced for deterministic output
             token=token
         )
 
-        # Existing tools
+        # Curated toolset - remove redundant/conflicting tools
         search_tool = DuckDuckGoSearchTool()
         wiki_search_tool = WikiSearchTool()
-        str_reverse_tool = StringReverseTool()
-        keywords_extract_tool = KeywordsExtractorTool()
-        speech_to_text_tool = SpeechToTextTool()
-        visit_webpage_tool = VisitWebpageTool()
-        final_answer_tool = FinalAnswerTool()
-        video_transcription_tool = VideoTranscriptionTool()
-
-        # ✅ New Llama Tool
-        code_llama_tool = CodeLlamaTool()
-        # ✅ Add Hugging Face default tools
         arxiv_search_tool = ArxivSearchTool()
         doc_qa_tool = HuggingFaceDocumentQATool()
-        image_qa_tool = HuggingFaceImageQATool()
-        translation_tool = HuggingFaceTranslationTool()
         python_tool = PythonCodeExecutionTool()
+        final_answer_tool = FinalAnswerTool()
 
-        system_prompt = f"""
-You are my general AI assistant. Your task is to answer the question I asked.
-First, provide an explanation of your reasoning, step by step, to arrive at the answer.
-Then, return your final answer in a single line, formatted as follows: "FINAL ANSWER: [YOUR FINAL ANSWER]".
-[YOUR FINAL ANSWER] should be a number, a string, or a comma-separated list of numbers and/or strings, depending on the question.
-If the answer is a number, do not use commas or units (e.g., $, %) unless specified.
-If the answer is a string, do not use articles or abbreviations (e.g., for cities), and write digits in plain text unless specified.
-If the answer is a comma-separated list, apply the above rules for each element based on whether it is a number or a string.
+        # Strategic tool selection
+        tools = [
+            search_tool, 
+            wiki_search_tool,
+            arxiv_search_tool,
+            doc_qa_tool,
+            python_tool,
+            final_answer_tool
+        ]
+
+        # Enhanced system prompt
+        system_prompt = """
+You are a precision question-answering AI. Follow this protocol:
+1. Analyze the question type: factual, computational, or multi-step
+2. Select the optimal tool:
+   - Use Search/Wiki/Arxiv for factual queries
+   - Use Python tool for calculations
+   - Use DocQA for document-based questions
+3. Execute necessary actions
+4. Verify answer matches question requirements
+5. Output FINAL ANSWER using this format:
+   "FINAL ANSWER: [EXACT_RESULT]"
+
+Answer rules:
+- Numbers: Plain format (e.g., 1000000)
+- Strings: No articles/abbreviations (e.g., "Paris" not "city of Paris")
+- Lists: Comma-separated (e.g., "red,blue,green")
+- Never include units ($, kg, etc.) unless explicitly required
+- For true/false: Use "true" or "false" lowercase
 """
 
         self.agent = CodeAgent(
             model=model,
-            tools=[
-                search_tool, wiki_search_tool, str_reverse_tool,
-                keywords_extract_tool, speech_to_text_tool,
-                visit_webpage_tool, final_answer_tool,
-                parse_excel_to_json, video_transcription_tool,
-                arxiv_search_tool,
-                doc_qa_tool, image_qa_tool,
-                translation_tool, python_tool, 
-                code_llama_tool  # 🔧 Add here
-            ],
-            add_base_tools=True
+            tools=tools,
+            add_base_tools=False  # Prevent tool conflicts
         )
-        self.agent.prompt_templates["system_prompt"] = self.agent.prompt_templates["system_prompt"] + system_prompt
+        # Force strict prompt template
+        self.agent.prompt_templates["system_prompt"] = system_prompt
 
     def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        answer = self.agent.run(question)
-        print(f"Agent returning answer: {answer}")
-        return answer
+        print(f"Processing: {question[:50]}...")
+        try:
+            result = self.agent.run(question)
+            # Extract final answer using regex
+            import re
+            match = re.search(r"FINAL ANSWER:\s*(.+)", result, re.IGNORECASE)
+            return match.group(1).strip() if match else result
+        except Exception as e:
+            print(f"Error: {str(e)}")
+            return "Unable to determine answer"