Very initial MCP server PoC for TrustGraph

trustgraph-mcp/README.md

@@ -0,0 +1 @@
+See https://trustgraph.ai/

trustgraph-mcp/scripts/mcp-server

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

trustgraph-mcp/setup.py

@@ -0,0 +1,44 @@
+import setuptools
+import os
+import importlib
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+# Load a version number module
+spec = importlib.util.spec_from_file_location(
+    'version', 'trustgraph/mcp_version.py'
+)
+version_module = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(version_module)
+
+version = version_module.__version__
+
+setuptools.setup(
+    name="trustgraph-mcp",
+    version=version,
+    author="trustgraph.ai",
+    author_email="security@trustgraph.ai",
+    description="TrustGraph provides a means to run a pipeline of flexible AI processing components in a flexible means to achieve a processing pipeline.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/trustgraph-ai/trustgraph",
+    packages=setuptools.find_namespace_packages(
+        where='./',
+    ),
+    classifiers=[ 
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: GNU General Public License v3 or later (GPLv3+)",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='>=3.8',
+    download_url = "https://github.com/trustgraph-ai/trustgraph/archive/refs/tags/v" + version + ".tar.gz",
+    install_requires=[
+        "trustgraph-base>=1.1,<1.2",
+        "mcp",
+        "websockets",
+    ],
+    scripts=[
+        "scripts/mcp-server",
+    ]
+)

trustgraph-mcp/trustgraph/mcp_server/__init__.py

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

trustgraph-mcp/trustgraph/mcp_server/__main__.py

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

trustgraph-mcp/trustgraph/mcp_server/tg_socket.py

@@ -0,0 +1,129 @@
+
+from dataclasses import dataclass
+from websockets.asyncio.client import connect
+import asyncio
+import logging
+import json
+import uuid
+import time
+
+class WebSocketManager:
+
+    def __init__(self, url):
+        self.url = url
+        self.socket = None
+
+    async def start(self):
+        self.socket = await connect(self.url)
+        self.pending_requests = {}
+        self.running = True
+        self.reader_task = asyncio.create_task(self.reader())
+
+    async def stop(self):
+        self.running = False
+        await self.reader_task
+
+    async def reader(self):
+        """
+        Background task to read websocket responses and route to correct
+        request
+        """
+
+        while self.running:
+            try:
+
+                try:
+                    response_text = await asyncio.wait_for(
+                        self.socket.recv(), 0.5
+                    )
+                except TimeoutError:
+                    continue
+
+                response = json.loads(response_text)
+
+                request_id = response.get("id")
+                if request_id and request_id in self.pending_requests:
+                    # Put the response in the queue
+                    queue = self.pending_requests[request_id]
+                    await queue.put(response)
+                else:
+                    logging.warning(
+                        f"Response for unknown request ID: {request_id}"
+                    )
+
+            except Exception as e:
+
+                logging.error(f"Error in websocket reader: {e}")
+
+                # Put error in all pending queues
+                for queue in self.pending_requests.values():
+                    try:
+                        await queue.put({"error": str(e)})
+                    except:
+                        pass
+
+                self.pending_requests.clear()
+                break
+
+        await self.socket.close()
+        self.socket = None
+
+    async def request(
+            self, service, request_data, flow_id="default",
+    ):
+        """
+        Send a request via websocket and handle single or streaming responses
+        """
+
+        # Generate unique request ID
+        request_id = f"{uuid.uuid4()}"
+
+        # Determine if this service streams responses
+        streaming_services = {"agent"}
+        is_streaming = service in streaming_services
+
+        # Create a queue for all responses (streaming and single)
+        response_queue = asyncio.Queue()
+        self.pending_requests[request_id] = response_queue
+
+        try:
+
+            # Build request message
+            message = {
+                "id": request_id,
+                "service": service,
+                "request": request_data,
+            }
+
+            if flow_id is not None:
+                message["flow"] = flow_id
+
+            # Send request
+            await self.socket.send(json.dumps(message))
+
+            while self.running:
+
+                try:
+                    response = await asyncio.wait_for(
+                        response_queue.get(), 0.5
+                    )
+                except TimeoutError:
+                    continue
+
+                if "error" in response:
+                    if "message" in response["error"]:
+                        raise RuntimeError(response["error"]["text"])
+                    else:
+                        raise RuntimeError(str(response["error"]))
+
+                yield response["response"]
+
+                if "complete" in response:
+                    if response["complete"]:
+                        break
+
+        except Exception as e:
+            # Clean up on error
+            self.pending_requests.pop(request_id, None)
+            raise e
+

trustgraph-mcp/trustgraph/mcp_version.py

@@ -0,0 +1 @@
+__version__ = "1.1.0"