mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-06 20:02:11 +02:00
init
This commit is contained in:
elpresidank
parent
c386f68743
commit
b6536eca38
100 changed files with 17680 additions and 377 deletions
|
|
@ -1,2 +1 @@
|
|||
export { createMcpServer, run } from "./server.js";
|
||||
export { SocketManager, type SocketManagerConfig } from "./socket-manager.js";
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
* TrustGraph MCP server.
|
||||
*
|
||||
* Exposes TrustGraph capabilities as MCP tools for AI assistants.
|
||||
* Communicates with the TrustGraph gateway via WebSocket.
|
||||
* Uses the vendored @trustgraph/client for all gateway communication.
|
||||
*
|
||||
* Python reference: trustgraph-mcp/trustgraph/mcp_server/mcp.py
|
||||
*/
|
||||
|
|
@ -10,10 +10,11 @@
|
|||
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||||
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
||||
import { z } from "zod";
|
||||
import { SocketManager } from "./socket-manager.js";
|
||||
import { createTrustGraphSocket, type BaseApi, type Term } from "@trustgraph/client";
|
||||
|
||||
export function createMcpServer(config: {
|
||||
gatewayUrl: string;
|
||||
user?: string;
|
||||
token?: string;
|
||||
flowId?: string;
|
||||
}) {
|
||||
|
|
@ -22,13 +23,17 @@ export function createMcpServer(config: {
|
|||
version: "0.1.0",
|
||||
});
|
||||
|
||||
const socket = new SocketManager({
|
||||
gatewayUrl: config.gatewayUrl,
|
||||
token: config.token,
|
||||
});
|
||||
const user = config.user ?? "mcp";
|
||||
const socket: BaseApi = createTrustGraphSocket(
|
||||
user,
|
||||
config.token,
|
||||
config.gatewayUrl,
|
||||
);
|
||||
|
||||
const flowId = config.flowId ?? "default";
|
||||
|
||||
// ===================== Flow-scoped tools =====================
|
||||
|
||||
// --- Text Completion ---
|
||||
server.tool(
|
||||
"text_completion",
|
||||
|
|
@ -38,8 +43,9 @@ export function createMcpServer(config: {
|
|||
prompt: z.string().describe("User prompt"),
|
||||
},
|
||||
async ({ system, prompt }) => {
|
||||
const resp = await socket.request("text-completion", { system, prompt }, { flowId }) as Record<string, unknown>;
|
||||
return { content: [{ type: "text" as const, text: String(resp.response ?? resp) }] };
|
||||
const flow = socket.flow(flowId);
|
||||
const response = await flow.textCompletion(system, prompt);
|
||||
return { content: [{ type: "text" as const, text: response }] };
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -51,14 +57,32 @@ export function createMcpServer(config: {
|
|||
query: z.string().describe("Natural language query"),
|
||||
entity_limit: z.number().optional().describe("Max entities to retrieve"),
|
||||
triple_limit: z.number().optional().describe("Max triples per entity"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ query, entity_limit, triple_limit }) => {
|
||||
const resp = await socket.request(
|
||||
"graph-rag",
|
||||
{ query, entity_limit, triple_limit },
|
||||
{ flowId },
|
||||
) as Record<string, unknown>;
|
||||
return { content: [{ type: "text" as const, text: String(resp.response ?? resp) }] };
|
||||
async ({ query, entity_limit, triple_limit, collection }) => {
|
||||
const flow = socket.flow(flowId);
|
||||
const response = await flow.graphRag(
|
||||
query,
|
||||
{ entityLimit: entity_limit, tripleLimit: triple_limit },
|
||||
collection,
|
||||
);
|
||||
return { content: [{ type: "text" as const, text: response }] };
|
||||
},
|
||||
);
|
||||
|
||||
// --- Document RAG ---
|
||||
server.tool(
|
||||
"document_rag",
|
||||
"Query documents using RAG",
|
||||
{
|
||||
query: z.string().describe("Natural language query"),
|
||||
doc_limit: z.number().optional().describe("Max documents to retrieve"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ query, doc_limit, collection }) => {
|
||||
const flow = socket.flow(flowId);
|
||||
const response = await flow.documentRag(query, doc_limit, collection);
|
||||
return { content: [{ type: "text" as const, text: response }] };
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -70,8 +94,23 @@ export function createMcpServer(config: {
|
|||
question: z.string().describe("Question for the agent"),
|
||||
},
|
||||
async ({ question }) => {
|
||||
const resp = await socket.request("agent", { question }, { flowId }) as Record<string, unknown>;
|
||||
return { content: [{ type: "text" as const, text: String(resp.answer ?? resp) }] };
|
||||
const flow = socket.flow(flowId);
|
||||
let fullAnswer = "";
|
||||
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
flow.agent(
|
||||
question,
|
||||
() => {}, // think — ignore for MCP
|
||||
() => {}, // observe — ignore for MCP
|
||||
(chunk, complete) => {
|
||||
fullAnswer += chunk;
|
||||
if (complete) resolve();
|
||||
},
|
||||
(err) => reject(new Error(err)),
|
||||
);
|
||||
});
|
||||
|
||||
return { content: [{ type: "text" as const, text: fullAnswer }] };
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -83,8 +122,9 @@ export function createMcpServer(config: {
|
|||
text: z.array(z.string()).describe("Texts to embed"),
|
||||
},
|
||||
async ({ text }) => {
|
||||
const resp = await socket.request("embeddings", { text }, { flowId }) as Record<string, unknown>;
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
const flow = socket.flow(flowId);
|
||||
const vectors = await flow.embeddings(text);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(vectors) }] };
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -97,15 +137,15 @@ export function createMcpServer(config: {
|
|||
p: z.string().optional().describe("Predicate IRI"),
|
||||
o: z.string().optional().describe("Object IRI or literal"),
|
||||
limit: z.number().optional().describe("Max results"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ s, p, o, limit }) => {
|
||||
const request: Record<string, unknown> = { limit };
|
||||
if (s) request.s = { type: "IRI", iri: s };
|
||||
if (p) request.p = { type: "IRI", iri: p };
|
||||
if (o) request.o = { type: "IRI", iri: o };
|
||||
|
||||
const resp = await socket.request("triples-query", request, { flowId }) as Record<string, unknown>;
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
async ({ s, p, o, limit, collection }) => {
|
||||
const flow = socket.flow(flowId);
|
||||
const sTerm: Term | undefined = s ? { t: "i", i: s } : undefined;
|
||||
const pTerm: Term | undefined = p ? { t: "i", i: p } : undefined;
|
||||
const oTerm: Term | undefined = o ? { t: "i", i: o } : undefined;
|
||||
const triples = await flow.triplesQuery(sTerm, pTerm, oTerm, limit, collection);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(triples, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
|
|
@ -116,28 +156,48 @@ export function createMcpServer(config: {
|
|||
{
|
||||
query: z.string().describe("Text to find similar entities for"),
|
||||
limit: z.number().optional().describe("Max results"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ query, limit }) => {
|
||||
async ({ query, limit, collection }) => {
|
||||
const flow = socket.flow(flowId);
|
||||
// First embed the query, then search
|
||||
const embResp = await socket.request("embeddings", { text: [query] }, { flowId }) as { vectors: number[][] };
|
||||
const resp = await socket.request(
|
||||
"graph-embeddings-query",
|
||||
{ vectors: embResp.vectors, limit: limit ?? 10 },
|
||||
{ flowId },
|
||||
) as Record<string, unknown>;
|
||||
const vectors = await flow.embeddings([query]);
|
||||
const entities = await flow.graphEmbeddingsQuery(
|
||||
vectors[0],
|
||||
limit ?? 10,
|
||||
collection,
|
||||
);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(entities, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// ===================== Config tools =====================
|
||||
|
||||
server.tool(
|
||||
"get_config_all",
|
||||
"Get all configuration values",
|
||||
{},
|
||||
async () => {
|
||||
const cfg = socket.config();
|
||||
const resp = await cfg.getConfigAll();
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// --- Config ---
|
||||
server.tool(
|
||||
"get_config",
|
||||
"Get configuration values",
|
||||
"Get specific configuration values",
|
||||
{
|
||||
keys: z.array(z.string()).describe("Config keys to retrieve"),
|
||||
keys: z.array(
|
||||
z.object({
|
||||
type: z.string().describe("Config type"),
|
||||
key: z.string().describe("Config key"),
|
||||
}),
|
||||
).describe("Config keys to retrieve"),
|
||||
},
|
||||
async ({ keys }) => {
|
||||
const resp = await socket.request("config", { operation: "get", keys }) as Record<string, unknown>;
|
||||
const cfg = socket.config();
|
||||
const resp = await cfg.getConfig(keys);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
|
@ -146,10 +206,206 @@ export function createMcpServer(config: {
|
|||
"put_config",
|
||||
"Set configuration values",
|
||||
{
|
||||
values: z.record(z.unknown()).describe("Key-value pairs to set"),
|
||||
values: z.array(
|
||||
z.object({
|
||||
type: z.string().describe("Config type"),
|
||||
key: z.string().describe("Config key"),
|
||||
value: z.string().describe("Config value (JSON-encoded)"),
|
||||
}),
|
||||
).describe("Key-value entries to set"),
|
||||
},
|
||||
async ({ values }) => {
|
||||
const resp = await socket.request("config", { operation: "put", values }) as Record<string, unknown>;
|
||||
const cfg = socket.config();
|
||||
const resp = await cfg.putConfig(values);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"delete_config",
|
||||
"Delete a configuration entry",
|
||||
{
|
||||
type: z.string().describe("Config type"),
|
||||
key: z.string().describe("Config key"),
|
||||
},
|
||||
async ({ type, key }) => {
|
||||
const cfg = socket.config();
|
||||
const resp = await cfg.deleteConfig({ type, key });
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// ===================== Flow management tools =====================
|
||||
|
||||
server.tool(
|
||||
"get_flows",
|
||||
"List all available flows",
|
||||
{},
|
||||
async () => {
|
||||
const flows = socket.flows();
|
||||
const ids = await flows.getFlows();
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(ids, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"get_flow",
|
||||
"Get a specific flow definition",
|
||||
{
|
||||
flow_id: z.string().describe("Flow ID to retrieve"),
|
||||
},
|
||||
async ({ flow_id }) => {
|
||||
const flows = socket.flows();
|
||||
const def = await flows.getFlow(flow_id);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(def, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"start_flow",
|
||||
"Start a flow instance",
|
||||
{
|
||||
flow_id: z.string().describe("Flow ID"),
|
||||
blueprint_name: z.string().describe("Blueprint name"),
|
||||
description: z.string().describe("Flow description"),
|
||||
parameters: z.record(z.unknown()).optional().describe("Optional flow parameters"),
|
||||
},
|
||||
async ({ flow_id, blueprint_name, description, parameters }) => {
|
||||
const flows = socket.flows();
|
||||
const resp = await flows.startFlow(flow_id, blueprint_name, description, parameters);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"stop_flow",
|
||||
"Stop a running flow",
|
||||
{
|
||||
flow_id: z.string().describe("Flow ID to stop"),
|
||||
},
|
||||
async ({ flow_id }) => {
|
||||
const flows = socket.flows();
|
||||
const resp = await flows.stopFlow(flow_id);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// ===================== Library (document) tools =====================
|
||||
|
||||
server.tool(
|
||||
"get_documents",
|
||||
"List all documents in the library",
|
||||
{},
|
||||
async () => {
|
||||
const lib = socket.librarian();
|
||||
const docs = await lib.getDocuments();
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(docs, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"load_document",
|
||||
"Upload a document to the library",
|
||||
{
|
||||
document: z.string().describe("Base64-encoded document content"),
|
||||
mime_type: z.string().describe("Document MIME type"),
|
||||
title: z.string().describe("Document title"),
|
||||
comments: z.string().optional().describe("Additional comments"),
|
||||
tags: z.array(z.string()).optional().describe("Document tags"),
|
||||
id: z.string().optional().describe("Optional document ID"),
|
||||
},
|
||||
async ({ document, mime_type, title, comments, tags, id }) => {
|
||||
const lib = socket.librarian();
|
||||
const resp = await lib.loadDocument(
|
||||
document,
|
||||
mime_type,
|
||||
title,
|
||||
comments ?? "",
|
||||
tags ?? [],
|
||||
id,
|
||||
);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"remove_document",
|
||||
"Remove a document from the library",
|
||||
{
|
||||
id: z.string().describe("Document ID to remove"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ id, collection }) => {
|
||||
const lib = socket.librarian();
|
||||
const resp = await lib.removeDocument(id, collection);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// ===================== Prompt tools =====================
|
||||
|
||||
server.tool(
|
||||
"get_prompts",
|
||||
"List available prompt templates",
|
||||
{},
|
||||
async () => {
|
||||
const cfg = socket.config();
|
||||
const prompts = await cfg.getPrompts();
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(prompts, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"get_prompt",
|
||||
"Get a specific prompt template",
|
||||
{
|
||||
id: z.string().describe("Prompt template ID"),
|
||||
},
|
||||
async ({ id }) => {
|
||||
const cfg = socket.config();
|
||||
const prompt = await cfg.getPrompt(id);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(prompt, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
// ===================== Knowledge core tools =====================
|
||||
|
||||
server.tool(
|
||||
"get_knowledge_cores",
|
||||
"List available knowledge graph cores",
|
||||
{},
|
||||
async () => {
|
||||
const knowledge = socket.knowledge();
|
||||
const cores = await knowledge.getKnowledgeCores();
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(cores, null, 2) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"delete_kg_core",
|
||||
"Delete a knowledge graph core",
|
||||
{
|
||||
id: z.string().describe("Knowledge core ID"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ id, collection }) => {
|
||||
const knowledge = socket.knowledge();
|
||||
const resp = await knowledge.deleteKgCore(id, collection);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
},
|
||||
);
|
||||
|
||||
server.tool(
|
||||
"load_kg_core",
|
||||
"Load a knowledge graph core",
|
||||
{
|
||||
id: z.string().describe("Knowledge core ID"),
|
||||
flow: z.string().describe("Flow to use for loading"),
|
||||
collection: z.string().optional().describe("Collection name"),
|
||||
},
|
||||
async ({ id, flow, collection }) => {
|
||||
const knowledge = socket.knowledge();
|
||||
const resp = await knowledge.loadKgCore(id, flow, collection);
|
||||
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
|
||||
},
|
||||
);
|
||||
|
|
@ -160,6 +416,7 @@ export function createMcpServer(config: {
|
|||
export async function run(): Promise<void> {
|
||||
const { server, socket } = createMcpServer({
|
||||
gatewayUrl: process.env.GATEWAY_URL ?? "ws://localhost:8088/api/v1/socket",
|
||||
user: process.env.USER_ID ?? "mcp",
|
||||
token: process.env.GATEWAY_SECRET,
|
||||
flowId: process.env.FLOW_ID ?? "default",
|
||||
});
|
||||
|
|
@ -167,8 +424,8 @@ export async function run(): Promise<void> {
|
|||
const transport = new StdioServerTransport();
|
||||
await server.connect(transport);
|
||||
|
||||
process.on("SIGINT", async () => {
|
||||
await socket.close();
|
||||
process.on("SIGINT", () => {
|
||||
socket.close();
|
||||
process.exit(0);
|
||||
});
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,147 +0,0 @@
|
|||
/**
|
||||
* WebSocket manager for communicating with the TrustGraph gateway.
|
||||
*
|
||||
* Maintains a persistent connection per user and handles request/response
|
||||
* correlation via UUIDs.
|
||||
*
|
||||
* Python reference: trustgraph-mcp/trustgraph/mcp_server/tg_socket.py
|
||||
*/
|
||||
|
||||
import WebSocket from "ws";
|
||||
import { randomUUID } from "node:crypto";
|
||||
|
||||
export interface SocketManagerConfig {
|
||||
gatewayUrl: string;
|
||||
token?: string;
|
||||
}
|
||||
|
||||
interface PendingRequest {
|
||||
resolve: (value: unknown) => void;
|
||||
reject: (error: Error) => void;
|
||||
responses: unknown[];
|
||||
streaming: boolean;
|
||||
onChunk?: (chunk: unknown) => void;
|
||||
}
|
||||
|
||||
export class SocketManager {
|
||||
private ws: WebSocket | null = null;
|
||||
private pending = new Map<string, PendingRequest>();
|
||||
private connected = false;
|
||||
|
||||
constructor(private readonly config: SocketManagerConfig) {}
|
||||
|
||||
async connect(): Promise<void> {
|
||||
if (this.connected) return;
|
||||
|
||||
const url = new URL(this.config.gatewayUrl);
|
||||
if (this.config.token) {
|
||||
url.searchParams.set("token", this.config.token);
|
||||
}
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
this.ws = new WebSocket(url.toString());
|
||||
|
||||
this.ws.on("open", () => {
|
||||
this.connected = true;
|
||||
resolve();
|
||||
});
|
||||
|
||||
this.ws.on("error", (err) => {
|
||||
if (!this.connected) reject(err);
|
||||
else console.error("[SocketManager] WebSocket error:", err);
|
||||
});
|
||||
|
||||
this.ws.on("message", (data) => {
|
||||
try {
|
||||
const msg = JSON.parse(data.toString());
|
||||
const { id, response, error, complete } = msg;
|
||||
|
||||
const req = this.pending.get(id);
|
||||
if (!req) return;
|
||||
|
||||
if (error) {
|
||||
req.reject(new Error(`${error.type}: ${error.message}`));
|
||||
this.pending.delete(id);
|
||||
return;
|
||||
}
|
||||
|
||||
if (req.streaming && req.onChunk) {
|
||||
req.onChunk(response);
|
||||
}
|
||||
|
||||
req.responses.push(response);
|
||||
|
||||
if (complete) {
|
||||
req.resolve(req.streaming ? req.responses : response);
|
||||
this.pending.delete(id);
|
||||
}
|
||||
} catch (err) {
|
||||
console.error("[SocketManager] Failed to parse message:", err);
|
||||
}
|
||||
});
|
||||
|
||||
this.ws.on("close", () => {
|
||||
this.connected = false;
|
||||
// Reject all pending requests
|
||||
for (const [id, req] of this.pending) {
|
||||
req.reject(new Error("WebSocket closed"));
|
||||
}
|
||||
this.pending.clear();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async request(
|
||||
service: string,
|
||||
requestData: Record<string, unknown>,
|
||||
options?: {
|
||||
flowId?: string;
|
||||
timeoutMs?: number;
|
||||
onChunk?: (chunk: unknown) => void;
|
||||
},
|
||||
): Promise<unknown> {
|
||||
await this.connect();
|
||||
if (!this.ws) throw new Error("Not connected");
|
||||
|
||||
const id = randomUUID();
|
||||
const timeoutMs = options?.timeoutMs ?? 300_000;
|
||||
|
||||
const msg = {
|
||||
id,
|
||||
service,
|
||||
flow: options?.flowId ?? "default",
|
||||
request: requestData,
|
||||
};
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const timer = setTimeout(() => {
|
||||
this.pending.delete(id);
|
||||
reject(new Error(`Request timed out after ${timeoutMs}ms`));
|
||||
}, timeoutMs);
|
||||
|
||||
this.pending.set(id, {
|
||||
resolve: (value) => {
|
||||
clearTimeout(timer);
|
||||
resolve(value);
|
||||
},
|
||||
reject: (err) => {
|
||||
clearTimeout(timer);
|
||||
reject(err);
|
||||
},
|
||||
responses: [],
|
||||
streaming: !!options?.onChunk,
|
||||
onChunk: options?.onChunk,
|
||||
});
|
||||
|
||||
this.ws!.send(JSON.stringify(msg));
|
||||
});
|
||||
}
|
||||
|
||||
async close(): Promise<void> {
|
||||
if (this.ws) {
|
||||
this.ws.close();
|
||||
this.ws = null;
|
||||
this.connected = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue