Update app.py

app.py

@@ -1,54 +1,199 @@
 import gradio as gr
 import requests
 from smolagents import CodeAgent
+import json
+import re
 
-# Your HF Space URL - Replace with your actual space URL
 HF_SPACE_URL = "https://Manavraj-Troubleshoot_MCP.hf.space"
 
-def call_mcp_server(message):
+def call_mcp_server(message, tool_type="knowledge_base"):
     """Call your MCP server and let it handle tool selection"""
     try:
-        # Call your server's main interface or create a unified endpoint
+        # 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)
+        
         response = requests.post(f"{HF_SPACE_URL}/api/predict", 
-                               json={"data": [message], "fn_index": 0})
+                               json={"data": [message], "fn_index": fn_index})
         return response.json().get("data", ["No response"])[0]
     except Exception as e:
         return f"MCP server error: {str(e)}"
 
-# Initialize CodeAgent 
+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 connected to an MCP server.
+    system_prompt="""You are a Technical Support AI Agent that uses a structured Thought-Action-Observation cycle to solve user queries effectively.
 
-Your goal is to help users with technical problems by routing their requests to the appropriate server tools and providing clear, helpful responses.
+WORKFLOW:
+For every user query, you MUST follow this exact structure:
 
-The MCP server has knowledge base search, web search, and response formatting capabilities. Simply process user queries and provide helpful technical support.
+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
 
-{{managed_agents_descriptions}}
+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):
-    """Route user message to MCP server"""
-    mcp_response = call_mcp_server(message)
+    """Enhanced chat interface with thought-action-observation cycle"""
+    
+    # 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.
+    """
+    
     try:
-        final_response = agent.run(f"User asked: {message}\nMCP Server response: {mcp_response}\nProvide a helpful response:")
-        return final_response
-    except:
-        return mcp_response  # Fallback to direct MCP response
+        # Get agent's initial thought process
+        agent_response = agent.run(thinking_prompt)
+        
+        # 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)
+        
+        # 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:
+        # 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", 
-    description="Technical support powered by MCP server tools",
+    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 tech news"
-    ]
+        "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__":