Update app.py

app.py

@@ -10,14 +10,56 @@ 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
+
+
+
 class BasicAgent:
     def __init__(self):
         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
+   
     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
+        #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)
+        else:
+            return "Didn't understand you."
+
+        
+        
+    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):
     """