This commit is contained in:
elpresidank 2026-04-05 21:09:33 -05:00
parent 9e9307a2aa
commit e26caa0b12
123 changed files with 3478 additions and 10078 deletions

4
ts/.gitignore vendored Normal file
View file

@ -0,0 +1,4 @@
node_modules/
dist/
*.tsbuildinfo
.turbo/

16
ts/package.json Normal file
View file

@ -0,0 +1,16 @@
{
"name": "trustgraph-ts",
"private": true,
"scripts": {
"build": "turbo build",
"dev": "turbo dev",
"lint": "turbo lint",
"test": "turbo test",
"clean": "turbo clean"
},
"devDependencies": {
"turbo": "^2.5.0",
"typescript": "^5.8.0"
},
"packageManager": "pnpm@9.15.0"
}

View file

@ -0,0 +1,21 @@
{
"name": "@trustgraph/base",
"version": "0.1.0",
"type": "module",
"main": "dist/index.js",
"types": "dist/index.d.ts",
"scripts": {
"build": "tsc",
"dev": "tsc --watch",
"clean": "rm -rf dist",
"test": "vitest run"
},
"dependencies": {
"nats": "^2.29.0",
"prom-client": "^15.1.0"
},
"devDependencies": {
"typescript": "^5.8.0",
"vitest": "^3.1.0"
}
}

View file

@ -0,0 +1,12 @@
export type {
Message,
BackendProducer,
BackendConsumer,
PubSubBackend,
CreateProducerOptions,
CreateConsumerOptions,
ConsumerType,
InitialPosition,
} from "./types.js";
export { NatsBackend } from "./nats.js";

View file

@ -0,0 +1,196 @@
/**
* NATS JetStream backend implementation.
*
* Replaces Pulsar as the message broker. NATS JetStream provides
* at-least-once delivery, consumer groups, and replay matching
* the QoS levels used by the Python Pulsar backend.
*
* Python reference: trustgraph-base/trustgraph/base/pulsar_backend.py
*/
import {
connect,
type NatsConnection,
type JetStreamClient,
type JetStreamManager,
type ConsumerMessages,
type JsMsg,
StringCodec,
AckPolicy,
} from "nats";
import type {
PubSubBackend,
BackendProducer,
BackendConsumer,
CreateProducerOptions,
CreateConsumerOptions,
Message,
} from "./types.js";
const sc = StringCodec();
class NatsMessage<T> implements Message<T> {
constructor(
private readonly msg: JsMsg,
private readonly decoded: T,
) {}
value(): T {
return this.decoded;
}
properties(): Record<string, string> {
const headers = this.msg.headers;
const props: Record<string, string> = {};
if (headers) {
for (const [key, values] of headers) {
props[key] = values[0];
}
}
return props;
}
}
class NatsProducer<T> implements BackendProducer<T> {
constructor(
private readonly js: JetStreamClient,
private readonly subject: string,
) {}
async send(message: T, properties?: Record<string, string>): Promise<void> {
const data = sc.encode(JSON.stringify(message));
const opts: Record<string, unknown> = {};
if (properties && Object.keys(properties).length > 0) {
const { headers } = await import("nats");
const hdrs = headers();
for (const [key, val] of Object.entries(properties)) {
hdrs.append(key, val);
}
opts.headers = hdrs;
}
await this.js.publish(this.subject, data, opts);
}
async flush(): Promise<void> {
// NATS publishes are flushed on the connection level
}
async close(): Promise<void> {
// No per-producer cleanup needed for NATS
}
}
class NatsConsumer<T> implements BackendConsumer<T> {
private messages: ConsumerMessages | null = null;
constructor(
private readonly js: JetStreamClient,
private readonly jsm: JetStreamManager,
private readonly subject: string,
private readonly subscription: string,
private readonly initialPosition: "latest" | "earliest",
) {}
async init(): Promise<void> {
// Ensure stream exists
const streamName = this.streamNameFromSubject(this.subject);
try {
await this.jsm.streams.info(streamName);
} catch {
await this.jsm.streams.add({
name: streamName,
subjects: [this.subject],
});
}
// Create or bind to durable consumer
const consumer = await this.js.consumers.get(streamName, this.subscription);
this.messages = await consumer.consume();
}
async receive(timeoutMs = 2000): Promise<Message<T> | null> {
if (!this.messages) throw new Error("Consumer not initialized");
const deadline = Date.now() + timeoutMs;
for await (const msg of this.messages) {
const decoded = JSON.parse(sc.decode(msg.data)) as T;
return new NatsMessage(msg, decoded);
}
if (Date.now() >= deadline) return null;
return null;
}
async acknowledge(message: Message<T>): Promise<void> {
const natsMsg = message as NatsMessage<T>;
// Access internal JsMsg for ack — in practice we'd store the ref
// This is a simplified version; real impl tracks msg refs
void natsMsg;
}
async negativeAcknowledge(message: Message<T>): Promise<void> {
void message;
}
async unsubscribe(): Promise<void> {
// Drain and close consumer
}
async close(): Promise<void> {
if (this.messages) {
this.messages.stop();
}
}
private streamNameFromSubject(subject: string): string {
// Convert topic like "tg.flow.text-completion" to stream name "tg_flow"
const parts = subject.split(".");
return parts.slice(0, 2).join("_");
}
}
export class NatsBackend implements PubSubBackend {
private connection: NatsConnection | null = null;
private js: JetStreamClient | null = null;
private jsm: JetStreamManager | null = null;
constructor(private readonly url: string = "nats://localhost:4222") {}
private async ensureConnected(): Promise<void> {
if (!this.connection) {
this.connection = await connect({ servers: this.url });
this.js = this.connection.jetstream();
this.jsm = await this.connection.jetstreamManager();
}
}
async createProducer<T>(options: CreateProducerOptions): Promise<BackendProducer<T>> {
await this.ensureConnected();
return new NatsProducer<T>(this.js!, options.topic);
}
async createConsumer<T>(options: CreateConsumerOptions): Promise<BackendConsumer<T>> {
await this.ensureConnected();
const consumer = new NatsConsumer<T>(
this.js!,
this.jsm!,
options.topic,
options.subscription,
options.initialPosition ?? "latest",
);
await consumer.init();
return consumer;
}
async close(): Promise<void> {
if (this.connection) {
await this.connection.drain();
this.connection = null;
this.js = null;
this.jsm = null;
}
}
}

View file

@ -0,0 +1,45 @@
/**
* Core pub/sub backend abstraction.
*
* Mirrors Python's backend.py Protocol classes. Any message broker
* (NATS, Pulsar, Redis Streams) implements these interfaces.
*/
export interface Message<T = unknown> {
value(): T;
properties(): Record<string, string>;
}
export interface BackendProducer<T = unknown> {
send(message: T, properties?: Record<string, string>): Promise<void>;
flush(): Promise<void>;
close(): Promise<void>;
}
export interface BackendConsumer<T = unknown> {
receive(timeoutMs?: number): Promise<Message<T> | null>;
acknowledge(message: Message<T>): Promise<void>;
negativeAcknowledge(message: Message<T>): Promise<void>;
unsubscribe(): Promise<void>;
close(): Promise<void>;
}
export type ConsumerType = "shared" | "exclusive" | "failover";
export type InitialPosition = "latest" | "earliest";
export interface CreateProducerOptions {
topic: string;
}
export interface CreateConsumerOptions {
topic: string;
subscription: string;
initialPosition?: InitialPosition;
consumerType?: ConsumerType;
}
export interface PubSubBackend {
createProducer<T>(options: CreateProducerOptions): Promise<BackendProducer<T>>;
createConsumer<T>(options: CreateConsumerOptions): Promise<BackendConsumer<T>>;
close(): Promise<void>;
}

View file

@ -0,0 +1,29 @@
/**
* Custom error types.
*
* Python reference: trustgraph-base/trustgraph/exceptions.py
*/
export class TooManyRequestsError extends Error {
constructor(message = "Rate limit exceeded") {
super(message);
this.name = "TooManyRequestsError";
}
}
export class LlmError extends Error {
constructor(
message: string,
public readonly errorType: string = "llm-error",
) {
super(message);
this.name = "LlmError";
}
}
export class ParseError extends Error {
constructor(message: string) {
super(message);
this.name = "ParseError";
}
}

View file

@ -0,0 +1,10 @@
// @trustgraph/base — core abstractions for the TrustGraph TypeScript port
export * from "./backend/index.js";
export * from "./messaging/index.js";
export * from "./processor/index.js";
export * from "./schema/index.js";
export * from "./spec/index.js";
export * from "./services/index.js";
export * from "./metrics/index.js";
export * from "./errors.js";

View file

@ -0,0 +1,105 @@
/**
* High-level consumer with concurrency, retry, and rate-limit handling.
*
* Python reference: trustgraph-base/trustgraph/base/consumer.py
*/
import type { PubSubBackend, BackendConsumer, Message } from "../backend/types.js";
import { TooManyRequestsError } from "../errors.js";
export type MessageHandler<T> = (
message: T,
properties: Record<string, string>,
flow: FlowContext,
) => Promise<void>;
export interface FlowContext {
id: string;
name: string;
}
export interface ConsumerOptions<T> {
pubsub: PubSubBackend;
topic: string;
subscription: string;
handler: MessageHandler<T>;
concurrency?: number;
initialPosition?: "latest" | "earliest";
rateLimitRetryMs?: number;
rateLimitTimeoutMs?: number;
}
export class Consumer<T> {
private backend: BackendConsumer<T> | null = null;
private running = false;
private abortController = new AbortController();
private readonly concurrency: number;
private readonly rateLimitRetryMs: number;
constructor(private readonly options: ConsumerOptions<T>) {
this.concurrency = options.concurrency ?? 1;
this.rateLimitRetryMs = options.rateLimitRetryMs ?? 10_000;
}
async start(flow: FlowContext): Promise<void> {
this.backend = await this.options.pubsub.createConsumer<T>({
topic: this.options.topic,
subscription: this.options.subscription,
initialPosition: this.options.initialPosition ?? "latest",
});
this.running = true;
// Spawn concurrent consumer tasks
const tasks = Array.from({ length: this.concurrency }, () =>
this.consumeLoop(flow),
);
// Run all concurrently — first rejection stops all
await Promise.all(tasks);
}
async stop(): Promise<void> {
this.running = false;
this.abortController.abort();
if (this.backend) {
await this.backend.close();
this.backend = null;
}
}
private async consumeLoop(flow: FlowContext): Promise<void> {
while (this.running) {
try {
const msg = await this.backend!.receive(2000);
if (!msg) continue;
await this.handleWithRetry(msg, flow);
await this.backend!.acknowledge(msg);
} catch (err) {
if (!this.running) break;
console.error("[Consumer] Error in consume loop:", err);
// Brief pause before retry
await sleep(1000);
}
}
}
private async handleWithRetry(msg: Message<T>, flow: FlowContext): Promise<void> {
try {
await this.options.handler(msg.value(), msg.properties(), flow);
} catch (err) {
if (err instanceof TooManyRequestsError) {
console.warn(`[Consumer] Rate limited, retrying in ${this.rateLimitRetryMs}ms`);
await sleep(this.rateLimitRetryMs);
await this.options.handler(msg.value(), msg.properties(), flow);
} else {
throw err;
}
}
}
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}

View file

@ -0,0 +1,4 @@
export { Producer } from "./producer.js";
export { Consumer, type MessageHandler, type FlowContext, type ConsumerOptions } from "./consumer.js";
export { Subscriber, AsyncQueue } from "./subscriber.js";
export { RequestResponse, type RequestResponseOptions } from "./request-response.js";

View file

@ -0,0 +1,40 @@
/**
* High-level async producer with queue buffering and retry.
*
* Python reference: trustgraph-base/trustgraph/base/producer.py
*/
import type { PubSubBackend, BackendProducer } from "../backend/types.js";
import type { ProducerMetrics } from "../metrics/prometheus.js";
export class Producer<T> {
private backend: BackendProducer<T> | null = null;
private running = false;
constructor(
private readonly pubsub: PubSubBackend,
private readonly topic: string,
private readonly metrics?: ProducerMetrics,
) {}
async start(): Promise<void> {
this.backend = await this.pubsub.createProducer<T>({ topic: this.topic });
this.running = true;
}
async send(id: string, message: T): Promise<void> {
if (!this.backend) throw new Error("Producer not started");
await this.backend.send(message, { id });
this.metrics?.inc();
}
async stop(): Promise<void> {
this.running = false;
if (this.backend) {
await this.backend.flush();
await this.backend.close();
this.backend = null;
}
}
}

View file

@ -0,0 +1,91 @@
/**
* Request/response pattern over pub/sub.
*
* Sends a request with a unique ID, subscribes for matching responses.
* Supports streaming (multiple responses per request) via a recipient callback.
*
* Python reference: trustgraph-base/trustgraph/base/request_response_spec.py
*/
import { randomUUID } from "node:crypto";
import { Producer } from "./producer.js";
import { Subscriber } from "./subscriber.js";
import type { PubSubBackend } from "../backend/types.js";
export interface RequestResponseOptions {
pubsub: PubSubBackend;
requestTopic: string;
responseTopic: string;
subscription: string;
}
export class RequestResponse<TReq, TRes> {
private producer: Producer<TReq>;
private subscriber: Subscriber<TRes>;
constructor(private readonly options: RequestResponseOptions) {
this.producer = new Producer<TReq>(options.pubsub, options.requestTopic);
this.subscriber = new Subscriber<TRes>(
options.pubsub,
options.responseTopic,
options.subscription,
);
}
async start(): Promise<void> {
await this.producer.start();
await this.subscriber.start();
}
async stop(): Promise<void> {
await this.producer.stop();
await this.subscriber.stop();
}
/**
* Send a request and wait for responses.
*
* @param request - The request payload
* @param options.timeoutMs - Total timeout in milliseconds (default: 300s)
* @param options.recipient - Optional callback for streaming responses.
* Return `true` to indicate the final response has been received.
* If omitted, returns the first response.
*/
async request(
request: TReq,
options?: {
timeoutMs?: number;
recipient?: (response: TRes) => Promise<boolean>;
},
): Promise<TRes> {
const id = randomUUID();
const timeoutMs = options?.timeoutMs ?? 300_000;
const recipient = options?.recipient;
const queue = this.subscriber.subscribe(id);
try {
await this.producer.send(id, request);
const deadline = Date.now() + timeoutMs;
while (true) {
const remaining = deadline - Date.now();
if (remaining <= 0) {
throw new Error(`Request timed out after ${timeoutMs}ms`);
}
const response = await queue.pop(remaining);
if (recipient) {
const isFinal = await recipient(response);
if (isFinal) return response;
} else {
return response;
}
}
} finally {
this.subscriber.unsubscribe(id);
}
}
}

View file

@ -0,0 +1,143 @@
/**
* Fan-out subscriber: routes responses to waiting callers by request ID.
*
* Python reference: trustgraph-base/trustgraph/base/subscriber.py
*/
import type { PubSubBackend, BackendConsumer } from "../backend/types.js";
type Resolver<T> = {
queue: AsyncQueue<T>;
};
/**
* Simple async queue for inter-task communication (replaces asyncio.Queue).
*/
export class AsyncQueue<T> {
private buffer: T[] = [];
private waiters: Array<(value: T) => void> = [];
push(item: T): void {
const waiter = this.waiters.shift();
if (waiter) {
waiter(item);
} else {
this.buffer.push(item);
}
}
async pop(timeoutMs?: number): Promise<T> {
const buffered = this.buffer.shift();
if (buffered !== undefined) return buffered;
return new Promise<T>((resolve, reject) => {
let timer: ReturnType<typeof setTimeout> | undefined;
const waiter = (value: T) => {
if (timer) clearTimeout(timer);
resolve(value);
};
this.waiters.push(waiter);
if (timeoutMs !== undefined) {
timer = setTimeout(() => {
const idx = this.waiters.indexOf(waiter);
if (idx !== -1) this.waiters.splice(idx, 1);
reject(new Error(`Queue.pop timed out after ${timeoutMs}ms`));
}, timeoutMs);
}
});
}
get length(): number {
return this.buffer.length;
}
}
export class Subscriber<T> {
private backend: BackendConsumer<T> | null = null;
private running = false;
// ID-specific subscriptions (request/response correlation)
private idSubscribers = new Map<string, Resolver<T>>();
// Wildcard subscribers (receive all messages)
private allSubscribers = new Map<string, Resolver<T>>();
constructor(
private readonly pubsub: PubSubBackend,
private readonly topic: string,
private readonly subscription: string,
) {}
async start(): Promise<void> {
this.backend = await this.pubsub.createConsumer<T>({
topic: this.topic,
subscription: this.subscription,
});
this.running = true;
// Start the dispatch loop (fire and forget — runs until stop)
this.dispatchLoop().catch((err) => {
if (this.running) console.error("[Subscriber] dispatch loop error:", err);
});
}
async stop(): Promise<void> {
this.running = false;
if (this.backend) {
await this.backend.close();
this.backend = null;
}
}
subscribe(id: string): AsyncQueue<T> {
const queue = new AsyncQueue<T>();
this.idSubscribers.set(id, { queue });
return queue;
}
subscribeAll(id: string): AsyncQueue<T> {
const queue = new AsyncQueue<T>();
this.allSubscribers.set(id, { queue });
return queue;
}
unsubscribe(id: string): void {
this.idSubscribers.delete(id);
}
unsubscribeAll(id: string): void {
this.allSubscribers.delete(id);
}
private async dispatchLoop(): Promise<void> {
while (this.running) {
try {
const msg = await this.backend!.receive(2000);
if (!msg) continue;
const props = msg.properties();
const id = props.id;
const value = msg.value();
// Route to ID-specific subscriber
if (id) {
const sub = this.idSubscribers.get(id);
if (sub) {
sub.queue.push(value);
}
}
// Broadcast to all-subscribers
for (const sub of this.allSubscribers.values()) {
sub.queue.push(value);
}
await this.backend!.acknowledge(msg);
} catch (err) {
if (!this.running) break;
console.error("[Subscriber] Error:", err);
}
}
}
}

View file

@ -0,0 +1 @@
export { ConsumerMetrics, ProducerMetrics, registry } from "./prometheus.js";

View file

@ -0,0 +1,68 @@
/**
* Prometheus metrics wrappers.
*
* Python reference: trustgraph-base/trustgraph/base/metrics.py
*/
import { Counter, Histogram, Registry, collectDefaultMetrics } from "prom-client";
export const registry = new Registry();
collectDefaultMetrics({ register: registry });
export class ConsumerMetrics {
private requestHistogram: Histogram;
private processingCounter: Counter;
private rateLimitCounter: Counter;
constructor(processor: string, flow: string, name: string) {
this.requestHistogram = new Histogram({
name: "tg_consumer_request_duration_seconds",
help: "Consumer request processing time",
labelNames: ["processor", "flow", "name"],
registers: [registry],
});
this.processingCounter = new Counter({
name: "tg_consumer_processing_total",
help: "Consumer processing outcomes",
labelNames: ["processor", "flow", "name", "status"],
registers: [registry],
});
this.rateLimitCounter = new Counter({
name: "tg_consumer_rate_limit_total",
help: "Consumer rate limit events",
labelNames: ["processor", "flow", "name"],
registers: [registry],
});
}
recordTime(seconds: number): void {
this.requestHistogram.observe(seconds);
}
process(status: "success" | "error"): void {
this.processingCounter.inc({ status });
}
rateLimit(): void {
this.rateLimitCounter.inc();
}
}
export class ProducerMetrics {
private counter: Counter;
constructor(processor: string, flow: string, name: string) {
this.counter = new Counter({
name: "tg_producer_items_total",
help: "Producer items sent",
labelNames: ["processor", "flow", "name"],
registers: [registry],
});
}
inc(): void {
this.counter.inc();
}
}

View file

@ -0,0 +1,83 @@
/**
* Base async processor foundation for all TrustGraph services.
*
* Handles pub/sub lifecycle, configuration subscription, and graceful shutdown.
*
* Python reference: trustgraph-base/trustgraph/base/async_processor.py
*/
import type { PubSubBackend } from "../backend/types.js";
import { NatsBackend } from "../backend/nats.js";
import { topics } from "../schema/topics.js";
export interface ProcessorConfig {
id: string;
pubsubUrl?: string;
metricsPort?: number;
}
export type ConfigHandler = (
config: Record<string, unknown>,
version: number,
) => Promise<void>;
export abstract class AsyncProcessor {
protected pubsub: PubSubBackend;
protected running = false;
private configHandlers: ConfigHandler[] = [];
private shutdownCallbacks: Array<() => Promise<void>> = [];
constructor(protected readonly config: ProcessorConfig) {
this.pubsub = new NatsBackend(config.pubsubUrl ?? "nats://localhost:4222");
}
registerConfigHandler(handler: ConfigHandler): void {
this.configHandlers.push(handler);
}
async start(): Promise<void> {
this.running = true;
// Set up graceful shutdown
const shutdown = async () => {
console.log(`[${this.config.id}] Shutting down...`);
await this.stop();
process.exit(0);
};
process.on("SIGINT", shutdown);
process.on("SIGTERM", shutdown);
await this.run();
}
async stop(): Promise<void> {
this.running = false;
for (const cb of this.shutdownCallbacks) {
await cb();
}
await this.pubsub.close();
}
protected onShutdown(callback: () => Promise<void>): void {
this.shutdownCallbacks.push(callback);
}
protected abstract run(): Promise<void>;
/**
* Static launch helper parses env/args and starts the processor.
* Subclasses call: `MyProcessor.launch("my-service")`
*/
static async launch<T extends AsyncProcessor>(
this: new (config: ProcessorConfig) => T,
id: string,
): Promise<void> {
const config: ProcessorConfig = {
id,
pubsubUrl: process.env.NATS_URL ?? process.env.PULSAR_HOST,
metricsPort: parseInt(process.env.METRICS_PORT ?? "8000", 10),
};
const processor = new this(config);
await processor.start();
}
}

View file

@ -0,0 +1,69 @@
/**
* Flow-aware processor that manages dynamic flow instances.
*
* Python reference: trustgraph-base/trustgraph/base/flow_processor.py
*/
import { AsyncProcessor, type ProcessorConfig } from "./async-processor.js";
import type { Spec } from "../spec/types.js";
import { Flow, type FlowDefinition } from "./flow.js";
export abstract class FlowProcessor extends AsyncProcessor {
private specifications: Spec[] = [];
private flows = new Map<string, Flow>();
constructor(config: ProcessorConfig) {
super(config);
this.registerConfigHandler(this.onConfigureFlows.bind(this));
}
registerSpecification(spec: Spec): void {
this.specifications.push(spec);
}
protected async run(): Promise<void> {
// The processor sits idle waiting for flow configurations
// to arrive via the config push topic. In the meantime,
// the consumer loop runs in the background.
await new Promise<void>((resolve) => {
const check = () => {
if (!this.running) resolve();
else setTimeout(check, 1000);
};
check();
});
}
private async onConfigureFlows(
config: Record<string, unknown>,
version: number,
): Promise<void> {
const flowDefs = config.flows as Record<string, FlowDefinition> | undefined;
if (!flowDefs) return;
// Stop removed flows
for (const [name, flow] of this.flows) {
if (!(name in flowDefs)) {
await flow.stop();
this.flows.delete(name);
}
}
// Start new flows
for (const [name, defn] of Object.entries(flowDefs)) {
if (!this.flows.has(name)) {
const flow = new Flow(name, this.config.id, this.pubsub, defn, this.specifications);
await flow.start();
this.flows.set(name, flow);
}
}
}
override async stop(): Promise<void> {
for (const flow of this.flows.values()) {
await flow.stop();
}
this.flows.clear();
await super.stop();
}
}

View file

@ -0,0 +1,83 @@
/**
* Runtime flow instance created by FlowProcessor for each configured flow.
*
* Python reference: trustgraph-base/trustgraph/base/flow.py
*/
import type { PubSubBackend } from "../backend/types.js";
import type { Spec } from "../spec/types.js";
import type { Producer } from "../messaging/producer.js";
import type { Consumer } from "../messaging/consumer.js";
export interface FlowDefinition {
/** Topic overrides keyed by spec name */
topics?: Record<string, string>;
/** Parameter values keyed by spec name */
parameters?: Record<string, unknown>;
}
export class Flow {
private producers = new Map<string, Producer<unknown>>();
private consumers = new Map<string, Consumer<unknown>>();
private parameters = new Map<string, unknown>();
constructor(
public readonly name: string,
public readonly processorId: string,
private readonly pubsub: PubSubBackend,
private readonly definition: FlowDefinition,
private readonly specifications: Spec[],
) {}
async start(): Promise<void> {
for (const spec of this.specifications) {
await spec.add(this, this.pubsub, this.definition);
}
// Start all consumers
for (const consumer of this.consumers.values()) {
consumer.start({ id: this.processorId, name: this.name }).catch((err) => {
console.error(`[Flow:${this.name}] Consumer error:`, err);
});
}
}
async stop(): Promise<void> {
for (const consumer of this.consumers.values()) {
await consumer.stop();
}
for (const producer of this.producers.values()) {
await producer.stop();
}
}
registerProducer(name: string, producer: Producer<unknown>): void {
this.producers.set(name, producer);
}
registerConsumer(name: string, consumer: Consumer<unknown>): void {
this.consumers.set(name, consumer);
}
setParameter(name: string, value: unknown): void {
this.parameters.set(name, value);
}
producer<T>(name: string): Producer<T> {
const p = this.producers.get(name);
if (!p) throw new Error(`Producer "${name}" not found in flow "${this.name}"`);
return p as Producer<T>;
}
consumer<T>(name: string): Consumer<T> {
const c = this.consumers.get(name);
if (!c) throw new Error(`Consumer "${name}" not found in flow "${this.name}"`);
return c as Consumer<T>;
}
parameter<T>(name: string): T {
const v = this.parameters.get(name);
if (v === undefined) throw new Error(`Parameter "${name}" not found in flow "${this.name}"`);
return v as T;
}
}

View file

@ -0,0 +1,3 @@
export { AsyncProcessor, type ProcessorConfig, type ConfigHandler } from "./async-processor.js";
export { FlowProcessor } from "./flow-processor.js";
export { Flow, type FlowDefinition } from "./flow.js";

View file

@ -0,0 +1,3 @@
export * from "./primitives.js";
export * from "./topics.js";
export * from "./messages.js";

View file

@ -0,0 +1,135 @@
/**
* Message types for service communication.
*
* Python reference: trustgraph-base/trustgraph/schema/services/
*/
import type { TgError, Triple, Term, RowSchema } from "./primitives.js";
// Text completion
export interface TextCompletionRequest {
system: string;
prompt: string;
model?: string;
temperature?: number;
streaming?: boolean;
}
export interface TextCompletionResponse {
response: string;
model?: string;
inToken?: number;
outToken?: number;
error?: TgError;
endOfStream?: boolean;
}
// Embeddings
export interface EmbeddingsRequest {
text: string[];
model?: string;
}
export interface EmbeddingsResponse {
vectors: number[][];
error?: TgError;
}
// Graph RAG
export interface GraphRagRequest {
query: string;
collection?: string;
entityLimit?: number;
tripleLimit?: number;
maxSubgraphSize?: number;
maxPathLength?: number;
streaming?: boolean;
}
export interface GraphRagResponse {
response: string;
error?: TgError;
endOfStream?: boolean;
}
// Document RAG
export interface DocumentRagRequest {
query: string;
collection?: string;
streaming?: boolean;
}
export interface DocumentRagResponse {
response: string;
error?: TgError;
endOfStream?: boolean;
}
// Agent
export interface AgentRequest {
question: string;
collection?: string;
streaming?: boolean;
}
export interface AgentResponse {
answer: string;
error?: TgError;
endOfStream?: boolean;
endOfSession?: boolean;
}
// Triples query
export interface TriplesQueryRequest {
s?: Term;
p?: Term;
o?: Term;
collection?: string;
limit?: number;
}
export interface TriplesQueryResponse {
triples: Triple[];
error?: TgError;
}
// Graph embeddings query
export interface GraphEmbeddingsRequest {
vectors: number[][];
limit?: number;
collection?: string;
}
export interface GraphEmbeddingsResponse {
entities: Term[];
error?: TgError;
}
// Config
export type ConfigOperation = "get" | "list" | "delete" | "put" | "config";
export interface ConfigRequest {
operation: ConfigOperation;
keys?: string[];
values?: Record<string, unknown>;
}
export interface ConfigResponse {
version?: number;
values?: Record<string, unknown>;
directory?: string[];
config?: Record<string, unknown>;
error?: TgError;
}
// Prompt
export interface PromptRequest {
name: string;
variables?: Record<string, string>;
}
export interface PromptResponse {
system: string;
prompt: string;
error?: TgError;
}

View file

@ -0,0 +1,72 @@
/**
* Core data types mirroring the Python schema primitives.
*
* Python reference: trustgraph-base/trustgraph/schema/core/primitives.py
*/
export interface TgError {
type: string;
message: string;
}
// RDF Term types — discriminated union
export type TermType = "IRI" | "BLANK" | "LITERAL" | "TRIPLE";
export interface IriTerm {
type: "IRI";
iri: string;
}
export interface BlankTerm {
type: "BLANK";
id: string;
}
export interface LiteralTerm {
type: "LITERAL";
value: string;
datatype?: string;
language?: string;
}
export interface TripleTerm {
type: "TRIPLE";
triple: Triple;
}
export type Term = IriTerm | BlankTerm | LiteralTerm | TripleTerm;
export interface Triple {
s: Term;
p: Term;
o: Term;
g?: Term; // Named graph (optional quad)
}
export interface Field {
name: string;
type: string;
description?: string;
}
export interface RowSchema {
name: string;
description?: string;
fields: Field[];
}
// LLM-related types
export interface LlmResult {
text: string;
inToken: number;
outToken: number;
model: string;
}
export interface LlmChunk {
text: string;
inToken: number | null;
outToken: number | null;
model: string;
isFinal: boolean;
}

View file

@ -0,0 +1,62 @@
/**
* Topic naming conventions.
*
* Python reference: trustgraph-base/trustgraph/schema/core/topic.py
*
* The Python version uses Pulsar URI format: "q1/tg/flow/queue-name"
* We use NATS subject format: "tg.flow.queue-name"
*/
export type QoS = "q0" | "q1" | "q2";
export function topic(
name: string,
tenant = "tg",
namespace = "flow",
): string {
return `${tenant}.${namespace}.${name}`;
}
// Well-known topics from the Python schema
export const topics = {
// Config
configRequest: topic("config-request"),
configResponse: topic("config-response"),
configPush: topic("config-push"),
// Text completion
textCompletionRequest: topic("text-completion-request"),
textCompletionResponse: topic("text-completion-response"),
// Embeddings
embeddingsRequest: topic("embeddings-request"),
embeddingsResponse: topic("embeddings-response"),
// Graph RAG
graphRagRequest: topic("graph-rag-request"),
graphRagResponse: topic("graph-rag-response"),
// Document RAG
documentRagRequest: topic("document-rag-request"),
documentRagResponse: topic("document-rag-response"),
// Agent
agentRequest: topic("agent-request"),
agentResponse: topic("agent-response"),
// Triples
triplesRequest: topic("triples-request"),
triplesResponse: topic("triples-response"),
// Graph embeddings
graphEmbeddingsRequest: topic("graph-embeddings-request"),
graphEmbeddingsResponse: topic("graph-embeddings-response"),
// Document embeddings
docEmbeddingsRequest: topic("doc-embeddings-request"),
docEmbeddingsResponse: topic("doc-embeddings-response"),
// Prompt
promptRequest: topic("prompt-request"),
promptResponse: topic("prompt-response"),
} as const;

View file

@ -0,0 +1,46 @@
/**
* Base embeddings service.
*
* Python reference: trustgraph-base/trustgraph/base/embeddings_service.py
*/
import { FlowProcessor } from "../processor/flow-processor.js";
import { ConsumerSpec } from "../spec/consumer-spec.js";
import { ProducerSpec } from "../spec/producer-spec.js";
import { ParameterSpec } from "../spec/parameter-spec.js";
import type { ProcessorConfig } from "../processor/async-processor.js";
import type { FlowContext } from "../messaging/consumer.js";
import type { EmbeddingsRequest, EmbeddingsResponse } from "../schema/messages.js";
export abstract class EmbeddingsService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
new ConsumerSpec<EmbeddingsRequest>(
"request",
this.onRequest.bind(this),
),
);
this.registerSpecification(new ProducerSpec<EmbeddingsResponse>("response"));
this.registerSpecification(new ParameterSpec("model"));
}
private async onRequest(
msg: EmbeddingsRequest,
properties: Record<string, string>,
_flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (!requestId) return;
try {
const vectors = await this.onEmbeddings(msg.text, msg.model);
void vectors; // Producer send would go here
} catch (err) {
console.error(`[EmbeddingsService] Error processing request:`, err);
}
}
abstract onEmbeddings(texts: string[], model?: string): Promise<number[][]>;
}

View file

@ -0,0 +1,2 @@
export { LlmService } from "./llm-service.js";
export { EmbeddingsService } from "./embeddings-service.js";

View file

@ -0,0 +1,88 @@
/**
* Base LLM service handles message plumbing, subclasses implement the LLM call.
*
* Python reference: trustgraph-base/trustgraph/base/llm_service.py
*/
import { FlowProcessor } from "../processor/flow-processor.js";
import { ConsumerSpec } from "../spec/consumer-spec.js";
import { ProducerSpec } from "../spec/producer-spec.js";
import { ParameterSpec } from "../spec/parameter-spec.js";
import type { ProcessorConfig } from "../processor/async-processor.js";
import type { FlowContext } from "../messaging/consumer.js";
import type { Flow } from "../processor/flow.js";
import type {
TextCompletionRequest,
TextCompletionResponse,
} from "../schema/messages.js";
import type { LlmResult, LlmChunk } from "../schema/primitives.js";
export abstract class LlmService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
new ConsumerSpec<TextCompletionRequest>(
"request",
this.onRequest.bind(this),
),
);
this.registerSpecification(new ProducerSpec<TextCompletionResponse>("response"));
this.registerSpecification(new ParameterSpec("model"));
this.registerSpecification(new ParameterSpec("temperature"));
}
private async onRequest(
msg: TextCompletionRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
// We need the actual flow instance to access producers/parameters.
// In the full implementation, FlowContext would carry a flow reference.
// For now this shows the pattern.
const requestId = properties.id;
if (!requestId) return;
try {
if (msg.streaming && this.supportsStreaming()) {
for await (const chunk of this.generateContentStream(
msg.system,
msg.prompt,
msg.model,
msg.temperature,
)) {
// Send each chunk as a response with the same request ID
void chunk; // Producer send would go here
}
} else {
const result = await this.generateContent(
msg.system,
msg.prompt,
msg.model,
msg.temperature,
);
void result; // Producer send would go here
}
} catch (err) {
console.error(`[LlmService] Error processing request:`, err);
}
}
abstract generateContent(
system: string,
prompt: string,
model?: string,
temperature?: number,
): Promise<LlmResult>;
abstract generateContentStream(
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk>;
supportsStreaming(): boolean {
return false;
}
}

View file

@ -0,0 +1,32 @@
/**
* Consumer specification declares a message consumer for a flow.
*
* Python reference: trustgraph-base/trustgraph/base/consumer_spec.py
*/
import type { Spec } from "./types.js";
import type { PubSubBackend } from "../backend/types.js";
import type { Flow, FlowDefinition } from "../processor/flow.js";
import { Consumer, type MessageHandler } from "../messaging/consumer.js";
export class ConsumerSpec<T> implements Spec {
constructor(
public readonly name: string,
private readonly handler: MessageHandler<T>,
private readonly concurrency = 1,
) {}
async add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise<void> {
const topic = definition.topics?.[this.name] ?? this.name;
const consumer = new Consumer<T>({
pubsub,
topic,
subscription: `${flow.processorId}-${flow.name}-${this.name}`,
handler: this.handler,
concurrency: this.concurrency,
});
flow.registerConsumer(this.name, consumer as Consumer<unknown>);
}
}

View file

@ -0,0 +1,4 @@
export type { Spec } from "./types.js";
export { ConsumerSpec } from "./consumer-spec.js";
export { ProducerSpec } from "./producer-spec.js";
export { ParameterSpec } from "./parameter-spec.js";

View file

@ -0,0 +1,18 @@
/**
* Parameter specification declares a configuration parameter for a flow.
*
* Python reference: trustgraph-base/trustgraph/base/parameter_spec.py
*/
import type { Spec } from "./types.js";
import type { PubSubBackend } from "../backend/types.js";
import type { Flow, FlowDefinition } from "../processor/flow.js";
export class ParameterSpec implements Spec {
constructor(public readonly name: string) {}
async add(flow: Flow, _pubsub: PubSubBackend, definition: FlowDefinition): Promise<void> {
const value = definition.parameters?.[this.name];
flow.setParameter(this.name, value);
}
}

View file

@ -0,0 +1,21 @@
/**
* Producer specification declares a message producer for a flow.
*
* Python reference: trustgraph-base/trustgraph/base/producer_spec.py
*/
import type { Spec } from "./types.js";
import type { PubSubBackend } from "../backend/types.js";
import type { Flow, FlowDefinition } from "../processor/flow.js";
import { Producer } from "../messaging/producer.js";
export class ProducerSpec<T> implements Spec {
constructor(public readonly name: string) {}
async add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise<void> {
const topic = definition.topics?.[this.name] ?? this.name;
const producer = new Producer<T>(pubsub, topic);
await producer.start();
flow.registerProducer(this.name, producer as Producer<unknown>);
}
}

View file

@ -0,0 +1,13 @@
/**
* Specification types for declarative flow configuration.
*
* Python reference: trustgraph-base/trustgraph/base/spec.py and siblings
*/
import type { PubSubBackend } from "../backend/types.js";
import type { Flow, FlowDefinition } from "../processor/flow.js";
export interface Spec {
name: string;
add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise<void>;
}

View file

@ -0,0 +1,8 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"outDir": "dist",
"rootDir": "src"
},
"include": ["src"]
}

View file

@ -0,0 +1,25 @@
{
"name": "@trustgraph/cli",
"version": "0.1.0",
"type": "module",
"bin": {
"tg": "dist/index.js"
},
"scripts": {
"build": "tsc",
"dev": "tsc --watch",
"clean": "rm -rf dist",
"test": "vitest run"
},
"dependencies": {
"@trustgraph/base": "workspace:*",
"@trustgraph/mcp": "workspace:*",
"commander": "^13.1.0",
"ws": "^8.18.0"
},
"devDependencies": {
"@types/ws": "^8.5.0",
"typescript": "^5.8.0",
"vitest": "^3.1.0"
}
}

View file

@ -0,0 +1,34 @@
/**
* Agent CLI commands.
*
* Python reference: trustgraph-cli/trustgraph/cli/invoke_agent.py
*/
import type { Command } from "commander";
import { createSocket, getOpts } from "./util.js";
export function registerAgentCommands(program: Command): void {
program
.command("agent")
.description("Ask the TrustGraph agent a question")
.argument("<question>", "Question to ask")
.action(async (question: string, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("agent", { question }, {
flowId: opts.flow,
onChunk: (chunk) => {
const c = chunk as { answer?: string };
if (c.answer) process.stdout.write(c.answer);
},
});
const r = resp as { answer?: string };
if (r.answer) console.log(r.answer);
} finally {
await socket.close();
}
});
}

View file

@ -0,0 +1,81 @@
/**
* Config CLI commands.
*
* Python reference: trustgraph-cli/trustgraph/cli/show_config.py etc.
*/
import type { Command } from "commander";
import { createSocket, getOpts } from "./util.js";
export function registerConfigCommands(program: Command): void {
const config = program
.command("config")
.description("Configuration management");
config
.command("show")
.description("Show current configuration")
.action(async (_opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("config", { operation: "config" });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
config
.command("get")
.description("Get a configuration value")
.argument("<key>", "Config key")
.action(async (key: string, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("config", { operation: "get", keys: [key] });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
config
.command("set")
.description("Set a configuration value")
.argument("<key>", "Config key")
.argument("<value>", "Config value (JSON)")
.action(async (key: string, value: string, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const parsed = JSON.parse(value);
const resp = await socket.request("config", {
operation: "put",
values: { [key]: parsed },
});
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
config
.command("list")
.description("List configuration keys")
.action(async (_opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("config", { operation: "list" });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
}

View file

@ -0,0 +1,80 @@
/**
* Flow management CLI commands.
*
* Python reference: trustgraph-cli/trustgraph/cli/start_flow.py, stop_flow.py, etc.
*/
import type { Command } from "commander";
import { createSocket, getOpts } from "./util.js";
export function registerFlowCommands(program: Command): void {
const flow = program
.command("flow")
.description("Flow management");
flow
.command("list")
.description("List active flows")
.action(async (_opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("flow", { operation: "list" });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
flow
.command("start")
.description("Start a flow")
.argument("<name>", "Flow name")
.action(async (name: string, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("flow", { operation: "start", name });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
flow
.command("stop")
.description("Stop a flow")
.argument("<name>", "Flow name")
.action(async (name: string, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("flow", { operation: "stop", name });
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
flow
.command("status")
.description("Show flow status")
.argument("[name]", "Flow name (all if omitted)")
.action(async (name: string | undefined, _opts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request("flow", {
operation: "status",
...(name ? { name } : {}),
});
console.log(JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
}

View file

@ -0,0 +1,58 @@
/**
* Graph RAG CLI commands.
*
* Python reference: trustgraph-cli/trustgraph/cli/invoke_graph_rag.py
*/
import type { Command } from "commander";
import { createSocket, getOpts } from "./util.js";
export function registerGraphRagCommands(program: Command): void {
program
.command("graph-rag")
.description("Query the knowledge graph using RAG")
.argument("<query>", "Natural language query")
.option("--entity-limit <n>", "Max entities", "50")
.option("--triple-limit <n>", "Max triples per entity", "30")
.action(async (query: string, cmdOpts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request(
"graph-rag",
{
query,
entity_limit: parseInt(cmdOpts.entityLimit, 10),
triple_limit: parseInt(cmdOpts.tripleLimit, 10),
},
{ flowId: opts.flow },
) as { response?: string };
console.log(resp.response ?? JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
program
.command("document-rag")
.description("Query documents using RAG")
.argument("<query>", "Natural language query")
.action(async (query: string, _cmdOpts, cmd) => {
const opts = getOpts(cmd);
const socket = await createSocket(opts);
try {
const resp = await socket.request(
"document-rag",
{ query },
{ flowId: opts.flow },
) as { response?: string };
console.log(resp.response ?? JSON.stringify(resp, null, 2));
} finally {
await socket.close();
}
});
}

View file

@ -0,0 +1,28 @@
/**
* Shared CLI utilities.
*/
import type { Command } from "commander";
import { SocketManager } from "@trustgraph/mcp";
export interface CliOpts {
gateway: string;
token?: string;
flow: string;
}
export function getOpts(cmd: Command): CliOpts {
// Walk up to root command to get global options
let root = cmd;
while (root.parent) root = root.parent;
return root.opts() as CliOpts;
}
export async function createSocket(opts: CliOpts): Promise<SocketManager> {
const socket = new SocketManager({
gatewayUrl: opts.gateway,
token: opts.token,
});
await socket.connect();
return socket;
}

View file

@ -0,0 +1,33 @@
#!/usr/bin/env node
/**
* Unified TrustGraph CLI.
*
* Replaces the 60+ individual Python CLI scripts with a single
* `tg` command using subcommands.
*
* Python reference: trustgraph-cli/trustgraph/cli/
*/
import { Command } from "commander";
import { registerAgentCommands } from "./commands/agent.js";
import { registerGraphRagCommands } from "./commands/graph-rag.js";
import { registerConfigCommands } from "./commands/config.js";
import { registerFlowCommands } from "./commands/flow.js";
const program = new Command();
program
.name("tg")
.description("TrustGraph CLI — interact with TrustGraph services")
.version("0.1.0")
.option("-g, --gateway <url>", "Gateway WebSocket URL", "ws://localhost:8088/api/v1/socket")
.option("-t, --token <token>", "Authentication token")
.option("-f, --flow <id>", "Flow ID", "default");
registerAgentCommands(program);
registerGraphRagCommands(program);
registerConfigCommands(program);
registerFlowCommands(program);
program.parse();

View file

@ -0,0 +1,12 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"outDir": "dist",
"rootDir": "src"
},
"include": ["src"],
"references": [
{ "path": "../base" },
{ "path": "../mcp" }
]
}

View file

@ -0,0 +1,26 @@
{
"name": "@trustgraph/flow",
"version": "0.1.0",
"type": "module",
"main": "dist/index.js",
"types": "dist/index.d.ts",
"scripts": {
"build": "tsc",
"dev": "tsc --watch",
"clean": "rm -rf dist",
"test": "vitest run"
},
"dependencies": {
"@trustgraph/base": "workspace:*",
"openai": "^4.85.0",
"@anthropic-ai/sdk": "^0.39.0",
"@qdrant/js-client-rest": "^1.13.0",
"neo4j-driver": "^5.28.0",
"fastify": "^5.2.0",
"@fastify/websocket": "^11.0.0"
},
"devDependencies": {
"typescript": "^5.8.0",
"vitest": "^3.1.0"
}
}

View file

@ -0,0 +1,110 @@
/**
* Dispatcher manager routes requests to backend services via pub/sub.
*
* Python reference: trustgraph-flow/trustgraph/gateway/dispatch/manager.py
*/
import { NatsBackend, RequestResponse, type PubSubBackend } from "@trustgraph/base";
import type { GatewayConfig } from "../server.js";
export type Responder = (response: unknown, complete: boolean) => Promise<void>;
export class DispatcherManager {
private pubsub: PubSubBackend;
private requestors = new Map<string, RequestResponse<unknown, unknown>>();
constructor(private readonly config: GatewayConfig) {
this.pubsub = new NatsBackend(config.natsUrl ?? "nats://localhost:4222");
}
async start(): Promise<void> {
// Pre-create requestors for known global services
// Flow-specific requestors are created on demand
}
async stop(): Promise<void> {
for (const rr of this.requestors.values()) {
await rr.stop();
}
await this.pubsub.close();
}
private async getRequestor(
requestTopic: string,
responseTopic: string,
key: string,
): Promise<RequestResponse<unknown, unknown>> {
let rr = this.requestors.get(key);
if (!rr) {
rr = new RequestResponse({
pubsub: this.pubsub,
requestTopic,
responseTopic,
subscription: `gateway-${key}`,
});
await rr.start();
this.requestors.set(key, rr);
}
return rr;
}
async dispatchGlobalService(
kind: string,
request: Record<string, unknown>,
): Promise<unknown> {
const requestTopic = `tg.flow.${kind}-request`;
const responseTopic = `tg.flow.${kind}-response`;
const rr = await this.getRequestor(requestTopic, responseTopic, `global:${kind}`);
return rr.request(request);
}
async dispatchFlowService(
flow: string,
kind: string,
request: Record<string, unknown>,
): Promise<unknown> {
const requestTopic = `tg.flow.${kind}-request`;
const responseTopic = `tg.flow.${kind}-response`;
const rr = await this.getRequestor(requestTopic, responseTopic, `flow:${flow}:${kind}`);
return rr.request(request);
}
async dispatchGlobalServiceStreaming(
kind: string,
request: Record<string, unknown>,
responder: Responder,
): Promise<void> {
const requestTopic = `tg.flow.${kind}-request`;
const responseTopic = `tg.flow.${kind}-response`;
const rr = await this.getRequestor(requestTopic, responseTopic, `global:${kind}`);
await rr.request(request, {
recipient: async (response) => {
const res = response as Record<string, unknown>;
const complete = !!res.complete || !!res.endOfStream || !!res.endOfSession;
await responder(res, complete);
return complete;
},
});
}
async dispatchFlowServiceStreaming(
flow: string,
kind: string,
request: Record<string, unknown>,
responder: Responder,
): Promise<void> {
const requestTopic = `tg.flow.${kind}-request`;
const responseTopic = `tg.flow.${kind}-response`;
const rr = await this.getRequestor(requestTopic, responseTopic, `flow:${flow}:${kind}`);
await rr.request(request, {
recipient: async (response) => {
const res = response as Record<string, unknown>;
const complete = !!res.complete || !!res.endOfStream || !!res.endOfSession;
await responder(res, complete);
return complete;
},
});
}
}

View file

@ -0,0 +1,86 @@
/**
* WebSocket multiplexer handles concurrent requests over a single connection.
*
* Python reference: trustgraph-flow/trustgraph/gateway/dispatch/mux.py
*/
import { AsyncQueue } from "@trustgraph/base";
const MAX_OUTSTANDING = 15;
const MAX_QUEUE_SIZE = 10;
export interface MuxRequest {
id: string;
service: string;
flow?: string;
request: Record<string, unknown>;
}
export type MuxHandler = (
request: MuxRequest,
respond: (response: unknown, complete: boolean) => Promise<void>,
) => Promise<void>;
export class Mux {
private queue = new AsyncQueue<MuxRequest>();
private outstanding = 0;
private running = true;
constructor(private readonly handler: MuxHandler) {}
receive(request: MuxRequest): void {
if (this.queue.length >= MAX_QUEUE_SIZE) {
console.warn("[Mux] Queue full, dropping request:", request.id);
return;
}
this.queue.push(request);
}
async run(send: (data: string) => void): Promise<void> {
while (this.running) {
if (this.outstanding >= MAX_OUTSTANDING) {
await sleep(50);
continue;
}
try {
const request = await this.queue.pop(1000);
this.outstanding++;
// Fire and forget — error handling inside
this.processRequest(request, send).finally(() => {
this.outstanding--;
});
} catch {
// Timeout on queue pop — just loop
}
}
}
stop(): void {
this.running = false;
}
private async processRequest(
request: MuxRequest,
send: (data: string) => void,
): Promise<void> {
try {
await this.handler(request, async (response, complete) => {
send(JSON.stringify({ id: request.id, response, complete }));
});
} catch (err) {
send(
JSON.stringify({
id: request.id,
error: { type: "internal", message: String(err) },
complete: true,
}),
);
}
}
}
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}

View file

@ -0,0 +1,3 @@
export { createGateway, run, type GatewayConfig } from "./server.js";
export { DispatcherManager } from "./dispatch/manager.js";
export { Mux, type MuxRequest, type MuxHandler } from "./dispatch/mux.js";

View file

@ -0,0 +1,136 @@
/**
* API Gateway HTTP + WebSocket server.
*
* Replaces the Python aiohttp gateway with Fastify.
*
* Python reference: trustgraph-flow/trustgraph/gateway/service.py
*/
import Fastify from "fastify";
import websocketPlugin from "@fastify/websocket";
import { DispatcherManager } from "./dispatch/manager.js";
export interface GatewayConfig {
port: number;
metricsPort: number;
secret?: string;
natsUrl?: string;
}
export async function createGateway(config: GatewayConfig) {
const app = Fastify({ logger: true });
await app.register(websocketPlugin);
const dispatcher = new DispatcherManager(config);
await dispatcher.start();
// Authentication middleware
app.addHook("onRequest", async (request, reply) => {
if (request.url === "/api/v1/metrics") return;
if (request.url === "/api/v1/socket") return; // Socket auth via query param
if (config.secret) {
const auth = request.headers.authorization;
if (!auth || auth !== `Bearer ${config.secret}`) {
reply.code(401).send({ error: "Unauthorized" });
}
}
});
// REST endpoint: POST /api/v1/:kind
app.post<{ Params: { kind: string } }>("/api/v1/:kind", async (request, reply) => {
const { kind } = request.params;
const body = request.body as Record<string, unknown>;
try {
const result = await dispatcher.dispatchGlobalService(kind, body);
return result;
} catch (err) {
reply.code(500).send({ error: { type: "internal", message: String(err) } });
}
});
// REST endpoint: POST /api/v1/flow/:flow/service/:kind
app.post<{ Params: { flow: string; kind: string } }>(
"/api/v1/flow/:flow/service/:kind",
async (request, reply) => {
const { flow, kind } = request.params;
const body = request.body as Record<string, unknown>;
try {
const result = await dispatcher.dispatchFlowService(flow, kind, body);
return result;
} catch (err) {
reply.code(500).send({ error: { type: "internal", message: String(err) } });
}
},
);
// WebSocket endpoint: /api/v1/socket
app.get("/api/v1/socket", { websocket: true }, (socket, request) => {
// Auth via query param
const url = new URL(request.url, `http://${request.headers.host}`);
const token = url.searchParams.get("token");
if (config.secret && token !== config.secret) {
socket.close(4001, "Unauthorized");
return;
}
socket.on("message", async (data) => {
try {
const msg = JSON.parse(data.toString());
const { id, service, flow, request: req } = msg;
const responder = async (response: unknown, complete: boolean) => {
socket.send(JSON.stringify({ id, response, complete }));
};
if (flow) {
await dispatcher.dispatchFlowServiceStreaming(flow, service, req, responder);
} else {
await dispatcher.dispatchGlobalServiceStreaming(service, req, responder);
}
} catch (err) {
const msg = JSON.parse(data.toString());
socket.send(
JSON.stringify({
id: msg.id,
error: { type: "internal", message: String(err) },
complete: true,
}),
);
}
});
socket.on("close", () => {
// Cleanup
});
});
// Metrics endpoint
app.get("/api/v1/metrics", async () => {
const { registry } = await import("@trustgraph/base");
return registry.metrics();
});
return {
start: () => app.listen({ port: config.port, host: "0.0.0.0" }),
stop: async () => {
await app.close();
await dispatcher.stop();
},
};
}
export async function run(): Promise<void> {
const config: GatewayConfig = {
port: parseInt(process.env.GATEWAY_PORT ?? "8088", 10),
metricsPort: parseInt(process.env.METRICS_PORT ?? "8000", 10),
secret: process.env.GATEWAY_SECRET,
natsUrl: process.env.NATS_URL,
};
const gateway = await createGateway(config);
await gateway.start();
console.log(`[Gateway] Listening on port ${config.port}`);
}

View file

@ -0,0 +1,7 @@
// @trustgraph/flow — processing services
export { createGateway, type GatewayConfig } from "./gateway/index.js";
export { OpenAIProcessor } from "./model/text-completion/openai.js";
export { ClaudeProcessor } from "./model/text-completion/claude.js";
export { GraphRag, type GraphRagConfig, type GraphRagClients } from "./retrieval/graph-rag.js";
export { DocumentRag, type DocumentRagClients } from "./retrieval/document-rag.js";

View file

@ -0,0 +1,129 @@
/**
* Anthropic Claude text completion service.
*
* Python reference: trustgraph-flow/trustgraph/model/text_completion/claude/llm.py
*/
import Anthropic from "@anthropic-ai/sdk";
import { LlmService, type ProcessorConfig, type LlmResult, type LlmChunk, TooManyRequestsError } from "@trustgraph/base";
export class ClaudeProcessor extends LlmService {
private client: Anthropic;
private defaultModel: string;
private defaultTemperature: number;
private maxOutput: number;
constructor(config: ProcessorConfig & {
model?: string;
apiKey?: string;
temperature?: number;
maxOutput?: number;
}) {
super(config);
this.defaultModel = config.model ?? "claude-sonnet-4-20250514";
this.defaultTemperature = config.temperature ?? 0.0;
this.maxOutput = config.maxOutput ?? 8192;
const apiKey = config.apiKey ?? process.env.CLAUDE_KEY;
if (!apiKey) throw new Error("Claude API key not specified");
this.client = new Anthropic({ apiKey });
console.log("[Claude] LLM service initialized");
}
async generateContent(
system: string,
prompt: string,
model?: string,
temperature?: number,
): Promise<LlmResult> {
const modelName = model ?? this.defaultModel;
const temp = temperature ?? this.defaultTemperature;
try {
const response = await this.client.messages.create({
model: modelName,
max_tokens: this.maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
});
const text = response.content[0].type === "text"
? response.content[0].text
: "";
return {
text,
inToken: response.usage.input_tokens,
outToken: response.usage.output_tokens,
model: modelName,
};
} catch (err) {
if (err instanceof Anthropic.RateLimitError) {
throw new TooManyRequestsError();
}
throw err;
}
}
override supportsStreaming(): boolean {
return true;
}
async *generateContentStream(
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
const modelName = model ?? this.defaultModel;
const temp = temperature ?? this.defaultTemperature;
try {
const stream = this.client.messages.stream({
model: modelName,
max_tokens: this.maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
});
for await (const event of stream) {
if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
yield {
text: event.delta.text,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
}
const finalMessage = await stream.finalMessage();
yield {
text: "",
inToken: finalMessage.usage.input_tokens,
outToken: finalMessage.usage.output_tokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if (err instanceof Anthropic.RateLimitError) {
throw new TooManyRequestsError();
}
throw err;
}
}
}
export async function run(): Promise<void> {
await ClaudeProcessor.launch("text-completion");
}

View file

@ -0,0 +1,138 @@
/**
* OpenAI text completion service.
*
* Python reference: trustgraph-flow/trustgraph/model/text_completion/openai/llm.py
*/
import OpenAI from "openai";
import { LlmService, type ProcessorConfig, type LlmResult, type LlmChunk, TooManyRequestsError } from "@trustgraph/base";
export class OpenAIProcessor extends LlmService {
private client: OpenAI;
private defaultModel: string;
private defaultTemperature: number;
private maxOutput: number;
constructor(config: ProcessorConfig & {
model?: string;
apiKey?: string;
baseUrl?: string;
temperature?: number;
maxOutput?: number;
}) {
super(config);
this.defaultModel = config.model ?? "gpt-4o";
this.defaultTemperature = config.temperature ?? 0.0;
this.maxOutput = config.maxOutput ?? 4096;
const apiKey = config.apiKey ?? process.env.OPENAI_TOKEN;
if (!apiKey) throw new Error("OpenAI API key not specified");
this.client = new OpenAI({
apiKey,
baseURL: config.baseUrl ?? process.env.OPENAI_BASE_URL,
});
console.log("[OpenAI] LLM service initialized");
}
async generateContent(
system: string,
prompt: string,
model?: string,
temperature?: number,
): Promise<LlmResult> {
const modelName = model ?? this.defaultModel;
const temp = temperature ?? this.defaultTemperature;
try {
const resp = await this.client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: this.maxOutput,
});
return {
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
};
} catch (err) {
if (err instanceof OpenAI.RateLimitError) {
throw new TooManyRequestsError();
}
throw err;
}
}
override supportsStreaming(): boolean {
return true;
}
async *generateContentStream(
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
const modelName = model ?? this.defaultModel;
const temp = temperature ?? this.defaultTemperature;
try {
const stream = await this.client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: this.maxOutput,
stream: true,
stream_options: { include_usage: true },
});
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
if (chunk.choices?.[0]?.delta?.content) {
yield {
text: chunk.choices[0].delta.content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
if (chunk.usage) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
}
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if (err instanceof OpenAI.RateLimitError) {
throw new TooManyRequestsError();
}
throw err;
}
}
}
export async function run(): Promise<void> {
await OpenAIProcessor.launch("text-completion");
}

View file

@ -0,0 +1,66 @@
/**
* Document RAG retrieval pipeline.
*
* Simpler than Graph RAG embeds the query, finds similar document chunks,
* and synthesizes an answer from the chunk content.
*
* Python reference: trustgraph-flow/trustgraph/retrieval/document_rag/
*/
import type {
RequestResponse,
TextCompletionRequest,
TextCompletionResponse,
EmbeddingsRequest,
EmbeddingsResponse,
PromptRequest,
PromptResponse,
} from "@trustgraph/base";
export interface DocumentRagClients {
llm: RequestResponse<TextCompletionRequest, TextCompletionResponse>;
embeddings: RequestResponse<EmbeddingsRequest, EmbeddingsResponse>;
docEmbeddings: RequestResponse<unknown, unknown>; // Doc embedding query
prompt: RequestResponse<PromptRequest, PromptResponse>;
}
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
export class DocumentRag {
constructor(private readonly clients: DocumentRagClients) {}
async query(
queryText: string,
options?: {
collection?: string;
streaming?: boolean;
chunkCallback?: ChunkCallback;
},
): Promise<string> {
// Step 1: Embed the query
const embResp = await this.clients.embeddings.request({ text: [queryText] });
const vectors = (embResp as EmbeddingsResponse).vectors;
// Step 2: Find similar document chunks
const docResp = await this.clients.docEmbeddings.request({ vectors, limit: 10 });
const chunks = docResp as { chunks: Array<{ content: string; document: string }> };
// Step 3: Build context from chunks
const context = (chunks.chunks ?? [])
.map((c) => c.content)
.join("\n\n---\n\n");
// Step 4: Synthesize answer
const promptResp = await this.clients.prompt.request({
name: "document-rag-synthesize",
variables: { query: queryText, context },
});
const resp = await this.clients.llm.request({
system: (promptResp as PromptResponse).system,
prompt: (promptResp as PromptResponse).prompt,
});
return (resp as TextCompletionResponse).response;
}
}

View file

@ -0,0 +1,207 @@
/**
* Graph RAG retrieval pipeline.
*
* This is the core RAG pipeline that:
* 1. Extracts concepts from the query
* 2. Embeds concepts to find matching entities
* 3. Traverses the knowledge graph from those entities
* 4. Scores and filters edges
* 5. Synthesizes an answer with the selected context
*
* Python reference: trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
*/
import type {
RequestResponse,
TextCompletionRequest,
TextCompletionResponse,
EmbeddingsRequest,
EmbeddingsResponse,
GraphEmbeddingsRequest,
GraphEmbeddingsResponse,
TriplesQueryRequest,
TriplesQueryResponse,
PromptRequest,
PromptResponse,
Term,
Triple,
} from "@trustgraph/base";
export interface GraphRagConfig {
entityLimit?: number;
tripleLimit?: number;
maxSubgraphSize?: number;
maxPathLength?: number;
edgeScoreLimit?: number;
edgeLimit?: number;
}
export interface GraphRagClients {
llm: RequestResponse<TextCompletionRequest, TextCompletionResponse>;
embeddings: RequestResponse<EmbeddingsRequest, EmbeddingsResponse>;
graphEmbeddings: RequestResponse<GraphEmbeddingsRequest, GraphEmbeddingsResponse>;
triples: RequestResponse<TriplesQueryRequest, TriplesQueryResponse>;
prompt: RequestResponse<PromptRequest, PromptResponse>;
}
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
export class GraphRag {
private config: Required<GraphRagConfig>;
constructor(
private readonly clients: GraphRagClients,
config: GraphRagConfig = {},
) {
this.config = {
entityLimit: config.entityLimit ?? 50,
tripleLimit: config.tripleLimit ?? 30,
maxSubgraphSize: config.maxSubgraphSize ?? 1000,
maxPathLength: config.maxPathLength ?? 2,
edgeScoreLimit: config.edgeScoreLimit ?? 30,
edgeLimit: config.edgeLimit ?? 25,
};
}
async query(
queryText: string,
options?: {
collection?: string;
streaming?: boolean;
chunkCallback?: ChunkCallback;
},
): Promise<string> {
// Step 1: Extract concepts from the query via prompt + LLM
const concepts = await this.extractConcepts(queryText);
// Step 2: Embed concepts concurrently
const vectors = await this.getVectors(concepts);
// Step 3: Find matching entities via graph embeddings
const entities = await this.getEntities(vectors);
// Step 4: Traverse the knowledge graph from entities
const subgraph = await this.followEdges(entities);
// Step 5: Score and filter edges via LLM
const scoredEdges = await this.scoreEdges(queryText, subgraph);
// Step 6: Synthesize answer
const answer = await this.synthesize(queryText, scoredEdges, options?.chunkCallback);
return answer;
}
private async extractConcepts(query: string): Promise<string[]> {
const promptResp = await this.clients.prompt.request({
name: "extract-concepts",
variables: { query },
});
const llmResp = await this.clients.llm.request({
system: (promptResp as PromptResponse).system,
prompt: (promptResp as PromptResponse).prompt,
});
// Parse concepts from LLM response (newline-separated)
return (llmResp as TextCompletionResponse).response
.split("\n")
.map((c) => c.trim())
.filter(Boolean);
}
private async getVectors(concepts: string[]): Promise<number[][]> {
const resp = await this.clients.embeddings.request({ text: concepts });
return (resp as EmbeddingsResponse).vectors;
}
private async getEntities(vectors: number[][]): Promise<Term[]> {
const resp = await this.clients.graphEmbeddings.request({
vectors,
limit: this.config.entityLimit,
});
return (resp as GraphEmbeddingsResponse).entities;
}
private async followEdges(entities: Term[]): Promise<Triple[]> {
// Batch triple queries for all entities
const allTriples: Triple[] = [];
const queries = entities.map((entity) =>
this.clients.triples.request({ s: entity, limit: this.config.tripleLimit }),
);
const results = await Promise.all(queries);
for (const result of results) {
allTriples.push(...(result as TriplesQueryResponse).triples);
}
// TODO: Multi-hop traversal up to maxPathLength
return allTriples.slice(0, this.config.maxSubgraphSize);
}
private async scoreEdges(query: string, triples: Triple[]): Promise<Triple[]> {
// TODO: LLM-based edge scoring and filtering
// For now, return top N edges
return triples.slice(0, this.config.edgeLimit);
}
private async synthesize(
query: string,
edges: Triple[],
chunkCallback?: ChunkCallback,
): Promise<string> {
// Format edges as context
const context = edges
.map((t) => `${termToString(t.s)} -> ${termToString(t.p)} -> ${termToString(t.o)}`)
.join("\n");
const promptResp = await this.clients.prompt.request({
name: "graph-rag-synthesize",
variables: { query, context },
});
if (chunkCallback) {
// Streaming response
let fullText = "";
await this.clients.llm.request(
{
system: (promptResp as PromptResponse).system,
prompt: (promptResp as PromptResponse).prompt,
streaming: true,
},
{
recipient: async (resp) => {
const r = resp as TextCompletionResponse;
if (r.response) {
fullText += r.response;
await chunkCallback(r.response, !!r.endOfStream);
}
return !!r.endOfStream;
},
},
);
return fullText;
}
const resp = await this.clients.llm.request({
system: (promptResp as PromptResponse).system,
prompt: (promptResp as PromptResponse).prompt,
});
return (resp as TextCompletionResponse).response;
}
}
function termToString(term: Term): string {
switch (term.type) {
case "IRI":
return term.iri;
case "LITERAL":
return term.value;
case "BLANK":
return `_:${term.id}`;
case "TRIPLE":
return `(${termToString(term.triple.s)} ${termToString(term.triple.p)} ${termToString(term.triple.o)})`;
}
}

View file

@ -0,0 +1,11 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"outDir": "dist",
"rootDir": "src"
},
"include": ["src"],
"references": [
{ "path": "../base" }
]
}

View file

@ -0,0 +1,23 @@
{
"name": "@trustgraph/mcp",
"version": "0.1.0",
"type": "module",
"main": "dist/index.js",
"types": "dist/index.d.ts",
"scripts": {
"build": "tsc",
"dev": "tsc --watch",
"clean": "rm -rf dist",
"test": "vitest run"
},
"dependencies": {
"@trustgraph/base": "workspace:*",
"@modelcontextprotocol/sdk": "^1.8.0",
"ws": "^8.18.0"
},
"devDependencies": {
"@types/ws": "^8.5.0",
"typescript": "^5.8.0",
"vitest": "^3.1.0"
}
}

View file

@ -0,0 +1,2 @@
export { createMcpServer, run } from "./server.js";
export { SocketManager, type SocketManagerConfig } from "./socket-manager.js";

View file

@ -0,0 +1,174 @@
/**
* TrustGraph MCP server.
*
* Exposes TrustGraph capabilities as MCP tools for AI assistants.
* Communicates with the TrustGraph gateway via WebSocket.
*
* Python reference: trustgraph-mcp/trustgraph/mcp_server/mcp.py
*/
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";
export function createMcpServer(config: {
gatewayUrl: string;
token?: string;
flowId?: string;
}) {
const server = new McpServer({
name: "trustgraph",
version: "0.1.0",
});
const socket = new SocketManager({
gatewayUrl: config.gatewayUrl,
token: config.token,
});
const flowId = config.flowId ?? "default";
// --- Text Completion ---
server.tool(
"text_completion",
"Run a text completion using the configured LLM",
{
system: z.string().describe("System prompt"),
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) }] };
},
);
// --- Graph RAG ---
server.tool(
"graph_rag",
"Query the knowledge graph using RAG",
{
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"),
},
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) }] };
},
);
// --- Agent ---
server.tool(
"agent",
"Ask the TrustGraph agent a question",
{
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) }] };
},
);
// --- Embeddings ---
server.tool(
"embeddings",
"Generate text embeddings",
{
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) }] };
},
);
// --- Triples Query ---
server.tool(
"triples_query",
"Query the knowledge graph for triples matching a pattern",
{
s: z.string().optional().describe("Subject IRI"),
p: z.string().optional().describe("Predicate IRI"),
o: z.string().optional().describe("Object IRI or literal"),
limit: z.number().optional().describe("Max results"),
},
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) }] };
},
);
// --- Graph Embeddings Query ---
server.tool(
"graph_embeddings_query",
"Find entities similar to a text query using vector embeddings",
{
query: z.string().describe("Text to find similar entities for"),
limit: z.number().optional().describe("Max results"),
},
async ({ query, limit }) => {
// 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>;
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
},
);
// --- Config ---
server.tool(
"get_config",
"Get configuration values",
{
keys: z.array(z.string()).describe("Config keys to retrieve"),
},
async ({ keys }) => {
const resp = await socket.request("config", { operation: "get", keys }) as Record<string, unknown>;
return { content: [{ type: "text" as const, text: JSON.stringify(resp, null, 2) }] };
},
);
server.tool(
"put_config",
"Set configuration values",
{
values: z.record(z.unknown()).describe("Key-value pairs to set"),
},
async ({ values }) => {
const resp = await socket.request("config", { operation: "put", values }) as Record<string, unknown>;
return { content: [{ type: "text" as const, text: JSON.stringify(resp) }] };
},
);
return { server, socket };
}
export async function run(): Promise<void> {
const { server, socket } = createMcpServer({
gatewayUrl: process.env.GATEWAY_URL ?? "ws://localhost:8088/api/v1/socket",
token: process.env.GATEWAY_SECRET,
flowId: process.env.FLOW_ID ?? "default",
});
const transport = new StdioServerTransport();
await server.connect(transport);
process.on("SIGINT", async () => {
await socket.close();
process.exit(0);
});
}

View file

@ -0,0 +1,147 @@
/**
* 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;
}
}
}

View file

@ -0,0 +1,11 @@
{
"extends": "../../tsconfig.base.json",
"compilerOptions": {
"outDir": "dist",
"rootDir": "src"
},
"include": ["src"],
"references": [
{ "path": "../base" }
]
}

2
ts/pnpm-workspace.yaml Normal file
View file

@ -0,0 +1,2 @@
packages:
- "packages/*"

18
ts/tsconfig.base.json Normal file
View file

@ -0,0 +1,18 @@
{
"compilerOptions": {
"target": "ES2022",
"module": "NodeNext",
"moduleResolution": "NodeNext",
"lib": ["ES2022"],
"strict": true,
"esModuleInterop": true,
"skipLibCheck": true,
"forceConsistentCasingInFileNames": true,
"resolveJsonModule": true,
"declaration": true,
"declarationMap": true,
"sourceMap": true,
"outDir": "dist",
"rootDir": "src"
}
}

12
ts/tsconfig.json Normal file
View file

@ -0,0 +1,12 @@
{
"extends": "./tsconfig.base.json",
"compilerOptions": {
"noEmit": true
},
"references": [
{ "path": "packages/base" },
{ "path": "packages/flow" },
{ "path": "packages/cli" },
{ "path": "packages/mcp" }
]
}

22
ts/turbo.json Normal file
View file

@ -0,0 +1,22 @@
{
"$schema": "https://turbo.build/schema.json",
"tasks": {
"build": {
"dependsOn": ["^build"],
"outputs": ["dist/**"]
},
"dev": {
"cache": false,
"persistent": true
},
"lint": {
"dependsOn": ["^build"]
},
"test": {
"dependsOn": ["build"]
},
"clean": {
"cache": false
}
}
}