Update app.py

app.py

@@ -1,70 +1,346 @@
-import gradio as gr
-from huggingface_hub import InferenceClient
-
-
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+from huggingface_hub import InferenceClient
+import json
+import asyncio
+from typing import Optional, Dict, Any, List
+
+
+class MCPContextManager:
+    """Model Context Protocol Manager for handling context and tools"""
+    
+    def __init__(self):
+        self.context = []
+        self.tools = {}
+        self.resources = {}
+        
+    def add_context(self, role: str, content: str, metadata: Optional[Dict] = None):
+        """Add context entry following MCP specification"""
+        entry = {
+            "role": role,
+            "content": content,
+            "timestamp": None,
+            "metadata": metadata or {}
+        }
+        self.context.append(entry)
+        return entry
+    
+    def register_tool(self, name: str, description: str, parameters: Dict, handler):
+        """Register a tool following MCP tool specification"""
+        self.tools[name] = {
+            "name": name,
+            "description": description,
+            "parameters": parameters,
+            "handler": handler
+        }
+        
+    def register_resource(self, uri: str, name: str, mime_type: str, content: Any):
+        """Register a resource following MCP resource specification"""
+        self.resources[uri] = {
+            "uri": uri,
+            "name": name,
+            "mimeType": mime_type,
+            "content": content
+        }
+    
+    def get_context_window(self, max_tokens: int = 4096) -> List[Dict]:
+        """Get context window within token limits"""
+        return self.context[-10:]  # Simple implementation
+    
+    async def call_tool(self, tool_name: str, arguments: Dict) -> Any:
+        """Execute a registered tool"""
+        if tool_name not in self.tools:
+            raise ValueError(f"Tool {tool_name} not found")
+        
+        tool = self.tools[tool_name]
+        return await tool["handler"](arguments)
+    
+    def get_available_tools(self) -> List[Dict]:
+        """Get list of available tools in MCP format"""
+        return [
+            {
+                "name": tool["name"],
+                "description": tool["description"],
+                "parameters": tool["parameters"]
+            }
+            for tool in self.tools.values()
+        ]
+    
+    def export_context(self) -> str:
+        """Export context in MCP-compatible JSON format"""
+        return json.dumps({
+            "context": self.context,
+            "tools": self.get_available_tools(),
+            "resources": list(self.resources.keys())
+        }, indent=2)
+
+
+# Initialize MCP Manager
+mcp_manager = MCPContextManager()
+
+
+# Register example MCP tools
+async def calculator_handler(args):
+    """Calculator tool handler"""
+    operation = args.get("operation")
+    a = float(args.get("a", 0))
+    b = float(args.get("b", 0))
+    
+    operations = {
+        "add": a + b,
+        "subtract": a - b,
+        "multiply": a * b,
+        "divide": a / b if b != 0 else "Error: Division by zero"
+    }
+    return operations.get(operation, "Invalid operation")
+
+
+async def memory_handler(args):
+    """Memory/context storage handler"""
+    action = args.get("action")
+    key = args.get("key")
+    value = args.get("value")
+    
+    if action == "store":
+        mcp_manager.register_resource(
+            uri=f"memory://{key}",
+            name=key,
+            mime_type="text/plain",
+            content=value
+        )
+        return f"Stored {key}"
+    elif action == "retrieve":
+        resource = mcp_manager.resources.get(f"memory://{key}")
+        return resource["content"] if resource else "Not found"
+    return "Invalid action"
+
+
+async def web_search_handler(args):
+    """Simulated web search handler"""
+    query = args.get("query")
+    return f"Search results for: {query} (MCP tool simulation)"
+
+
+# Register tools with MCP manager
+mcp_manager.register_tool(
+    name="calculator",
+    description="Perform basic arithmetic operations",
+    parameters={
+        "type": "object",
+        "properties": {
+            "operation": {"type": "string", "enum": ["add", "subtract", "multiply", "divide"]},
+            "a": {"type": "number"},
+            "b": {"type": "number"}
+        },
+        "required": ["operation", "a", "b"]
+    },
+    handler=calculator_handler
+)
+
+mcp_manager.register_tool(
+    name="memory",
+    description="Store and retrieve information in memory",
+    parameters={
+        "type": "object",
+        "properties": {
+            "action": {"type": "string", "enum": ["store", "retrieve"]},
+            "key": {"type": "string"},
+            "value": {"type": "string"}
+        },
+        "required": ["action", "key"]
+    },
+    handler=memory_handler
+)
+
+mcp_manager.register_tool(
+    name="web_search",
+    description="Search the web for information",
+    parameters={
+        "type": "object",
+        "properties": {
+            "query": {"type": "string"}
+        },
+        "required": ["query"]
+    },
+    handler=web_search_handler
+)
+
+
+def respond(
+    message,
+    history: list[dict[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    enable_mcp,
+    hf_token: gr.OAuthToken,
+):
+    """
+    Enhanced respond function with MCP support
+    """
+    # Add message to MCP context if enabled
+    if enable_mcp:
+        mcp_manager.add_context("user", message)
+    
+    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+
+    # Build messages with MCP context
+    messages = [{"role": "system", "content": system_message}]
+    
+    if enable_mcp:
+        # Add MCP tools information to system message
+        tools_info = "\n\nAvailable MCP Tools:\n" + json.dumps(mcp_manager.get_available_tools(), indent=2)
+        messages[0]["content"] += tools_info
+
+    messages.extend(history)
+    messages.append({"role": "user", "content": message})
+
+    response = ""
+
+    for msg in client.chat_completion(
+        messages,
+        max_tokens=max_tokens,
+        stream=True,
+        temperature=temperature,
+        top_p=top_p,
+    ):
+        choices = msg.choices
+        token = ""
+        if len(choices) and choices[0].delta.content:
+            token = choices[0].delta.content
+
+        response += token
+        yield response
+    
+    # Add response to MCP context if enabled
+    if enable_mcp:
+        mcp_manager.add_context("assistant", response)
+
+
+def call_mcp_tool(tool_name: str, arguments_json: str):
+    """Interface for calling MCP tools from UI"""
+    try:
+        arguments = json.loads(arguments_json)
+        result = asyncio.run(mcp_manager.call_tool(tool_name, arguments))
+        return json.dumps({"success": True, "result": result}, indent=2)
+    except Exception as e:
+        return json.dumps({"success": False, "error": str(e)}, indent=2)
+
+
+def export_mcp_context():
+    """Export current MCP context"""
+    return mcp_manager.export_context()
+
+
+def get_mcp_tools_list():
+    """Get formatted list of available MCP tools"""
+    tools = mcp_manager.get_available_tools()
+    return json.dumps(tools, indent=2)
+
+
+"""
+Enhanced Gradio interface with Model Context Protocol (MCP) support
+"""
+chatbot = gr.ChatInterface(
+    respond,
+    type="messages",
+    additional_inputs=[
+        gr.Textbox(
+            value="You are a friendly AI assistant with access to Model Context Protocol (MCP) tools. You can use tools like calculator, memory storage, and web search to help users.",
+            label="System message",
+            lines=3
+        ),
+        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.95,
+            step=0.05,
+            label="Top-p (nucleus sampling)",
+        ),
+        gr.Checkbox(value=True, label="Enable MCP (Model Context Protocol)"),
+    ],
+)
+
+with gr.Blocks(title="AI Chat with MCP Support") as demo:
+    gr.Markdown("# 🤖 AI Chatbot with Model Context Protocol (MCP)")
+    gr.Markdown("Enhanced with MCP for context management, tool usage, and resource handling")
+    
+    with gr.Tabs():
+        with gr.Tab("💬 Chat"):
+            with gr.Sidebar():
+                gr.LoginButton()
+            chatbot.render()
+        
+        with gr.Tab("🛠️ MCP Tools"):
+            gr.Markdown("### Available MCP Tools")
+            tools_display = gr.JSON(label="Registered Tools")
+            refresh_tools_btn = gr.Button("Refresh Tools List")
+            
+            gr.Markdown("### Call MCP Tool")
+            with gr.Row():
+                tool_name_input = gr.Textbox(label="Tool Name", placeholder="calculator")
+                tool_args_input = gr.Textbox(
+                    label="Arguments (JSON)",
+                    placeholder='{"operation": "add", "a": 5, "b": 3}',
+                    lines=3
+                )
+            call_tool_btn = gr.Button("Call Tool")
+            tool_result = gr.JSON(label="Tool Result")
+            
+            refresh_tools_btn.click(
+                fn=get_mcp_tools_list,
+                outputs=tools_display
+            )
+            
+            call_tool_btn.click(
+                fn=call_mcp_tool,
+                inputs=[tool_name_input, tool_args_input],
+                outputs=tool_result
+            )
+        
+        with gr.Tab("📋 MCP Context"):
+            gr.Markdown("### Export MCP Context")
+            gr.Markdown("View and export the current context window managed by MCP")
+            export_btn = gr.Button("Export Context")
+            context_display = gr.Code(language="json", label="MCP Context")
+            
+            export_btn.click(
+                fn=export_mcp_context,
+                outputs=context_display
+            )
+        
+        with gr.Tab("ℹ️ About MCP"):
+            gr.Markdown("""
+            ### Model Context Protocol (MCP)
+            
+            MCP is a protocol for managing context, tools, and resources in AI applications.
+            
+            **Features:**
+            - 🔧 **Tool Integration**: Register and call tools dynamically
+            - 💾 **Context Management**: Maintain conversation history and context
+            - 📦 **Resource Handling**: Store and retrieve resources (files, data, etc.)
+            - 🔌 **Extensible**: Easy to add new tools and capabilities
+            
+            **Available Tools:**
+            1. **Calculator**: Perform arithmetic operations
+            2. **Memory**: Store and retrieve information
+            3. **Web Search**: Search for information (simulated)
+            
+            **Example Tool Call:**
+            ```json
+            {
+              "tool": "calculator",
+              "arguments": {
+                "operation": "add",
+                "a": 10,
+                "b": 5
+              }
+            }
+            ```
+            """)
+
+
+if __name__ == "__main__":
+    demo.launch()