Create app.py

app.py

@@ -0,0 +1,468 @@
+import asyncio
+import os
+import json
+from typing import List, Dict, Any, Union
+from contextlib import AsyncExitStack
+
+import gradio as gr
+from gradio.components.chatbot import ChatMessage
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.client.sse import sse_client
+from anthropic import Anthropic
+from datasets import load_dataset
+import pandas as pd
+
+loop = asyncio.new_event_loop()
+asyncio.set_event_loop(loop)
+
+class MCPClientWrapper:
+    def __init__(self):
+        self.session = None
+        self.exit_stack = None
+        self.anthropic = None
+        self.tools = []
+        self.dataset = None
+        self.validation_results = []
+    
+    def set_api_key(self, api_key: str) -> str:
+        """Set the Anthropic API key and initialize the client"""
+        if not api_key or not api_key.strip():
+            return "Please enter a valid Anthropic API key"
+        
+        try:
+            self.anthropic = Anthropic(api_key=api_key.strip())
+            return "API key set successfully ✅"
+        except Exception as e:
+            return f"Failed to set API key: {str(e)}"
+    
+    def connect(self, server_input: str) -> str:
+        if not self.anthropic:
+            return "Please set your Anthropic API key first"
+        return loop.run_until_complete(self._connect(server_input))
+    
+    async def _connect(self, server_input: str) -> str:
+        if self.exit_stack:
+            await self.exit_stack.aclose()
+        
+        self.exit_stack = AsyncExitStack()
+        
+        try:
+            # Check if input is a URL (starts with http:// or https://)
+            if server_input.startswith(('http://', 'https://')):
+                # Connect via SSE
+                read, write = await self.exit_stack.enter_async_context(
+                    sse_client(server_input)
+                )
+                connection_type = "SSE URL"
+            else:
+                # Connect via stdio (local file)
+                is_python = server_input.endswith('.py')
+                command = "python" if is_python else "node"
+                
+                server_params = StdioServerParameters(
+                    command=command,
+                    args=[server_input],
+                    env={"PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"}
+                )
+                
+                read, write = await self.exit_stack.enter_async_context(
+                    stdio_client(server_params)
+                )
+                connection_type = "Local script"
+            
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(read, write)
+            )
+            await self.session.initialize()
+            
+            response = await self.session.list_tools()
+            self.tools = [{ 
+                "name": tool.name,
+                "description": tool.description,
+                "input_schema": tool.inputSchema
+            } for tool in response.tools]
+            
+            tool_names = [tool["name"] for tool in self.tools]
+            return f"Connected to MCP server via {connection_type}. Available tools: {', '.join(tool_names)}"
+            
+        except Exception as e:
+            return f"Connection failed: {str(e)}"
+    
+    def load_dataset(self) -> tuple:
+        """Load the TAAIC Phase1 validation dataset"""
+        try:
+            self.dataset = load_dataset("aitxchallenge/Phase1_Model_Validator", split="train")
+            dataset_info = f"Dataset loaded successfully! {len(self.dataset)} validation cases available."
+            
+            # Create a preview of the dataset
+            df = pd.DataFrame(self.dataset)
+            preview = df.head().to_string()
+            
+            return (
+                dataset_info,
+                gr.Button("🔍 Validate", interactive=True),
+                gr.Textbox(value=f"Dataset Preview:\n{preview}", visible=True)
+            )
+        except Exception as e:
+            return (
+                f"Failed to load dataset: {str(e)}",
+                gr.Button("📥 Load Dataset", interactive=True),
+                gr.Textbox(visible=False)
+            )
+    
+    def validate_tools(self) -> str:
+        """Run validation on all dataset cases"""
+        if not self.anthropic:
+            return "Please set your Anthropic API key first."
+            
+        if not self.dataset:
+            return "Please load the dataset first."
+        
+        if not self.session:
+            return "Please connect to an MCP server first."
+        
+        return loop.run_until_complete(self._run_validation())
+    
+    async def _run_validation(self) -> str:
+        """Async validation runner"""
+        self.validation_results = []
+        total_cases = len(self.dataset)
+        passed = 0
+        failed = 0
+        
+        for i, case in enumerate(self.dataset):
+            try:
+                # Extract test case information
+                query = case.get('query', case.get('question', ''))
+                expected_output = case.get('expected_output', case.get('expected', ''))
+                test_id = case.get('id', f'test_{i}')
+                
+                # Run the query through the MCP tools
+                result = await self._validate_single_case(query, expected_output, test_id)
+                self.validation_results.append(result)
+                
+                if result['passed']:
+                    passed += 1
+                else:
+                    failed += 1
+                    
+            except Exception as e:
+                failed += 1
+                self.validation_results.append({
+                    'test_id': test_id,
+                    'query': query,
+                    'error': str(e),
+                    'passed': False
+                })
+        
+        # Generate validation report
+        report = f"""
+VALIDATION COMPLETE
+==================
+Total Cases: {total_cases}
+Passed: {passed}
+Failed: {failed}
+Success Rate: {(passed/total_cases)*100:.1f}%
+
+DETAILED RESULTS:
+"""
+        
+        for result in self.validation_results:
+            status = "✅ PASS" if result['passed'] else "❌ FAIL"
+            report += f"\n{status} [{result['test_id']}] {result['query'][:50]}..."
+            if not result['passed'] and 'error' in result:
+                report += f"\n   Error: {result['error']}"
+        
+        return report
+    
+    async def _validate_single_case(self, query: str, expected_output: str, test_id: str) -> Dict[str, Any]:
+        """Validate a single test case"""
+        try:
+            # Send query to Claude with MCP tools
+            claude_messages = [{"role": "user", "content": query}]
+            
+            response = self.anthropic.messages.create(
+                model="claude-3-5-sonnet-20241022",
+                max_tokens=1000,
+                messages=claude_messages,
+                tools=self.tools
+            )
+            
+            # Process tool calls if any
+            actual_output = ""
+            for content in response.content:
+                if content.type == 'text':
+                    actual_output += content.text
+                elif content.type == 'tool_use':
+                    tool_result = await self.session.call_tool(content.name, content.input)
+                    actual_output += str(tool_result.content)
+            
+            # Simple validation logic - you may want to customize this
+            passed = self._validate_output(actual_output, expected_output)
+            
+            return {
+                'test_id': test_id,
+                'query': query,
+                'expected': expected_output,
+                'actual': actual_output,
+                'passed': passed
+            }
+            
+        except Exception as e:
+            return {
+                'test_id': test_id,
+                'query': query,
+                'error': str(e),
+                'passed': False
+            }
+    
+    def _validate_output(self, actual: str, expected: str) -> bool:
+        """Basic output validation - customize based on your needs"""
+        # This is a simple implementation - you may want more sophisticated validation
+        if not expected:
+            return True  # If no expected output specified, consider it passed
+        
+        # You can implement more sophisticated matching here
+        # For now, using simple substring matching
+        return expected.lower() in actual.lower()
+    
+    def process_message(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]) -> tuple:
+        if not self.anthropic:
+            return history + [
+                {"role": "user", "content": message}, 
+                {"role": "assistant", "content": "Please set your Anthropic API key first."}
+            ], gr.Textbox(value="")
+            
+        if not self.session:
+            return history + [
+                {"role": "user", "content": message}, 
+                {"role": "assistant", "content": "Please connect to an MCP server first."}
+            ], gr.Textbox(value="")
+        
+        new_messages = loop.run_until_complete(self._process_query(message, history))
+        return history + [{"role": "user", "content": message}] + new_messages, gr.Textbox(value="")
+    
+    async def _process_query(self, message: str, history: List[Union[Dict[str, Any], ChatMessage]]):
+        claude_messages = []
+        for msg in history:
+            if isinstance(msg, ChatMessage):
+                role, content = msg.role, msg.content
+            else:
+                role, content = msg.get("role"), msg.get("content")
+            
+            if role in ["user", "assistant", "system"]:
+                claude_messages.append({"role": role, "content": content})
+        
+        claude_messages.append({"role": "user", "content": message})
+        
+        response = self.anthropic.messages.create(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=1000,
+            messages=claude_messages,
+            tools=self.tools
+        )
+
+        result_messages = []
+        
+        for content in response.content:
+            if content.type == 'text':
+                result_messages.append({
+                    "role": "assistant", 
+                    "content": content.text
+                })
+                
+            elif content.type == 'tool_use':
+                tool_name = content.name
+                tool_args = content.input
+                
+                result_messages.append({
+                    "role": "assistant",
+                    "content": f"I'll only use the {tool_name} tool to help answer your question.",
+                    "metadata": {
+                        "title": f"Using tool: {tool_name}",
+                        "log": f"Parameters: {json.dumps(tool_args, ensure_ascii=True)}",
+                        "status": "pending",
+                        "id": f"tool_call_{tool_name}"
+                    }
+                })
+                
+                result_messages.append({
+                    "role": "assistant",
+                    "content": "```json\n" + json.dumps(tool_args, indent=2, ensure_ascii=True) + "\n```",
+                    "metadata": {
+                        "parent_id": f"tool_call_{tool_name}",
+                        "id": f"params_{tool_name}",
+                        "title": "Tool Parameters"
+                    }
+                })
+                
+                try:
+                    result = await self.session.call_tool(tool_name, tool_args)
+                    
+                    if result_messages and "metadata" in result_messages[-2]:
+                        result_messages[-2]["metadata"]["status"] = "done"
+                    
+                    result_messages.append({
+                        "role": "assistant",
+                        "content": "Here are the results from the tool:",
+                        "metadata": {
+                            "title": f"Tool Result for {tool_name}",
+                            "status": "done",
+                            "id": f"result_{tool_name}"
+                        }
+                    })
+                    
+                    result_content = result.content
+                    if isinstance(result_content, list):
+                        result_content = "\n".join(str(item) for item in result_content)
+                    
+                    try:
+                        result_json = json.loads(result_content)
+                        if isinstance(result_json, dict) and "type" in result_json:
+                            if result_json["type"] == "image" and "url" in result_json:
+                                result_messages.append({
+                                    "role": "assistant",
+                                    "content": {"path": result_json["url"], "alt_text": result_json.get("message", "Generated image")},
+                                    "metadata": {
+                                        "parent_id": f"result_{tool_name}",
+                                        "id": f"image_{tool_name}",
+                                        "title": "Generated Image"
+                                    }
+                                })
+                            else:
+                                result_messages.append({
+                                    "role": "assistant",
+                                    "content": "```\n" + result_content + "\n```",
+                                    "metadata": {
+                                        "parent_id": f"result_{tool_name}",
+                                        "id": f"raw_result_{tool_name}",
+                                        "title": "Raw Output"
+                                    }
+                                })
+                    except:
+                        result_messages.append({
+                            "role": "assistant",
+                            "content": "```\n" + result_content + "\n```",
+                            "metadata": {
+                                "parent_id": f"result_{tool_name}",
+                                "id": f"raw_result_{tool_name}",
+                                "title": "Raw Output"
+                            }
+                        })
+                    
+                    claude_messages.append({"role": "user", "content": f"Tool result for {tool_name}: {result_content}"})
+                    next_response = self.anthropic.messages.create(
+                        model="claude-3-5-sonnet-20241022",
+                        max_tokens=1000,
+                        messages=claude_messages,
+                    )
+                    
+                    if next_response.content and next_response.content[0].type == 'text':
+                        result_messages.append({
+                            "role": "assistant",
+                            "content": next_response.content[0].text
+                        })
+                
+                except Exception as e:
+                    result_messages.append({
+                        "role": "assistant",
+                        "content": f"Error calling tool {tool_name}: {str(e)}",
+                        "metadata": {
+                            "title": f"Error - {tool_name}",
+                            "status": "error",
+                            "id": f"error_{tool_name}"
+                        }
+                    })
+
+        return result_messages
+
+client = MCPClientWrapper()
+
+def gradio_interface():
+    with gr.Blocks(title="TAAIC Tool Validation") as demo:
+        gr.Markdown("# TAAIC Tool Validation")
+        gr.Markdown("Connect your Gradio MCP Tool for validation for the TAAIC challenge.")
+
+        # API Key input section
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=4):
+                api_key_input = gr.Textbox(
+                    label="Anthropic API Key",
+                    placeholder="Enter your Anthropic API key (sk-ant-...)",
+                    type="password"
+                )
+            with gr.Column(scale=1):
+                api_key_btn = gr.Button("Set API Key")
+        
+        api_key_status = gr.Textbox(label="API Key Status", interactive=False)
+        
+        # MCP Server connection section
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=4):
+                server_input = gr.Textbox(
+                    label="MCP Server URL or Script Path",
+                    placeholder="Enter URL (e.g., https://cyrilzakka-clinical-trials.hf.space/gradio_api/mcp/sse) or local script path (e.g., weather.py)",
+                    value="https://cyrilzakka-clinical-trials.hf.space/gradio_api/mcp/sse"
+                )
+            with gr.Column(scale=1):
+                connect_btn = gr.Button("Connect")
+        
+        status = gr.Textbox(label="Connection Status", interactive=False)
+        
+        # Dataset loading section
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=3):
+                dataset_status = gr.Textbox(
+                    label="Dataset Status",
+                    value="Click 'Load Dataset' to load validation cases",
+                    interactive=False
+                )
+            with gr.Column(scale=1):
+                dataset_btn = gr.Button("📥 Load Dataset", interactive=True)
+        
+        dataset_preview = gr.Textbox(
+            label="Dataset Preview",
+            visible=False,
+            interactive=False,
+            max_lines=10
+        )
+        
+        # Validation results
+        validation_results = gr.Textbox(
+            label="Validation Results",
+            visible=False,
+            interactive=False,
+            max_lines=20
+        )
+        
+        # Event handlers
+        api_key_btn.click(client.set_api_key, inputs=api_key_input, outputs=api_key_status)
+        connect_btn.click(client.connect, inputs=server_input, outputs=status)
+        
+        dataset_btn.click(
+            client.load_dataset, 
+            outputs=[dataset_status, dataset_btn, dataset_preview]
+        )
+        
+        def run_validation():
+            results = client.validate_tools()
+            return gr.Textbox(value=results, visible=True)
+        
+        dataset_btn.click(
+            lambda: client.validate_tools() if client.dataset else "Please load dataset first.",
+            outputs=validation_results,
+            show_progress=True
+        ).then(
+            lambda: gr.Textbox(visible=True),
+            outputs=validation_results
+        )
+        
+        # msg.submit(client.process_message, [msg, chatbot], [chatbot, msg])
+        # clear_btn.click(lambda: [], None, chatbot)
+        
+    return demo
+
+if __name__ == "__main__":
+    interface = gradio_interface()
+    interface.launch(debug=True)