app.py

import gradio as gr
import requests
from smolagents import CodeAgent
import json
import re
import logging
from tenacity import retry, stop_after_attempt, wait_exponential

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Use lowercase and verify your exact Hugging Face Space URL
HF_SPACE_URL = "https://manavraj-troubleshoot-mcp.hf.space"  # Note: Changed underscores to hyphens

@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def call_mcp_server(message, tool_type="knowledge_base"):
    """Call your MCP server with retries and proper error handling"""
    try:
        # Determine which tool to use based on tool_type
        fn_index_map = {
            "knowledge_base": 0,
            "web_search": 1, 
            "formatter": 2
        }
        
        fn_index = fn_index_map.get(tool_type, 0)
        
        logger.info(f"Calling MCP server with message: {message[:50]}...")
        
        response = requests.post(
            f"{HF_SPACE_URL}/api/predict",
            json={"data": [message], "fn_index": fn_index},
            verify=False,  # Temporarily disable SSL verification
            timeout=30
        )
        
        if response.status_code != 200:
            raise Exception(f"MCP server returned {response.status_code}")
            
        data = response.json()
        if not isinstance(data, dict) or 'data' not in data:
            raise Exception("Invalid MCP response format")
            
        return data['data'][0]
        
    except Exception as e:
        logger.error(f"MCP server call failed: {str(e)}")
        return f"MCP server error: {str(e)}"

def extract_thought_action_observation(response):
    """Extract and format the thought-action-observation cycle from agent response"""
    sections = {
        'thought': '',
        'action': '',
        'observation': ''
    }
    
    # Try to extract sections using regex
    thought_match = re.search(r'(?:THOUGHT|Thought):\s*(.*?)(?=(?:ACTION|Action|OBSERVATION|Observation|$))', response, re.DOTALL | re.IGNORECASE)
    action_match = re.search(r'(?:ACTION|Action):\s*(.*?)(?=(?:OBSERVATION|Observation|FINAL|Final|$))', response, re.DOTALL | re.IGNORECASE)
    observation_match = re.search(r'(?:OBSERVATION|Observation):\s*(.*?)(?=(?:FINAL|Final|$))', response, re.DOTALL | re.IGNORECASE)
    
    if thought_match:
        sections['thought'] = thought_match.group(1).strip()
    if action_match:
        sections['action'] = action_match.group(1).strip()
    if observation_match:
        sections['observation'] = observation_match.group(1).strip()
    
    return sections

# Initialize CodeAgent with enhanced system prompt
agent = CodeAgent(
    tools=[],
    model="microsoft/DialoGPT-medium",
    system_prompt="""You are a Technical Support AI Agent that uses a structured Thought-Action-Observation cycle to solve user queries effectively.
WORKFLOW:
For every user query, you MUST follow this exact structure:
THOUGHT: Analyze the user's problem and decide what action to take
- Identify the type of issue (technical problem, information request, formatting need)
- Determine which MCP server tool would be most appropriate
- Consider what information you need to provide a complete answer
ACTION: Specify exactly what action you will take
- "search_knowledge_base" for known technical issues (wifi, screen, hardware problems)
- "web_search" for current information, news, or unknown topics
- "format_response" for organizing or structuring information
- "direct_response" if you can answer without tools
OBSERVATION: Analyze the results and determine next steps
- Evaluate if the action provided sufficient information
- Identify if additional actions are needed
- Note any gaps in the information
FINAL RESPONSE: Provide a clear, helpful answer to the user
EXAMPLE:
User: "My wifi keeps disconnecting"
THOUGHT: The user has a wifi connectivity issue. This is a common technical problem that likely has solutions in our knowledge base. I should search the knowledge base first for wifi-related solutions.
ACTION: search_knowledge_base with query "wifi disconnecting problem"
OBSERVATION: The knowledge base provided basic wifi troubleshooting steps. These steps are helpful but I should also check for any recent wifi issues or advanced solutions via web search to provide comprehensive help.
FINAL RESPONSE: [Provide complete solution combining knowledge base and any additional insights]
Remember:
- Always use the THOUGHT-ACTION-OBSERVATION-FINAL RESPONSE structure
- Be specific about which tool you're using and why
- Provide actionable, clear solutions
- If one action isn't sufficient, explain what additional actions you'll take
{{managed_agents_descriptions}}
{{authorized_imports}}"""
)

def determine_tool_type(message):
    """Determine which MCP tool to use based on user message"""
    message_lower = message.lower()
    
    # Technical issues - use knowledge base
    tech_keywords = ["wifi", "screen", "computer", "laptop", "error", "problem", "issue", "broken", "not working"]
    if any(keyword in message_lower for keyword in tech_keywords):
        return "knowledge_base"
    
    # Search queries - use web search
    search_keywords = ["search", "find", "news", "latest", "current", "what is", "how to", "information about"]
    if any(keyword in message_lower for keyword in search_keywords):
        return "web_search"
    
    # Formatting requests - use formatter
    format_keywords = ["format", "organize", "list", "steps", "number"]
    if any(keyword in message_lower for keyword in format_keywords):
        return "formatter"
    
    # Default to knowledge base for technical support
    return "knowledge_base"

def chat_interface(message, history):
    """Enhanced chat interface with thought-action-observation cycle"""
    
    try:
        # Step 1: Let agent think about the problem
        thinking_prompt = f"""
        User Query: "{message}"
        
        Follow the THOUGHT-ACTION-OBSERVATION cycle:
        
        THOUGHT: Analyze this query and determine the best approach.
        """
        
        # Get agent's initial thought process
        agent_response = agent.run(thinking_prompt)
        if not isinstance(agent_response, str):
            raise ValueError("Agent returned non-string response")
        
        # Extract thought and determine action
        cycle_parts = extract_thought_action_observation(agent_response)
        
        # Determine tool type based on message content and agent's thought
        tool_type = determine_tool_type(message)
        
        # Step 2: Execute the action via MCP server
        mcp_response = call_mcp_server(message, tool_type)
        
        # Step 3: Let agent observe and provide final response
        final_prompt = f"""
        User Query: "{message}"
        
        THOUGHT: {cycle_parts.get('thought', 'Analyzing user query for appropriate technical support response.')}
        
        ACTION: Used {tool_type} tool with query: "{message}"
        
        OBSERVATION: The MCP server returned: "{mcp_response}"
        
        FINAL RESPONSE: Based on the observation, provide a complete, helpful response to the user. If the MCP server response is insufficient, explain what additional steps might be needed.
        """
        
        final_response = agent.run(final_prompt)
        if not isinstance(final_response, str):
            final_response = str(final_response)
        
        # Format the response to show the thinking process
        formatted_response = f"""🤔 **THOUGHT:** {cycle_parts.get('thought', 'Analyzing your query...')}
⚡ **ACTION:** Used {tool_type.replace('_', ' ')} tool
👁️ **OBSERVATION:** {mcp_response[:200]}{'...' if len(mcp_response) > 200 else ''}
✅ **SOLUTION:**
{final_response}"""
        
        return formatted_response
        
    except Exception as e:
        logger.error(f"Error in chat interface: {str(e)}")
        # Fallback to direct MCP response with simple structure
        tool_type = determine_tool_type(message)
        mcp_response = call_mcp_server(message, tool_type)
        
        return f"""🤔 **THOUGHT:** You have a {tool_type.replace('_', ' ')} related query.
⚡ **ACTION:** Consulted the MCP server using {tool_type} tool.
👁️ **OBSERVATION:** Found relevant information.
✅ **SOLUTION:**
{mcp_response}"""

# Create Gradio interface
demo = gr.ChatInterface(
    fn=chat_interface,
    title="🔧 MCP Technical Support Agent (TAO Cycle)", 
    description="Technical support powered by MCP server tools using Thought-Action-Observation methodology",
    examples=[
        "I have a wifi connection problem",
        "My screen keeps flickering", 
        "Search for latest cybersecurity news",
        "Format these steps: Install driver. Restart computer. Test connection"
    ],
    chatbot=gr.Chatbot(height=600)
)

if __name__ == "__main__":
    demo.launch()