Update app.py

app.py

@@ -6,6 +6,18 @@ from typing import Dict, List, Optional, Any, Generator, Tuple
 from dataclasses import dataclass
 from pathlib import Path
 import hashlib
+import re
+import requests
+import pandas as pd
+import tempfile
+import gradio as gr
+from transformers import pipeline
+from PIL import Image
+import soundfile as sf
+import numpy as np
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+from huggingface_hub import InferenceClient, get_inference_client, Agent
 
 # Setup log file path
 log_path = os.path.join(os.getcwd(), f"agent_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
@@ -33,36 +45,11 @@ logging.getLogger().addHandler(console_handler)
 logging.info(f"===== Application Startup at {datetime.now().isoformat()} =====")
 logging.info(f"📂 Log file created at: {log_path}")
 
-import requests
-import pandas as pd
-import tempfile
-import gradio as gr
-from transformers import pipeline
-from PIL import Image
-import soundfile as sf
-import numpy as np
-from requests.adapters import HTTPAdapter
-from urllib3.util.retry import Retry
-
-# Configuration constants
-API_URL = "https://agents-course-unit4-scoring.hf.space"
-DEFAULT_TIMEOUT = 30  # seconds
-MAX_RETRIES = 3
-BACKOFF_FACTOR = 0.3
-MAX_FILE_SIZE = 50 * 1024 * 1024  # 50MB limit
-TEMP_DIR = tempfile.gettempdir()
+# --- Constants and Data Classes ---
+API_URL = "https://huggingface.co/api/agent-course"
 
-@dataclass
-class ProcessingResult:
-    """Structured result for processing operations"""
-    success: bool
-    answer: str
-    error: Optional[str] = None
-    processing_time: float = 0.0
-    
 @dataclass
 class QuestionLog:
-    """Structured log entry for questions"""
     question_num: int
     question_preview: str
     question_type: str
@@ -70,677 +57,479 @@ class QuestionLog:
     processing_time: float
     status: str
 
-class RobustHTTPSession:
-    """HTTP session with retry logic and proper error handling"""
-    
-    def __init__(self, timeout: int = DEFAULT_TIMEOUT):
-        self.session = requests.Session()
-        self.timeout = timeout
-        
-        # Configure retry strategy
-        retry_strategy = Retry(
-            total=MAX_RETRIES,
-            backoff_factor=BACKOFF_FACTOR,
-            status_forcelist=[429, 500, 502, 503, 504],
-            allowed_methods=["HEAD", "GET", "POST"]
-        )
-        
-        adapter = HTTPAdapter(max_retries=retry_strategy)
-        self.session.mount("http://", adapter)
-        self.session.mount("https://", adapter)
-        
-        # Set headers
-        self.session.headers.update({
-            'User-Agent': 'SmartAgent/1.0',
-            'Accept': 'application/json, */*'
-        })
-    
-    def get(self, url: str, **kwargs) -> requests.Response:
-        """GET request with error handling"""
+# Define some common patterns for cleaning up agent output
+ANSWER_PATTERNS = [
+    r"the answer is:\s*(.*)",
+    r"final_answer\((.*)\)",
+    r"final_answer\(\s*\"(.*?)\"\s*\)",
+    r"final_answer\(\s*'(.*?)'\s*\)",
+    r"final_answer\(\s*```(.*?)```\s*\)",
+    r"final_answer\(\s*`(.*?)`\s*\)",
+    r"final_answer\(\s*\[(.*?)\]\s*\)", # For list-like answers
+    r"final_answer\(\s*\{(.*?)\}\s*\)", # For dict-like answers
+    r"final_answer\(\s*(.*?)\s*\)", # Most general final_answer
+    r"answer is\s*(.*)",
+    r"final answer is\s*(.*)",
+    r"final answer:\s*(.*)",
+    r"final answer of the question is:\s*(.*)",
+    r"the final answer is:\s*(.*)",
+    r"final answer\s*:\s*(.*)",
+    r"final_answer\s*:\s*(.*)",
+    r"final_answer is\s*(.*)",
+    r"the value is:\s*(.*)",
+    r" (.*)",
+    r"The final answer is: (.*)",
+    r"final answer of the question is (.*)",
+    r"final answer to the question is (.*)",
+    r"final response: (.*)",
+    r"Final Response: (.*)",
+    r"the final response is (.*)",
+    r"Final Response is: (.*)",
+    r"Final Answer:\s*(.*)",
+    r"Final Answer is:\s*(.*)",
+    r"Answer:\s*(.*)",
+    r"The answer is:\s*(.*)",
+    r"Final Answer\s*\[([^\]]+)\]",
+    r"The final answer is\s*\[([^\]]+)\]",
+    r"The answer is\s*\[([^\]]+)\]",
+    r"Answer\s*\[([^\]]+)\]",
+    r"```json\n\{\"answer\":\s*\"(.*?)\"\n\}```", # For JSON encapsulated answers
+    r"```json\n\{\"answer\":\s*(.*?)\n\}```", # For JSON encapsulated answers (non-string)
+    r"\"answer\":\s*\"(.*?)\"",
+    r"\"answer\":\s*(.*)",
+    r"(\w[\w\s\.\-,\/]*)\s*$", # Catch any trailing string
+]
+
+# --- SmartAgent Class ---
+class SmartAgent:
+    def __init__(self, username: str, http_session: requests.Session):
+        self.username = username
+        self.http_session = http_session
+        self.agent_id = None
+        self.agent: Optional[Agent] = None # Type hint for clarity
+        self.pipelines = {} # Cache for pipelines
+
+        self.tool_code_cache: Dict[str, str] = {} # To cache tool code
+        logging.info("SmartAgent initialized.")
+
+    def setup(self):
+        logging.info("Setting up agent...")
         try:
-            return self.session.get(url, timeout=self.timeout, **kwargs)
-        except requests.exceptions.RequestException as e:
-            logging.error(f"HTTP GET error for {url}: {e}")
+            self.agent = self._initialize_agent()
+            logging.info("Agent setup complete.")
+        except Exception as e:
+            logging.exception(f"Error during agent setup: {e}")
             raise
-    
-    def post(self, url: str, **kwargs) -> requests.Response:
-        """POST request with error handling"""
+
+    def _initialize_agent(self) -> Agent:
+        logging.info("Initializing Hugging Face Agent...")
         try:
-            return self.session.post(url, timeout=self.timeout * 2, **kwargs)  # Double timeout for posts
-        except requests.exceptions.RequestException as e:
-            logging.error(f"HTTP POST error for {url}: {e}")
+            # Ensure HF_TOKEN is available in the environment or via login
+            hf_token = os.environ.get("HF_TOKEN")
+            if not hf_token:
+                # Fallback to get_inference_client's default token handling (e.g., from ~/.cache/huggingface/token)
+                logging.warning("HF_TOKEN environment variable not set. Relying on Hugging Face Hub client token.")
+                client = get_inference_client() # This will pick up token from env or cache
+            else:
+                client = InferenceClient(token=hf_token)
+
+            # Check if agent already exists for this username
+            list_agents_resp = self.http_session.get(f"{API_URL}/agents")
+            list_agents_resp.raise_for_status()
+            existing_agents = list_agents_resp.json()
+            logging.debug(f"Existing agents: {existing_agents}")
+
+            for agent_info in existing_agents:
+                if agent_info.get("username") == self.username:
+                    self.agent_id = agent_info["agent_id"]
+                    logging.info(f"Re-using existing agent with ID: {self.agent_id}")
+                    return Agent(id=self.agent_id, client=client)
+
+            # If no existing agent, create a new one
+            create_agent_resp = self.http_session.post(f"{API_URL}/agents", json={"username": self.username})
+            create_agent_resp.raise_for_status()
+            created_agent_info = create_agent_resp.json()
+            self.agent_id = created_agent_info["agent_id"]
+            logging.info(f"Created new agent with ID: {self.agent_id}")
+            return Agent(id=self.agent_id, client=client)
+
+        except requests.exceptions.RequestException as req_e:
+            logging.error(f"Network or API error during agent initialization: {req_e}")
+            raise
+        except Exception as e:
+            logging.error(f"Unexpected error during agent initialization: {e}")
             raise
 
-class SmartAgent:
-    """Enhanced AI agent with robust error handling and performance optimizations"""
-    
-    def __init__(self):
-        self.models_loaded = {"text": False, "image": False, "audio": False}
-        self.http_session = RobustHTTPSession()
-        self.temp_files = set()  # Track temp files for cleanup
-        
-        # Initialize models with better error handling
-        self._initialize_models()
-        
-        # Add a simple knowledge base for common patterns
-        self.knowledge_patterns = {
-            "opposite": {"left": "right", "right": "left", "up": "down", "down": "up", "in": "out", "out": "in"},
-            "math": {"plus": "+", "minus": "-", "times": "*", "divided by": "/"},
-            "common_responses": {
-                "youtube": "I cannot access YouTube videos",
-                "wikipedia": "I cannot access Wikipedia articles", 
-                "image": "I cannot process images in text questions",
-                "video": "I cannot process video content",
-                "excel": "I cannot process Excel files",
-                "python": "I cannot execute Python code"
-            }
-        }
-    
-    def _initialize_models(self) -> None:
-        """Initialize all ML models with comprehensive error handling"""
-        model_configs = [
-            ("text", "text-generation", "distilgpt2", "text_gen"),
-            ("image", "image-classification", "google/vit-base-patch16-224", "image_cls"),
-            ("audio", "automatic-speech-recognition", "openai/whisper-tiny", "audio_stt")
-        ]
-        
-        for model_type, task, model_name, attr_name in model_configs:
-            try:
-                start_time = time.time()
-                model = pipeline(task, model=model_name, device=-1)
-                load_time = time.time() - start_time
-                
-                setattr(self, attr_name, model)
-                self.models_loaded[model_type] = True
-                logging.info(f"✅ {model_type.title()} model loaded successfully in {load_time:.2f}s")
-                
-            except Exception as e:
-                logging.exception(f"❌ {model_type.title()} model load failed: {e}")
-                setattr(self, attr_name, None)
-    
-    def _simple_pattern_matching(self, question: str) -> Optional[str]:
-        """Simple pattern matching for common question types"""
-        question_lower = question.lower()
-        
-        # Handle reversed text questions
-        if question_lower.startswith(".rewsna eht sa"):
-            # This is asking for the opposite of "left" 
-            return "right"
-        
-        # Handle opposite questions
-        if "opposite" in question_lower:
-            for word, opposite in self.knowledge_patterns["opposite"].items():
-                if word in question_lower:
-                    return opposite
-        
-        # Handle grocery/shopping lists
-        if "grocery" in question_lower or "shopping" in question_lower:
-            return "I cannot create specific shopping lists without more context"
-        
-        # Handle specific platform references
-        for platform, response in self.knowledge_patterns["common_responses"].items():
-            if platform in question_lower:
-                return response
-        
-        return None
-    
-    def _generate_file_hash(self, task_id: str) -> str:
-        """Generate a unique hash for file naming"""
-        return hashlib.md5(f"{task_id}_{datetime.now().isoformat()}".encode()).hexdigest()[:8]
-    
-    def get_file_for_task(self, task_id: str) -> Optional[str]:
-        """Download file with enhanced error handling and validation"""
+    def _get_tool_code(self, tool_code_hash: str) -> str:
+        if tool_code_hash in self.tool_code_cache:
+            return self.tool_code_cache[tool_code_hash]
+
+        logging.info(f"Fetching tool code for hash: {tool_code_hash}")
         try:
-            logging.debug(f"📥 Downloading file for task {task_id}")
-            resp = self.http_session.get(f"{API_URL}/files/{task_id}")
-            
-            if resp.status_code != 200:
-                logging.error(f"❌ File download failed with status {resp.status_code} for task {task_id}")
-                return None
-            
-            # Check file size
-            content_length = resp.headers.get('content-length')
-            if content_length and int(content_length) > MAX_FILE_SIZE:
-                logging.error(f"❌ File too large ({content_length} bytes) for task {task_id}")
-                return None
-            
-            # Create secure temp file
-            file_hash = self._generate_file_hash(task_id)
-            file_path = os.path.join(TEMP_DIR, f"agent_{file_hash}_{task_id}.tmp")
-            
-            with open(file_path, "wb") as f:
-                f.write(resp.content)
-            
-            # Verify file was written correctly
-            if not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
-                logging.error(f"❌ File write verification failed for task {task_id}")
-                return None
-            
-            self.temp_files.add(file_path)
-            logging.debug(f"📁 File downloaded successfully for task {task_id}: {file_path}")
-            return file_path
-            
+            resp = self.http_session.get(f"{API_URL}/tool_code/{tool_code_hash}")
+            resp.raise_for_status()
+            tool_code = resp.json().get("tool_code", "")
+            if not tool_code:
+                raise ValueError(f"Tool code for hash {tool_code_hash} is empty.")
+            self.tool_code_cache[tool_code_hash] = tool_code
+            return tool_code
+        except requests.exceptions.RequestException as req_e:
+            logging.error(f"Failed to fetch tool code for {tool_code_hash}: {req_e}")
+            raise
         except Exception as e:
-            logging.exception(f"❌ Download error for task {task_id}: {e}")
-            return None
-    
-    def _cleanup_temp_file(self, file_path: str) -> None:
-        """Safely cleanup temporary files"""
+            logging.error(f"Error getting tool code: {e}")
+            raise
+
+    def _execute_agent(self, question: str, question_type: str, tools_code: Optional[List[Dict]] = None) -> str:
+        if not self.agent:
+            raise ValueError("Agent not initialized. Call setup() first.")
+
+        logging.info(f"Executing agent for question type '{question_type}': {question[:50]}...")
         try:
-            if file_path and os.path.exists(file_path):
-                os.remove(file_path)
-                self.temp_files.discard(file_path)
-                logging.debug(f"🗑️ Cleaned up temp file: {file_path}")
+            # Prepare special tools if provided by the API
+            special_tools = []
+            if tools_code:
+                for tool_def in tools_code:
+                    tool_code_hash = tool_def.get("tool_code_hash")
+                    tool_name = tool_def.get("tool_name")
+                    if tool_code_hash and tool_name:
+                        tool_code_str = self._get_tool_code(tool_code_hash)
+                        # Create a unique function name to avoid conflicts if `exec` is used multiple times
+                        unique_func_name = f"dynamic_tool_func_{hashlib.md5(tool_code_str.encode()).hexdigest()}"
+                        tool_code_str = tool_code_str.replace("def run_tool", f"def {unique_func_name}")
+                        
+                        global_vars = {}
+                        local_vars = {"inputs": None, "tool_code_hash": tool_code_hash} # inputs will be passed by agent
+                        exec(tool_code_str, global_vars, local_vars)
+                        
+                        if unique_func_name not in global_vars:
+                            raise ValueError(f"Function {unique_func_name} not found after executing tool code.")
+
+                        special_tools.append(
+                            Agent.create_tool(
+                                name=tool_name,
+                                description=f"Dynamically loaded tool for {tool_name}",
+                                function=global_vars[unique_func_name]
+                            )
+                        )
+                        logging.debug(f"Added dynamic tool: {tool_name}")
+
+            # Run the agent
+            agent_output = self.agent.run(question, additional_tools=special_tools if special_tools else None)
+
+            # Extract the last response from chat history
+            raw_answer = agent_output.chat_history[-1].response
+            logging.debug(f"Agent raw output: {raw_answer}")
+            return raw_answer
         except Exception as e:
-            logging.warning(f"⚠️ Failed to cleanup temp file {file_path}: {e}")
-    
-    def _process_text_question(self, question: str, task_id: str) -> ProcessingResult:
-        """Process text-based questions with enhanced validation"""
-        if not self.text_gen:
-            return ProcessingResult(False, "SKIPPED", "Text model not available")
-        
-        start_time = time.time()
-        try:
-            # First try simple pattern matching
-            simple_answer = self._simple_pattern_matching(question)
-            if simple_answer:
-                processing_time = time.time() - start_time
-                logging.debug(f"📝 Pattern-matched answer for {task_id}: '{simple_answer}'")
-                return ProcessingResult(True, simple_answer, None, processing_time)
-            
-            # For these complex questions, we need a different approach
-            # Many of these questions require web search or specific knowledge
-            # Let's try a more direct approach with better prompting
-            
-            # Check if question seems to need web search/external data
-            web_indicators = ["youtube.com", "wikipedia", "article", "published", "http", "www"]
-            if any(indicator in question.lower() for indicator in web_indicators):
-                # Return a generic response indicating web search is needed
-                return ProcessingResult(True, "I need access to web search to answer this question accurately.", None, time.time() - start_time)
-            
-            # For other questions, use a simple conversational approach
-            try:
-                # Use a cleaner prompt that's more likely to work
-                prompt = f"{question.strip()}\n\nAnswer:"
-                
-                result = self.text_gen(
-                    prompt,
-                    max_new_tokens=30,  # Even shorter responses
-                    do_sample=False,     # Use greedy decoding for consistency
-                    temperature=0.1,     # Very low temperature
-                    top_p=0.5,
-                    repetition_penalty=1.2,
-                    pad_token_id=self.text_gen.tokenizer.eos_token_id if hasattr(self.text_gen.tokenizer, 'eos_token_id') else None
-                )
-                
-                processing_time = time.time() - start_time
-                logging.debug(f"🧪 Raw text-gen output for {task_id}: {result}")
-                
-                # Extract the generated text
-                generated_text = result[0]["generated_text"]
-                
-                # Extract only the answer part
-                if "\n\nAnswer:" in generated_text:
-                    answer = generated_text.split("\n\nAnswer:")[-1].strip()
-                elif "Answer:" in generated_text:
-                    answer = generated_text.split("Answer:")[-1].strip()
-                else:
-                    # Take everything after the original question
-                    if question.strip() in generated_text:
-                        answer = generated_text.replace(question.strip(), "").strip()
-                    else:
-                        answer = generated_text.strip()
-                
-                # Clean the answer more aggressively
-                answer = ' '.join(answer.split())  # Normalize whitespace
-                answer = answer.replace('\n', ' ').replace('\r', ' ').strip()
-                
-                # Remove common artifacts and prefixes
-                artifacts = ["Answer:", "answer:", "A:", "a:", "Response:", "response:", "The answer is", "the answer is"]
-                for artifact in artifacts:
-                    if answer.startswith(artifact):
-                        answer = answer[len(artifact):].strip()
-                
-                # Remove trailing punctuation that might be artifacts
-                if answer.endswith(('.', '!', '?', ':')):
-                    answer = answer[:-1].strip()
-                
-                # Final validation
-                if not answer or len(answer) < 2:
-                    logging.warning(f"⚠️ Generated answer too short for {task_id}: '{answer}'")
-                    return ProcessingResult(True, "Unable to determine answer", None, processing_time)
-                
-                # Check for gibberish patterns more aggressively
-                words = answer.split()
-                if len(words) > 3:
-                    # Check for excessive repetition
-                    unique_words = set(words)
-                    if len(unique_words) / len(words) < 0.4:  # Less than 40% unique words
-                        logging.warning(f"⚠️ Detected repetitive output for {task_id}: '{answer}'")
-                        return ProcessingResult(True, "Unable to determine answer", None, processing_time)
-                
-                # Check for specific nonsensical patterns from tiny-gpt2
-                nonsense_indicators = [
-                    "Observ", "Motorola", "Hancock", "ESV", "Rh", "Prob", "antibiotic", "trilogy", "dispatch",
-                    "credibility", "autonomy", "intermittent", "circumcised", "hauled", "vendors", "scalp",
-                    "Participation", "conservation", "reborn", "stairs", "pawn", "Brew", "Jr", "Amph"
-                ]
-                if any(indicator in answer for indicator in nonsense_indicators):
-                    logging.warning(f"⚠️ Detected nonsensical output for {task_id}: '{answer}'")
-                    return ProcessingResult(True, "Unable to determine answer", None, processing_time)
-                
-                # Check for very long answers that are likely gibberish
-                if len(words) > 20:
-                    logging.warning(f"⚠️ Answer too long, likely gibberish for {task_id}: '{answer[:100]}...'")
-                    return ProcessingResult(True, "Unable to determine answer", None, processing_time)
-                
-                logging.debug(f"📝 Text answer for {task_id}: '{answer}' (processed in {processing_time:.2f}s)")
-                return ProcessingResult(True, answer, None, processing_time)
+            logging.error(f"Error during agent execution: {e}")
+            raise
+
+    def _extract_answer(self, raw_answer: str, question_type: str) -> str:
+        logging.debug(f"Extracting answer from raw_answer: {raw_answer}")
+        answer = "ERROR"
+        # Try to find a match using the defined patterns
+        for pattern in ANSWER_PATTERNS:
+            match = re.search(pattern, raw_answer, re.IGNORECASE | re.DOTALL)
+            if match:
+                extracted_content = match.group(1).strip()
+                # Basic cleanup
+                extracted_content = extracted_content.replace("\\n", "").replace("\\", "")
+                # Remove quotes if the answer is just a quoted string (e.g., "answer")
+                if extracted_content.startswith('"') and extracted_content.endswith('"'):
+                    extracted_content = extracted_content[1:-1]
+                if extracted_content.startswith("'") and extracted_content.endswith("'"):
+                    extracted_content = extracted_content[1:-1]
                 
-            except Exception as inner_e:
-                logging.warning(f"⚠️ Text generation failed for {task_id}: {inner_e}")
-                return ProcessingResult(True, "Unable to determine answer", None, time.time() - start_time)
-            
-        except Exception as e:
-            processing_time = time.time() - start_time
-            logging.exception(f"❌ Text processing error for {task_id}: {e}")
-            return ProcessingResult(False, "ERROR", str(e), processing_time)
-    
-    def _process_image_question(self, task_id: str) -> ProcessingResult:
-        """Process image-based questions with enhanced validation"""
-        if not self.image_cls:
-            return ProcessingResult(False, "SKIPPED", "Image model not available")
+                answer = extracted_content
+                logging.debug(f"Extracted answer using pattern '{pattern}': {answer}")
+                break # Stop after the first successful extraction
         
-        file_path = self.get_file_for_task(task_id)
-        if not file_path:
-            return ProcessingResult(False, "SKIPPED", "File download failed")
-        
-        start_time = time.time()
+        if answer == "ERROR" and raw_answer:
+            # If no pattern matches, try to use the raw answer if it looks reasonable
+            # (e.g., if the agent just directly gave the answer without a wrapper)
+            if len(raw_answer) < 200 and not any(kw in raw_answer.lower() for kw in ["thought", "tool", "action", "observation"]):
+                answer = raw_answer.strip()
+                logging.debug(f"No pattern matched, using raw answer directly: {answer}")
+
+        if not answer: # Ensure answer is not empty if extraction failed
+            answer = "ERROR"
+
+        return answer
+
+    def _load_pipeline(self, pipeline_name: str, **kwargs):
+        if pipeline_name not in self.pipelines:
+            logging.info(f"Loading pipeline: {pipeline_name}")
+            self.pipelines[pipeline_name] = pipeline(pipeline_name, **kwargs)
+        return self.pipelines[pipeline_name]
+
+    def _run_tool(self, tool_name: str, inputs: Dict[str, Any]) -> Any:
+        logging.info(f"Running tool: {tool_name} with inputs: {inputs}")
+        result = None
+        temp_file_paths = [] # To keep track of temp files for cleanup
+
         try:
-            # Validate image file
-            with Image.open(file_path) as img:
-                # Convert to RGB and validate
-                img = img.convert("RGB")
-                
-                # Check image dimensions
-                width, height = img.size
-                if width < 10 or height < 10:
-                    raise ValueError(f"Image too small: {width}x{height}")
-                
-                # Classify image
-                results = self.image_cls(img)
-                
-                processing_time = time.time() - start_time
-                
-                if not results or len(results) == 0:
-                    logging.warning(f"⚠️ No classification results for {task_id}")
-                    return ProcessingResult(False, "ERROR", "No classification results", processing_time)
-                
-                answer = results[0]["label"]
-                confidence = results[0].get("score", 0.0)
-                
-                logging.debug(f"🖼️ Image classification for {task_id}: '{answer}' (confidence: {confidence:.3f}, processed in {processing_time:.2f}s)")
-                return ProcessingResult(True, answer, None, processing_time)
+            if tool_name == "image-to-text":
+                image_bytes = self.http_session.get(inputs["image"]).content
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_img_file:
+                    tmp_img_file.write(image_bytes)
+                    image_path = tmp_img_file.name
+                    temp_file_paths.append(image_path)
                 
+                image = Image.open(image_path)
+                image_to_text_pipeline = self._load_pipeline("image-to-text")
+                result = image_to_text_pipeline(image)[0]["generated_text"]
+                logging.info(f"Image-to-text result: {result[:50]}...")
+
+            elif tool_name == "text-to-image":
+                text_to_image_pipeline = self._load_pipeline("text-to-image", model="runwayml/stable-diffusion-v1-5")
+                images = text_to_image_pipeline(inputs["text"])
+                if images and images.images:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_output_img_file:
+                        images.images[0].save(tmp_output_img_file.name)
+                        result = tmp_output_img_file.name # Return path or bytes as needed by agent
+                        temp_file_paths.append(result)
+                    logging.info(f"Text-to-image result saved to: {result}")
+                else:
+                    logging.warning("Text-to-image pipeline returned no images.")
+                    result = None
+
+            elif tool_name == "speech-to-text":
+                audio_bytes = self.http_session.get(inputs["audio"]).content
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".flac") as tmp_audio_file:
+                    tmp_audio_file.write(audio_bytes)
+                    audio_path = tmp_audio_file.name
+                    temp_file_paths.append(audio_path)
+
+                speech_to_text_pipeline = self._load_pipeline("automatic-speech-recognition", model="openai/whisper-tiny.en")
+                result = speech_to_text_pipeline(audio_path)["text"]
+                logging.info(f"Speech-to-text result: {result[:50]}...")
+
+            elif tool_name == "text-to-speech":
+                text_to_speech_pipeline = self._load_pipeline("text-to-speech", model="suno/bark-small")
+                speech = text_to_speech_pipeline(inputs["text"])
+                if speech and speech.audio is not None:
+                    with tempfile.NamedTemporaryFile(delete=False, suffix=".flac") as tmp_output_audio_file:
+                        sf.write(tmp_output_audio_file.name, speech.audio.numpy(), samplerate=speech.sampling_rate)
+                        result = tmp_output_audio_file.name # Return path
+                        temp_file_paths.append(result)
+                    logging.info(f"Text-to-speech result saved to: {result}")
+                else:
+                    logging.warning("Text-to-speech pipeline returned no audio.")
+                    result = None
+            else:
+                logging.warning(f"Unknown tool: {tool_name}. Skipping execution.")
+                return f"Error: Unknown tool {tool_name}"
+
+            return result
+
         except Exception as e:
-            processing_time = time.time() - start_time
-            logging.exception(f"❌ Image processing error for {task_id}: {e}")
-            return ProcessingResult(False, "ERROR", str(e), processing_time)
-        finally:
-            self._cleanup_temp_file(file_path)
-    
-    def _process_audio_question(self, task_id: str) -> ProcessingResult:
-        """Process audio-based questions with enhanced validation"""
-        if not self.audio_stt:
-            return ProcessingResult(False, "SKIPPED", "Audio model not available")
-        
-        file_path = self.get_file_for_task(task_id)
-        if not file_path:
-            return ProcessingResult(False, "SKIPPED", "File download failed")
-        
-        start_time = time.time()
-        try:
-            # Load and validate audio
-            audio_data, sample_rate = sf.read(file_path)
-            
-            # Validate audio data
-            if len(audio_data) == 0:
-                raise ValueError("Empty audio file")
-            
-            # Convert to mono if stereo
-            if len(audio_data.shape) > 1:
-                audio_data = np.mean(audio_data, axis=1)
-            
-            # Normalize audio
-            if np.max(np.abs(audio_data)) > 0:
-                audio_data = audio_data / np.max(np.abs(audio_data))
-            
-            # Transcribe audio
-            result = self.audio_stt({
-                "array": audio_data.astype(np.float32),
-                "sampling_rate": sample_rate
-            })
-            
-            processing_time = time.time() - start_time
-            logging.debug(f"🧠 Whisper raw output for {task_id}: {result}")
-            
-            text = result.get("text", "").strip()
-            
-            # Enhanced validation
-            if not text:
-                logging.warning(f"⚠️ Empty transcription for {task_id}")
-                return ProcessingResult(False, "ERROR", "Empty transcription", processing_time)
-            
-            # Check for common transcription artifacts
-            invalid_patterns = ["\x00", "�", "653", "[BLANK_AUDIO]", "Thank you.", "Thanks for watching."]
-            if any(pattern in text for pattern in invalid_patterns):
-                logging.warning(f"⚠️ Invalid transcription pattern detected for {task_id}: '{text}'")
-                return ProcessingResult(False, "ERROR", "Invalid transcription pattern", processing_time)
-            
-            # Additional validation for very short or repetitive text
-            if len(text) < 2 or len(set(text.split())) == 1:
-                logging.warning(f"⚠️ Suspicious transcription for {task_id}: '{text}'")
-                return ProcessingResult(False, "ERROR", "Suspicious transcription", processing_time)
-            
-            logging.debug(f"🎵 Audio transcription for {task_id}: '{text}' (processed in {processing_time:.2f}s)")
-            return ProcessingResult(True, text, None, processing_time)
-            
-        except Exception as e:
-            processing_time = time.time() - start_time
-            logging.exception(f"❌ Audio processing error for {task_id}: {e}")
-            return ProcessingResult(False, "ERROR", str(e), processing_time)
+            logging.error(f"Error running tool '{tool_name}': {e}", exc_info=True)
+            return f"Error running tool {tool_name}: {e}"
         finally:
-            self._cleanup_temp_file(file_path)
-    
-    def __call__(self, data: Dict[str, Any]) -> str:
-        """Process a question with comprehensive error handling and logging"""
-        question_type = data.get("type", "text")
-        question = data.get("question", "")
-        task_id = data.get("task_id", "UNKNOWN")
-        
-        logging.info(f"🔄 Processing {question_type} question for task {task_id}")
-        
-        # Route to appropriate processor
-        if question_type == "text":
-            result = self._process_text_question(question, task_id)
-        elif question_type == "image":
-            result = self._process_image_question(task_id)
-        elif question_type == "audio":
-            result = self._process_audio_question(task_id)
-        else:
-            logging.error(f"❌ Unknown question type: {question_type}")
-            return "ERROR"
-        
-        # Log result
-        if result.success:
-            logging.info(f"✅ Successfully processed {question_type} task {task_id} in {result.processing_time:.2f}s")
-        else:
-            logging.warning(f"⚠️ Failed to process {question_type} task {task_id}: {result.error}")
-        
-        return result.answer
-    
-    def cleanup(self) -> None:
-        """Cleanup all temporary files"""
-        for file_path in list(self.temp_files):
-            self._cleanup_temp_file(file_path)
-        logging.info("🧹 Cleanup completed")
-
-def run_and_submit(username: str) -> Generator[Tuple[str, str, str, str, Optional[pd.DataFrame]], None, None]:
-    """Enhanced main execution function with better error handling"""
+            # Clean up temporary files
+            for fp in temp_file_paths:
+                if os.path.exists(fp):
+                    try:
+                        os.unlink(fp)
+                        logging.debug(f"Cleaned up temporary file: {fp}")
+                    except OSError as e:
+                        logging.warning(f"Could not delete temporary file {fp}: {e}")
+
+    def cleanup(self):
+        logging.info("Cleaning up agent resources...")
+        if self.agent and self.agent_id:
+            try:
+                # Assuming the Hugging Face Agent object has a delete method or similar
+                # For this course, you might need to call a specific API endpoint or rely on
+                # the agent session expiring if no direct delete method is provided by the lib.
+                # If there's a specific API endpoint for deleting agents, use it here.
+                # For now, we'll log its ID as deleted.
+                # If agent.delete() exists, use it:
+                # self.agent.delete()
+                # For the course, simply logging and clearing internal state might be enough.
+                logging.info(f"Agent with ID {self.agent_id} is conceptually deleted (or will expire).")
+            except Exception as e:
+                logging.warning(f"Failed to delete agent or clean up its remote state: {e}")
+        # Clear cached pipelines and tool code
+        self.pipelines.clear()
+        self.tool_code_cache.clear()
+        logging.info("SmartAgent resources cleaned up.")
+
+
+# --- Main Run and Submit Function ---
+def run_and_submit(username: str) -> Generator[Tuple[str, str, str, str, pd.DataFrame], None, None]:
+    logging.info(f"Starting run and submit for user: {username}")
     
-    def generate_updates():
-        yield f"✅ Logged in as {username}", "Initializing...", "", "", None
-        
-        # Initialize agent
-        agent = SmartAgent()
-        
-        # Report model loading status
-        loaded_models = [k for k, v in agent.models_loaded.items() if v]
-        if loaded_models:
-            yield f"🤖 Models loaded: {', '.join(loaded_models)}", f"Models ready: {', '.join(loaded_models)}", "", "", None
-        else:
-            yield "❌ No models loaded successfully", "No models available", "", "", None
-            return
+    # Configure HTTP session with retries
+    retry_strategy = Retry(
+        total=3,
+        backoff_factor=1,
+        status_forcelist=[429, 500, 502, 503, 504],
+        allowed_methods=["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"] # Added "POST" for retries on agent API calls
+    )
+    adapter = HTTPAdapter(max_retries=retry_strategy)
+    http_session = requests.Session()
+    http_session.mount("https://", adapter)
+    http_session.mount("http://", adapter)
+
+    agent = SmartAgent(username=username, http_session=http_session)
+    answers: List[Dict[str, str]] = []
+    logs: List[Dict[str, Any]] = [] # List of dicts for DataFrame
+
+    try:
+        agent.setup()
         
         # Fetch questions
-        try:
-            logging.info("📥 Fetching questions from API...")
-            resp = agent.http_session.get(f"{API_URL}/questions")
-            resp.raise_for_status()
-            questions = resp.json()
-            logging.info(f"✅ Successfully fetched {len(questions)} questions")
-        except Exception as e:
-            error_msg = f"❌ Error fetching questions: {e}"
-            logging.exception(error_msg)
-            yield error_msg, "Failed to fetch questions", "", "", None
-            return
-        
-        # Process questions
-        answers = []
-        logs = []
-        total_questions = len(questions)
-        total_processing_time = 0
-        
-        logging.info(f"🚀 Starting to process {total_questions} questions")
-        yield f"📋 Processing {total_questions} questions...", f"0/{total_questions} questions processed", "", "", None
-        
-        for i, question_data in enumerate(questions, 1):
-            start_time = time.time()
-            
-            # Extract question details
-            question_preview = question_data.get('question', 'UNKNOWN')[:80] + "..." if len(question_data.get('question', '')) > 80 else question_data.get('question', 'UNKNOWN')
-            full_question = question_data.get('question', 'UNKNOWN')
-            question_type = question_data.get('type', 'UNKNOWN')
-            task_id = question_data.get('task_id', 'UNKNOWN')
-            
-            logging.info(f"📦 Processing Q{i}/{total_questions} (ID: {task_id}): {question_preview}")
-            
-            # Update UI
+        logging.info("Fetching questions from API...")
+        questions_resp = agent.http_session.get(f"{API_URL}/questions")
+        questions_resp.raise_for_status()
+        questions_data = questions_resp.json()
+        total_questions = len(questions_data)
+        logging.info(f"Received {total_questions} questions.")
+
+        yield (
+            "🚀 Starting agent run...",
+            f"0/{total_questions} questions processed",
+            "Initializing...",
+            "",
+            pd.DataFrame(logs)
+        )
+
+        for i, question_data in enumerate(questions_data, 1):
+            task_id = question_data.get("task_id")
+            question_type = question_data.get("question_type")
+            question_preview = question_data.get("question_preview")
+            full_question = question_data.get("question")
+            tools_code = question_data.get("tools_code") # List of tool definitions
+
+            if not all([task_id, question_type, question_preview, full_question]):
+                logging.warning(f"Skipping malformed question data: {question_data}")
+                logs.append(QuestionLog(
+                    question_num=i,
+                    question_preview="Malformed Data",
+                    question_type="N/A",
+                    answer="ERROR: Malformed Question",
+                    processing_time=0.0,
+                    status="ERROR"
+                ).__dict__)
+                continue
+
+            logging.info(f"Processing Q{i}/{total_questions} [Type: {question_type}]: {question_preview}")
             yield (
-                f"📦 Processing question {i}/{total_questions}: {question_preview}...",
+                f"Working on Q{i}/{total_questions}...",
                 f"{i-1}/{total_questions} questions processed",
                 f"Q{i} [{question_type}]: {full_question}",
                 "Processing...",
-                pd.DataFrame(logs) if logs else None
+                pd.DataFrame(logs)
             )
             
-            # Process question
+            start_time = time.time()
+            answer = "ERROR" # Default answer in case of unhandled error
+            status = "ERROR" # Default status
+
             try:
-                answer = agent(question_data)
+                agent_answer = agent._execute_agent(full_question, question_type, tools_code)
+                answer = agent._extract_answer(agent_answer, question_type)
                 processing_time = time.time() - start_time
-                total_processing_time += processing_time
-                
-                logging.info(f"✅ Q{i} completed in {processing_time:.2f}s: {answer}")
-                
+                logging.info(f"Q{i} [{question_type}] Answer: {answer} (Took {processing_time:.2f}s)")
+                status = "OK" if answer != "ERROR" else "ERROR" # Update status based on extracted answer
+
             except Exception as e:
-                processing_time = time.time() - start_time
-                logging.exception(f"❌ Unexpected error processing Q{i}: {e}")
-                answer = "ERROR"
-            
-            # Store results
-            answers.append({
-                "task_id": task_id,
-                "submitted_answer": answer
-            })
-            
-            # Update logs
-            log_entry = QuestionLog(
+                logging.exception(f"❌ Q{i} [{question_type}] Error processing question: {e}")
+                answer = "ERROR" # Ensure answer is "ERROR" on exception
+                # Calculate processing_time even on error if start_time was set
+                processing_time = time.time() - start_time if 'start_time' in locals() else 0.0
+                status = "ERROR" # Explicitly set status to ERROR
+
+            logs.append(QuestionLog(
                 question_num=i,
                 question_preview=question_preview,
                 question_type=question_type,
-                answer=answer or "SKIPPED",
+                answer=answer,
                 processing_time=processing_time,
-                status="✅" if answer not in ["ERROR", "SKIPPED"] else "❌" if answer == "ERROR" else "⏭️"
-            )
-            logs.append(log_entry.__dict__)
+                status=status
+            ).__dict__)
             
-            # Update UI with completed question
+            # Only append to answers list if not a critical error preventing an answer
+            # The course might expect "ERROR" as a valid answer for submission in some cases
+            answers.append({
+                "task_id": task_id,
+                "answer": answer
+            })
+
             yield (
-                f"✅ Completed question {i}/{total_questions} in {processing_time:.2f}s",
+                f"✅ Q{i}/{total_questions} processed",
                 f"{i}/{total_questions} questions processed",
                 f"Q{i} [{question_type}]: {full_question}",
-                answer or "SKIPPED",
+                f"Answer: {answer}",
                 pd.DataFrame(logs)
             )
-        
-        # Submit results
-        logging.info(f"📤 Submitting {len(answers)} answers (total processing time: {total_processing_time:.2f}s)")
-        yield f"📤 Submitting {len(answers)} answers...", f"{total_questions}/{total_questions} questions processed", "Submitting results...", "", pd.DataFrame(logs)
-        
-        try:
-            space_id = "dmfelder/DF_Final_Assignment_Template"
-            submission_body = {
-                "username": username,
-                "agent_code": f"https://huggingface.co/spaces/{space_id}/tree/main",
-                "answers": answers
-            }
-            
-            logging.debug(f"📤 Submitting payload: {len(answers)} answers")
-            response = agent.http_session.post(f"{API_URL}/submit", json=submission_body)
-            response.raise_for_status()
-            
-            result = response.json()
-            score = result.get("score", "N/A")
-            
-            success_msg = f"🏁 Final Score: {score}% (processed in {total_processing_time:.2f}s total)"
-            logging.info(f"✅ Submission successful: {score}%")
-            
-            yield success_msg, f"Complete! Score: {score}%", "Evaluation complete", f"Final score: {score}%", pd.DataFrame(logs)
-            
-        except Exception as e:
-            error_msg = f"❌ Submission error: {e}"
-            logging.exception(error_msg)
-            yield error_msg, "Submission failed", "Error occurred", str(e), pd.DataFrame(logs)
-        finally:
-            # Cleanup
-            agent.cleanup()
-    
-    return generate_updates()
-
-# Enhanced Gradio interface
-with gr.Blocks(
-    title="Enhanced Smart Agent",
-    theme=gr.themes.Soft(),
-    css="""
-    .gradio-container {
-        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-    }
-    .status-box {
-        background-color: #f8f9fa;
-        border-radius: 8px;
-        padding: 10px;
-    }
-    """
-) as app:
-    gr.Markdown("""
-    # 🤖 Enhanced Smart Agent Evaluation
-    
-    **Features:**
-    - Robust error handling and recovery
-    - Enhanced model performance
-    - Comprehensive logging and monitoring
-    - Improved file handling and cleanup
-    - Better answer validation and processing
-    """)
-    
-    with gr.Row():
-        login = gr.LoginButton(value="Login Required")
-        run = gr.Button("🚀 Run Enhanced Agent", variant="primary", size="lg")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            status = gr.Textbox(
-                lines=12,
-                label="📊 Status & Progress",
-                placeholder="Ready to start...",
-                elem_classes=["status-box"]
-            )
-            progress = gr.Textbox(
-                label="🔄 Current Progress",
-                value="Ready to start...",
-                interactive=False
-            )
-        
-        with gr.Column(scale=1):
-            current_question = gr.Textbox(
-                lines=4,
-                label="❓ Current Question",
-                placeholder="No question being processed...",
-                interactive=False
-            )
-            current_answer = gr.Textbox(
-                label="💬 Current Answer",
-                placeholder="No answer yet...",
-                interactive=False
-            )
-    
-    results = gr.DataFrame(
-        label="📈 Processing Results",
-        headers=["Q", "Question", "Type", "Answer", "Processing Time", "Status"],
-        wrap=True
-    )
-    
-    with gr.Row():
-        log_file = gr.File(label="📋 Download Detailed Log")
-        gr.Markdown("*Log file will be available after processing completes*")
-    
-    def start_processing(profile):
-        """Start the agent processing with enhanced error handling"""
-        if not profile:
-            yield "❌ You must log in first.", "Not logged in", "", "", pd.DataFrame(), None
-            return
-        
-        updates = []
-        final_df = None
-        
-        try:
-            for msg, prog, question, answer, df in run_and_submit(profile):
-                updates.append(msg)
-                if df is not None:
-                    final_df = df
-                yield (
-                    "\n".join(updates[-10:]),  # Show last 10 updates
-                    prog,
-                    question,
-                    answer,
-                    final_df if final_df is not None else pd.DataFrame(),
-                    None
-                )
-            
-            # Final yield includes log file
-            yield (
-                "\n".join(updates[-10:]),
-                prog,
-                question,
-                answer,
-                final_df if final_df is not None else pd.DataFrame(),
-                log_path
-            )
-            
-        except Exception as e:
-            error_msg = f"❌ Critical error during processing: {e}"
-            logging.exception(error_msg)
-            yield (
-                error_msg,
-                "Critical error occurred",
-                "Error",
-                str(e),
-                pd.DataFrame(),
-                log_path
-            )
-    
-    run.click(
-        start_processing,
-        inputs=[login],
-        outputs=[status, progress, current_question, current_answer, results, log_file]
+
+        # Submit answers
+        logging.info("📤 Submitting answers to API")
+        resp = agent.http_session.post(f"{API_URL}/submit", json={
+            "username": username,
+            "answers": answers
+        })
+        resp.raise_for_status()
+        result = resp.json()
+        logging.info(f"✅ Submission successful: {result}")
+        yield (
+            "🎉 Submission completed successfully!",
+            f"Final Score: {result.get('score', 'Unknown')}",
+            f"Details: {result.get('message', '')}",
+            "",
+            pd.DataFrame(logs)
+        )
+    except requests.exceptions.RequestException as req_e:
+        logging.exception(f"❌ API Request failed during run or submission: {req_e}")
+        yield (
+            "❌ API Error",
+            "Submission failed due to network/API error",
+            str(req_e),
+            "",
+            pd.DataFrame(logs)
+        )
+    except Exception as e:
+        logging.exception(f"❌ An unexpected error occurred during run or submission: {e}")
+        yield (
+            "❌ Run Failed",
+            "An unexpected error occurred",
+            str(e),
+            "",
+            pd.DataFrame(logs)
+        )
+    finally:
+        agent.cleanup()
+
+
+# --- Gradio UI ---
+with gr.Blocks() as demo:
+    username_input = gr.Textbox(label="Username")
+    status_display = gr.Textbox(label="Status")
+    progress_display = gr.Textbox(label="Progress")
+    question_display = gr.Textbox(label="Current Question")
+    answer_display = gr.Textbox(label="Answer")
+    logs_display = gr.Dataframe(label="Log", headers=["question_num", "question_preview", "question_type", "answer", "processing_time", "status"])
+
+    run_button = gr.Button("Run Agent")
+
+    def start_run(username):
+        # This function acts as a simple wrapper to call the generator
+        yield from run_and_submit(username)
+
+    run_button.click(
+        fn=start_run,
+        inputs=[username_input],
+        outputs=[
+            status_display,
+            progress_display,
+            question_display,
+            answer_display,
+            logs_display
+        ]
     )
 
 if __name__ == "__main__":
-    logging.info("🚀 Starting Enhanced Smart Agent application")
-    app.launch(
-        share=False,
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        quiet=False
-        )
+    demo.launch()