Update smolagents_agent.py

smolagents_agent.py

@@ -581,7 +581,6 @@ class OptimizedSmolagentsGAIAgent:
             print(f"Error initializing HuggingFace model: {e}")
             return None
 
-
     def classify_question(self, question: str) -> Dict[str, Any]:
         """Enhanced question classification with confidence scores"""
         q_lower = question.lower()
@@ -667,27 +666,37 @@ class OptimizedSmolagentsGAIAgent:
             result_str = str(result)
             
             # Format response for HuggingFace compatibility
-            # Wrap any code blocks in <code> tags as expected by the parsing system
-            if "```" in result_str or "def " in result_str or "import " in result_str:
-                # Extract code blocks and wrap them properly
+            # Only wrap actual code blocks, not explanatory text
+            if "```" in result_str:
+                # Extract markdown code blocks and wrap them properly
                 import re
-                # Find Python code blocks (either fenced with ```python or ``` or plain code)
-                code_pattern = r'```(?:python)?\n(.*?)\n```|```\n(.*?)\n```|([a-zA-Z_][a-zA-Z0-9_]*\s*=\s*.*|def\s+\w+.*|import\s+\w+.*|from\s+\w+.*import.*)'
+                # Find Python code blocks (either fenced with ```python or ```)
+                code_pattern = r'```(?:python)?\n(.*?)\n```|```\n(.*?)\n```'
                 matches = re.findall(code_pattern, result_str, re.DOTALL)
                 
                 if matches:
                     formatted_result = result_str
                     for match in matches:
-                        code_content = match[0] if match[0] else match[1] if match[1] else match[2]
-                        if code_content:
+                        code_content = match[0] if match[0] else match[1]
+                        if code_content and code_content.strip():
                             # Wrap in <code> tags as expected by HuggingFace parser
                             wrapped_code = f"<code>{code_content.strip()}</code>"
-                            formatted_result = formatted_result.replace(code_content.strip(), wrapped_code)
+                            # Replace the original code block with wrapped version
+                            formatted_result = formatted_result.replace(f"```{code_content.strip()}```", wrapped_code)
+                            # Handle ```python variant
+                            formatted_result = formatted_result.replace(f"```python\n{code_content.strip()}```", wrapped_code)
                     result_str = formatted_result
-                else:
-                    # If no code blocks found but response contains code-like patterns, wrap entire response
-                    if any(keyword in result_str.lower() for keyword in ['def ', 'import ', 'return ', 'print(', '# ']):
-                        result_str = f"<code>{result_str}</code>"
+            
+            # Only wrap if it's clearly actual code (function definitions, imports, etc.)
+            # and NOT explanatory text or thoughts
+            if (result_str.startswith(('def ', 'import ', 'from ', 'class ', '# ')) or 
+                (result_str.count('\n') == 0 and ('=' in result_str or '(' in result_str))):
+                # Make sure it's actual code, not explanatory text
+                if not any(phrase in result_str.lower() for phrase in [
+                    'first thought', 'i need to', 'let me', 'here\'s', 'here is',
+                    'to solve', 'the answer is', 'result is', 'response is'
+                ]):
+                    result_str = f"<code>{result_str}</code>"
             
             # Cache the result
             self.cache.set(cache_key, result_str)
@@ -695,6 +704,10 @@ class OptimizedSmolagentsGAIAgent:
             return result_str
             
         except Exception as e:
+            # Handle timeout errors gracefully
+            error_msg = str(e)
+            if "Gateway Time-out" in error_msg or "timeout" in error_msg.lower():
+                return f"Model timeout occurred. Please try again or simplify the question. Error: {error_msg}"
             return f"Agent processing error: {str(e)}"
 
     def get_tool_recommendations(self, question: str) -> List[str]:
@@ -729,7 +742,6 @@ if __name__ == "__main__":
 
     print("=== OPTIMIZED SMOLAGENTS AGENT TEST ===\n")
     
-    
     for i, question in enumerate(test_questions, 1):
         print(f"Test {i}: {question}")
         recommendations = agent.get_tool_recommendations(question)