Update app.py

app.py

@@ -3,6 +3,7 @@ import requests
 from smolagents import CodeAgent
 import logging
 from tenacity import retry, stop_after_attempt, wait_exponential
+from typing import Optional
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -12,7 +13,7 @@ logger = logging.getLogger(__name__)
 HF_SPACE_URL = "https://manavraj-troubleshoot-mcp.hf.space"
 
 @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
-def call_mcp_server(message, tool_type="knowledge_base"):
+def call_mcp_server(message: str, tool_type: str = "knowledge_base") -> str:
     """Call MCP server with proper endpoint routing"""
     try:
         endpoint_map = {
@@ -37,7 +38,7 @@ def call_mcp_server(message, tool_type="knowledge_base"):
         if not isinstance(data, dict) or 'data' not in data:
             raise Exception("Invalid MCP response format")
             
-        return data['data'][0]
+        return str(data['data'][0])
         
     except Exception as e:
         logger.error(f"MCP call failed: {str(e)}")
@@ -45,20 +46,18 @@ def call_mcp_server(message, tool_type="knowledge_base"):
 
 # Initialize CodeAgent with proper system prompt including required tags
 agent = CodeAgent(
-    tools=[],  # No tools defined here as they're in MCP
+    tools=[call_mcp_server],
     model="microsoft/DialoGPT-medium",
     system_prompt="""{{authorized_imports}}
 - requests for API calls
 - standard Python libraries
-
 {{managed_agents_descriptions}}
-
 You are an advanced Technical Support Assistant with these capabilities:
 1. Troubleshooting technical issues (wifi, hardware, software)
 2. Finding information through web search
 3. Formatting instructions into clear steps
 
-Tools available via MCP server:
+Available Tools:
 - knowledge_base: For technical issues (wifi, screen, sound problems)
 - web_search: For finding latest information or non-technical queries
 - formatter: To organize raw steps into numbered lists
@@ -72,19 +71,25 @@ Response Guidelines:
 
 Example workflow:
 THOUGHT: User has wifi issue, should check knowledge base
-ACTION: knowledge_base
+ACTION: knowledge_base("wifi keeps disconnecting")
 OBSERVATION: Found 4 troubleshooting steps
 FINAL RESPONSE: Here's how to fix your wifi: [steps]
 
 Never show THOUGHT/ACTION/OBSERVATION to user - only final response."""
 )
 
-def chat_interface(message, history):
+def chat_interface(message: str, history: Optional[list] = None) -> str:
     """Simplified chat interface"""
     try:
         # Let agent handle the complete process
         response = agent.run(message)
-        return str(response).split("FINAL RESPONSE:")[-1].strip()
+        
+        # Extract final response and clean it up
+        if isinstance(response, str):
+            if "FINAL RESPONSE:" in response:
+                return response.split("FINAL RESPONSE:")[-1].strip()
+            return response
+        return str(response)
     except Exception as e:
         logger.error(f"Chat error: {str(e)}")
         return f"Error processing request: {str(e)}"
@@ -97,7 +102,8 @@ demo = gr.ChatInterface(
         "My wifi keeps disconnecting",
         "Find the latest Windows 11 update",
         "Format: Restart. Check connections. Update drivers"
-    ]
+    ],
+    description="A technical support assistant that can troubleshoot issues, search the web, and format instructions."
 )
 
 if __name__ == "__main__":