MCP client support (#427)

- MCP client service
- Tool request/response schema
- API gateway support for mcp-tool
- Message translation for tool request & response
- Make mcp-tool using configuration service for information
  about where the MCP services are.

trustgraph-flow/scripts/mcp-tool

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from trustgraph.agent.mcp_tool import run
+
+run()
+

trustgraph-flow/setup.py

@@ -49,6 +49,7 @@ setuptools.setup(
         "langchain-community",
         "langchain-core",
         "langchain-text-splitters",
+        "mcp",
         "minio",
         "mistralai",
         "neo4j",
@@ -99,6 +100,7 @@ setuptools.setup(
         "scripts/kg-store",
         "scripts/kg-manager",
         "scripts/librarian",
+        "scripts/mcp-tool",
         "scripts/metering",
         "scripts/object-extract-row",
         "scripts/oe-write-milvus",

trustgraph-flow/trustgraph/agent/mcp_tool/__init__.py

@@ -0,0 +1,3 @@
+
+from . service import *
+

trustgraph-flow/trustgraph/agent/mcp_tool/__main__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from . service import run
+
+if __name__ == '__main__':
+    run()
+

trustgraph-flow/trustgraph/agent/mcp_tool/service.py

@@ -0,0 +1,105 @@
+
+"""
+MCP tool-calling service, calls an external MCP tool.  Input is
+name + parameters, output is the response, either a string or an object.
+"""
+
+import json
+from mcp.client.streamable_http import streamablehttp_client
+from mcp import ClientSession
+
+from ... base import ToolService
+
+default_ident = "mcp-tool"
+
+class Service(ToolService):
+
+    def __init__(self, **params):
+
+        super(Service, self).__init__(
+            **params
+        )
+
+        self.register_config_handler(self.on_mcp_config)
+
+        self.mcp_services = {}
+
+    async def on_mcp_config(self, config, version):
+
+        print("Got config version", version)
+
+        if "mcp" not in config: return
+
+        self.mcp_services = {
+            k: json.loads(v)
+            for k, v in config["mcp"].items()
+        }
+
+    async def invoke_tool(self, name, parameters):
+
+        try:
+
+            if name not in self.mcp_services:
+                raise RuntimeError(f"MCP service {name} not known")
+
+            if "url" not in self.mcp_services[name]:
+                raise RuntimeError(f"MCP service {name} URL not defined")
+
+            url = self.mcp_services[name]["url"]
+
+            if "name" in self.mcp_services[name]:
+                remote_name = self.mcp_services[name]["name"]
+            else:
+                remote_name = name
+
+            print("Invoking", remote_name, "at", url, flush=True)
+
+            # Connect to a streamable HTTP server
+            async with streamablehttp_client(url) as (
+                    read_stream,
+                    write_stream,
+                    _,
+            ):
+
+                # Create a session using the client streams
+                async with ClientSession(read_stream, write_stream) as session:
+
+                    # Initialize the connection
+                    await session.initialize()
+
+                    # Call a tool
+                    result = await session.call_tool(
+                        remote_name,
+                        parameters
+                    )
+
+                    if result.structuredContent:
+                        return result.structuredContent
+                    elif hasattr(result, "content"):
+                            return "".join([
+                                x.text
+                                for x in result.content
+                            ])
+                    else:
+                        return "No content"
+
+        except BaseExceptionGroup as e:
+
+            for child in e.exceptions:
+                print(child)
+
+            raise e.exceptions[0]
+
+        except Exception as e:
+
+            print(e)
+            raise e
+            
+    @staticmethod
+    def add_args(parser):
+
+        ToolService.add_args(parser)
+
+def run():
+    Service.launch(default_ident, __doc__)
+

trustgraph-flow/trustgraph/gateway/dispatch/manager.py

@@ -17,7 +17,7 @@ from . document_rag import DocumentRagRequestor
 from . triples_query import TriplesQueryRequestor
 from . embeddings import EmbeddingsRequestor
 from . graph_embeddings_query import GraphEmbeddingsQueryRequestor
-from . prompt import PromptRequestor
+from . mcp_tool import McpToolRequestor
 from . text_load import TextLoad
 from . document_load import DocumentLoad
 
@@ -40,6 +40,7 @@ request_response_dispatchers = {
     "agent": AgentRequestor,
     "text-completion": TextCompletionRequestor,
     "prompt": PromptRequestor,
+    "mcp-tool": McpToolRequestor,
     "graph-rag": GraphRagRequestor,
     "document-rag": DocumentRagRequestor,
     "embeddings": EmbeddingsRequestor,

trustgraph-flow/trustgraph/gateway/dispatch/mcp_tool.py

@@ -0,0 +1,32 @@
+
+from ... schema import ToolRequest, ToolResponse
+from ... messaging import TranslatorRegistry
+
+from . requestor import ServiceRequestor
+
+class McpToolRequestor(ServiceRequestor):
+    def __init__(
+            self, pulsar_client, request_queue, response_queue, timeout,
+            consumer, subscriber,
+    ):
+
+        super(McpToolRequestor, self).__init__(
+            pulsar_client=pulsar_client,
+            request_queue=request_queue,
+            response_queue=response_queue,
+            request_schema=ToolRequest,
+            response_schema=ToolResponse,
+            subscription = subscriber,
+            consumer_name = consumer,
+            timeout=timeout,
+        )
+
+        self.request_translator = TranslatorRegistry.get_request_translator("tool")
+        self.response_translator = TranslatorRegistry.get_response_translator("tool")
+
+    def to_request(self, body):
+        return self.request_translator.to_pulsar(body)
+
+    def from_response(self, message):
+        return self.response_translator.from_response_with_completion(message)
+