PoC MCP server (#419)

* Very initial MCP server PoC for TrustGraph

* Put service on port 8000

* Add MCP container and packages to buildout

Makefile

@@ -17,6 +17,7 @@ wheels:
 	pip3 wheel --no-deps --wheel-dir dist trustgraph-embeddings-hf/
 	pip3 wheel --no-deps --wheel-dir dist trustgraph-cli/
 	pip3 wheel --no-deps --wheel-dir dist trustgraph-ocr/
+	pip3 wheel --no-deps --wheel-dir dist trustgraph-mcp/
 
 packages: update-package-versions
 	rm -rf dist/
@@ -28,6 +29,7 @@ packages: update-package-versions
 	cd trustgraph-embeddings-hf && python3 setup.py sdist --dist-dir ../dist/
 	cd trustgraph-cli && python3 setup.py sdist --dist-dir ../dist/
 	cd trustgraph-ocr && python3 setup.py sdist --dist-dir ../dist/
+	cd trustgraph-mcp && python3 setup.py sdist --dist-dir ../dist/
 
 pypi-upload:
 	twine upload dist/*-${VERSION}.*
@@ -45,6 +47,7 @@ update-package-versions:
 	echo __version__ = \"${VERSION}\" > trustgraph-cli/trustgraph/cli_version.py
 	echo __version__ = \"${VERSION}\" > trustgraph-ocr/trustgraph/ocr_version.py
 	echo __version__ = \"${VERSION}\" > trustgraph/trustgraph/trustgraph_version.py
+	echo __version__ = \"${VERSION}\" > trustgraph-mcp/trustgraph/mcp_version.py
 
 container: update-package-versions
 	${DOCKER} build -f containers/Containerfile.base \
@@ -59,12 +62,16 @@ container: update-package-versions
 	    -t ${CONTAINER_BASE}/trustgraph-hf:${VERSION} .
 	${DOCKER} build -f containers/Containerfile.ocr \
 	    -t ${CONTAINER_BASE}/trustgraph-ocr:${VERSION} .
+	${DOCKER} build -f containers/Containerfile.mcp \
+	    -t ${CONTAINER_BASE}/trustgraph-mcp:${VERSION} .
 
 some-containers:
 	${DOCKER} build -f containers/Containerfile.base \
 	    -t ${CONTAINER_BASE}/trustgraph-base:${VERSION} .
 	${DOCKER} build -f containers/Containerfile.flow \
 	    -t ${CONTAINER_BASE}/trustgraph-flow:${VERSION} .
+	${DOCKER} build -f containers/Containerfile.mcp \
+	    -t ${CONTAINER_BASE}/trustgraph-mcp:${VERSION} .
 #	${DOCKER} build -f containers/Containerfile.vertexai \
 #	    -t ${CONTAINER_BASE}/trustgraph-vertexai:${VERSION} .
 #	${DOCKER} build -f containers/Containerfile.bedrock \

containers/Containerfile.mcp

@@ -0,0 +1,46 @@
+
+# ----------------------------------------------------------------------------
+# Build an AI container.  This does the torch install which is huge, and I
+# like to avoid re-doing this.
+# ----------------------------------------------------------------------------
+
+FROM docker.io/fedora:42 AS base
+
+ENV PIP_BREAK_SYSTEM_PACKAGES=1
+
+RUN dnf install -y python3.12 && \
+  alternatives --install /usr/bin/python python /usr/bin/python3.12 1 && \
+  python -m ensurepip --upgrade && \
+  pip3 install --no-cache-dir mcp websockets && \
+  dnf clean all
+
+# ----------------------------------------------------------------------------
+# Build a container which contains the built Python packages.  The build
+# creates a bunch of left-over cruft, a separate phase means this is only
+# needed to support package build
+# ----------------------------------------------------------------------------
+
+FROM base AS build
+
+COPY trustgraph-mcp/ /root/build/trustgraph-mcp/
+
+WORKDIR /root/build/
+
+RUN pip3 wheel -w /root/wheels/ --no-deps ./trustgraph-mcp/
+
+RUN ls /root/wheels
+
+# ----------------------------------------------------------------------------
+# Finally, the target container.  Start with base and add the package.
+# ----------------------------------------------------------------------------
+
+FROM base
+
+COPY --from=build /root/wheels /root/wheels
+
+RUN \
+    pip3 install --no-cache-dir /root/wheels/trustgraph_mcp-* && \
+    rm -rf /root/wheels
+
+WORKDIR /
+

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,43 @@
+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=[
+        "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"