Update code_agent.py

code_agent.py

@@ -3,9 +3,13 @@ Real `smolagents` CodeAgent integration for the final project.
 
 This module initializes a `CodeAgent` with standard tools and exposes a 
 `run_agent(prompt)` function returning the final answer string.
-"""
-
 
+Notes:
+- Requires `smolagents` installed in the environment.
+- For serverless inference via Hugging Face, you must be logged in 
+  (`huggingface-cli login`) or have `HF_TOKEN` environment variable set,
+  and have sufficient provider credits.
+"""
 from typing import List, Any
 import logging
 import os
@@ -22,12 +26,29 @@ from smolagents import (
 )
 
 # Get HF token from environment (set by huggingface-cli login or HF_TOKEN env var)
-HF_TOKEN = os.environ.get("HF_TOKEN")
-if not HF_TOKEN:
+def _get_hf_token() -> str | None:
+    """Get HF token from environment or huggingface_hub cache."""
+    token = os.environ.get("HF_TOKEN")
+    if token:
+        return token
+    
+    # Fallback: try huggingface_hub.get_token() if available
+    try:
+        from huggingface_hub import get_token
+        token = get_token()
+        if token:
+            logger.info("✓ Using HF token from huggingface_hub cache")
+            return token
+    except ImportError:
+        pass
+    
     logger.warning(
-        "HF_TOKEN not found in environment. Run 'huggingface-cli login' or set HF_TOKEN env var. "
+        "⚠️  HF_TOKEN not found in environment. Run 'huggingface-cli login' or set HF_TOKEN env var. "
         "CodeAgent initialization will fail without valid credentials."
     )
+    return None
+
+HF_TOKEN = _get_hf_token()
 
 
 def make_code_agent(
@@ -72,14 +93,30 @@ def make_code_agent(
         if model_name:
             try:
                 model = InferenceClientModel(model_id=model_name, token=HF_TOKEN)
-            except Exception:
+                logger.info(f"InferenceClientModel initialized with custom model: {model_name}")
+            except Exception as e:
+                logger.warning(f"Failed to initialize custom model {model_name}: {e}. Falling back to OLMo-Instruct")
                 try:
                     model = InferenceClientModel(model_id='allenai/Olmo-3-7B-Instruct', token=HF_TOKEN)
-                except Exception:
+                    logger.info("InferenceClientModel initialized with fallback model: allenai/Olmo-3-7B-Instruct")
+                except Exception as e2:
+                    logger.warning(f"Failed to initialize OLMo-Instruct: {e2}. Falling back to default model")
                     model = InferenceClientModel(token=HF_TOKEN)
+                    logger.info("InferenceClientModel initialized with default model")
         else:
-            model = InferenceClientModel(token=HF_TOKEN)
-        logger.info("InferenceClientModel initialized successfully with HF_TOKEN")
+            # Default to OLMo-Think model
+            try:
+                model = InferenceClientModel(model_id='allenai/Olmo-3-7B-Think', token=HF_TOKEN)
+                logger.info("InferenceClientModel initialized with default model: allenai/Olmo-3-7B-Think")
+            except Exception as e:
+                logger.warning(f"Failed to initialize OLMo-Think: {e}. Falling back to OLMo-Instruct")
+                try:
+                    model = InferenceClientModel(model_id='allenai/Olmo-3-7B-Instruct', token=HF_TOKEN)
+                    logger.info("InferenceClientModel initialized with fallback model: allenai/Olmo-3-7B-Instruct")
+                except Exception as e2:
+                    logger.warning(f"Failed to initialize OLMo-Instruct: {e2}. Using default model")
+                    model = InferenceClientModel(token=HF_TOKEN)
+                    logger.info("InferenceClientModel initialized with default model")
     except Exception as e:
         logger.error(
             "InferenceClientModel initialization failed (ensure HF_TOKEN is set and has credits): %s", e
@@ -90,6 +127,55 @@ def make_code_agent(
     return agent
 
 
+def _extract_clean_answer(text: str) -> str:
+    """Clean agent output by removing markdown, code blocks, and formatting.
+    
+    Handles:
+    - Code fences (```python ... ```)
+    - Code keywords (import, def, class, from)
+    - Markdown formatting (*#-`)
+    - Excess whitespace
+    
+    Args:
+        text: raw agent output (may contain code, markdown)
+    
+    Returns:
+        Clean plain-text answer string.
+    """
+    lines = text.strip().split('\n')
+    cleaned_lines = []
+    in_code_block = False
+    
+    for line in lines:
+        stripped = line.strip()
+        
+        # Track code blocks
+        if '```' in stripped:
+            in_code_block = not in_code_block
+            continue
+        
+        # Skip code blocks
+        if in_code_block:
+            continue
+        
+        # Skip lines that look like code
+        if stripped.startswith(('import ', 'from ', 'def ', 'class ', '>>>')):
+            continue
+        
+        # Skip empty lines at boundaries
+        if stripped:
+            cleaned_lines.append(stripped)
+    
+    # Join and clean markdown formatting
+    answer = ' '.join(cleaned_lines)
+    
+    # Remove markdown bold/italic/strikethrough
+    for char in ('*', '**', '_', '~~'):
+        answer = answer.replace(char, '')
+    
+    return answer.strip()
+
+
 _AGENT_SINGLETON = None
 
 
@@ -97,76 +183,97 @@ def get_agent():
     """Get or create the singleton CodeAgent instance."""
     global _AGENT_SINGLETON
     if _AGENT_SINGLETON is None:
-        _AGENT_SINGLETON = make_code_agent(model_name="allenai/Olmo-3-7B-Think")
+        logger.info("⏳ Initializing CodeAgent singleton with model: allenai/Olmo-3-7B-Think")
+        try:
+            _AGENT_SINGLETON = make_code_agent(model_name="allenai/Olmo-3-7B-Think")
+            logger.info("✓ CodeAgent singleton initialized successfully")
+        except Exception as e:
+            logger.error("❌ Failed to initialize CodeAgent: %s", e)
+            raise
     return _AGENT_SINGLETON
 
 
 def run_agent(prompt: str) -> str:
-    """
-    Run the CodeAgent and return the final answer string.
-        
+    """Run the CodeAgent and return the final answer string.
+    
     Args:
         prompt: the reasoning task/question for the agent.
     
     Returns:
-        The agent's final answer as a string.
-
+        The agent's final answer as a string (plain-text, no code/markdown).
+    
     Raises:
         Exception: if CodeAgent.run fails (e.g., no HF credentials or credits).
     """
     agent = get_agent()
 
-    # System-level instruction to avoid returning raw code or chained tool
-    # calls in a single code block. We require a single-line plain-text
-    # final answer with no surrounding code fences or explanation.
+    # System-level instruction to force plain-text output only
     system_instruction = (
-        "SYSTEM_INSTRUCTION: Return ONLY the concise final answer as plain text on a single line. "
-        "Do NOT produce code blocks, do NOT chain tool calls inside the same code block, and do NOT include explanations."
+        "IMPORTANT: Your final answer MUST be plain text only. "
+        "Do NOT produce code blocks (```), do NOT return Python code, and do NOT include explanations. "
+        "Return ONLY the direct answer as a single line or paragraph of plain text."
     )
 
     full_prompt = f"{system_instruction}\n\n{prompt}"
 
     try:
         res = agent.run(full_prompt)
+        logger.debug("Agent raw response type: %s", type(res).__name__)
     except Exception as e:
-        # Surface the error (likely HF auth/credits); re-raise for the caller to handle
-        logger.error("CodeAgent.run failed: %s", e)
+        # Surface the error (likely HF auth/credits)
+        logger.error("❌ CodeAgent.run failed: %s", e)
         raise
 
-    # If the agent produced a dict-like response, prefer canonical keys
+    # Handle dict-like responses
     if isinstance(res, dict):
-        for key in ("answer", "final_answer", "final"):
+        for key in ("answer", "final_answer", "final", "output"):
             if key in res and isinstance(res[key], str):
-                return res[key].strip()
-        # fallback: stringify
-        return str(res).strip()
+                answer = _extract_clean_answer(res[key])
+                logger.debug("Extracted answer from key '%s': %s", key, answer[:50])
+                return answer
+        # Fallback: stringify
+        return _extract_clean_answer(str(res))
 
-    # If model output contains code fences or appears to be code, try a single-pass retry
+    # Handle string responses
     if isinstance(res, str):
         text = res.strip()
-        # naive detection of code fences or '```' or 'import ' which suggests code
-        if "```" in text or text.lstrip().startswith("import ") or "def " in text:
-            logger.warning("Agent returned code-like output; retrying with stricter instruction")
+        
+        # Check for code-like output
+        has_code_fence = "```" in text
+        has_code_keyword = any(text.lstrip().startswith(kw) for kw in ["import ", "def ", "class ", ">>>"])
+        looks_like_code = has_code_fence or has_code_keyword
+        
+        if looks_like_code:
+            logger.warning("⚠️  Agent returned code-like output; retrying with stricter instruction")
             stricter = (
-                "STRICT_INSTRUCTION: You must return only the final answer. No code, no markdown, no explanations. "
-                "If the answer is a number, return only the number."
+                "FINAL INSTRUCTION: Return ONLY the plain-text answer. No code blocks, no code, no explanations. "
+                "If the answer is a number, return only the number. If it's a date, return only the date. "
+                "Do not explain your answer."
             )
             retry_prompt = f"{stricter}\n\n{prompt}"
             try:
                 retry_res = agent.run(retry_prompt)
                 if isinstance(retry_res, str):
-                    return retry_res.strip()
+                    answer = _extract_clean_answer(retry_res)
+                    logger.info("✓ Retry successful, extracted clean answer")
+                    return answer
                 if isinstance(retry_res, dict):
-                    for key in ("answer", "final_answer", "final"):
+                    for key in ("answer", "final_answer", "final", "output"):
                         if key in retry_res and isinstance(retry_res[key], str):
-                            return retry_res[key].strip()
-                    return str(retry_res).strip()
+                            answer = _extract_clean_answer(retry_res[key])
+                            return answer
+                    return _extract_clean_answer(str(retry_res))
             except Exception as e:
-                logger.error("Retry run failed: %s", e)
-                # fallthrough to return original text
-        return text
+                logger.error("❌ Retry failed: %s", e)
+                # Fallthrough to clean original response
+        
+        # Clean and return
+        answer = _extract_clean_answer(text)
+        logger.debug("✓ Final answer (cleaned): %s", answer[:80])
+        return answer
 
-    return str(res)
+    # Fallback for other types
+    return _extract_clean_answer(str(res))
 
 
 if __name__ == "__main__":