Update app.py

app.py

@@ -10,56 +10,61 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # --- Basic Agent Definition ---
 # ----- THIS IS WERE YOU CAN BUILD WHAT YOU WANT ------
-from transformers import pipeline
-
-
+from tools import (
+    web_search_tool,
+    image_analysis_tool,
+    audio_transcription_tool,
+    document_analysis_tool,
+    llm_tool
+)
 
 class BasicAgent:
-    def __init__(self):
+    def __init__(self, llm):
         print("BasicAgent initialized.")
-        llm_pipeline = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2")
-        from tools import (
-            web_search_tool,
-            analyze_image_tool,
-            ocr_tool,
-            video_ocr_tool,
-            document_tool,
-            chat_tool,
-        )
-        self.web_search_tool = web_search_tool
-        self.analyze_image_tool = analyze_image_tool
-        self.ocr_tool = ocr_tool
-        self.video_ocr_tool = video_ocr_tool
-        self.document_tool = document_tool
-        self.chat_tool = chat_tool
+        self.llm = llm
+
+    def classify_question(self, question: str) -> str:
+        prompt = f"""
+        You are a question classifier.
+        Classify the question into one of the following categories:
+        
+        - web: for questions that require searching current information online
+        - image: when user uploads an image or asks about visual content
+        - audio: when user uploads an audio or asks what is said
+        - video: for YouTube or video questions
+        - document: for PDFs or documents
+        - text: for general logic, reasoning, or knowledge questions
+        
+        Question: "{question}"
+        
+        Return only the label.
+        """
+        return self.llm(prompt).strip().lower()
    
-    def __call__(self, question: str) -> str:
-        print(f"Agent received question (first 50 chars): {question[:50]}...")
-        #fixed_answer = "This is a default answer."
-        #print(f"Agent returning fixed answer: {fixed_answer}")
-        #return fixed_answer
-
-        classification = llm(f"Classify the question: '{question}' into one of: web_search, image_analysis, ocr, video_ocr, document, chat")
-        if classification == "web_search":
-            return self.web_search_tool(question)
-        elif classification == "image_analysis":
-            return self.analyze_image_tool(question)
-        elif classification == "ocr":
-            return self.ocr_tool(question)
-        elif classification == "video_ocr":
-            return self.video_ocr_tool(question)
-        elif classification == "document":
-            return self.document_tool(question)
-        elif classification == "chat":
-            return self.chat_tool(question)
+    def __call__(self, question: str, file_path: str = None) -> str:
+        category = self.classify_question(question)
+
+        if category == "web":
+            return web_search_tool(question)
+
+        elif category == "image" and file_path:
+            return image_analysis_tool(file_path)
+
+        elif category == "audio" and file_path:
+            return audio_transcription_tool(file_path)
+
+        elif category == "document" and file_path:
+            with open(file_path, 'r', encoding='utf-8') as f:
+                content = f.read()
+            return document_analysis_tool(content)
+
+        elif category == "text":
+            return llm_tool(question, self.llm)
+
         else:
-            return "Didn't understand you."
+            return f"Sorry, I couldn't handle this category: {category}"
 
         
-        
-    def llm(prompt: str) -> str:
-        result = llm_pipeline(prompt, max_new_tokens=100, do_sample=False)
-        return result[0]['generated_text']
 
 def run_and_submit_all( profile: gr.OAuthProfile | None):
     """