mcp

mcp.py

from .utility.strings import extract_json
from .extensions import NewelleExtension
import threading
import os 
import json 

class MCPExtension(NewelleExtension):
    id = "mcp"
    name = "MCP Extension"


    def get_replace_codeblocks_langs(self) -> list:
        return ["tool"]

    @staticmethod
    def get_extra_requirements() -> list:
        return ["mcp"]

    def preprocess_history(self, history: list, prompts: list) -> tuple[list, list]:        
        import asyncio
        from mcp import ClientSession, StdioServerParameters
        query = history[-1]["Message"]
        tools = self.sync_get_tools(None)
        prompt = self.get_prompt_to_identify_tool_and_arguments(query, tools)
        prompts.append(prompt)
        return history, prompts

    def get_answer(self, codeblock: str, lang: str) -> str | None:
        print(codeblock) 
        js = codeblock
        print(js)
        call = json.loads(js)
        print(call)
        if "tool" not in call:
            return "Missing tool name"
        tool_name = call["tool"]
        args = call["arguments"]
        result = self.sync_call_tool(tool_name, args)
        return result

    def get_prompt_to_identify_tool_and_arguments(self,query,tools):
        tools_description = "\n".join([f"- {tool.name}, {tool.description}, {tool.inputSchema} " for tool in tools])
        return  ("- You are a helpful assistant with access to these tools:\n\n"
                    f"{tools_description}\n"
                    "Choose the appropriate tool based on the user's question. \n"
                    f"User's Question: {query}\n"                
                    "If no tool is needed, reply directly.\n\n"
                    "IMPORTANT: When you need to use a tool, you must ONLY respond with "                
                    "the exact JSON object format below, nothing else:\n"
                    "Keep the values in str "
                    "```tool\n"
                    "{\n"
                    '    "tool": "tool-name",\n'
                    '    "arguments": {\n'
                    '        "argument-name": "value"\n'
                    "    }\n"
                    "}\n```\n")

    def sync_get_tools(self, server_params):
        """Synchronous wrapper to get available tools"""
        import asyncio
        from mcp.client.stdio import stdio_client
        from mcp import ClientSession
        from mcp.client.streamable_http import streamablehttp_client
        url = "https://server.smithery.ai/@smithery/toolbox/mcp?api_key=N&profile=extreme-lark-TgkZVO"
        async def _async_get_tools():
            async with streamablehttp_client(url) as (read, write, _):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    tools = await session.list_tools()
                    return tools.tools 
        return asyncio.run(_async_get_tools())

    def sync_call_tool(self,tool_name, arguments):
        """Synchronous wrapper to call a tool"""
        import asyncio
        from mcp import ClientSession
        from mcp.client.streamable_http import streamablehttp_client
        url = "https://server.smithery.ai/@smithery/toolbox/mcp?api_key=N&profile=extreme-lark-TgkZVO"
        async def _async_call_tool():
            async with streamablehttp_client(url) as (read, write, _):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    result = await session.call_tool(tool_name, arguments=arguments)
                    return result
        
        return asyncio.run(_async_call_tool())