Advance TS port Effect workbench

This commit is contained in:
elpresidank 2026-06-01 16:22:25 -05:00
parent 92dae8c374
commit 3515106670
116 changed files with 12286 additions and 9584 deletions

View file

@ -0,0 +1,131 @@
import { describe, expect, it } from "@effect/vitest";
import { Effect } from "effect";
import type {
DocumentEmbeddingsRequest,
DocumentEmbeddingsResponse,
EmbeddingsRequest,
EmbeddingsResponse,
FlowRequestor,
GraphEmbeddingsRequest,
GraphEmbeddingsResponse,
PromptRequest,
PromptResponse,
TextCompletionRequest,
TextCompletionResponse,
TriplesQueryRequest,
TriplesQueryResponse,
} from "@trustgraph/base";
import { makeDocumentRagEngine, type DocumentRagClients } from "../retrieval/document-rag.js";
import { makeGraphRagEngine, type GraphRagClients } from "../retrieval/graph-rag.js";
const requestor = <TReq, TRes>(
handler: (request: TReq) => TRes | Promise<TRes>,
): FlowRequestor<TReq, TRes> => ({
request: async (request) => handler(request),
stop: async () => undefined,
});
describe("RAG engines", () => {
it.effect(
"runs Graph RAG without per-request service objects",
Effect.fnUntraced(function* () {
const prompts: Array<PromptRequest> = [];
const triplesRequests: Array<TriplesQueryRequest> = [];
let synthesisContext = "";
const clients: GraphRagClients = {
prompt: requestor<PromptRequest, PromptResponse>((request) => {
prompts.push(request);
if (request.name === "extract-concepts") {
return { system: "extract-system", prompt: "extract-prompt" };
}
synthesisContext = String(request.variables?.context ?? "");
return { system: "synth-system", prompt: "synth-prompt" };
}),
llm: requestor<TextCompletionRequest, TextCompletionResponse>((request) => {
if (request.prompt === "extract-prompt") {
return { response: "alpha\nbeta", endOfStream: true };
}
return { response: `answer:${request.prompt}`, endOfStream: true };
}),
embeddings: requestor<EmbeddingsRequest, EmbeddingsResponse>((request) => {
expect(request.text).toEqual(["alpha", "beta"]);
return { vectors: [[1], [2]] };
}),
graphEmbeddings: requestor<GraphEmbeddingsRequest, GraphEmbeddingsResponse>((request) => {
expect(request.collection).toBe("project");
return {
entities: [{ type: "IRI", iri: "https://example.test/entity/a" }],
};
}),
triples: requestor<TriplesQueryRequest, TriplesQueryResponse>((request) => {
triplesRequests.push(request);
return {
triples: [
{
s: { type: "IRI", iri: "https://example.test/entity/a" },
p: { type: "IRI", iri: "https://example.test/relation" },
o: { type: "LITERAL", value: "related value" },
},
],
};
}),
};
const engine = makeGraphRagEngine();
const result = yield* engine.query(
clients,
"who is related?",
{ collection: "project" },
{ maxPathLength: 1 },
);
expect(result.answer).toBe("answer:synth-prompt");
expect(result.subgraph).toHaveLength(1);
expect(prompts.map((prompt) => prompt.name)).toEqual([
"extract-concepts",
"graph-rag-synthesize",
]);
expect(triplesRequests).toHaveLength(1);
expect(synthesisContext).toContain("https://example.test/entity/a");
expect(synthesisContext).toContain("related value");
}),
);
it.effect(
"builds Document RAG synthesis context from returned chunks",
Effect.fnUntraced(function* () {
let synthesisContext = "";
const clients: DocumentRagClients = {
embeddings: requestor<EmbeddingsRequest, EmbeddingsResponse>((request) => {
expect(request.text).toEqual(["explain docs"]);
return { vectors: [[0.1, 0.2]] };
}),
docEmbeddings: requestor<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>((request) => {
expect(request.collection).toBe("docs");
return {
chunks: [
{ chunkId: "1", score: 0.9, content: "first chunk" },
{ chunkId: "2", score: 0.8, content: "" },
{ chunkId: "3", score: 0.7, content: "second chunk" },
],
};
}),
prompt: requestor<PromptRequest, PromptResponse>((request) => {
synthesisContext = String(request.variables?.context ?? "");
return { system: "doc-system", prompt: "doc-prompt" };
}),
llm: requestor<TextCompletionRequest, TextCompletionResponse>((request) => ({
response: `doc-answer:${request.prompt}`,
endOfStream: true,
})),
};
const engine = makeDocumentRagEngine();
const response = yield* engine.query(clients, "explain docs", { collection: "docs" });
expect(response).toBe("doc-answer:doc-prompt");
expect(synthesisContext).toBe("first chunk\n\n---\n\nsecond chunk");
}),
);
});

View file

@ -1 +1 @@
export { McpToolService } from "./service.js";
export { McpToolService, run } from "./service.js";

View file

@ -17,152 +17,312 @@ import {
FlowProcessor,
ConsumerSpec,
ProducerSpec,
makeFlowProcessorProgram,
errorMessage,
type ProcessorConfig,
type FlowContext,
type ToolRequest,
type ToolResponse,
type EffectConfigHandler,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Context, Effect, Layer, Ref } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
interface McpServiceConfig {
url: string;
"remote-name"?: string;
"auth-token"?: string;
const McpServiceConfig = S.Struct({
url: S.String,
"remote-name": S.optionalKey(S.String),
"auth-token": S.optionalKey(S.String),
});
type McpServiceConfig = typeof McpServiceConfig.Type;
const decodeRawMcpConfig = S.decodeUnknownOption(S.Record(S.String, S.String));
const decodeMcpServiceConfig = S.decodeUnknownOption(McpServiceConfig.pipe(S.fromJsonString));
const decodeToolParameters = S.decodeUnknownOption(S.Record(S.String, S.Unknown).pipe(S.fromJsonString));
const encodeJson = S.encodeUnknownOption(S.UnknownFromJsonString);
export class McpToolError extends S.TaggedErrorClass<McpToolError>()(
"McpToolError",
{
message: S.String,
operation: S.String,
tool: S.optionalKey(S.String),
},
) {}
export interface McpToolRuntimeService {
readonly configure: (
config: Record<string, unknown>,
version: number,
) => Effect.Effect<void>;
readonly invokeTool: (
name: string,
parameters: Record<string, unknown>,
) => Effect.Effect<string | unknown, McpToolError>;
}
export class McpToolService extends FlowProcessor {
private mcpServices: Record<string, McpServiceConfig> = {};
export class McpToolRuntime extends Context.Service<
McpToolRuntime,
McpToolRuntimeService
>()("@trustgraph/flow/agent/mcp-tool/service/McpToolRuntime") {}
const mcpToolError = (
operation: string,
cause: unknown,
tool?: string,
): McpToolError =>
new McpToolError({
operation,
message: errorMessage(cause),
...(tool === undefined ? {} : { tool }),
});
const closeTransport = (
transport: StreamableHTTPClientTransport,
tool: string,
) =>
Effect.tryPromise({
try: () => transport.close(),
catch: (cause) => mcpToolError("close-transport", cause, tool),
}).pipe(
Effect.catch((error) =>
Effect.logError("[McpToolService] Failed to close MCP transport", {
error: error.message,
tool: error.tool ?? tool,
}),
),
);
const loadMcpServices = Effect.fn("McpToolRuntime.loadMcpServices")(function* (
config: Record<string, unknown>,
version: number,
) {
yield* Effect.log(`[McpToolService] Got config version ${version}`);
if (!("mcp" in config) || typeof config.mcp !== "object" || config.mcp === null) {
return {};
}
const rawConfig = decodeRawMcpConfig(config.mcp);
if (O.isNone(rawConfig)) {
yield* Effect.logError("[McpToolService] MCP config must be an object of JSON strings");
return {};
}
const services: Record<string, McpServiceConfig> = {};
for (const [name, value] of Object.entries(rawConfig.value)) {
const decoded = decodeMcpServiceConfig(value);
if (O.isNone(decoded)) {
yield* Effect.logError(`[McpToolService] Failed to parse MCP config for ${name}`);
continue;
}
services[name] = decoded.value;
yield* Effect.log(`[McpToolService] Registered MCP service: ${name}`);
}
yield* Effect.log(
`[McpToolService] ${Object.keys(services).length} MCP services configured`,
);
return services;
});
const invokeConfiguredTool = Effect.fn("McpToolRuntime.invokeTool")(function* (
services: Record<string, McpServiceConfig>,
name: string,
parameters: Record<string, unknown>,
) {
if (!(name in services)) {
return yield* mcpToolError("lookup-service", `MCP service "${name}" not known`, name);
}
const svcConfig = services[name];
if (svcConfig.url.length === 0) {
return yield* mcpToolError("validate-service", `MCP service "${name}" URL not defined`, name);
}
const remoteName = svcConfig["remote-name"] ?? name;
const headers: Record<string, string> = {};
if (svcConfig["auth-token"] !== undefined && svcConfig["auth-token"].length > 0) {
headers.Authorization = `Bearer ${svcConfig["auth-token"]}`;
}
yield* Effect.log(`[McpToolService] Invoking ${remoteName} at ${svcConfig.url}`);
const url = yield* Effect.try({
try: () => new URL(svcConfig.url),
catch: (cause) => mcpToolError("validate-url", cause, name),
});
const transport = new StreamableHTTPClientTransport(
url,
{ requestInit: { headers } },
);
const client = new Client({ name: "trustgraph-mcp-client", version: "1.0.0" });
const result = yield* Effect.acquireUseRelease(
Effect.tryPromise({
try: async () => {
await client.connect(transport as unknown as Parameters<Client["connect"]>[0]);
return client;
},
catch: (cause) => mcpToolError("connect", cause, name),
}),
(connectedClient) =>
Effect.tryPromise({
try: () =>
connectedClient.callTool({
name: remoteName,
arguments: parameters,
}),
catch: (cause) => mcpToolError("call-tool", cause, name),
}),
() => closeTransport(transport, name),
);
if (result.structuredContent !== undefined && result.structuredContent !== null) {
return result.structuredContent;
}
if (result.content !== undefined && Array.isArray(result.content)) {
return result.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("");
}
return "No content";
});
export const makeMcpToolRuntime = Effect.gen(function* () {
const servicesRef = yield* Ref.make<Record<string, McpServiceConfig>>({});
return McpToolRuntime.of({
configure: Effect.fn("McpToolRuntime.configure")(function* (config, version) {
const services = yield* loadMcpServices(config, version);
yield* Ref.set(servicesRef, services);
}),
invokeTool: Effect.fn("McpToolRuntime.invokeToolFromRef")(function* (name, parameters) {
const services = yield* Ref.get(servicesRef);
return yield* invokeConfiguredTool(services, name, parameters);
}),
});
});
export const McpToolRuntimeLive = Layer.effect(McpToolRuntime, makeMcpToolRuntime);
const onMcpConfig = Effect.fn("McpToolService.onConfig")(function* (
config: Record<string, unknown>,
version: number,
) {
const runtime = yield* McpToolRuntime;
yield* runtime.configure(config, version);
});
type McpToolHandlerError =
| FlowResourceNotFoundError
| MessagingDeliveryError;
const parametersFromJson = (
name: string,
parameters: string,
): Effect.Effect<Record<string, unknown>, McpToolError> => {
if (parameters.length === 0) return Effect.succeed({});
const decoded = decodeToolParameters(parameters);
if (O.isNone(decoded)) {
return Effect.fail(mcpToolError("decode-parameters", "Tool parameters must be a JSON object", name));
}
return Effect.succeed(decoded.value);
};
const onMcpToolRequest = Effect.fn("McpToolService.onRequest")(function* (
msg: ToolRequest,
properties: Record<string, string>,
flowCtx: FlowContext<McpToolRuntime>,
): Effect.fn.Return<void, McpToolHandlerError, McpToolRuntime> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = yield* flowCtx.flow.producerEffect<ToolResponse>("mcp-tool-response");
const runtime = yield* McpToolRuntime;
const result = yield* parametersFromJson(msg.name, msg.parameters).pipe(
Effect.flatMap((parameters) => runtime.invokeTool(msg.name, parameters)),
Effect.catch((error) =>
Effect.logError(`[McpToolService] Error invoking tool ${msg.name}`, {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
responseProducer.send(requestId, {
error: { type: "tool-error", message: error.message },
}),
),
Effect.as(undefined),
),
),
);
if (result === undefined) return;
if (typeof result === "string") {
yield* responseProducer.send(requestId, { text: result });
return;
}
const encoded = encodeJson(result);
yield* responseProducer.send(requestId, {
object: O.isSome(encoded) ? encoded.value : String(result),
});
});
export const makeMcpToolSpecs = (): ReadonlyArray<Spec<McpToolRuntime>> => [
new ConsumerSpec<ToolRequest, McpToolHandlerError, McpToolRuntime>(
"mcp-tool-request",
onMcpToolRequest,
),
new ProducerSpec<ToolResponse>("mcp-tool-response"),
];
export const makeMcpToolConfigHandlers = (): ReadonlyArray<
EffectConfigHandler<never, McpToolRuntime>
> => [onMcpConfig];
export class McpToolService extends FlowProcessor<McpToolRuntime> {
private readonly runtime = Effect.runSync(makeMcpToolRuntime);
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
ConsumerSpec.fromPromise<ToolRequest>("mcp-tool-request", this.onRequest.bind(this)),
);
this.registerSpecification(new ProducerSpec<ToolResponse>("mcp-tool-response"));
this.registerConfigHandler(this.onMcpConfig.bind(this));
}
private async onMcpConfig(
config: Record<string, unknown>,
version: number,
): Promise<void> {
console.log(`[McpToolService] Got config version ${version}`);
if (!("mcp" in config) || typeof config.mcp !== "object" || config.mcp === null) {
this.mcpServices = {};
return;
for (const spec of makeMcpToolSpecs()) {
this.registerSpecification(spec);
}
const mcpConfig = config.mcp as Record<string, string>;
this.mcpServices = {};
for (const [name, value] of Object.entries(mcpConfig)) {
try {
this.mcpServices[name] = JSON.parse(value) as McpServiceConfig;
console.log(`[McpToolService] Registered MCP service: ${name}`);
} catch (err) {
console.error(`[McpToolService] Failed to parse MCP config for ${name}:`, err);
}
}
console.log(
`[McpToolService] ${Object.keys(this.mcpServices).length} MCP services configured`,
this.registerConfigHandler((config, version) =>
Effect.runPromise(onMcpConfig(config, version).pipe(
Effect.provideService(McpToolRuntime, this.runtime),
)),
);
}
private async onRequest(
msg: ToolRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = flowCtx.flow.producer<ToolResponse>("mcp-tool-response");
try {
const result = await this.invokeTool(
msg.name,
msg.parameters !== undefined && msg.parameters.length > 0
? JSON.parse(msg.parameters) as Record<string, unknown>
: {},
);
if (typeof result === "string") {
await responseProducer.send(requestId, { text: result });
} else {
await responseProducer.send(requestId, { object: JSON.stringify(result) });
}
} catch (err) {
console.error(`[McpToolService] Error invoking tool ${msg.name}:`, err);
const message = err instanceof Error ? err.message : String(err);
await responseProducer.send(requestId, {
error: { type: "tool-error", message },
});
}
}
private async invokeTool(
name: string,
parameters: Record<string, unknown>,
): Promise<string | unknown> {
if (!(name in this.mcpServices)) {
throw new Error(`MCP service "${name}" not known`);
}
const svcConfig = this.mcpServices[name];
if (svcConfig.url.length === 0) {
throw new Error(`MCP service "${name}" URL not defined`);
}
const remoteName = svcConfig["remote-name"] ?? name;
// Build headers with optional bearer token
const headers: Record<string, string> = {};
if (svcConfig["auth-token"] !== undefined && svcConfig["auth-token"].length > 0) {
headers["Authorization"] = `Bearer ${svcConfig["auth-token"]}`;
}
console.log(`[McpToolService] Invoking ${remoteName} at ${svcConfig.url}`);
// Connect to streamable HTTP MCP server
const transport = new StreamableHTTPClientTransport(
new URL(svcConfig.url),
{ requestInit: { headers } },
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(McpToolRuntime, this.runtime),
);
const client = new Client({ name: "trustgraph-mcp-client", version: "1.0.0" });
try {
await client.connect(transport as unknown as Parameters<Client["connect"]>[0]);
const result = await client.callTool({
name: remoteName,
arguments: parameters,
});
// Extract response — prefer structured content, fall back to text
if (result.structuredContent !== undefined && result.structuredContent !== null) {
return result.structuredContent;
}
if (result.content !== undefined && Array.isArray(result.content)) {
return result.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("");
}
return "No content";
} finally {
await transport.close();
}
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, McpToolRuntime>({
id: "mcp-tool",
make: (config) => new McpToolService(config),
specs: () => makeMcpToolSpecs(),
configHandlers: () => makeMcpToolConfigHandlers(),
layer: () => McpToolRuntimeLive,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
}

View file

@ -21,6 +21,8 @@ import {
ConsumerSpec,
ProducerSpec,
RequestResponseSpec,
makeFlowProcessorProgram,
errorMessage,
type ProcessorConfig,
type FlowContext,
type AgentRequest,
@ -35,8 +37,18 @@ import {
type TriplesQueryResponse,
type ToolRequest,
type ToolResponse,
type EffectConfigHandler,
type EffectRequestOptions,
type EffectRequestResponse,
type FlowRequestOptions,
type FlowRequestor,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Context, Effect, Layer, Ref } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
import {
createKnowledgeQueryTool,
@ -51,398 +63,490 @@ import type { AgentTool, ToolArg } from "./types.js";
const MAX_ITERATIONS = 10;
export class AgentService extends FlowProcessor {
/** Config-driven tools; null means "use hardcoded fallback". */
private configuredTools: AgentTool[] | null = null;
class AgentToolExecutionError extends S.TaggedErrorClass<AgentToolExecutionError>()(
"AgentToolExecutionError",
{
message: S.String,
},
) {}
const UnknownRecord = S.Record(S.String, S.Unknown);
const ToolArgumentConfig = S.StructWithRest(
S.Struct({
name: S.optionalKey(S.String),
type: S.optionalKey(S.String),
description: S.optionalKey(S.String),
}),
[UnknownRecord],
);
const ToolConfigEntry = S.StructWithRest(
S.Struct({
type: S.optionalKey(S.String),
name: S.optionalKey(S.String),
description: S.optionalKey(S.String),
arguments: ToolArgumentConfig.pipe(S.Array, S.optionalKey),
}),
[UnknownRecord],
);
type ToolConfigEntry = typeof ToolConfigEntry.Type;
const decodeRawToolConfig = S.decodeUnknownOption(S.Record(S.String, S.String));
const decodeToolConfigEntry = S.decodeUnknownOption(ToolConfigEntry.pipe(S.fromJsonString));
export interface AgentRuntimeService {
readonly configureTools: (
config: Record<string, unknown>,
version: number,
) => Effect.Effect<void>;
readonly getConfiguredTools: Effect.Effect<ReadonlyArray<AgentTool> | null>;
}
export class AgentRuntime extends Context.Service<AgentRuntime, AgentRuntimeService>()(
"@trustgraph/flow/agent/react/service/AgentRuntime",
) {}
const toEffectRequestOptions = <TRes>(
options: FlowRequestOptions<TRes> | undefined,
): EffectRequestOptions<TRes> | undefined => {
if (options === undefined) return undefined;
return {
...(options.timeoutMs === undefined ? {} : { timeoutMs: options.timeoutMs }),
...(options.recipient === undefined
? {}
: {
recipient: (response: TRes) =>
Effect.promise(() => options.recipient?.(response) ?? Promise.resolve(true)),
}),
};
};
const toPromiseRequestor = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
): FlowRequestor<TReq, TRes> => ({
request: (request, options) =>
Effect.runPromise(requestor.request(request, toEffectRequestOptions(options))),
stop: () => Effect.runPromise(requestor.stop),
});
const buildConfiguredTool = (
toolId: string,
data: ToolConfigEntry,
): AgentTool | null => {
const implType = data.type ?? "";
const name = data.name ?? "";
const description = data.description ?? "";
const config = { ...data } as Record<string, unknown>;
if (name.length === 0) {
console.warn(`[AgentService] Skipping tool with no name: ${toolId}`);
return null;
}
switch (implType) {
case "knowledge-query":
return {
name,
description:
description.length > 0
? description
: "Query the knowledge graph for information about entities and their relationships.",
args: [{ name: "question", type: "string", description: "The question to ask" }],
config,
execute: async () => "",
};
case "document-query":
return {
name,
description:
description.length > 0
? description
: "Search documents for relevant information.",
args: [{ name: "question", type: "string", description: "The question to search for" }],
config,
execute: async () => "",
};
case "triples-query":
return {
name,
description:
description.length > 0
? description
: "Query for specific triples in the knowledge graph.",
args: [
{ name: "subject", type: "string", description: "Subject entity (optional)" },
{ name: "predicate", type: "string", description: "Predicate/relationship (optional)" },
{ name: "object", type: "string", description: "Object entity (optional)" },
],
config,
execute: async () => "",
};
case "mcp-tool": {
const args: ToolArg[] = (data.arguments ?? []).map((arg) => ({
name: arg.name ?? "",
type: arg.type ?? "string",
description: arg.description ?? "",
}));
return {
name,
description,
args,
config,
execute: async () => "",
};
}
default:
console.warn(`[AgentService] Unknown tool type "${implType}" for ${name}`);
return null;
}
};
const loadConfiguredTools = Effect.fn("AgentRuntime.loadConfiguredTools")(function* (
config: Record<string, unknown>,
version: number,
) {
yield* Effect.log(`[AgentService] Loading tool configuration version ${version}`);
if (!("tool" in config) || typeof config.tool !== "object" || config.tool === null) {
yield* Effect.log("[AgentService] No tool config found, using default tools");
return null;
}
const rawConfig = decodeRawToolConfig(config.tool);
if (O.isNone(rawConfig)) {
yield* Effect.logError("[AgentService] Tool config must be an object of JSON strings");
return null;
}
const tools: AgentTool[] = [];
for (const [toolId, toolValue] of Object.entries(rawConfig.value)) {
const decoded = decodeToolConfigEntry(toolValue);
if (O.isNone(decoded)) {
yield* Effect.logError(`[AgentService] Failed to parse tool config ${toolId}`);
continue;
}
const tool = buildConfiguredTool(toolId, decoded.value);
if (tool === null) continue;
tools.push(tool);
yield* Effect.log(`[AgentService] Registered tool: ${tool.name} (${tool.config?.type ?? "unknown"})`);
}
yield* Effect.log(`[AgentService] ${tools.length} tools loaded from config`);
return tools.length > 0 ? tools : null;
});
export const makeAgentRuntime = Effect.gen(function* () {
const configuredToolsRef = yield* Ref.make<ReadonlyArray<AgentTool> | null>(null);
return AgentRuntime.of({
configureTools: Effect.fn("AgentRuntime.configureTools")(function* (config, version) {
const tools = yield* loadConfiguredTools(config, version);
yield* Ref.set(configuredToolsRef, tools);
}),
getConfiguredTools: Ref.get(configuredToolsRef),
});
});
export const AgentRuntimeLive = Layer.effect(AgentRuntime, makeAgentRuntime);
const onToolsConfig = Effect.fn("AgentService.onToolsConfig")(function* (
config: Record<string, unknown>,
version: number,
) {
const runtime = yield* AgentRuntime;
yield* runtime.configureTools(config, version);
});
const wireTools = Effect.fn("AgentService.wireTools")(function* (
tools: ReadonlyArray<AgentTool>,
flowCtx: FlowContext<AgentRuntime>,
collection: string | undefined,
onExplain: (data: ExplainData) => void,
) {
const graphRag = yield* flowCtx.flow.requestorEffect<GraphRagRequest, GraphRagResponse>("graph-rag");
const docRag = yield* flowCtx.flow.requestorEffect<DocumentRagRequest, DocumentRagResponse>("doc-rag");
const triples = yield* flowCtx.flow.requestorEffect<TriplesQueryRequest, TriplesQueryResponse>("triples");
const mcpTool = yield* flowCtx.flow.requestorEffect<ToolRequest, ToolResponse>("mcp-tool");
return tools.map((tool) => {
const implType = tool.config?.type as string | undefined;
switch (implType) {
case "knowledge-query": {
const live = createKnowledgeQueryTool(
toPromiseRequestor(graphRag),
collection,
onExplain,
);
return { ...tool, execute: live.execute };
}
case "document-query": {
const live = createDocumentQueryTool(
toPromiseRequestor(docRag),
collection,
);
return { ...tool, execute: live.execute };
}
case "triples-query": {
const live = createTriplesQueryTool(
toPromiseRequestor(triples),
collection,
);
return { ...tool, execute: live.execute };
}
case "mcp-tool": {
const live = createMcpTool(
toPromiseRequestor(mcpTool),
tool.name,
tool.description,
tool.args,
);
return { ...tool, execute: live.execute };
}
default:
return tool;
}
});
});
const defaultTools = Effect.fn("AgentService.defaultTools")(function* (
flowCtx: FlowContext<AgentRuntime>,
collection: string | undefined,
onExplain: (data: ExplainData) => void,
) {
const graphRag = yield* flowCtx.flow.requestorEffect<GraphRagRequest, GraphRagResponse>("graph-rag");
const docRag = yield* flowCtx.flow.requestorEffect<DocumentRagRequest, DocumentRagResponse>("doc-rag");
const triples = yield* flowCtx.flow.requestorEffect<TriplesQueryRequest, TriplesQueryResponse>("triples");
return [
createKnowledgeQueryTool(
toPromiseRequestor(graphRag),
collection,
onExplain,
),
createDocumentQueryTool(
toPromiseRequestor(docRag),
collection,
),
createTriplesQueryTool(
toPromiseRequestor(triples),
collection,
),
];
});
const executeTool = (
tool: AgentTool,
input: string,
): Effect.Effect<string> =>
Effect.tryPromise({
try: () => tool.execute(input),
catch: (cause) => new AgentToolExecutionError({ message: errorMessage(cause) }),
}).pipe(
Effect.catch((error: AgentToolExecutionError) =>
Effect.succeed(`Error executing tool: ${error.message}`),
),
);
type AgentHandlerError =
| FlowResourceNotFoundError
| MessagingDeliveryError;
const onAgentRequest = Effect.fn("AgentService.onRequest")(function* (
msg: AgentRequest,
properties: Record<string, string>,
flowCtx: FlowContext<AgentRuntime>,
): Effect.fn.Return<void, AgentHandlerError, AgentRuntime> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = yield* flowCtx.flow.producerEffect<AgentResponse>("agent-response");
yield* Effect.gen(function* () {
const runtime = yield* AgentRuntime;
const explainEvents: ExplainData[] = [];
const onExplain = (data: ExplainData) => {
explainEvents.push(data);
};
const configuredTools = yield* runtime.getConfiguredTools;
let tools = configuredTools !== null
? yield* wireTools(configuredTools, flowCtx, msg.collection, onExplain)
: yield* defaultTools(flowCtx, msg.collection, onExplain);
tools = filterToolsByGroupAndState(tools, msg.group, msg.state);
const { system, prompt: initialPrompt } = buildReActPrompt(
tools,
msg.question,
);
const llmClient = yield* flowCtx.flow.requestorEffect<
TextCompletionRequest,
TextCompletionResponse
>("llm");
let conversation = initialPrompt;
for (let iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
yield* Effect.log(
`[AgentService] Iteration ${iteration + 1}/${MAX_ITERATIONS} for request ${requestId}`,
);
const llmResponse = yield* llmClient.request({
system,
prompt: conversation,
});
if (llmResponse.error !== undefined) {
yield* responseProducer.send(requestId, {
chunk_type: "error",
content: `LLM error: ${llmResponse.error.message}`,
end_of_dialog: true,
});
return;
}
const text = llmResponse.response;
const parsed = parseReActResponse(text);
if (parsed.thought.length > 0) {
yield* responseProducer.send(requestId, {
chunk_type: "thought",
content: parsed.thought,
end_of_message: true,
});
}
if (parsed.finalAnswer.length > 0) {
for (const explain of explainEvents) {
yield* responseProducer.send(requestId, {
chunk_type: "explain",
content: "",
explain_id: explain.explainId,
explain_triples: explain.triples,
} as AgentResponse);
}
yield* responseProducer.send(requestId, {
chunk_type: "answer",
content: parsed.finalAnswer,
end_of_message: true,
end_of_dialog: true,
});
return;
}
if (parsed.action.length > 0 && parsed.actionInput.length > 0) {
const tool = tools.find((candidate) => candidate.name === parsed.action);
const observation = tool === undefined
? `Unknown tool: ${parsed.action}. Available tools: ${tools.map((candidate) => candidate.name).join(", ")}`
: yield* executeTool(tool, parsed.actionInput);
yield* responseProducer.send(requestId, {
chunk_type: "observation",
content: observation,
end_of_message: true,
});
conversation += `\n${text}\nObservation: ${observation}\n`;
} else if (parsed.finalAnswer.length === 0) {
conversation += `\n${text}\nObservation: You must either use a tool (Action + Action Input) or provide a Final Answer.\n`;
}
}
yield* responseProducer.send(requestId, {
chunk_type: "error",
content:
"Maximum reasoning iterations reached without a final answer. " +
"The agent was unable to complete the task within the allowed steps.",
end_of_message: true,
end_of_dialog: true,
});
}).pipe(
Effect.catch((error: unknown) =>
Effect.logError(`[AgentService] Error processing request ${requestId}`, {
error: error instanceof Error ? error.message : String(error),
}).pipe(
Effect.flatMap(() =>
responseProducer.send(requestId, {
chunk_type: "error",
content: `Agent error: ${error instanceof Error ? error.message : String(error)}`,
end_of_message: true,
end_of_dialog: true,
}),
),
),
),
);
});
export const makeAgentSpecs = (): ReadonlyArray<Spec<AgentRuntime>> => [
new ConsumerSpec<AgentRequest, AgentHandlerError, AgentRuntime>(
"agent-request",
onAgentRequest,
),
new ProducerSpec<AgentResponse>("agent-response"),
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
new RequestResponseSpec<GraphRagRequest, GraphRagResponse>(
"graph-rag",
"graph-rag-request",
"graph-rag-response",
),
new RequestResponseSpec<DocumentRagRequest, DocumentRagResponse>(
"doc-rag",
"document-rag-request",
"document-rag-response",
),
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
),
new RequestResponseSpec<ToolRequest, ToolResponse>(
"mcp-tool",
"mcp-tool-request",
"mcp-tool-response",
),
];
export const makeAgentConfigHandlers = (): ReadonlyArray<
EffectConfigHandler<never, AgentRuntime>
> => [onToolsConfig];
export class AgentService extends FlowProcessor<AgentRuntime> {
private readonly runtime = Effect.runSync(makeAgentRuntime);
constructor(config: ProcessorConfig) {
super(config);
// Consumer: agent requests
this.registerSpecification(
ConsumerSpec.fromPromise<AgentRequest>("agent-request", this.onRequest.bind(this)),
);
for (const spec of makeAgentSpecs()) {
this.registerSpecification(spec);
}
// Producer: agent responses (streaming chunks)
this.registerSpecification(new ProducerSpec<AgentResponse>("agent-response"));
// Request-response clients for tool execution
this.registerSpecification(
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
this.registerConfigHandler((config, version) =>
Effect.runPromise(onToolsConfig(config, version).pipe(
Effect.provideService(AgentRuntime, this.runtime),
)),
);
this.registerSpecification(
new RequestResponseSpec<GraphRagRequest, GraphRagResponse>(
"graph-rag",
"graph-rag-request",
"graph-rag-response",
),
);
this.registerSpecification(
new RequestResponseSpec<DocumentRagRequest, DocumentRagResponse>(
"doc-rag",
"document-rag-request",
"document-rag-response",
),
);
this.registerSpecification(
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
),
);
// MCP tool invocation client
this.registerSpecification(
new RequestResponseSpec<ToolRequest, ToolResponse>(
"mcp-tool",
"mcp-tool-request",
"mcp-tool-response",
),
);
// Register for config-push to build tools dynamically
this.registerConfigHandler(this.onToolsConfig.bind(this));
console.log("[AgentService] Service initialized");
}
// ---------- Config-driven tool registration ----------
private async onToolsConfig(
config: Record<string, unknown>,
version: number,
): Promise<void> {
console.log(`[AgentService] Loading tool configuration version ${version}`);
try {
if (!("tool" in config) || typeof config.tool !== "object" || config.tool === null) {
// No tool config — keep using hardcoded fallback
this.configuredTools = null;
console.log("[AgentService] No tool config found, using default tools");
return;
}
const toolConfig = config.tool as Record<string, string>;
const tools: AgentTool[] = [];
for (const [_toolId, toolValue] of Object.entries(toolConfig)) {
try {
const data = JSON.parse(toolValue) as Record<string, unknown>;
const implType = typeof data["type"] === "string" ? data["type"] : "";
const name = typeof data["name"] === "string" ? data["name"] : "";
const description =
typeof data["description"] === "string" ? data["description"] : "";
if (name.length === 0) {
console.warn(`[AgentService] Skipping tool with no name: ${_toolId}`);
continue;
}
let tool: AgentTool | null = null;
switch (implType) {
case "knowledge-query":
// Will be wired to requestor at request time
tool = {
name,
description:
description.length > 0
? description
: "Query the knowledge graph for information about entities and their relationships.",
args: [{ name: "question", type: "string", description: "The question to ask" }],
config: data,
execute: async () => "", // placeholder — wired at request time
};
break;
case "document-query":
tool = {
name,
description:
description.length > 0
? description
: "Search documents for relevant information.",
args: [{ name: "question", type: "string", description: "The question to search for" }],
config: data,
execute: async () => "",
};
break;
case "triples-query":
tool = {
name,
description:
description.length > 0
? description
: "Query for specific triples in the knowledge graph.",
args: [
{ name: "subject", type: "string", description: "Subject entity (optional)" },
{ name: "predicate", type: "string", description: "Predicate/relationship (optional)" },
{ name: "object", type: "string", description: "Object entity (optional)" },
],
config: data,
execute: async () => "",
};
break;
case "mcp-tool": {
const configArgs = (data["arguments"] as Array<Record<string, string>>) ?? [];
const args: ToolArg[] = configArgs.map((a) => ({
name: a.name ?? "",
type: a.type ?? "string",
description: a.description ?? "",
}));
// Create a placeholder — will be wired to the MCP requestor at request time
tool = {
name,
description,
args,
config: data,
execute: async () => "", // placeholder
};
break;
}
default:
console.warn(`[AgentService] Unknown tool type "${implType}" for ${name}`);
continue;
}
if (tool !== null) {
tools.push(tool);
console.log(`[AgentService] Registered tool: ${name} (${implType})`);
}
} catch (err) {
console.error(`[AgentService] Failed to parse tool config ${_toolId}:`, err);
}
}
this.configuredTools = tools.length > 0 ? tools : null;
console.log(`[AgentService] ${tools.length} tools loaded from config`);
} catch (err) {
console.error("[AgentService] Config reload failed:", err);
}
}
/**
* Wire up tool execute functions with live requestors from the flow context.
* Config-driven tools store placeholders; this replaces them with real impls.
*/
private wireTools(
tools: AgentTool[],
flowCtx: FlowContext,
collection?: string,
onExplain?: (data: ExplainData) => void,
): AgentTool[] {
return tools.map((tool) => {
const implType = tool.config?.["type"] as string | undefined;
switch (implType) {
case "knowledge-query": {
const live = createKnowledgeQueryTool(
flowCtx.flow.requestor<GraphRagRequest, GraphRagResponse>("graph-rag"),
collection,
onExplain,
);
return { ...tool, execute: live.execute };
}
case "document-query": {
const live = createDocumentQueryTool(
flowCtx.flow.requestor<DocumentRagRequest, DocumentRagResponse>("doc-rag"),
collection,
);
return { ...tool, execute: live.execute };
}
case "triples-query": {
const live = createTriplesQueryTool(
flowCtx.flow.requestor<TriplesQueryRequest, TriplesQueryResponse>("triples"),
collection,
);
return { ...tool, execute: live.execute };
}
case "mcp-tool": {
const live = createMcpTool(
flowCtx.flow.requestor<ToolRequest, ToolResponse>("mcp-tool"),
tool.name,
tool.description,
tool.args,
);
return { ...tool, execute: live.execute };
}
default:
return tool;
}
});
}
private async onRequest(
msg: AgentRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = flowCtx.flow.producer<AgentResponse>("agent-response");
try {
// Accumulate explain data from tool calls for emission after completion
const explainEvents: ExplainData[] = [];
const onExplain = (data: ExplainData) => {
explainEvents.push(data);
};
// Build tools — config-driven or hardcoded fallback
let tools: AgentTool[];
if (this.configuredTools !== null) {
tools = this.wireTools(this.configuredTools, flowCtx, msg.collection, onExplain);
} else {
// Hardcoded fallback (backward compat)
tools = [
createKnowledgeQueryTool(
flowCtx.flow.requestor<GraphRagRequest, GraphRagResponse>("graph-rag"),
msg.collection,
onExplain,
),
createDocumentQueryTool(
flowCtx.flow.requestor<DocumentRagRequest, DocumentRagResponse>("doc-rag"),
msg.collection,
),
createTriplesQueryTool(
flowCtx.flow.requestor<TriplesQueryRequest, TriplesQueryResponse>("triples"),
msg.collection,
),
];
}
// Apply tool filtering by group and state
tools = filterToolsByGroupAndState(tools, msg.group, msg.state);
// Build the ReAct prompt
const { system, prompt: initialPrompt } = buildReActPrompt(
tools,
msg.question,
);
const llmClient = flowCtx.flow.requestor<
TextCompletionRequest,
TextCompletionResponse
>("llm");
// Conversation accumulates the full exchange for multi-turn reasoning
let conversation = initialPrompt;
for (let iteration = 0; iteration < MAX_ITERATIONS; iteration++) {
console.log(
`[AgentService] Iteration ${iteration + 1}/${MAX_ITERATIONS} for request ${requestId}`,
);
// Call LLM (non-streaming for MVP)
const llmResponse = await llmClient.request({
system,
prompt: conversation,
});
if (llmResponse.error !== undefined) {
await responseProducer.send(requestId, {
chunk_type: "error",
content: `LLM error: ${llmResponse.error.message}`,
end_of_dialog: true,
});
return;
}
const text = llmResponse.response;
// Parse the LLM response with simple line-based parsing
const parsed = parseReActResponse(text);
// Send thought chunk
if (parsed.thought.length > 0) {
await responseProducer.send(requestId, {
chunk_type: "thought",
content: parsed.thought,
end_of_message: true,
});
}
// If we got a final answer, emit explain events then send the answer
if (parsed.finalAnswer.length > 0) {
// Emit explain events collected from tool calls
for (const explain of explainEvents) {
await responseProducer.send(requestId, {
chunk_type: "explain",
content: "",
explain_id: explain.explainId,
explain_triples: explain.triples,
} as AgentResponse);
}
await responseProducer.send(requestId, {
chunk_type: "answer",
content: parsed.finalAnswer,
end_of_message: true,
end_of_dialog: true,
});
return;
}
// Execute tool if action was specified
if (parsed.action.length > 0 && parsed.actionInput.length > 0) {
const tool = tools.find((t) => t.name === parsed.action);
let observation: string;
if (tool !== undefined) {
try {
observation = await tool.execute(parsed.actionInput);
} catch (err) {
observation = `Error executing tool: ${err instanceof Error ? err.message : String(err)}`;
}
} else {
observation = `Unknown tool: ${parsed.action}. Available tools: ${tools.map((t) => t.name).join(", ")}`;
}
// Send observation chunk
await responseProducer.send(requestId, {
chunk_type: "observation",
content: observation,
end_of_message: true,
});
// Append the full exchange to conversation for the next iteration
conversation += `\n${text}\nObservation: ${observation}\n`;
} else if (parsed.finalAnswer.length === 0) {
// LLM didn't produce a valid action or final answer -- nudge it
conversation += `\n${text}\nObservation: You must either use a tool (Action + Action Input) or provide a Final Answer.\n`;
}
}
// Max iterations reached without a final answer
await responseProducer.send(requestId, {
chunk_type: "error",
content:
"Maximum reasoning iterations reached without a final answer. " +
"The agent was unable to complete the task within the allowed steps.",
end_of_message: true,
end_of_dialog: true,
});
} catch (err) {
console.error(`[AgentService] Error processing request ${requestId}:`, err);
await responseProducer.send(requestId, {
chunk_type: "error",
content: `Agent error: ${err instanceof Error ? err.message : String(err)}`,
end_of_message: true,
end_of_dialog: true,
});
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(AgentRuntime, this.runtime),
);
}
}
@ -524,11 +628,13 @@ function parseReActResponse(text: string): {
};
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, AgentRuntime>({
id: "agent",
make: (config) => new AgentService(config),
specs: () => makeAgentSpecs(),
configHandlers: () => makeAgentConfigHandlers(),
layer: () => AgentRuntimeLive,
});
export async function run(): Promise<void> {
await AgentService.launch("agent");
await Effect.runPromise(program);
}

View file

@ -21,81 +21,86 @@ import {
type TextDocument,
type Chunk,
type Triples,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import { recursiveSplit } from "./recursive-splitter.js";
const DEFAULT_CHUNK_SIZE = 2000;
const DEFAULT_CHUNK_OVERLAP = 100;
const onChunkMessage = Effect.fn("ChunkingService.onMessage")(function* (
msg: TextDocument,
properties: Record<string, string>,
flowCtx: FlowContext,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const chunkSize = yield* flowCtx.flow.parameterEffect<number>("chunk-size").pipe(
Effect.orElseSucceed(() => DEFAULT_CHUNK_SIZE),
);
const chunkOverlap = yield* flowCtx.flow.parameterEffect<number>("chunk-overlap").pipe(
Effect.orElseSucceed(() => DEFAULT_CHUNK_OVERLAP),
);
const text = msg.text;
if (text.trim().length === 0) {
yield* Effect.logWarning(`[ChunkingService] Empty text received for document ${msg.documentId}`);
return;
}
const chunks = recursiveSplit(text, chunkSize, chunkOverlap);
yield* Effect.log(
`[ChunkingService] Split document ${msg.documentId} into ${chunks.length} chunks (size=${chunkSize}, overlap=${chunkOverlap})`,
);
const outputProducer = yield* flowCtx.flow.producerEffect<Chunk>("chunk-output");
yield* Effect.forEach(
chunks,
(chunkText) =>
outputProducer.send(requestId, {
metadata: msg.metadata,
chunk: chunkText,
documentId: msg.documentId,
}),
{ discard: true },
);
});
export const makeChunkingSpecs = (): ReadonlyArray<
Spec<never>
> => [
new ConsumerSpec<TextDocument, FlowResourceNotFoundError | MessagingDeliveryError>(
"chunk-input",
onChunkMessage,
),
new ProducerSpec<Chunk>("chunk-output"),
new ProducerSpec<Triples>("chunk-triples"),
new ParameterSpec("chunk-size"),
new ParameterSpec("chunk-overlap"),
];
export class ChunkingService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
new ConsumerSpec<TextDocument, FlowResourceNotFoundError | MessagingDeliveryError>(
"chunk-input",
this.onMessageEffect.bind(this),
),
);
this.registerSpecification(new ProducerSpec<Chunk>("chunk-output"));
this.registerSpecification(new ProducerSpec<Triples>("chunk-triples"));
this.registerSpecification(new ParameterSpec("chunk-size"));
this.registerSpecification(new ParameterSpec("chunk-overlap"));
for (const spec of makeChunkingSpecs()) {
this.registerSpecification(spec);
}
console.log("[ChunkingService] Service initialized");
}
private onMessageEffect(
msg: TextDocument,
properties: Record<string, string>,
flowCtx: FlowContext,
) {
return Effect.gen(function* () {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const chunkSize = yield* flowCtx.flow.parameterEffect<number>("chunk-size").pipe(
Effect.catch(() => Effect.succeed(DEFAULT_CHUNK_SIZE)),
);
const chunkOverlap = yield* flowCtx.flow.parameterEffect<number>("chunk-overlap").pipe(
Effect.catch(() => Effect.succeed(DEFAULT_CHUNK_OVERLAP)),
);
const text = msg.text;
if (text.trim().length === 0) {
yield* Effect.logWarning(`[ChunkingService] Empty text received for document ${msg.documentId}`);
return;
}
const chunks = recursiveSplit(text, chunkSize, chunkOverlap);
yield* Effect.log(
`[ChunkingService] Split document ${msg.documentId} into ${chunks.length} chunks (size=${chunkSize}, overlap=${chunkOverlap})`,
);
const outputProducer = yield* flowCtx.flow.producerEffect<Chunk>("chunk-output");
yield* Effect.forEach(
chunks,
(chunkText) =>
outputProducer.send(requestId, {
metadata: msg.metadata,
chunk: chunkText,
documentId: msg.documentId,
}),
{ discard: true },
);
});
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram({
id: "chunking",
make: (config) => new ChunkingService(config),
specs: () => makeChunkingSpecs(),
});
export async function run(): Promise<void> {
await ChunkingService.launch("chunking");
await Effect.runPromise(program);
}

View file

@ -44,8 +44,33 @@ const ConfigPushSchema = S.Struct({
config: S.Record(S.String, S.Unknown),
});
const DEFAULT_WORKSPACE = "default";
interface ConfigKeyLike {
type: string;
key?: string;
}
interface ConfigValueLike {
workspace?: string;
type: string;
key: string;
value: unknown;
}
type NamespaceStore = Map<string, unknown>;
type WorkspaceStore = Map<string, NamespaceStore>;
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
function optionalString(value: unknown): string | undefined {
return typeof value === "string" && value.length > 0 ? value : undefined;
}
export class ConfigService extends AsyncProcessor {
private store = new Map<string, Map<string, unknown>>();
private store = new Map<string, WorkspaceStore>();
private version = 0;
private readonly persistPath: string | null;
private consumer: BackendConsumer<ConfigRequest> | null = null;
@ -137,36 +162,146 @@ export class ConfigService extends AsyncProcessor {
switch (op) {
case "get":
return this.handleGet(request.keys ?? []);
return this.handleGet(request);
case "put":
return await this.handlePut(request.keys ?? [], request.values ?? {});
return await this.handlePut(request);
case "delete":
return await this.handleDelete(request.keys ?? []);
return await this.handleDelete(request);
case "list":
return this.handleList(request.keys ?? []);
return this.handleList(request);
case "config":
return this.handleConfigDump();
return this.handleConfigDump(request);
case "getvalues":
return this.handleGetValues(request);
case "getvalues-all-ws":
return this.handleGetValuesAllWorkspaces(request);
default:
throw new Error(`Unknown config operation: ${op as string}`);
}
}
private handleGet(keys: string[]): ConfigResponse {
private requestRecord(request: ConfigRequest): Record<string, unknown> {
return request as Record<string, unknown>;
}
private workspaceFor(request: ConfigRequest): string {
return optionalString(this.requestRecord(request).workspace) ?? DEFAULT_WORKSPACE;
}
private workspaceStore(workspace: string, create: boolean): WorkspaceStore | undefined {
let store = this.store.get(workspace);
if (store === undefined && create) {
store = new Map<string, NamespaceStore>();
this.store.set(workspace, store);
}
return store;
}
private namespaceStore(
workspace: string,
namespace: string,
create: boolean,
): NamespaceStore | undefined {
const ws = this.workspaceStore(workspace, create);
if (ws === undefined) return undefined;
let ns = ws.get(namespace);
if (ns === undefined && create) {
ns = new Map<string, unknown>();
ws.set(namespace, ns);
}
return ns;
}
private rawKeys(request: ConfigRequest): unknown[] {
const keys = this.requestRecord(request).keys;
return Array.isArray(keys) ? keys : [];
}
private stringKeys(request: ConfigRequest): string[] {
return this.rawKeys(request).filter((key): key is string => typeof key === "string");
}
private objectKeys(request: ConfigRequest): ConfigKeyLike[] {
return this.rawKeys(request).flatMap((key) => {
if (!isRecord(key)) return [];
const type = optionalString(key.type);
if (type === undefined) return [];
const item: ConfigKeyLike = { type };
const keyValue = optionalString(key.key);
if (keyValue !== undefined) item.key = keyValue;
return [item];
});
}
private requestType(request: ConfigRequest): string | undefined {
return optionalString(this.requestRecord(request).type) ?? this.stringKeys(request)[0];
}
private configValues(request: ConfigRequest): ConfigValueLike[] {
const req = this.requestRecord(request);
const rawValues = req.values;
const workspace = this.workspaceFor(request);
if (Array.isArray(rawValues)) {
return rawValues.flatMap((value) => {
if (!isRecord(value)) return [];
const type = optionalString(value.type);
const key = optionalString(value.key);
if (type === undefined || key === undefined) return [];
return [{
workspace: optionalString(value.workspace) ?? workspace,
type,
key,
value: value.value,
}];
});
}
if (isRecord(rawValues)) {
const namespace = this.requestType(request);
if (namespace === undefined) return [];
return Object.entries(rawValues).map(([key, value]) => ({
workspace,
type: namespace,
key,
value,
}));
}
return [];
}
private handleGet(request: ConfigRequest): ConfigResponse {
const workspace = this.workspaceFor(request);
const objectKeys = this.objectKeys(request);
if (objectKeys.length > 0) {
const values = objectKeys.map((key) => ({
type: key.type,
key: key.key ?? "",
value: key.key !== undefined
? this.namespaceStore(workspace, key.type, false)?.get(key.key)
: undefined,
}));
return { version: this.version, values };
}
const keys = this.stringKeys(request);
if (keys.length === 0) {
return { version: this.version, values: {} };
}
const values: Record<string, unknown> = {};
const namespace = keys[0];
const subMap = this.store.get(namespace);
const subMap = this.namespaceStore(workspace, namespace, false);
if (subMap !== undefined) {
if (keys.length === 1) {
@ -188,23 +323,12 @@ export class ConfigService extends AsyncProcessor {
return { version: this.version, values };
}
private async handlePut(
keys: string[],
values: Record<string, unknown>,
): Promise<ConfigResponse> {
if (keys.length === 0) {
throw new Error("Put requires at least one key (namespace)");
}
private async handlePut(request: ConfigRequest): Promise<ConfigResponse> {
const values = this.configValues(request);
if (values.length === 0) throw new Error("Put requires config values");
const namespace = keys[0];
let subMap = this.store.get(namespace);
if (subMap === undefined) {
subMap = new Map<string, unknown>();
this.store.set(namespace, subMap);
}
for (const [k, v] of Object.entries(values)) {
subMap.set(k, v);
for (const item of values) {
this.namespaceStore(item.workspace ?? DEFAULT_WORKSPACE, item.type, true)?.set(item.key, item.value);
}
this.version++;
@ -214,25 +338,49 @@ export class ConfigService extends AsyncProcessor {
return { version: this.version };
}
private async handleDelete(keys: string[]): Promise<ConfigResponse> {
private async handleDelete(request: ConfigRequest): Promise<ConfigResponse> {
const workspace = this.workspaceFor(request);
const objectKeys = this.objectKeys(request);
if (objectKeys.length > 0) {
for (const key of objectKeys) {
const ws = this.workspaceStore(workspace, false);
if (ws === undefined) continue;
if (key.key === undefined) {
ws.delete(key.type);
} else {
const ns = ws.get(key.type);
ns?.delete(key.key);
if (ns !== undefined && ns.size === 0) ws.delete(key.type);
}
}
this.version++;
await this.persist();
await this.pushConfig();
return { version: this.version };
}
const keys = this.stringKeys(request);
if (keys.length === 0) {
throw new Error("Delete requires at least one key");
}
const namespace = keys[0];
const ws = this.workspaceStore(workspace, false);
if (ws === undefined) return { version: this.version };
if (keys.length === 1) {
// Delete entire namespace
this.store.delete(namespace);
ws.delete(namespace);
} else {
// Delete specific keys within namespace
const subMap = this.store.get(namespace);
const subMap = ws.get(namespace);
if (subMap !== undefined) {
for (let i = 1; i < keys.length; i++) {
subMap.delete(keys[i]);
}
if (subMap.size === 0) {
this.store.delete(namespace);
ws.delete(namespace);
}
}
}
@ -244,17 +392,20 @@ export class ConfigService extends AsyncProcessor {
return { version: this.version };
}
private handleList(keys: string[]): ConfigResponse {
if (keys.length === 0) {
private handleList(request: ConfigRequest): ConfigResponse {
const workspace = this.workspaceFor(request);
const ws = this.workspaceStore(workspace, false);
const namespace = this.requestType(request);
if (namespace === undefined) {
// List all namespaces
return {
version: this.version,
directory: [...this.store.keys()],
directory: ws !== undefined ? [...ws.keys()] : [],
};
}
const namespace = keys[0];
const subMap = this.store.get(namespace);
const subMap = ws?.get(namespace);
return {
version: this.version,
@ -263,30 +414,48 @@ export class ConfigService extends AsyncProcessor {
}
private handleGetValues(request: ConfigRequest): ConfigResponse {
const type = request.type ?? "";
const workspace = this.workspaceFor(request);
const type = this.requestType(request) ?? "";
const ws = this.workspaceStore(workspace, false);
const values: { key: string; value: unknown }[] = [];
const values: { type: string; key: string; value: unknown }[] = [];
for (const [namespace, subMap] of this.store) {
for (const [namespace, subMap] of ws ?? new Map<string, NamespaceStore>()) {
if (
type.length === 0 ||
namespace === type ||
namespace.startsWith(`${type}.`) ||
namespace.startsWith(`${type}/`)
namespace === type
) {
for (const [k, v] of subMap) {
values.push({ key: `${namespace}.${k}`, value: v });
values.push({ type: namespace, key: k, value: v });
}
}
}
return { version: this.version, values: values as unknown as Record<string, unknown> };
return { version: this.version, values };
}
private handleConfigDump(): ConfigResponse {
private handleGetValuesAllWorkspaces(request: ConfigRequest): ConfigResponse {
const type = this.requestType(request) ?? "";
const values: { workspace: string; type: string; key: string; value: unknown }[] = [];
for (const [workspace, ws] of this.store) {
for (const [namespace, subMap] of ws) {
if (type.length > 0 && namespace !== type) continue;
for (const [key, value] of subMap) {
values.push({ workspace, type: namespace, key, value });
}
}
}
return { version: this.version, values };
}
private handleConfigDump(request: ConfigRequest): ConfigResponse {
const workspace = this.workspaceFor(request);
const ws = this.workspaceStore(workspace, false);
const config: Record<string, unknown> = {};
for (const [namespace, subMap] of this.store) {
for (const [namespace, subMap] of ws ?? new Map<string, NamespaceStore>()) {
const obj: Record<string, unknown> = {};
for (const [k, v] of subMap) {
obj[k] = v;
@ -305,7 +474,8 @@ export class ConfigService extends AsyncProcessor {
if (pushProducer === null) return;
const config: Record<string, unknown> = {};
for (const [namespace, subMap] of this.store) {
const ws = this.workspaceStore(DEFAULT_WORKSPACE, false);
for (const [namespace, subMap] of ws ?? new Map<string, NamespaceStore>()) {
const obj: Record<string, unknown> = {};
for (const [k, v] of subMap) {
obj[k] = v;
@ -326,18 +496,22 @@ export class ConfigService extends AsyncProcessor {
if (persistPath === null) return;
try {
const data: Record<string, Record<string, unknown>> = {};
const workspaces: Record<string, Record<string, Record<string, unknown>>> = {};
for (const [namespace, subMap] of this.store) {
const obj: Record<string, unknown> = {};
for (const [k, v] of subMap) {
obj[k] = v;
for (const [workspace, ws] of this.store) {
const workspaceData: Record<string, Record<string, unknown>> = {};
for (const [namespace, subMap] of ws) {
const obj: Record<string, unknown> = {};
for (const [k, v] of subMap) {
obj[k] = v;
}
workspaceData[namespace] = obj;
}
data[namespace] = obj;
workspaces[workspace] = workspaceData;
}
const json = JSON.stringify(
{ version: this.version, data },
{ version: this.version, workspaces },
null,
2,
);
@ -356,22 +530,39 @@ export class ConfigService extends AsyncProcessor {
const raw = await readTextFile(persistPath);
const parsed = JSON.parse(raw) as {
version: number;
data: Record<string, Record<string, unknown>>;
data?: Record<string, Record<string, unknown>>;
workspaces?: Record<string, Record<string, Record<string, unknown>>>;
};
this.version = parsed.version ?? 0;
this.store.clear();
for (const [namespace, obj] of Object.entries(parsed.data ?? {})) {
const subMap = new Map<string, unknown>();
for (const [k, v] of Object.entries(obj)) {
subMap.set(k, v);
if (parsed.workspaces !== undefined) {
for (const [workspace, namespaces] of Object.entries(parsed.workspaces)) {
const ws = new Map<string, NamespaceStore>();
for (const [namespace, obj] of Object.entries(namespaces)) {
const subMap = new Map<string, unknown>();
for (const [k, v] of Object.entries(obj)) {
subMap.set(k, v);
}
ws.set(namespace, subMap);
}
this.store.set(workspace, ws);
}
this.store.set(namespace, subMap);
} else {
const ws = new Map<string, NamespaceStore>();
for (const [namespace, obj] of Object.entries(parsed.data ?? {})) {
const subMap = new Map<string, unknown>();
for (const [k, v] of Object.entries(obj)) {
subMap.set(k, v);
}
ws.set(namespace, subMap);
}
this.store.set(DEFAULT_WORKSPACE, ws);
}
console.log(
`[ConfigService] Loaded persisted config (version=${this.version}, namespaces=${this.store.size})`,
`[ConfigService] Loaded persisted config (version=${this.version}, workspaces=${this.store.size})`,
);
} catch {
// File doesn't exist yet or is invalid — start fresh

View file

@ -21,7 +21,8 @@ import {
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import type { BackendProducer, BackendConsumer, Message } from "@trustgraph/base";
import { joinPath, readTextFile, writeTextFile } from "../runtime/effect-files.js";
import { Effect } from "effect";
import { ensureDirectory, joinPath, readTextFile, writeTextFile } from "../runtime/effect-files.js";
export interface KnowledgeCoreServiceConfig extends ProcessorConfig {
dataDir?: string;
@ -32,9 +33,17 @@ interface KnowledgeCore {
graphEmbeddings: { entity: Term; vectors: number[][] }[];
}
interface DocumentEmbeddingsCore {
metadata?: Record<string, unknown>;
chunks?: unknown[];
[key: string]: unknown;
}
export class KnowledgeCoreService extends AsyncProcessor {
/** Keyed by `${user}:${id}` */
private cores = new Map<string, KnowledgeCore>();
private deCores = new Map<string, DocumentEmbeddingsCore[]>();
private readonly dataDir: string;
private readonly persistPath: string;
private consumer: BackendConsumer<KnowledgeRequest> | null = null;
@ -43,6 +52,7 @@ export class KnowledgeCoreService extends AsyncProcessor {
constructor(config: KnowledgeCoreServiceConfig) {
super(config);
const dataDir = config.dataDir ?? process.env.KNOWLEDGE_DATA_DIR ?? "./data/knowledge";
this.dataDir = dataDir;
this.persistPath = joinPath(dataDir, "knowledge-state.json");
}
@ -51,6 +61,7 @@ export class KnowledgeCoreService extends AsyncProcessor {
}
protected override async run(): Promise<void> {
await ensureDirectory(this.dataDir);
// Load persisted state
await this.loadFromDisk();
@ -116,11 +127,40 @@ export class KnowledgeCoreService extends AsyncProcessor {
return this.putKgCore(request, requestId);
case "load-kg-core":
return this.loadKgCore(request, requestId);
case "unload-kg-core":
return this.unloadKgCore(request, requestId);
case "list-de-cores":
return this.listDeCores(request, requestId);
case "get-de-core":
return this.getDeCore(request, requestId);
case "delete-de-core":
return this.deleteDeCore(request, requestId);
case "put-de-core":
return this.putDeCore(request, requestId);
case "load-de-core":
return this.loadDeCore(request, requestId);
default:
throw new Error(`Unknown knowledge operation: ${request.operation as string}`);
}
}
private requestRecord(request: KnowledgeRequest): Record<string, unknown> {
return request as Record<string, unknown>;
}
private graphEmbeddings(request: KnowledgeRequest): { entity: Term; vectors: number[][] }[] {
const req = this.requestRecord(request);
const value = request.graphEmbeddings ?? req["graph-embeddings"];
return Array.isArray(value) ? value as { entity: Term; vectors: number[][] }[] : [];
}
private documentEmbeddings(request: KnowledgeRequest): DocumentEmbeddingsCore | undefined {
const req = this.requestRecord(request);
const value = request.documentEmbeddings ?? req["document-embeddings"];
if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined;
return value as DocumentEmbeddingsCore;
}
private async listKgCores(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const prefix = user.length > 0 ? `${user}:` : "";
@ -167,7 +207,7 @@ export class KnowledgeCoreService extends AsyncProcessor {
const isLast = i + BATCH_SIZE >= core.graphEmbeddings.length;
await this.responseProducer!.send(
{ graphEmbeddings: batch, eos: isLast },
{ graphEmbeddings: batch, "graph-embeddings": batch, eos: isLast } as KnowledgeResponse,
{ id: requestId },
);
}
@ -207,8 +247,9 @@ export class KnowledgeCoreService extends AsyncProcessor {
}
// Append graph embeddings if provided
if (request.graphEmbeddings !== undefined && request.graphEmbeddings.length > 0) {
core.graphEmbeddings.push(...request.graphEmbeddings);
const graphEmbeddings = this.graphEmbeddings(request);
if (graphEmbeddings.length > 0) {
core.graphEmbeddings.push(...graphEmbeddings);
}
await this.persist();
@ -229,22 +270,108 @@ export class KnowledgeCoreService extends AsyncProcessor {
throw new Error(`Knowledge core not found: ${key}`);
}
// MVP: just acknowledge. Full implementation would publish triples
// to flow storage topics via the flow config.
if (core.triples.length > 0) {
const producer = await this.pubsub.createProducer<unknown>({ topic: "tg.flow.triples" });
try {
await producer.send({
metadata: {
id: coreId,
root: coreId,
user,
collection: request.collection ?? "default",
},
triples: core.triples,
});
} finally {
await producer.close();
}
}
console.log(
`[KnowledgeCoreService] Load requested for core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length}) — returning success`,
`[KnowledgeCoreService] Loaded core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length})`,
);
await this.responseProducer!.send({}, { id: requestId });
}
private async unloadKgCore(_request: KnowledgeRequest, requestId: string): Promise<void> {
await this.responseProducer!.send({}, { id: requestId });
}
private async listDeCores(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const prefix = user.length > 0 ? `${user}:` : "";
const ids = [...this.deCores.keys()]
.filter((key) => prefix.length === 0 || key.startsWith(prefix))
.map((key) => key.slice(prefix.length));
await this.responseProducer!.send({ ids }, { id: requestId });
}
private async getDeCore(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
const core = this.deCores.get(key);
if (core === undefined) throw new Error(`Document embeddings core not found: ${key}`);
for (let i = 0; i < core.length; i++) {
const isLast = i === core.length - 1;
await this.responseProducer!.send(
{
documentEmbeddings: core[i],
"document-embeddings": core[i],
eos: isLast,
} as KnowledgeResponse,
{ id: requestId },
);
}
if (core.length === 0) {
await this.responseProducer!.send({ eos: true }, { id: requestId });
}
}
private async deleteDeCore(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
this.deCores.delete(this.coreKey(user, coreId));
await this.persist();
await this.responseProducer!.send({}, { id: requestId });
}
private async putDeCore(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
const item = this.documentEmbeddings(request);
if (item === undefined) throw new Error("put-de-core requires document-embeddings");
const core = this.deCores.get(key) ?? [];
core.push(item);
this.deCores.set(key, core);
await this.persist();
await this.responseProducer!.send({}, { id: requestId });
}
private async loadDeCore(request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
if (!this.deCores.has(key)) throw new Error(`Document embeddings core not found: ${key}`);
await this.responseProducer!.send({}, { id: requestId });
}
// ---------- Persistence ----------
private async persist(): Promise<void> {
try {
// Serialize Map to object
const data: Record<string, KnowledgeCore> = {};
const data: {
kg: Record<string, KnowledgeCore>;
de: Record<string, DocumentEmbeddingsCore[]>;
} = { kg: {}, de: {} };
for (const [key, core] of this.cores) {
data[key] = core;
data.kg[key] = core;
}
for (const [key, core] of this.deCores) {
data.de[key] = core;
}
const json = JSON.stringify(data, null, 2);
@ -257,14 +384,24 @@ export class KnowledgeCoreService extends AsyncProcessor {
private async loadFromDisk(): Promise<void> {
try {
const raw = await readTextFile(this.persistPath);
const parsed = JSON.parse(raw) as Record<string, KnowledgeCore>;
const parsed = JSON.parse(raw) as Record<string, KnowledgeCore> | {
kg?: Record<string, KnowledgeCore>;
de?: Record<string, DocumentEmbeddingsCore[]>;
};
this.cores.clear();
for (const [key, core] of Object.entries(parsed)) {
this.deCores.clear();
const kg = "kg" in parsed && parsed.kg !== undefined ? parsed.kg : parsed as Record<string, KnowledgeCore>;
for (const [key, core] of Object.entries(kg)) {
this.cores.set(key, core);
}
if ("de" in parsed && parsed.de !== undefined) {
for (const [key, core] of Object.entries(parsed.de)) {
this.deCores.set(key, core);
}
}
console.log(`[KnowledgeCoreService] Loaded persisted state (cores=${this.cores.size})`);
console.log(`[KnowledgeCoreService] Loaded persisted state (kg=${this.cores.size}, de=${this.deCores.size})`);
} catch {
console.log("[KnowledgeCoreService] No persisted state found, starting fresh");
}
@ -293,5 +430,5 @@ export const program = makeProcessorProgram({
});
export async function run(): Promise<void> {
await KnowledgeCoreService.launch("knowledge-svc");
await Effect.runPromise(program);
}

View file

@ -22,6 +22,7 @@ import {
RequestResponseSpec,
type ProcessorConfig,
type FlowContext,
type FlowResourceNotFoundError,
type Document,
type TextDocument,
type Triples,
@ -29,170 +30,205 @@ import {
type Term,
type LibrarianRequest,
type LibrarianResponse,
type MessagingDeliveryError,
type MessagingTimeoutError,
type Spec,
errorMessage,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import * as S from "effect/Schema";
export class PdfDecoderError extends S.TaggedErrorClass<PdfDecoderError>()(
"PdfDecoderError",
{
message: S.String,
operation: S.String,
documentId: S.String,
cause: S.DefectWithStack,
},
) {}
type PdfDecoderHandlerError =
| FlowResourceNotFoundError
| MessagingDeliveryError
| MessagingTimeoutError
| PdfDecoderError;
type PdfDocument = Awaited<ReturnType<typeof getDocument>["promise"]>;
const pdfDecoderError = (
operation: string,
documentId: string,
cause: unknown,
) =>
new PdfDecoderError({
operation,
documentId,
message: errorMessage(cause),
cause,
});
const loadPdf = (documentId: string, pdfBuffer: Buffer) =>
Effect.tryPromise({
try: () => getDocument({ data: new Uint8Array(pdfBuffer) }).promise,
catch: (cause) => pdfDecoderError("load-pdf", documentId, cause),
});
const loadPageText = (documentId: string, pageNumber: number, pdf: PdfDocument) =>
Effect.tryPromise({
try: async () => {
const page = await pdf.getPage(pageNumber);
const textContent = await page.getTextContent();
return textContent.items
.filter((item): item is TextItem => "str" in item)
.map((item) => item.str)
.join(" ");
},
catch: (cause) => pdfDecoderError("load-page-text", documentId, cause),
});
const onPdfDecodeMessage = Effect.fn("PdfDecoderService.onMessage")(function* (
msg: Document,
properties: Record<string, string>,
flowCtx: FlowContext,
): Effect.fn.Return<void, PdfDecoderHandlerError> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const { documentId } = msg;
const user = msg.metadata.user;
const librarian = yield* flowCtx.flow.requestorEffect<LibrarianRequest, LibrarianResponse>(
"librarian-client",
);
const metadataResp = yield* librarian.request({
operation: "get-document-metadata",
documentId,
user,
});
if (metadataResp.error !== undefined) {
yield* Effect.logError(`[PdfDecoder] Failed to get metadata for ${documentId}`, {
error: metadataResp.error.message,
});
return;
}
const kind = metadataResp.documentMetadata?.kind;
if (kind !== "application/pdf") {
yield* Effect.log(`[PdfDecoder] Skipping document ${documentId}: kind=${kind} (not PDF)`);
return;
}
const contentResp = yield* librarian.request({
operation: "get-document-content",
documentId,
user,
});
if (
contentResp.error !== undefined ||
contentResp.content === undefined ||
contentResp.content.length === 0
) {
yield* Effect.logError(`[PdfDecoder] Failed to get content for ${documentId}`, {
error: contentResp.error?.message ?? "no content",
});
return;
}
const pdfBuffer = Buffer.from(contentResp.content, "base64");
const pdf = yield* loadPdf(documentId, pdfBuffer);
yield* Effect.log(`[PdfDecoder] Document ${documentId}: ${pdf.numPages} pages`);
const outputProducer = yield* flowCtx.flow.producerEffect<TextDocument>("decode-output");
const triplesProducer = yield* flowCtx.flow.producerEffect<Triples>("decode-triples");
for (let i = 1; i <= pdf.numPages; i++) {
const pageText = yield* loadPageText(documentId, i, pdf);
if (pageText.trim().length === 0) {
yield* Effect.log(`[PdfDecoder] Skipping empty page ${i} of document ${documentId}`);
continue;
}
const childResp = yield* librarian.request({
operation: "add-child-document",
documentMetadata: {
id: "",
user,
kind: "text/plain",
title: `Page ${i}`,
parentId: documentId,
documentType: "page",
time: Date.now(),
comments: "",
tags: [],
},
content: Buffer.from(pageText).toString("base64"),
});
if (childResp.error !== undefined) {
yield* Effect.logError(`[PdfDecoder] Failed to save page ${i} of ${documentId}`, {
error: childResp.error.message,
});
continue;
}
const childDocId = childResp.documentMetadata?.id ?? "";
yield* outputProducer.send(requestId, {
metadata: msg.metadata,
text: pageText,
documentId: childDocId,
});
const triples: Triple[] = [
{
s: iriTerm(`urn:tg:page:${childDocId}`),
p: iriTerm("http://www.w3.org/ns/prov#wasDerivedFrom"),
o: iriTerm(`urn:tg:doc:${documentId}`),
},
{
s: iriTerm(`urn:tg:page:${childDocId}`),
p: iriTerm("http://www.w3.org/2000/01/rdf-schema#label"),
o: literalTerm(`Page ${i}`),
},
];
yield* triplesProducer.send(requestId, {
metadata: msg.metadata,
triples,
});
}
yield* Effect.log(`[PdfDecoder] Finished processing document ${documentId}`);
});
export const makePdfDecoderSpecs = (): ReadonlyArray<Spec<never>> => [
new ConsumerSpec<Document, PdfDecoderHandlerError>("decode-input", onPdfDecodeMessage),
new ProducerSpec<TextDocument>("decode-output"),
new ProducerSpec<Triples>("decode-triples"),
new RequestResponseSpec<LibrarianRequest, LibrarianResponse>(
"librarian-client",
"librarian-request",
"librarian-response",
),
];
export class PdfDecoderService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
ConsumerSpec.fromPromise<Document>("decode-input", this.onMessage.bind(this)),
);
this.registerSpecification(new ProducerSpec<TextDocument>("decode-output"));
this.registerSpecification(new ProducerSpec<Triples>("decode-triples"));
this.registerSpecification(
new RequestResponseSpec<LibrarianRequest, LibrarianResponse>(
"librarian-client",
"librarian-request",
"librarian-response",
),
);
for (const spec of makePdfDecoderSpecs()) {
this.registerSpecification(spec);
}
console.log("[PdfDecoder] Service initialized");
}
private async onMessage(
msg: Document,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const { documentId } = msg;
const user = msg.metadata.user;
const librarian = flowCtx.flow.requestor<LibrarianRequest, LibrarianResponse>(
"librarian-client",
);
// 1. Fetch document metadata to check MIME type
const metadataResp = await librarian.request({
operation: "get-document-metadata",
documentId,
user,
});
if (metadataResp.error !== undefined) {
console.error(
`[PdfDecoder] Failed to get metadata for ${documentId}:`,
metadataResp.error.message,
);
return;
}
const kind = metadataResp.documentMetadata?.kind;
if (kind !== "application/pdf") {
console.log(
`[PdfDecoder] Skipping document ${documentId}: kind=${kind} (not PDF)`,
);
return;
}
// 2. Fetch document content
const contentResp = await librarian.request({
operation: "get-document-content",
documentId,
user,
});
if (
contentResp.error !== undefined ||
contentResp.content === undefined ||
contentResp.content.length === 0
) {
console.error(
`[PdfDecoder] Failed to get content for ${documentId}:`,
contentResp.error?.message ?? "no content",
);
return;
}
// 3. Decode base64 content and extract text per page
const pdfBuffer = Buffer.from(contentResp.content, "base64");
const pdf = await getDocument({ data: new Uint8Array(pdfBuffer) }).promise;
console.log(
`[PdfDecoder] Document ${documentId}: ${pdf.numPages} pages`,
);
const outputProducer = flowCtx.flow.producer<TextDocument>("decode-output");
const triplesProducer = flowCtx.flow.producer<Triples>("decode-triples");
for (let i = 1; i <= pdf.numPages; i++) {
const page = await pdf.getPage(i);
const textContent = await page.getTextContent();
const pageText = textContent.items
.filter((item): item is TextItem => "str" in item)
.map((item) => item.str)
.join(" ");
if (pageText.trim().length === 0) {
console.log(
`[PdfDecoder] Skipping empty page ${i} of document ${documentId}`,
);
continue;
}
// 4. Save as child document in librarian
const childResp = await librarian.request({
operation: "add-child-document",
documentMetadata: {
id: "",
user,
kind: "text/plain",
title: `Page ${i}`,
parentId: documentId,
documentType: "page",
time: Date.now(),
comments: "",
tags: [],
},
content: Buffer.from(pageText).toString("base64"),
});
if (childResp.error !== undefined) {
console.error(
`[PdfDecoder] Failed to save page ${i} of ${documentId}:`,
childResp.error.message,
);
continue;
}
const childDocId = childResp.documentMetadata?.id ?? "";
// 5. Emit TextDocument for the chunking pipeline
await outputProducer.send(requestId, {
metadata: msg.metadata,
text: pageText,
documentId: childDocId,
});
// 6. Emit provenance triples
const triples: Triple[] = [
{
s: iriTerm(`urn:tg:page:${childDocId}`),
p: iriTerm("http://www.w3.org/ns/prov#wasDerivedFrom"),
o: iriTerm(`urn:tg:doc:${documentId}`),
},
{
s: iriTerm(`urn:tg:page:${childDocId}`),
p: iriTerm("http://www.w3.org/2000/01/rdf-schema#label"),
o: literalTerm(`Page ${i}`),
},
];
await triplesProducer.send(requestId, {
metadata: msg.metadata,
triples,
});
}
console.log(
`[PdfDecoder] Finished processing document ${documentId}`,
);
}
}
function iriTerm(iri: string): Term {
@ -203,11 +239,11 @@ function literalTerm(value: string): Term {
return { type: "LITERAL", value };
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram({
id: "pdf-decoder",
make: (config) => new PdfDecoderService(config),
specs: () => makePdfDecoderSpecs(),
});
export async function run(): Promise<void> {
await PdfDecoderService.launch("pdf-decoder");
await Effect.runPromise(program);
}

View file

@ -10,10 +10,11 @@ import {
Embeddings,
EmbeddingsService,
embeddingsError,
makeEmbeddingsSpecs,
type EmbeddingsServiceShape,
type ProcessorConfig,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { makeFlowProcessorProgram } from "@trustgraph/base";
export interface OllamaEmbeddingsConfig extends ProcessorConfig {
model?: string;
@ -102,11 +103,12 @@ export class OllamaEmbeddingsProcessor extends EmbeddingsService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<OllamaEmbeddingsConfig, never, Embeddings>({
id: "embeddings",
make: (config) => new OllamaEmbeddingsProcessor(config),
specs: () => makeEmbeddingsSpecs(),
layer: (config) => OllamaEmbeddingsLive(config),
});
export async function run(): Promise<void> {
await OllamaEmbeddingsProcessor.launch("embeddings");
await Effect.runPromise(program);
}

View file

@ -15,6 +15,7 @@ import {
ConsumerSpec,
ProducerSpec,
RequestResponseSpec,
makeFlowProcessorProgram,
type ProcessorConfig,
type FlowContext,
type Chunk,
@ -27,229 +28,270 @@ import {
type TextCompletionResponse,
type Triple,
type Term,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type EffectRequestResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
// Well-known RDF/SKOS IRIs
const RDFS_LABEL = "http://www.w3.org/2000/01/rdf-schema#label";
const SKOS_DEFINITION = "http://www.w3.org/2004/02/skos/core#definition";
interface ExtractedRelationship {
subject: string;
predicate: string;
object: string;
}
const ExtractedRelationship = S.Struct({
subject: S.String,
predicate: S.String,
object: S.String,
});
type ExtractedRelationship = typeof ExtractedRelationship.Type;
interface ExtractedDefinition {
entity: string;
definition: string;
}
const ExtractedRelationshipsFromJson = S.Array(ExtractedRelationship).pipe(S.fromJsonString);
const decodeExtractedRelationships = S.decodeUnknownOption(ExtractedRelationshipsFromJson);
const ExtractedDefinition = S.Struct({
entity: S.String,
definition: S.String,
});
type ExtractedDefinition = typeof ExtractedDefinition.Type;
const ExtractedDefinitionsFromJson = S.Array(ExtractedDefinition).pipe(S.fromJsonString);
const decodeExtractedDefinitions = S.decodeUnknownOption(ExtractedDefinitionsFromJson);
type KnowledgeExtractHandlerError =
| FlowResourceNotFoundError
| MessagingDeliveryError;
type PromptClient = EffectRequestResponse<PromptRequest, PromptResponse>;
type LlmClient = EffectRequestResponse<TextCompletionRequest, TextCompletionResponse>;
const requestPrompt = Effect.fn("KnowledgeExtract.requestPrompt")(function* (
promptClient: PromptClient,
name: string,
text: string,
) {
return yield* promptClient.request(
{ name, variables: { text } },
{ timeoutMs: 10_000 },
);
});
const requestCompletion = Effect.fn("KnowledgeExtract.requestCompletion")(function* (
llmClient: LlmClient,
prompt: PromptResponse,
) {
return yield* llmClient.request(
{ system: prompt.system, prompt: prompt.prompt },
{ timeoutMs: 120_000 },
);
});
const extractRelationships = Effect.fn("KnowledgeExtract.extractRelationships")(function* (
promptClient: PromptClient,
llmClient: LlmClient,
text: string,
) {
const relPrompt = yield* requestPrompt(promptClient, "extract-relationships", text);
if (relPrompt.error !== undefined) return null;
for (let attempt = 0; attempt < 3; attempt++) {
const relCompletion = yield* requestCompletion(llmClient, relPrompt);
if (relCompletion.error !== undefined || relCompletion.response.length === 0) {
break;
}
const relationships = parseRelationshipsResponse(relCompletion.response);
if (relationships !== null) return relationships;
yield* Effect.logWarning(
`[KnowledgeExtract] Relationship parse failed, attempt ${attempt + 1}/3`,
);
}
return null;
});
const extractDefinitions = Effect.fn("KnowledgeExtract.extractDefinitions")(function* (
promptClient: PromptClient,
llmClient: LlmClient,
text: string,
) {
const defPrompt = yield* requestPrompt(promptClient, "extract-definitions", text);
if (defPrompt.error !== undefined) return null;
for (let attempt = 0; attempt < 3; attempt++) {
const defCompletion = yield* requestCompletion(llmClient, defPrompt);
if (defCompletion.error !== undefined || defCompletion.response.length === 0) {
break;
}
const definitions = parseDefinitionsResponse(defCompletion.response);
if (definitions !== null) return definitions;
yield* Effect.logWarning(
`[KnowledgeExtract] Definition parse failed, attempt ${attempt + 1}/3`,
);
}
return null;
});
const onKnowledgeExtractMessage = Effect.fn("KnowledgeExtractService.onMessage")(function* (
msg: Chunk,
properties: Record<string, string>,
flowCtx: FlowContext,
): Effect.fn.Return<void, KnowledgeExtractHandlerError> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const text = msg.chunk;
if (text.trim().length === 0) return;
const promptClient = yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt-client");
const llmClient = yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm-client");
const triplesProducer = yield* flowCtx.flow.producerEffect<Triples>("extract-triples");
const entityContextsProducer = yield* flowCtx.flow.producerEffect<EntityContexts>("extract-entity-contexts");
const allTriples: Triple[] = [];
const allEntityContexts: EntityContext[] = [];
const relationships = yield* extractRelationships(promptClient, llmClient, text).pipe(
Effect.catch((error: unknown) =>
Effect.logError("[KnowledgeExtract] Relationship extraction failed", {
error: error instanceof Error ? error.message : String(error),
}).pipe(Effect.as(null)),
),
);
if (relationships !== null) {
for (const rel of relationships) {
if (
rel.subject.length === 0 ||
rel.predicate.length === 0 ||
rel.object.length === 0
) {
continue;
}
const subjectIri = toEntityIri(rel.subject);
const predicateIri = toEntityIri(rel.predicate);
const objectIri = toEntityIri(rel.object);
allTriples.push({ s: subjectIri, p: predicateIri, o: objectIri });
allTriples.push({
s: subjectIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.subject),
});
allTriples.push({
s: predicateIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.predicate),
});
allTriples.push({
s: objectIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.object),
});
allEntityContexts.push({
entity: subjectIri,
context: text,
chunkId: msg.documentId,
});
allEntityContexts.push({
entity: objectIri,
context: text,
chunkId: msg.documentId,
});
}
yield* Effect.log(`[KnowledgeExtract] Extracted ${relationships.length} relationships`);
}
const definitions = yield* extractDefinitions(promptClient, llmClient, text).pipe(
Effect.catch((error: unknown) =>
Effect.logError("[KnowledgeExtract] Definition extraction failed", {
error: error instanceof Error ? error.message : String(error),
}).pipe(Effect.as(null)),
),
);
if (definitions !== null) {
for (const def of definitions) {
if (def.entity.length === 0 || def.definition.length === 0) continue;
const entityIri = toEntityIri(def.entity);
allTriples.push({
s: entityIri,
p: iriTerm(SKOS_DEFINITION),
o: literalTerm(def.definition),
});
allTriples.push({
s: entityIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(def.entity),
});
allEntityContexts.push({
entity: entityIri,
context: text,
chunkId: msg.documentId,
});
}
yield* Effect.log(`[KnowledgeExtract] Extracted ${definitions.length} definitions`);
}
if (allTriples.length > 0) {
yield* triplesProducer.send(requestId, {
metadata: msg.metadata,
triples: allTriples,
});
}
if (allEntityContexts.length > 0) {
yield* entityContextsProducer.send(requestId, {
metadata: msg.metadata,
entities: allEntityContexts,
});
}
});
export const makeKnowledgeExtractSpecs = (): ReadonlyArray<Spec<never>> => [
new ConsumerSpec<Chunk, KnowledgeExtractHandlerError>(
"extract-input",
onKnowledgeExtractMessage,
),
new ProducerSpec<Triples>("extract-triples"),
new ProducerSpec<EntityContexts>("extract-entity-contexts"),
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt-client",
"prompt-request",
"prompt-response",
),
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm-client",
"text-completion-request",
"text-completion-response",
),
];
export class KnowledgeExtractService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
this.registerSpecification(
ConsumerSpec.fromPromise<Chunk>("extract-input", this.onMessage.bind(this)),
);
this.registerSpecification(new ProducerSpec<Triples>("extract-triples"));
this.registerSpecification(new ProducerSpec<EntityContexts>("extract-entity-contexts"));
this.registerSpecification(
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt-client",
"prompt-request",
"prompt-response",
),
);
this.registerSpecification(
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm-client",
"text-completion-request",
"text-completion-response",
),
);
for (const spec of makeKnowledgeExtractSpecs()) {
this.registerSpecification(spec);
}
console.log("[KnowledgeExtract] Service initialized");
}
private async onMessage(
msg: Chunk,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const text = msg.chunk;
if (text.trim().length === 0) return;
const promptClient = flowCtx.flow.requestor<PromptRequest, PromptResponse>("prompt-client");
const llmClient = flowCtx.flow.requestor<TextCompletionRequest, TextCompletionResponse>("llm-client");
const triplesProducer = flowCtx.flow.producer<Triples>("extract-triples");
const entityContextsProducer = flowCtx.flow.producer<EntityContexts>("extract-entity-contexts");
const allTriples: Triple[] = [];
const allEntityContexts: EntityContext[] = [];
// --- Extract relationships ---
try {
const relPrompt = await promptClient.request(
{ name: "extract-relationships", variables: { text } },
{ timeoutMs: 10_000 },
);
if (relPrompt.error === undefined) {
let relationships: ExtractedRelationship[] | null = null;
for (let attempt = 0; attempt < 3; attempt++) {
const relCompletion = await llmClient.request(
{ system: relPrompt.system, prompt: relPrompt.prompt },
{ timeoutMs: 120_000 },
);
if (
relCompletion.error === undefined &&
relCompletion.response.length > 0
) {
relationships = parseJsonResponse<ExtractedRelationship[]>(relCompletion.response);
if (relationships !== null) break;
console.warn(`[KnowledgeExtract] Relationship parse failed, attempt ${attempt + 1}/3`);
} else {
break; // LLM error, don't retry
}
}
if (relationships !== null) {
for (const rel of relationships) {
if (
rel.subject.length === 0 ||
rel.predicate.length === 0 ||
rel.object.length === 0
) {
continue;
}
const subjectIri = toEntityIri(rel.subject);
const predicateIri = toEntityIri(rel.predicate);
const objectIri = toEntityIri(rel.object);
// Main relationship triple
allTriples.push({ s: subjectIri, p: predicateIri, o: objectIri });
// rdfs:label triples for each entity
allTriples.push({
s: subjectIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.subject),
});
allTriples.push({
s: predicateIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.predicate),
});
allTriples.push({
s: objectIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(rel.object),
});
// Entity contexts for subject and object
allEntityContexts.push({
entity: subjectIri,
context: text,
chunkId: msg.documentId,
});
allEntityContexts.push({
entity: objectIri,
context: text,
chunkId: msg.documentId,
});
}
console.log(`[KnowledgeExtract] Extracted ${relationships.length} relationships`);
}
}
} catch (err) {
console.error("[KnowledgeExtract] Relationship extraction failed:", err);
}
// --- Extract definitions ---
try {
const defPrompt = await promptClient.request(
{ name: "extract-definitions", variables: { text } },
{ timeoutMs: 10_000 },
);
if (defPrompt.error === undefined) {
let definitions: ExtractedDefinition[] | null = null;
for (let attempt = 0; attempt < 3; attempt++) {
const defCompletion = await llmClient.request(
{ system: defPrompt.system, prompt: defPrompt.prompt },
{ timeoutMs: 120_000 },
);
if (
defCompletion.error === undefined &&
defCompletion.response.length > 0
) {
definitions = parseJsonResponse<ExtractedDefinition[]>(defCompletion.response);
if (definitions !== null) break;
console.warn(`[KnowledgeExtract] Definition parse failed, attempt ${attempt + 1}/3`);
} else {
break; // LLM error, don't retry
}
}
if (definitions !== null) {
for (const def of definitions) {
if (def.entity.length === 0 || def.definition.length === 0) continue;
const entityIri = toEntityIri(def.entity);
// Definition triple
allTriples.push({
s: entityIri,
p: iriTerm(SKOS_DEFINITION),
o: literalTerm(def.definition),
});
// Label triple
allTriples.push({
s: entityIri,
p: iriTerm(RDFS_LABEL),
o: literalTerm(def.entity),
});
// Entity context
allEntityContexts.push({
entity: entityIri,
context: text,
chunkId: msg.documentId,
});
}
console.log(`[KnowledgeExtract] Extracted ${definitions.length} definitions`);
}
}
} catch (err) {
console.error("[KnowledgeExtract] Definition extraction failed:", err);
}
// --- Emit results ---
if (allTriples.length > 0) {
await triplesProducer.send(requestId, {
metadata: msg.metadata,
triples: allTriples,
});
}
if (allEntityContexts.length > 0) {
await entityContextsProducer.send(requestId, {
metadata: msg.metadata,
entities: allEntityContexts,
});
}
}
}
// ---------- Helpers ----------
@ -275,53 +317,68 @@ function literalTerm(value: string): Term {
* Uses progressive fallback: direct parse, array extraction, truncated array repair, single object wrap.
*/
export function parseJsonResponse<T>(raw: string): T | null {
// Attempt 1: direct parse after stripping fences
let cleaned = raw.trim();
const fenceMatch = cleaned.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```$/);
if (fenceMatch !== null) {
cleaned = (fenceMatch[1] ?? "").trim();
}
try {
return JSON.parse(cleaned) as T;
} catch { /* fall through */ }
// Attempt 2: extract first JSON array from the text
const arrayMatch = cleaned.match(/\[[\s\S]*\]/);
if (arrayMatch !== null) {
try {
return JSON.parse(arrayMatch[0]) as T;
} catch { /* fall through */ }
// Attempt 3: try to fix truncated array by closing it after the last complete object
const partial = arrayMatch[0];
const lastBrace = partial.lastIndexOf('}');
if (lastBrace > 0) {
const truncated = partial.slice(0, lastBrace + 1) + ']';
try {
return JSON.parse(truncated) as T;
} catch { /* fall through */ }
}
}
// Attempt 4: extract first JSON object, wrap in array
const objMatch = cleaned.match(/\{[\s\S]*?\}/);
if (objMatch !== null) {
try {
const obj = JSON.parse(objMatch[0]);
return [obj] as unknown as T;
} catch { /* fall through */ }
const decodeJson = S.decodeUnknownOption(S.UnknownFromJsonString);
for (const candidate of jsonCandidates(raw)) {
const decoded = decodeJson(candidate);
if (O.isSome(decoded)) return decoded.value as T;
}
console.warn("[KnowledgeExtract] Failed to parse JSON from LLM response:", raw.slice(0, 300));
return null;
}
export const program = makeProcessorProgram({
function parseRelationshipsResponse(raw: string): ReadonlyArray<ExtractedRelationship> | null {
for (const candidate of jsonCandidates(raw)) {
const decoded = decodeExtractedRelationships(candidate);
if (O.isSome(decoded)) return decoded.value;
}
console.warn("[KnowledgeExtract] Failed to parse relationships from LLM response:", raw.slice(0, 300));
return null;
}
function parseDefinitionsResponse(raw: string): ReadonlyArray<ExtractedDefinition> | null {
for (const candidate of jsonCandidates(raw)) {
const decoded = decodeExtractedDefinitions(candidate);
if (O.isSome(decoded)) return decoded.value;
}
console.warn("[KnowledgeExtract] Failed to parse definitions from LLM response:", raw.slice(0, 300));
return null;
}
function jsonCandidates(raw: string): ReadonlyArray<string> {
const candidates: string[] = [];
let cleaned = raw.trim();
const fenceMatch = cleaned.match(/^```(?:json)?\s*\n?([\s\S]*?)\n?```$/);
if (fenceMatch !== null) {
cleaned = (fenceMatch[1] ?? "").trim();
}
candidates.push(cleaned);
const arrayMatch = cleaned.match(/\[[\s\S]*\]/);
if (arrayMatch !== null) {
candidates.push(arrayMatch[0]);
const partial = arrayMatch[0];
const lastBrace = partial.lastIndexOf("}");
if (lastBrace > 0) {
candidates.push(partial.slice(0, lastBrace + 1) + "]");
}
}
const objMatch = cleaned.match(/\{[\s\S]*?\}/);
if (objMatch !== null) {
candidates.push(`[${objMatch[0]}]`);
}
return candidates;
}
export const program = makeFlowProcessorProgram({
id: "knowledge-extract",
make: (config) => new KnowledgeExtractService(config),
specs: () => makeKnowledgeExtractSpecs(),
});
export async function run(): Promise<void> {
await KnowledgeExtractService.launch("knowledge-extract");
await Effect.runPromise(program);
}

View file

@ -28,6 +28,7 @@ import type {
BackendConsumer,
Message,
} from "@trustgraph/base";
import { Effect } from "effect";
// ---------- Internal state types ----------
@ -35,13 +36,48 @@ interface FlowInstance {
id: string;
blueprintName: string;
description: string;
parameters: Record<string, string>;
parameters: Record<string, unknown>;
status: "running" | "stopped";
}
interface Blueprint {
description: string;
topics: Record<string, string>;
parameters?: Record<string, unknown>;
[key: string]: unknown;
}
interface ConfigValueEntry {
key: string;
value: unknown;
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
function optionalString(value: unknown): string | undefined {
return typeof value === "string" && value.length > 0 ? value : undefined;
}
function configValues(response: ConfigResponse): ConfigValueEntry[] {
const values = response.values;
if (!Array.isArray(values)) return [];
return values.flatMap((value) => {
if (!isRecord(value)) return [];
const key = optionalString(value.key);
if (key === undefined) return [];
return [{ key, value: value.value }];
});
}
function parseConfigRecord(value: unknown): Record<string, unknown> | undefined {
try {
const parsed = typeof value === "string" ? JSON.parse(value) as unknown : value;
return isRecord(parsed) ? parsed : undefined;
} catch {
return undefined;
}
}
// ---------- Default blueprint ----------
@ -122,6 +158,8 @@ export class FlowManagerService extends AsyncProcessor {
subscription: `${this.config.id}-config-client`,
});
await this.configClient.start();
await this.ensureDefaultBlueprint();
await this.refreshBlueprintsFromConfig();
// Create producer for flow-response topic
this.responseProducer = await this.pubsub.createProducer<Record<string, unknown>>({
@ -178,15 +216,101 @@ export class FlowManagerService extends AsyncProcessor {
}
}
private async configRequest(request: ConfigRequest): Promise<ConfigResponse> {
if (this.configClient === null) throw new Error("Config client not started");
return this.configClient.request(request);
}
private async ensureDefaultBlueprint(): Promise<void> {
const response = await this.configRequest({
operation: "getvalues",
type: "flow-blueprint",
});
if (configValues(response).some((value) => value.key === "default")) {
return;
}
await this.configRequest({
operation: "put",
keys: ["flow-blueprint"],
values: {
default: JSON.stringify(DEFAULT_BLUEPRINT),
},
});
}
private async refreshBlueprintsFromConfig(): Promise<void> {
const response = await this.configRequest({
operation: "getvalues",
type: "flow-blueprint",
});
const next = new Map<string, Blueprint>();
for (const item of configValues(response)) {
const parsed = parseConfigRecord(item.value);
if (parsed === undefined) continue;
next.set(item.key, parsed as Blueprint);
}
if (!next.has("default")) {
next.set("default", DEFAULT_BLUEPRINT);
}
this.blueprints = next;
}
private async refreshFlowsFromConfig(): Promise<void> {
const response = await this.configRequest({
operation: "getvalues",
type: "flow",
});
const next = new Map<string, FlowInstance>();
for (const item of configValues(response)) {
const parsed = parseConfigRecord(item.value);
if (parsed === undefined) continue;
const parameters = isRecord(parsed.parameters) ? parsed.parameters : {};
next.set(item.key, {
id: item.key,
blueprintName: optionalString(parsed["blueprint-name"]) ?? optionalString(parsed.blueprintName) ?? "default",
description: optionalString(parsed.description) ?? "",
parameters,
status: "running",
});
}
if (next.size === 0) {
const flowsResponse = await this.configRequest({
operation: "getvalues",
type: "flows",
});
for (const item of configValues(flowsResponse)) {
next.set(item.key, {
id: item.key,
blueprintName: "default",
description: "",
parameters: {},
status: "running",
});
}
}
this.flows = next;
}
private async handleOperation(
request: Record<string, unknown>,
): Promise<Record<string, unknown>> {
const op = request.operation as string;
await this.refreshBlueprintsFromConfig();
await this.refreshFlowsFromConfig();
switch (op) {
case "list-blueprints":
return this.handleListBlueprints();
case "put-blueprint":
return await this.handlePutBlueprint(request);
case "get-blueprint":
return this.handleGetBlueprint(request);
@ -236,9 +360,33 @@ export class FlowManagerService extends AsyncProcessor {
};
}
private handleDeleteBlueprint(
private async handlePutBlueprint(
request: Record<string, unknown>,
): Record<string, unknown> {
): Promise<Record<string, unknown>> {
const name = request["blueprint-name"] as string | undefined;
if (name === undefined || name.length === 0) {
throw new Error("Missing blueprint-name");
}
const rawDefinition = request["blueprint-definition"];
if (rawDefinition === undefined) {
throw new Error("Missing blueprint-definition");
}
const definition = typeof rawDefinition === "string"
? rawDefinition
: JSON.stringify(rawDefinition);
await this.configRequest({
operation: "put",
keys: ["flow-blueprint"],
values: { [name]: definition },
});
await this.refreshBlueprintsFromConfig();
return {};
}
private async handleDeleteBlueprint(
request: Record<string, unknown>,
): Promise<Record<string, unknown>> {
const name = request["blueprint-name"] as string | undefined;
if (name === undefined || name.length === 0) {
throw new Error("Missing blueprint-name");
@ -248,10 +396,11 @@ export class FlowManagerService extends AsyncProcessor {
throw new Error("Cannot delete the default blueprint");
}
const existed = this.blueprints.delete(name);
if (!existed) {
throw new Error(`Blueprint not found: ${name}`);
}
await this.configRequest({
operation: "delete",
keys: ["flow-blueprint", name],
});
this.blueprints.delete(name);
return {};
}
@ -292,7 +441,7 @@ export class FlowManagerService extends AsyncProcessor {
const id = request["flow-id"] as string | undefined;
const blueprintName = (request["blueprint-name"] as string) ?? "default";
const description = (request["description"] as string) ?? "";
const parameters = (request["parameters"] as Record<string, string>) ?? {};
const parameters = (request["parameters"] as Record<string, unknown>) ?? {};
if (id === undefined || id.length === 0) {
throw new Error("Missing flow-id");
@ -342,13 +491,15 @@ export class FlowManagerService extends AsyncProcessor {
this.flows.delete(id);
console.log(`[FlowManager] Stopped flow "${id}"`);
console.log(`[FlowManager] Stopped flow "${id}"`);
// Push updated flows config (without the removed flow)
await this.pushFlowsConfig();
await this.deleteFlowConfig(id);
return {};
}
// Push updated flows config (without the removed flow)
await this.pushFlowsConfig();
return {};
}
// ---------- Config push ----------
@ -360,10 +511,16 @@ export class FlowManagerService extends AsyncProcessor {
if (this.configClient === null) return;
const flowsConfig: Record<string, { topics: Record<string, string> }> = {};
const flowRecords: Record<string, string> = {};
for (const [id, inst] of this.flows) {
const blueprint = this.blueprints.get(inst.blueprintName);
if (blueprint !== undefined) {
flowsConfig[id] = { topics: blueprint.topics };
flowRecords[id] = JSON.stringify({
"blueprint-name": inst.blueprintName,
description: inst.description,
parameters: inst.parameters,
});
}
}
@ -373,6 +530,11 @@ export class FlowManagerService extends AsyncProcessor {
keys: ["flows"],
values: flowsConfig,
});
await this.configClient.request({
operation: "put",
keys: ["flow"],
values: flowRecords,
});
console.log(
`[FlowManager] Pushed flows config (${this.flows.size} active flows)`,
);
@ -381,6 +543,18 @@ export class FlowManagerService extends AsyncProcessor {
}
}
private async deleteFlowConfig(id: string): Promise<void> {
if (this.configClient === null) return;
await this.configClient.request({
operation: "delete",
keys: ["flows", id],
});
await this.configClient.request({
operation: "delete",
keys: ["flow", id],
});
}
// ---------- Lifecycle ----------
override async stop(): Promise<void> {
@ -410,5 +584,5 @@ export const program = makeProcessorProgram({
});
export async function run(): Promise<void> {
await FlowManagerService.launch("flow-manager");
await Effect.runPromise(program);
}

View file

@ -1,89 +0,0 @@
/**
* 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;
private readonly handler: MuxHandler;
constructor(handler: MuxHandler) {
this.handler = handler;
}
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

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

View file

@ -0,0 +1,35 @@
import { Schema as S } from "effect";
import * as Rpc from "effect/unstable/rpc/Rpc";
import * as RpcGroup from "effect/unstable/rpc/RpcGroup";
export class DispatchPayload extends S.Class<DispatchPayload>("DispatchPayload")({
scope: S.Literals(["global", "flow"]),
service: S.String,
flow: S.optionalKey(S.String),
request: S.Record(S.String, S.Unknown),
}) {}
export class DispatchStreamChunk extends S.Class<DispatchStreamChunk>("DispatchStreamChunk")({
response: S.Unknown,
complete: S.Boolean,
}) {}
export class DispatchError extends S.ErrorClass<DispatchError>("DispatchError")({
_tag: S.tag("DispatchError"),
message: S.String,
}) {}
export class Dispatch extends Rpc.make("Dispatch", {
payload: DispatchPayload,
success: S.Unknown,
error: DispatchError,
}) {}
export class DispatchStream extends Rpc.make("DispatchStream", {
payload: DispatchPayload,
success: DispatchStreamChunk,
error: DispatchError,
stream: true,
}) {}
export const TrustGraphRpcs = RpcGroup.make(Dispatch, DispatchStream);

View file

@ -0,0 +1,92 @@
import { Effect, Queue, Scope } from "effect";
import * as RpcMessage from "effect/unstable/rpc/RpcMessage";
import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
import * as RpcServer from "effect/unstable/rpc/RpcServer";
import * as Socket from "effect/unstable/socket/Socket";
export const makeSocketRpcProtocol = Effect.gen(function* () {
const serialization = yield* RpcSerialization.RpcSerialization;
const disconnects = yield* Queue.make<number>();
let nextClientId = 0;
const clients = new Map<number, {
readonly write: (response: RpcMessage.FromServerEncoded) => Effect.Effect<void>;
}>();
const clientIds = new Set<number>();
let writeRequest!: (
clientId: number,
message: RpcMessage.FromClientEncoded,
) => Effect.Effect<void>;
const onSocket = function* (
socket: Socket.Socket,
headers?: ReadonlyArray<[string, string]>,
) {
const scope = yield* Effect.scope;
const parser = serialization.makeUnsafe();
const clientId = nextClientId++;
yield* Scope.addFinalizerExit(scope, () => {
clients.delete(clientId);
clientIds.delete(clientId);
return Queue.offer(disconnects, clientId);
});
const writeRaw = yield* socket.writer;
const write = (response: RpcMessage.FromServerEncoded) => {
try {
const encoded = parser.encode(response);
if (encoded === undefined) return Effect.void;
return Effect.orDie(writeRaw(encoded));
} catch (cause) {
return Effect.orDie(
writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!),
);
}
};
clients.set(clientId, { write });
clientIds.add(clientId);
yield* socket.runRaw((data) => {
try {
const decoded = parser.decode(data) as ReadonlyArray<RpcMessage.FromClientEncoded>;
return Effect.forEach(decoded, (message) => {
if (message._tag === "Request" && headers !== undefined) {
return writeRequest(clientId, {
...message,
headers: headers.concat(message.headers),
});
}
return writeRequest(clientId, message);
}, { discard: true });
} catch (cause) {
return writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!);
}
}).pipe(
Effect.catchReason("SocketError", "SocketCloseError", () => Effect.void),
Effect.orDie,
);
};
const protocol = yield* RpcServer.Protocol.make((writeRequest_) => {
writeRequest = writeRequest_;
return Effect.succeed({
disconnects,
send: (clientId, response) => {
const client = clients.get(clientId);
if (client === undefined) return Effect.void;
return Effect.orDie(client.write(response));
},
end: () => Effect.void,
clientIds: Effect.sync(() => clientIds),
initialMessage: Effect.succeedNone,
supportsAck: true,
supportsTransferables: false,
supportsSpanPropagation: true,
});
});
return { onSocket, protocol } as const;
});

View file

@ -0,0 +1,109 @@
import { Cause, Effect, Layer, Queue, Scope } from "effect";
import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
import * as RpcServer from "effect/unstable/rpc/RpcServer";
import type * as Socket from "effect/unstable/socket/Socket";
import { errorMessage } from "@trustgraph/base";
import type { DispatcherManager } from "./dispatch/manager.js";
import { DispatchError, DispatchPayload, DispatchStreamChunk, TrustGraphRpcs } from "./rpc-contract.js";
import { makeSocketRpcProtocol } from "./rpc-protocol.js";
export interface GatewayRpcServer {
readonly onSocket: (
socket: Socket.Socket,
headers?: ReadonlyArray<[string, string]>,
) => Effect.Effect<void, never, Scope.Scope>;
}
export const makeGatewayRpcServer = Effect.fn("makeGatewayRpcServer")(function* (
dispatcher: DispatcherManager,
) {
const { onSocket, protocol } = yield* makeSocketRpcProtocol;
const serverLayer = RpcServer.layer(TrustGraphRpcs, {
disableFatalDefects: true,
}).pipe(
Layer.provide(Layer.succeed(RpcServer.Protocol, protocol)),
Layer.provide(makeGatewayRpcHandlers(dispatcher)),
Layer.provide(RpcSerialization.layerNdjson),
);
yield* Layer.launch(serverLayer).pipe(Effect.forkScoped);
return {
onSocket: Effect.fn("GatewayRpc.onSocket")(function* (socket, headers) {
yield* onSocket(socket, headers);
}),
} satisfies GatewayRpcServer;
});
const makeGatewayRpcHandlers = (dispatcher: DispatcherManager) =>
TrustGraphRpcs.toLayer(Effect.succeed(
TrustGraphRpcs.of({
Dispatch: (payload) =>
Effect.tryPromise({
try: () => dispatchOne(dispatcher, payload),
catch: (cause) => new DispatchError({ message: errorMessage(cause) }),
}),
DispatchStream: Effect.fn("GatewayRpc.DispatchStream")(function* (payload) {
const context = yield* Effect.context<never>();
const runPromise = Effect.runPromiseWith(context);
const queue = yield* Queue.bounded<DispatchStreamChunk, DispatchError | Cause.Done>(16);
yield* Effect.addFinalizer(() => Queue.shutdown(queue));
yield* Effect.tryPromise({
try: () =>
dispatchStream(dispatcher, payload, async (response, complete) => {
await runPromise(Queue.offer(queue, new DispatchStreamChunk({ response, complete })));
return complete;
}),
catch: (cause) => new DispatchError({ message: errorMessage(cause) }),
}).pipe(
Effect.flatMap(() => Queue.end(queue)),
Effect.catch((error) => Queue.fail(queue, error)),
Effect.forkScoped,
);
return queue;
}),
}),
));
function dispatchOne(
dispatcher: DispatcherManager,
payload: DispatchPayload,
): Promise<unknown> {
if (payload.scope === "flow") {
return dispatcher.dispatchFlowService(
payload.flow ?? "default",
payload.service,
payload.request,
);
}
return dispatcher.dispatchGlobalService(payload.service, payload.request);
}
async function dispatchStream(
dispatcher: DispatcherManager,
payload: DispatchPayload,
responder: (response: unknown, complete: boolean) => Promise<boolean>,
): Promise<void> {
const send = async (response: unknown, complete: boolean) => {
await responder(response, complete);
};
if (payload.scope === "flow") {
await dispatcher.dispatchFlowServiceStreaming(
payload.flow ?? "default",
payload.service,
payload.request,
send,
);
return;
}
await dispatcher.dispatchGlobalServiceStreaming(
payload.service,
payload.request,
send,
);
}

View file

@ -2,19 +2,20 @@
* API Gateway HTTP + WebSocket server.
*
* Replaces the Python aiohttp gateway with Fastify.
* Uses the Mux class for WebSocket multiplexing (queue-based request
* buffering, concurrency control, proper task lifecycle).
* Uses Effect RPC over WebSocket for streaming client requests.
*
* Python reference: trustgraph-flow/trustgraph/gateway/service.py
*/
import Fastify from "fastify";
import websocketPlugin from "@fastify/websocket";
import { Config, Effect } from "effect";
import { Config, Effect, Exit, Scope } from "effect";
import * as O from "effect/Option";
import { errorMessage, optionalStringConfig, registry, toTgError } from "@trustgraph/base";
import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
import * as EffectSocket from "effect/unstable/socket/Socket";
import { optionalStringConfig, registry, toTgError } from "@trustgraph/base";
import { DispatcherManager } from "./dispatch/manager.js";
import { Mux, type MuxRequest, type MuxHandler } from "./dispatch/mux.js";
import { makeGatewayRpcServer } from "./rpc-server.js";
export interface GatewayConfig {
port: number;
@ -29,11 +30,18 @@ export async function createGateway(config: GatewayConfig) {
const dispatcher = new DispatcherManager(config);
await dispatcher.start();
const rpcScope = await Effect.runPromise(Scope.make());
const rpcServer = await Effect.runPromise(
makeGatewayRpcServer(dispatcher).pipe(
Effect.provideService(RpcSerialization.RpcSerialization, RpcSerialization.ndjson),
Scope.provide(rpcScope),
),
);
// Authentication middleware
app.addHook("onRequest", async (request, reply) => {
if (request.url === "/api/v1/metrics") return;
if (request.url.startsWith("/api/v1/socket")) return; // Socket auth via query param
if (request.url.startsWith("/api/v1/rpc")) return; // RPC socket auth via query param
if (config.secret !== undefined && config.secret.length > 0) {
const auth = request.headers.authorization;
@ -43,6 +51,38 @@ export async function createGateway(config: GatewayConfig) {
}
});
app.post<{
Body: {
scope?: string;
service?: string;
flow?: string;
request?: Record<string, unknown>;
};
}>("/api/v1/workbench/dispatch", async (request, reply) => {
const body = request.body;
const service = body.service;
const payload = body.request;
if (service === undefined || service.length === 0 || payload === undefined) {
return reply.code(400).send({
error: { type: "bad-request", message: "service and request are required" },
});
}
try {
const result = body.scope === "flow"
? await dispatcher.dispatchFlowService(body.flow ?? "default", service, payload)
: await dispatcher.dispatchGlobalService(service, payload);
const err = (result as Record<string, unknown>)?.error as { type?: string; message?: string } | undefined;
if (err !== undefined) {
const statusCode = err.type === "not-found" ? 404 : 400;
return reply.code(statusCode).send(result);
}
return result;
} catch (err) {
reply.code(500).send({ error: toTgError(err) });
}
});
// REST endpoint: POST /api/v1/:kind (global services)
app.post<{ Params: { kind: string } }>("/api/v1/:kind", async (request, reply) => {
const { kind } = request.params;
@ -124,10 +164,8 @@ export async function createGateway(config: GatewayConfig) {
},
);
// WebSocket endpoint: /api/v1/socket
// Uses Mux for queue-based request buffering and concurrency control.
app.get("/api/v1/socket", { websocket: true }, (socket, request) => {
// Auth via query param
// Effect RPC WebSocket endpoint: /api/v1/rpc
app.get("/api/v1/rpc", { websocket: true }, (socket, request) => {
const url = new URL(request.url, `http://${request.headers.host}`);
const token = url.searchParams.get("token");
if (config.secret !== undefined && config.secret.length > 0 && token !== config.secret) {
@ -135,91 +173,18 @@ export async function createGateway(config: GatewayConfig) {
return;
}
// Build the MuxHandler that dispatches to the DispatcherManager
const handler: MuxHandler = async (muxReq, respond) => {
if (muxReq.flow !== undefined && muxReq.flow.length > 0) {
await dispatcher.dispatchFlowServiceStreaming(
muxReq.flow,
muxReq.service,
muxReq.request,
respond,
const program = Effect.scoped(
Effect.gen(function* () {
const effectSocket = yield* EffectSocket.fromWebSocket(
Effect.succeed(socket as unknown as globalThis.WebSocket),
{ closeCodeIsError: (code) => code !== 1000 },
);
} else {
await dispatcher.dispatchGlobalServiceStreaming(
muxReq.service,
muxReq.request,
respond,
);
}
};
yield* rpcServer.onSocket(effectSocket, headersFrom(request.headers));
}),
);
const mux = new Mux(handler);
// Start the Mux run loop — sends responses back over the socket
const runPromise = mux.run((data) => {
// Only send if the socket is still open (readyState 1 = OPEN)
if (socket.readyState === 1) {
socket.send(data);
}
});
// Incoming messages get queued into the Mux
socket.on("message", (data) => {
try {
const msg = JSON.parse(data.toString()) as {
id?: string;
service?: string;
flow?: string;
request?: Record<string, unknown>;
};
if (
msg.id === undefined ||
msg.id.length === 0 ||
msg.service === undefined ||
msg.service.length === 0 ||
msg.request === undefined
) {
socket.send(
JSON.stringify({
id: msg.id ?? null,
error: { type: "bad-request", message: "Missing id, service, or request" },
complete: true,
}),
);
return;
}
const muxReq: MuxRequest = {
id: msg.id,
service: msg.service,
request: msg.request,
...(msg.flow !== undefined ? { flow: msg.flow } : {}),
};
mux.receive(muxReq);
} catch (err) {
socket.send(
JSON.stringify({
error: { type: "parse-error", message: errorMessage(err) },
complete: true,
}),
);
}
});
socket.on("close", () => {
mux.stop();
});
socket.on("error", () => {
mux.stop();
});
// Ensure runPromise errors don't go unhandled
runPromise.catch((err) => {
console.error("[Gateway] Mux run loop error:", err);
mux.stop();
Effect.runPromise(program.pipe(Scope.provide(rpcScope))).catch((err) => {
console.error("[Gateway] RPC WebSocket error:", err);
if (socket.readyState === 1) {
socket.close(1011, "Internal server error");
}
@ -236,11 +201,21 @@ export async function createGateway(config: GatewayConfig) {
start: () => app.listen({ port: config.port, host: "0.0.0.0" }),
stop: async () => {
await app.close();
await Effect.runPromise(Scope.close(rpcScope, Exit.void));
await dispatcher.stop();
},
};
}
function headersFrom(headers: Record<string, string | string[] | number | undefined>): ReadonlyArray<[string, string]> {
return Object.entries(headers).flatMap(([key, value]) => {
if (typeof value === "string") return [[key, value] satisfies [string, string]];
if (typeof value === "number") return [[key, String(value)] satisfies [string, string]];
if (Array.isArray(value)) return value.map((item) => [key, item] satisfies [string, string]);
return [];
});
}
export async function run(): Promise<void> {
await Effect.runPromise(program);
}

View file

@ -3,8 +3,28 @@
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";
export {
GraphRag,
GraphRagEngine,
GraphRagLive,
makeGraphRagEngine,
normalizeGraphRagConfig,
stringToTerm,
termToString,
type GraphRagConfig,
type GraphRagClients,
type GraphRagEngineShape,
type GraphRagQueryOptions,
} from "./retrieval/graph-rag.js";
export {
DocumentRag,
DocumentRagEngine,
DocumentRagLive,
makeDocumentRagEngine,
type DocumentRagClients,
type DocumentRagEngineShape,
type DocumentRagQueryOptions,
} from "./retrieval/document-rag.js";
export { FalkorDBTriplesStore, type FalkorDBConfig } from "./storage/triples/falkordb.js";
export { FalkorDBTriplesQuery, type FalkorDBQueryConfig } from "./query/triples/falkordb.js";

View file

@ -23,6 +23,7 @@ import {
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import type { BackendProducer, BackendConsumer, Message } from "@trustgraph/base";
import { Effect } from "effect";
import { CollectionManager } from "./collection-manager.js";
import {
ensureDirectory,
@ -38,9 +39,29 @@ export interface LibrarianServiceConfig extends ProcessorConfig {
dataDir?: string;
}
interface UploadSession {
id: string;
documentMetadata: DocumentMetadata;
totalSize: number;
chunkSize: number;
totalChunks: number;
createdAt: string;
chunks: Map<number, string>;
user: string;
}
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
function optionalString(value: unknown): string | undefined {
return typeof value === "string" && value.length > 0 ? value : undefined;
}
export class LibrarianService extends AsyncProcessor {
private documents = new Map<string, DocumentMetadata>();
private processing = new Map<string, ProcessingMetadata>();
private uploads = new Map<string, UploadSession>();
private collectionManager = new CollectionManager();
private readonly dataDir: string;
private readonly persistPath: string;
@ -112,6 +133,107 @@ export class LibrarianService extends AsyncProcessor {
// ---------- Librarian message handling ----------
private requestRecord(request: LibrarianRequest): Record<string, unknown> {
return request as Record<string, unknown>;
}
private documentId(request: LibrarianRequest): string | undefined {
const req = this.requestRecord(request);
return optionalString(req.documentId) ?? optionalString(req["document-id"]);
}
private processingId(request: LibrarianRequest): string | undefined {
const req = this.requestRecord(request);
return optionalString(req.processingId) ?? optionalString(req["processing-id"]);
}
private documentMetadata(request: LibrarianRequest): DocumentMetadata | undefined {
const req = this.requestRecord(request);
const value = req.documentMetadata ?? req["document-metadata"];
return isRecord(value) ? this.normaliseDocumentMetadata(value) : undefined;
}
private processingMetadata(request: LibrarianRequest): ProcessingMetadata | undefined {
const req = this.requestRecord(request);
const value = req.processingMetadata ?? req["processing-metadata"];
if (!isRecord(value)) return undefined;
const documentId = optionalString(value.documentId) ?? optionalString(value["document-id"]) ?? "";
return {
id: optionalString(value.id) ?? crypto.randomUUID(),
documentId,
"document-id": documentId,
time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000),
flow: optionalString(value.flow) ?? "default",
user: optionalString(value.user) ?? optionalString(this.requestRecord(request).user) ?? "default",
collection: optionalString(value.collection) ?? optionalString(this.requestRecord(request).collection) ?? "default",
tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [],
};
}
private normaliseDocumentMetadata(value: Record<string, unknown>): DocumentMetadata {
const id = optionalString(value.id) ?? crypto.randomUUID();
const parentId = optionalString(value.parentId) ?? optionalString(value["parent-id"]);
const documentType = optionalString(value.documentType) ?? optionalString(value["document-type"]) ?? "source";
return {
id,
time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000),
kind: optionalString(value.kind) ?? "application/octet-stream",
title: optionalString(value.title) ?? "",
comments: optionalString(value.comments) ?? "",
user: optionalString(value.user) ?? "default",
tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [],
...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}),
documentType,
"document-type": documentType,
...(Array.isArray(value.metadata) ? { metadata: value.metadata as NonNullable<DocumentMetadata["metadata"]> } : {}),
};
}
private publicDocument(doc: DocumentMetadata): DocumentMetadata {
const parentId = doc.parentId ?? doc["parent-id"];
const documentType = doc.documentType ?? doc["document-type"] ?? "source";
return {
...doc,
...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}),
documentType,
"document-type": documentType,
};
}
private publicProcessing(proc: ProcessingMetadata): ProcessingMetadata {
const documentId = proc.documentId ?? proc["document-id"] ?? "";
return {
...proc,
documentId,
"document-id": documentId,
};
}
private documentResponse(doc: DocumentMetadata): LibrarianResponse {
const publicDoc = this.publicDocument(doc);
return {
documentMetadata: publicDoc,
"document-metadata": publicDoc,
};
}
private documentsResponse(docs: DocumentMetadata[]): LibrarianResponse {
const publicDocs = docs.map((doc) => this.publicDocument(doc));
return {
documents: publicDocs,
"document-metadatas": publicDocs,
};
}
private processingResponse(records: ProcessingMetadata[]): LibrarianResponse {
const publicRecords = records.map((proc) => this.publicProcessing(proc));
return {
processing: publicRecords,
"processing-metadata": publicRecords,
"processing-metadatas": publicRecords,
};
}
private async handleLibrarianMessage(msg: Message<LibrarianRequest>): Promise<void> {
const request = msg.value();
const props = msg.properties();
@ -123,6 +245,12 @@ export class LibrarianService extends AsyncProcessor {
}
try {
if (request.operation === "stream-document") {
for (const response of await this.streamDocument(request)) {
await this.libProducer!.send(response, { id: requestId });
}
return;
}
const response = await this.handleLibrarianOperation(request);
await this.libProducer!.send(response, { id: requestId });
} catch (err) {
@ -140,6 +268,8 @@ export class LibrarianService extends AsyncProcessor {
return this.addDocument(request);
case "remove-document":
return this.removeDocument(request);
case "update-document":
return this.updateDocument(request);
case "list-documents":
return this.listDocuments(request);
case "get-document-metadata":
@ -156,17 +286,31 @@ export class LibrarianService extends AsyncProcessor {
return this.removeProcessing(request);
case "list-processing":
return this.listProcessing(request);
case "begin-upload":
return this.beginUpload(request);
case "upload-chunk":
return this.uploadChunk(request);
case "complete-upload":
return this.completeUpload(request);
case "get-upload-status":
return this.getUploadStatus(request);
case "abort-upload":
return this.abortUpload(request);
case "list-uploads":
return this.listUploads(request);
case "stream-document":
throw new Error("stream-document must be handled as a streaming operation");
default:
throw new Error(`Unknown librarian operation: ${request.operation as string}`);
}
}
private async addDocument(request: LibrarianRequest): Promise<LibrarianResponse> {
const meta = request.documentMetadata;
const meta = this.documentMetadata(request);
if (meta === undefined) throw new Error("add-document requires documentMetadata");
const id = crypto.randomUUID();
const now = Date.now();
const id = meta.id;
const now = Math.floor(Date.now() / 1000);
const doc: DocumentMetadata = {
...meta,
@ -186,11 +330,11 @@ export class LibrarianService extends AsyncProcessor {
await this.persist();
console.log(`[LibrarianService] Added document ${id}: ${doc.title}`);
return { documentMetadata: doc };
return this.documentResponse(doc);
}
private async removeDocument(request: LibrarianRequest): Promise<LibrarianResponse> {
const id = request.documentId;
const id = this.documentId(request);
if (id === undefined || id.length === 0) {
throw new Error("remove-document requires documentId");
}
@ -234,23 +378,45 @@ export class LibrarianService extends AsyncProcessor {
return {};
}
private async updateDocument(request: LibrarianRequest): Promise<LibrarianResponse> {
const id = this.documentId(request) ?? this.documentMetadata(request)?.id;
if (id === undefined || id.length === 0) {
throw new Error("update-document requires documentId");
}
const existing = this.documents.get(id);
if (existing === undefined) throw new Error(`Document not found: ${id}`);
const meta = this.documentMetadata(request);
if (meta === undefined) throw new Error("update-document requires documentMetadata");
const doc: DocumentMetadata = this.publicDocument({
...existing,
...meta,
id,
time: meta.time ?? existing.time,
});
this.documents.set(id, doc);
await this.persist();
return this.documentResponse(doc);
}
private listDocuments(request: LibrarianRequest): LibrarianResponse {
const user = request.user ?? "";
const includeChildren = this.requestRecord(request)["include-children"] === true;
const docs: DocumentMetadata[] = [];
for (const doc of this.documents.values()) {
// Filter by user
if (user.length > 0 && doc.user !== user) continue;
// Exclude children (only top-level documents) unless explicitly requested
if (doc.parentId !== undefined && doc.parentId.length > 0) continue;
if (!includeChildren && doc.parentId !== undefined && doc.parentId.length > 0) continue;
docs.push(doc);
}
return { documents: docs };
return this.documentsResponse(docs);
}
private getDocumentMetadata(request: LibrarianRequest): LibrarianResponse {
const id = request.documentId;
const id = this.documentId(request);
if (id === undefined || id.length === 0) {
throw new Error("get-document-metadata requires documentId");
}
@ -258,11 +424,11 @@ export class LibrarianService extends AsyncProcessor {
const doc = this.documents.get(id);
if (doc === undefined) throw new Error(`Document not found: ${id}`);
return { documentMetadata: doc };
return this.documentResponse(doc);
}
private async getDocumentContent(request: LibrarianRequest): Promise<LibrarianResponse> {
const id = request.documentId;
const id = this.documentId(request);
if (id === undefined || id.length === 0) {
throw new Error("get-document-content requires documentId");
}
@ -274,14 +440,14 @@ export class LibrarianService extends AsyncProcessor {
const filePath = joinPath(this.dataDir, "docs", `${id}.bin`);
const buf = await readBinaryFile(filePath);
const content = Buffer.from(buf).toString("base64");
return { documentMetadata: doc, content };
return { ...this.documentResponse(doc), content };
} catch {
throw new Error(`Document content not found on disk: ${id}`);
}
}
private async addChildDocument(request: LibrarianRequest): Promise<LibrarianResponse> {
const meta = request.documentMetadata;
const meta = this.documentMetadata(request);
if (meta === undefined) {
throw new Error("add-child-document requires documentMetadata");
}
@ -294,8 +460,8 @@ export class LibrarianService extends AsyncProcessor {
throw new Error(`Parent document not found: ${meta.parentId}`);
}
const id = crypto.randomUUID();
const now = Date.now();
const id = meta.id;
const now = Math.floor(Date.now() / 1000);
const doc: DocumentMetadata = {
...meta,
@ -315,11 +481,11 @@ export class LibrarianService extends AsyncProcessor {
await this.persist();
console.log(`[LibrarianService] Added child document ${id} (parent: ${meta.parentId})`);
return { documentMetadata: doc };
return this.documentResponse(doc);
}
private listChildren(request: LibrarianRequest): LibrarianResponse {
const parentId = request.documentId;
const parentId = this.documentId(request);
if (parentId === undefined || parentId.length === 0) {
throw new Error("list-children requires documentId");
}
@ -331,15 +497,15 @@ export class LibrarianService extends AsyncProcessor {
}
}
return { documents: children };
return this.documentsResponse(children);
}
private async addProcessing(request: LibrarianRequest): Promise<LibrarianResponse> {
const proc = request.processingMetadata;
const proc = this.processingMetadata(request);
if (proc === undefined) throw new Error("add-processing requires processingMetadata");
const id = crypto.randomUUID();
const now = Date.now();
const id = proc.id;
const now = Math.floor(Date.now() / 1000);
const record: ProcessingMetadata = {
...proc,
@ -351,11 +517,11 @@ export class LibrarianService extends AsyncProcessor {
await this.persist();
console.log(`[LibrarianService] Added processing ${id} for document ${proc.documentId}`);
return { processing: [record] };
return this.processingResponse([record]);
}
private async removeProcessing(request: LibrarianRequest): Promise<LibrarianResponse> {
const id = request.processingId;
const id = this.processingId(request);
if (id === undefined || id.length === 0) {
throw new Error("remove-processing requires processingId");
}
@ -367,17 +533,167 @@ export class LibrarianService extends AsyncProcessor {
}
private listProcessing(request: LibrarianRequest): LibrarianResponse {
const documentId = request.documentId;
const documentId = this.documentId(request);
const records: ProcessingMetadata[] = [];
for (const proc of this.processing.values()) {
if (documentId !== undefined && documentId.length > 0 && proc.documentId !== documentId) {
const procDocumentId = proc.documentId ?? proc["document-id"];
if (documentId !== undefined && documentId.length > 0 && procDocumentId !== documentId) {
continue;
}
records.push(proc);
}
return { processing: records };
return this.processingResponse(records);
}
private beginUpload(request: LibrarianRequest): LibrarianResponse {
const meta = this.documentMetadata(request);
if (meta === undefined) throw new Error("begin-upload requires documentMetadata");
const req = this.requestRecord(request);
const totalSize = typeof req["total-size"] === "number" ? req["total-size"] : 0;
if (totalSize <= 0) throw new Error("begin-upload requires total-size");
const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0
? req["chunk-size"]
: 3 * 1024 * 1024;
const totalChunks = Math.max(1, Math.ceil(totalSize / chunkSize));
const uploadId = crypto.randomUUID();
this.uploads.set(uploadId, {
id: uploadId,
documentMetadata: meta,
totalSize,
chunkSize,
totalChunks,
createdAt: new Date().toISOString(),
chunks: new Map<number, string>(),
user: meta.user ?? optionalString(req.user) ?? "default",
});
return {
"upload-id": uploadId,
"chunk-size": chunkSize,
"total-chunks": totalChunks,
} as LibrarianResponse;
}
private uploadChunk(request: LibrarianRequest): LibrarianResponse {
const req = this.requestRecord(request);
const uploadId = optionalString(req["upload-id"]);
if (uploadId === undefined) throw new Error("upload-chunk requires upload-id");
const session = this.uploads.get(uploadId);
if (session === undefined) throw new Error(`Upload not found: ${uploadId}`);
const chunkIndex = typeof req["chunk-index"] === "number" ? req["chunk-index"] : -1;
if (!Number.isInteger(chunkIndex) || chunkIndex < 0 || chunkIndex >= session.totalChunks) {
throw new Error("upload-chunk requires a valid chunk-index");
}
const content = optionalString(req.content);
if (content === undefined) throw new Error("upload-chunk requires content");
session.chunks.set(chunkIndex, content);
const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0);
return {
"upload-id": uploadId,
"chunk-index": chunkIndex,
"chunks-received": session.chunks.size,
"total-chunks": session.totalChunks,
"bytes-received": bytesReceived,
"total-bytes": session.totalSize,
} as LibrarianResponse;
}
private async completeUpload(request: LibrarianRequest): Promise<LibrarianResponse> {
const uploadId = optionalString(this.requestRecord(request)["upload-id"]);
if (uploadId === undefined) throw new Error("complete-upload requires upload-id");
const session = this.uploads.get(uploadId);
if (session === undefined) throw new Error(`Upload not found: ${uploadId}`);
if (session.chunks.size !== session.totalChunks) {
throw new Error(`Upload incomplete: ${session.chunks.size}/${session.totalChunks} chunks received`);
}
const content = Array.from({ length: session.totalChunks }, (_, i) => session.chunks.get(i) ?? "").join("");
const response = await this.addDocument({
operation: "add-document",
documentMetadata: session.documentMetadata,
"document-metadata": session.documentMetadata,
content,
user: session.user,
} as LibrarianRequest);
this.uploads.delete(uploadId);
const documentId = response.documentMetadata?.id ?? response["document-metadata"]?.id ?? session.documentMetadata.id;
return {
...response,
"document-id": documentId,
"object-id": documentId,
} as LibrarianResponse;
}
private getUploadStatus(request: LibrarianRequest): LibrarianResponse {
const uploadId = optionalString(this.requestRecord(request)["upload-id"]);
if (uploadId === undefined) throw new Error("get-upload-status requires upload-id");
const session = this.uploads.get(uploadId);
if (session === undefined) throw new Error(`Upload not found: ${uploadId}`);
const receivedChunks = [...session.chunks.keys()].sort((a, b) => a - b);
const receivedSet = new Set(receivedChunks);
const missingChunks = Array.from({ length: session.totalChunks }, (_, i) => i).filter((i) => !receivedSet.has(i));
const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0);
return {
"upload-id": uploadId,
"upload-state": "in-progress",
"chunks-received": session.chunks.size,
"total-chunks": session.totalChunks,
"received-chunks": receivedChunks,
"missing-chunks": missingChunks,
"bytes-received": bytesReceived,
"total-bytes": session.totalSize,
} as LibrarianResponse;
}
private abortUpload(request: LibrarianRequest): LibrarianResponse {
const uploadId = optionalString(this.requestRecord(request)["upload-id"]);
if (uploadId === undefined) throw new Error("abort-upload requires upload-id");
this.uploads.delete(uploadId);
return {};
}
private listUploads(request: LibrarianRequest): LibrarianResponse {
const user = optionalString(this.requestRecord(request).user);
const sessions = [...this.uploads.values()]
.filter((session) => user === undefined || session.user === user)
.map((session) => ({
"upload-id": session.id,
"document-id": session.documentMetadata.id,
"document-metadata-json": JSON.stringify(this.publicDocument(session.documentMetadata)),
"total-size": session.totalSize,
"chunk-size": session.chunkSize,
"total-chunks": session.totalChunks,
"chunks-received": session.chunks.size,
"created-at": session.createdAt,
}));
return { "upload-sessions": sessions } as LibrarianResponse;
}
private async streamDocument(request: LibrarianRequest): Promise<LibrarianResponse[]> {
const id = this.documentId(request);
if (id === undefined) throw new Error("stream-document requires documentId");
const req = this.requestRecord(request);
const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0
? req["chunk-size"]
: 1024 * 1024;
const filePath = joinPath(this.dataDir, "docs", `${id}.bin`);
const buf = await readBinaryFile(filePath);
const base64 = Buffer.from(buf).toString("base64");
const totalChunks = Math.max(1, Math.ceil(base64.length / chunkSize));
return Array.from({ length: totalChunks }, (_, index) => {
const start = index * chunkSize;
const content = base64.slice(start, start + chunkSize);
return {
content,
"chunk-index": index,
"total-chunks": totalChunks,
eos: index === totalChunks - 1,
} as LibrarianResponse;
});
}
// ---------- Collection management ----------
@ -471,14 +787,14 @@ export class LibrarianService extends AsyncProcessor {
this.documents.clear();
if (parsed.documents !== undefined) {
for (const [id, doc] of Object.entries(parsed.documents)) {
this.documents.set(id, doc);
this.documents.set(id, this.publicDocument(doc));
}
}
this.processing.clear();
if (parsed.processing !== undefined) {
for (const [id, proc] of Object.entries(parsed.processing)) {
this.processing.set(id, proc);
this.processing.set(id, this.publicProcessing(proc));
}
}
@ -525,5 +841,5 @@ export const program = makeProcessorProgram({
});
export async function run(): Promise<void> {
await LibrarianService.launch("librarian-svc");
await Effect.runPromise(program);
}

View file

@ -10,13 +10,17 @@
import { AzureOpenAI } from "openai";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class AzureOpenAIProcessor extends LlmService {
private client: AzureOpenAI;
@ -157,11 +161,16 @@ export class AzureOpenAIProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new AzureOpenAIProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new AzureOpenAIProcessor(config))),
),
});
export async function run(): Promise<void> {
await AzureOpenAIProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -5,8 +5,18 @@
*/
import Anthropic from "@anthropic-ai/sdk";
import { LlmService, type ProcessorConfig, type LlmResult, type LlmChunk, tooManyRequestsError } from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class ClaudeProcessor extends LlmService {
private client: Anthropic;
@ -127,11 +137,16 @@ export class ClaudeProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new ClaudeProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new ClaudeProcessor(config))),
),
});
export async function run(): Promise<void> {
await ClaudeProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -8,13 +8,17 @@
import { Mistral } from "@mistralai/mistralai";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class MistralProcessor extends LlmService {
private client: Mistral;
@ -143,11 +147,16 @@ export class MistralProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new MistralProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new MistralProcessor(config))),
),
});
export async function run(): Promise<void> {
await MistralProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -7,8 +7,17 @@
*/
import { Ollama } from "ollama";
import { LlmService, type ProcessorConfig, type LlmResult, type LlmChunk } from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class OllamaProcessor extends LlmService {
private client: Ollama;
@ -113,11 +122,16 @@ export class OllamaProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new OllamaProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new OllamaProcessor(config))),
),
});
export async function run(): Promise<void> {
await OllamaProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -11,12 +11,16 @@
import OpenAI from "openai";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class OpenAICompatibleProcessor extends LlmService {
private client: OpenAI;
@ -137,11 +141,16 @@ export class OpenAICompatibleProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new OpenAICompatibleProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new OpenAICompatibleProcessor(config))),
),
});
export async function run(): Promise<void> {
await OpenAICompatibleProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -5,8 +5,18 @@
*/
import OpenAI from "openai";
import { LlmService, type ProcessorConfig, type LlmResult, type LlmChunk, tooManyRequestsError } from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import {
Llm,
LlmService,
makeFlowProcessorProgram,
makeLlmServiceShape,
makeLlmSpecs,
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
export class OpenAIProcessor extends LlmService {
private client: OpenAI;
@ -137,11 +147,16 @@ export class OpenAIProcessor extends LlmService {
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
id: "text-completion",
make: (config) => new OpenAIProcessor(config),
specs: () => makeLlmSpecs(),
layer: (config) =>
Layer.succeed(
Llm,
Llm.of(makeLlmServiceShape(new OpenAIProcessor(config))),
),
});
export async function run(): Promise<void> {
await OpenAIProcessor.launch("text-completion");
await Effect.runPromise(program);
}

View file

@ -29,11 +29,17 @@ import {
ConsumerSpec,
ProducerSpec,
type ProcessorConfig,
type EffectConfigHandler,
type FlowContext,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type PromptRequest,
type PromptResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import * as S from "effect/Schema";
export interface PromptTemplate {
system: string;
@ -44,94 +50,129 @@ export interface PromptTemplateConfig extends ProcessorConfig {
configKey?: string;
}
const PromptTemplateEntry = S.Struct({
system: S.optionalKey(S.String),
prompt: S.optionalKey(S.String),
});
const PromptTemplateEntries = S.Record(S.String, PromptTemplateEntry);
interface PromptTemplateRuntime {
readonly specs: ReadonlyArray<Spec<never>>;
readonly configHandlers: ReadonlyArray<EffectConfigHandler>;
}
const programRuntimes = new WeakMap<PromptTemplateConfig, PromptTemplateRuntime>();
const makePromptTemplateRuntime = (config: PromptTemplateConfig): PromptTemplateRuntime => {
const templates = new Map<string, PromptTemplate>();
const configKey = config.configKey ?? "prompt";
const onPromptConfig = Effect.fn("PromptTemplateService.onConfig")(function* (
pushedConfig: Record<string, unknown>,
version: number,
) {
yield* Effect.log(`[PromptTemplate] Loading prompt configuration version ${version}`);
const promptConfig = pushedConfig[configKey];
if (promptConfig === undefined) {
yield* Effect.logWarning(`[PromptTemplate] No key "${configKey}" in config`);
return;
}
const decoded = yield* S.decodeUnknownEffect(PromptTemplateEntries)(promptConfig).pipe(
Effect.catch((error) =>
Effect.logError("[PromptTemplate] Failed to decode prompt configuration", {
error: error.message,
configKey,
}).pipe(Effect.as(null)),
),
);
if (decoded === null) return;
templates.clear();
for (const [name, template] of Object.entries(decoded)) {
templates.set(name, {
system: template.system ?? "",
prompt: template.prompt ?? "",
});
}
yield* Effect.log(
`[PromptTemplate] Loaded ${templates.size} template(s): ${[...templates.keys()].join(", ")}`,
);
});
const onRequest = Effect.fn("PromptTemplateService.onRequest")(function* (
msg: PromptRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = yield* flowCtx.flow.producerEffect<PromptResponse>("prompt-response");
const template = templates.get(msg.name);
if (template === undefined) {
yield* responseProducer.send(requestId, {
system: "",
prompt: "",
error: {
type: "prompt-error",
message: `Unknown prompt template: "${msg.name}"`,
},
});
return;
}
const variables = msg.variables ?? {};
yield* responseProducer.send(requestId, {
system: renderTemplate(template.system, variables),
prompt: renderTemplate(template.prompt, variables),
});
});
return {
specs: [
new ConsumerSpec<PromptRequest, FlowResourceNotFoundError | MessagingDeliveryError>(
"prompt-request",
onRequest,
),
new ProducerSpec<PromptResponse>("prompt-response"),
],
configHandlers: [onPromptConfig],
};
};
const promptTemplateRuntime = (config: PromptTemplateConfig): PromptTemplateRuntime => {
const existing = programRuntimes.get(config);
if (existing !== undefined) return existing;
const runtime = makePromptTemplateRuntime(config);
programRuntimes.set(config, runtime);
return runtime;
};
export class PromptTemplateService extends FlowProcessor {
private templates = new Map<string, PromptTemplate>();
private readonly configKey: string;
private readonly runtime: PromptTemplateRuntime;
constructor(config: PromptTemplateConfig) {
super(config);
this.configKey = config.configKey ?? "prompt";
this.runtime = makePromptTemplateRuntime(config);
this.registerSpecification(
ConsumerSpec.fromPromise<PromptRequest>(
"prompt-request",
this.onRequest.bind(this),
),
);
this.registerSpecification(new ProducerSpec<PromptResponse>("prompt-response"));
this.registerConfigHandler(this.onPromptConfig.bind(this));
for (const spec of this.runtime.specs) {
this.registerSpecification(spec);
}
for (const handler of this.runtime.configHandlers) {
this.registerConfigHandler((pushedConfig, version) =>
Effect.runPromise(handler(pushedConfig, version)),
);
}
console.log("[PromptTemplate] Service initialized");
}
private async onPromptConfig(
config: Record<string, unknown>,
version: number,
): Promise<void> {
console.log(`[PromptTemplate] Loading prompt configuration version ${version}`);
const promptConfig = config[this.configKey] as
| Record<string, { system?: string; prompt?: string }>
| undefined;
if (promptConfig === undefined) {
console.warn(`[PromptTemplate] No key "${this.configKey}" in config`);
return;
}
try {
this.templates.clear();
for (const [name, template] of Object.entries(promptConfig)) {
this.templates.set(name, {
system: template.system ?? "",
prompt: template.prompt ?? "",
});
}
console.log(
`[PromptTemplate] Loaded ${this.templates.size} template(s): ${[...this.templates.keys()].join(", ")}`,
);
} catch (err) {
console.error("[PromptTemplate] Failed to load prompt configuration:", err);
}
}
private async onRequest(
msg: PromptRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const responseProducer = flowCtx.flow.producer<PromptResponse>("prompt-response");
try {
const template = this.templates.get(msg.name);
if (template === undefined) {
throw new Error(`Unknown prompt template: "${msg.name}"`);
}
const variables = msg.variables ?? {};
const system = renderTemplate(template.system, variables);
const prompt = renderTemplate(template.prompt, variables);
await responseProducer.send(requestId, { system, prompt });
} catch (err) {
console.error(`[PromptTemplate] Error processing request:`, err);
const message = err instanceof Error ? err.message : String(err);
await responseProducer.send(requestId, {
system: "",
prompt: "",
error: { type: "prompt-error", message },
});
}
}
}
/**
@ -150,11 +191,12 @@ function renderTemplate(
});
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram({
id: "prompt",
make: (config) => new PromptTemplateService(config),
specs: (config: PromptTemplateConfig) => promptTemplateRuntime(config).specs,
configHandlers: (config: PromptTemplateConfig) => promptTemplateRuntime(config).configHandlers,
});
export async function run(): Promise<void> {
await PromptTemplateService.launch("prompt");
await Effect.runPromise(program);
}

View file

@ -13,79 +13,108 @@ import {
ProducerSpec,
type ProcessorConfig,
type FlowContext,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type DocumentEmbeddingsRequest,
type DocumentEmbeddingsResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { QdrantDocEmbeddingsQuery } from "./qdrant-doc.js";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import {
QdrantDocEmbeddingsQueryLive,
QdrantDocEmbeddingsQueryService,
makeQdrantDocEmbeddingsQueryService,
type QdrantDocQueryConfig,
} from "./qdrant-doc.js";
export class DocEmbeddingsQueryService extends FlowProcessor {
private query: QdrantDocEmbeddingsQuery;
const onDocEmbeddingsQueryMessage = Effect.fn("DocEmbeddingsQueryService.onMessage")(function* (
msg: DocumentEmbeddingsRequest,
properties: Record<string, string>,
flowCtx: FlowContext<QdrantDocEmbeddingsQueryService>,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<DocumentEmbeddingsResponse>("document-embeddings-response");
const query = yield* QdrantDocEmbeddingsQueryService;
const collection = msg.collection ?? "default";
const allChunks: DocumentEmbeddingsResponse["chunks"] = [];
for (const vector of msg.vectors ?? []) {
const matches = yield* query.query({
vector,
user: msg.user ?? "default",
collection,
limit: msg.limit ?? 10,
}).pipe(
Effect.catch((error) =>
Effect.logError("[DocEmbeddingsQuery] Query failed", {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
producer.send(requestId, {
chunks: [],
error: { type: "query-error", message: error.message },
})
),
Effect.as(null),
),
),
);
if (matches === null) return;
for (const match of matches) {
allChunks.push({
chunkId: match.chunkId,
score: match.score,
...(match.content !== undefined ? { content: match.content } : {}),
});
}
}
yield* producer.send(requestId, { chunks: allChunks });
});
export const makeDocEmbeddingsQuerySpecs = (): ReadonlyArray<Spec<QdrantDocEmbeddingsQueryService>> => [
new ConsumerSpec<
DocumentEmbeddingsRequest,
FlowResourceNotFoundError | MessagingDeliveryError,
QdrantDocEmbeddingsQueryService
>("document-embeddings-request", onDocEmbeddingsQueryMessage),
new ProducerSpec<DocumentEmbeddingsResponse>("document-embeddings-response"),
];
export class DocEmbeddingsQueryService extends FlowProcessor<QdrantDocEmbeddingsQueryService> {
private readonly query = makeQdrantDocEmbeddingsQueryService();
constructor(config: ProcessorConfig) {
super(config);
this.query = new QdrantDocEmbeddingsQuery();
this.registerSpecification(
ConsumerSpec.fromPromise<DocumentEmbeddingsRequest>(
"document-embeddings-request",
this.onMessage.bind(this),
),
);
this.registerSpecification(
new ProducerSpec<DocumentEmbeddingsResponse>("document-embeddings-response"),
);
for (const spec of makeDocEmbeddingsQuerySpecs()) {
this.registerSpecification(spec);
}
console.log("[DocEmbeddingsQuery] Service initialized");
}
private async onMessage(
msg: DocumentEmbeddingsRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = flowCtx.flow.producer<DocumentEmbeddingsResponse>("document-embeddings-response");
const collection = msg.collection ?? "default";
try {
const allChunks: DocumentEmbeddingsResponse["chunks"] = [];
for (const vector of msg.vectors ?? []) {
const matches = await this.query.query({
vector,
user: msg.user ?? "default",
collection,
limit: msg.limit ?? 10,
});
for (const match of matches) {
allChunks.push({
chunkId: match.chunkId,
score: match.score,
...(match.content !== undefined ? { content: match.content } : {}),
});
}
}
await producer.send(requestId, { chunks: allChunks });
} catch (err) {
console.error("[DocEmbeddingsQuery] Query failed:", err);
await producer.send(requestId, {
chunks: [],
error: { type: "query-error", message: String(err) },
});
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(
QdrantDocEmbeddingsQueryService,
QdrantDocEmbeddingsQueryService.of(this.query),
),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig & QdrantDocQueryConfig, never, QdrantDocEmbeddingsQueryService>({
id: "doc-embeddings-query",
make: (config) => new DocEmbeddingsQueryService(config),
specs: () => makeDocEmbeddingsQuerySpecs(),
layer: (config) => QdrantDocEmbeddingsQueryLive(config),
});
export async function run(): Promise<void> {
await DocEmbeddingsQueryService.launch("doc-embeddings-query");
await Effect.runPromise(program);
}

View file

@ -8,6 +8,9 @@
*/
import { QdrantClient } from "@qdrant/js-client-rest";
import { errorMessage } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface QdrantDocQueryConfig {
url?: string;
@ -83,3 +86,54 @@ export class QdrantDocEmbeddingsQuery {
return chunks;
}
}
export class QdrantDocEmbeddingsQueryError extends S.TaggedErrorClass<QdrantDocEmbeddingsQueryError>()(
"QdrantDocEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantDocEmbeddingsQueryServiceShape {
readonly query: (
request: DocEmbeddingsQueryRequest,
) => Effect.Effect<ReadonlyArray<ChunkMatch>, QdrantDocEmbeddingsQueryError>;
}
export class QdrantDocEmbeddingsQueryService extends Context.Service<
QdrantDocEmbeddingsQueryService,
QdrantDocEmbeddingsQueryServiceShape
>()(
"@trustgraph/flow/query/embeddings/qdrant-doc/QdrantDocEmbeddingsQueryService",
) {}
const qdrantDocEmbeddingsQueryError = (operation: string, cause: unknown) =>
new QdrantDocEmbeddingsQueryError({
operation,
message: errorMessage(cause),
cause,
});
export const makeQdrantDocEmbeddingsQueryService = (
config: QdrantDocQueryConfig = {},
): QdrantDocEmbeddingsQueryServiceShape => {
const query = new QdrantDocEmbeddingsQuery(config);
return {
query: Effect.fn("QdrantDocEmbeddingsQuery.query")(function* (request) {
return yield* Effect.tryPromise({
try: () => query.query(request),
catch: (cause) => qdrantDocEmbeddingsQueryError("query", cause),
});
}),
};
};
export const QdrantDocEmbeddingsQueryLive = (
config: QdrantDocQueryConfig = {},
): Layer.Layer<QdrantDocEmbeddingsQueryService> =>
Layer.succeed(
QdrantDocEmbeddingsQueryService,
QdrantDocEmbeddingsQueryService.of(makeQdrantDocEmbeddingsQueryService(config)),
);

View file

@ -13,78 +13,109 @@ import {
ProducerSpec,
type ProcessorConfig,
type FlowContext,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type GraphEmbeddingsRequest,
type GraphEmbeddingsResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { QdrantGraphEmbeddingsQuery } from "./qdrant-graph.js";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import {
QdrantGraphEmbeddingsQueryLive,
QdrantGraphEmbeddingsQueryService,
makeQdrantGraphEmbeddingsQueryService,
type QdrantGraphQueryConfig,
} from "./qdrant-graph.js";
export class GraphEmbeddingsQueryService extends FlowProcessor {
private query: QdrantGraphEmbeddingsQuery;
const onGraphEmbeddingsQueryMessage = Effect.fn("GraphEmbeddingsQueryService.onMessage")(function* (
msg: GraphEmbeddingsRequest,
properties: Record<string, string>,
flowCtx: FlowContext<QdrantGraphEmbeddingsQueryService>,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<GraphEmbeddingsResponse>("graph-embeddings-response");
const query = yield* QdrantGraphEmbeddingsQueryService;
const user = msg.user ?? "default";
const collection = msg.collection ?? "default";
yield* Effect.log(
`[GraphEmbeddingsQuery] Request: user=${user}, collection=${collection}, vectors=${msg.vectors?.length ?? 0}, limit=${msg.limit}`,
);
const allEntities: GraphEmbeddingsResponse["entities"] = [];
for (const vector of msg.vectors ?? []) {
const matches = yield* query.query({
vector,
user,
collection,
limit: msg.limit ?? 50,
}).pipe(
Effect.catch((error) =>
Effect.logError("[GraphEmbeddingsQuery] Query failed", {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
producer.send(requestId, {
entities: [],
error: { type: "query-error", message: error.message },
})
),
Effect.as(null),
),
),
);
if (matches === null) return;
for (const match of matches) {
allEntities.push(match.entity);
}
}
yield* producer.send(requestId, { entities: allEntities });
});
export const makeGraphEmbeddingsQuerySpecs = (): ReadonlyArray<Spec<QdrantGraphEmbeddingsQueryService>> => [
new ConsumerSpec<
GraphEmbeddingsRequest,
FlowResourceNotFoundError | MessagingDeliveryError,
QdrantGraphEmbeddingsQueryService
>("graph-embeddings-request", onGraphEmbeddingsQueryMessage),
new ProducerSpec<GraphEmbeddingsResponse>("graph-embeddings-response"),
];
export class GraphEmbeddingsQueryService extends FlowProcessor<QdrantGraphEmbeddingsQueryService> {
private readonly query = makeQdrantGraphEmbeddingsQueryService();
constructor(config: ProcessorConfig) {
super(config);
this.query = new QdrantGraphEmbeddingsQuery();
this.registerSpecification(
ConsumerSpec.fromPromise<GraphEmbeddingsRequest>(
"graph-embeddings-request",
this.onMessage.bind(this),
),
);
this.registerSpecification(
new ProducerSpec<GraphEmbeddingsResponse>("graph-embeddings-response"),
);
for (const spec of makeGraphEmbeddingsQuerySpecs()) {
this.registerSpecification(spec);
}
console.log("[GraphEmbeddingsQuery] Service initialized");
}
private async onMessage(
msg: GraphEmbeddingsRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = flowCtx.flow.producer<GraphEmbeddingsResponse>("graph-embeddings-response");
const user = msg.user ?? "default";
const collection = msg.collection ?? "default";
console.log(`[GraphEmbeddingsQuery] Request: user=${user}, collection=${collection}, vectors=${msg.vectors?.length ?? 0}, limit=${msg.limit}`);
try {
// Query for each vector and aggregate results
const allEntities: GraphEmbeddingsResponse["entities"] = [];
for (const vector of msg.vectors ?? []) {
const matches = await this.query.query({
vector,
user,
collection,
limit: msg.limit ?? 50,
});
for (const match of matches) {
allEntities.push(match.entity);
}
}
await producer.send(requestId, { entities: allEntities });
} catch (err) {
console.error("[GraphEmbeddingsQuery] Query failed:", err);
await producer.send(requestId, {
entities: [],
error: { type: "query-error", message: String(err) },
});
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(
QdrantGraphEmbeddingsQueryService,
QdrantGraphEmbeddingsQueryService.of(this.query),
),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig & QdrantGraphQueryConfig, never, QdrantGraphEmbeddingsQueryService>({
id: "graph-embeddings-query",
make: (config) => new GraphEmbeddingsQueryService(config),
specs: () => makeGraphEmbeddingsQuerySpecs(),
layer: (config) => QdrantGraphEmbeddingsQueryLive(config),
});
export async function run(): Promise<void> {
await GraphEmbeddingsQueryService.launch("graph-embeddings-query");
await Effect.runPromise(program);
}

View file

@ -11,7 +11,9 @@
*/
import { QdrantClient } from "@qdrant/js-client-rest";
import type { Term } from "@trustgraph/base";
import { errorMessage, type Term } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface QdrantGraphQueryConfig {
url?: string;
@ -104,3 +106,54 @@ export class QdrantGraphEmbeddingsQuery {
return entities;
}
}
export class QdrantGraphEmbeddingsQueryError extends S.TaggedErrorClass<QdrantGraphEmbeddingsQueryError>()(
"QdrantGraphEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantGraphEmbeddingsQueryServiceShape {
readonly query: (
request: GraphEmbeddingsQueryRequest,
) => Effect.Effect<ReadonlyArray<EntityMatch>, QdrantGraphEmbeddingsQueryError>;
}
export class QdrantGraphEmbeddingsQueryService extends Context.Service<
QdrantGraphEmbeddingsQueryService,
QdrantGraphEmbeddingsQueryServiceShape
>()(
"@trustgraph/flow/query/embeddings/qdrant-graph/QdrantGraphEmbeddingsQueryService",
) {}
const qdrantGraphEmbeddingsQueryError = (operation: string, cause: unknown) =>
new QdrantGraphEmbeddingsQueryError({
operation,
message: errorMessage(cause),
cause,
});
export const makeQdrantGraphEmbeddingsQueryService = (
config: QdrantGraphQueryConfig = {},
): QdrantGraphEmbeddingsQueryServiceShape => {
const query = new QdrantGraphEmbeddingsQuery(config);
return {
query: Effect.fn("QdrantGraphEmbeddingsQuery.query")(function* (request) {
return yield* Effect.tryPromise({
try: () => query.query(request),
catch: (cause) => qdrantGraphEmbeddingsQueryError("query", cause),
});
}),
};
};
export const QdrantGraphEmbeddingsQueryLive = (
config: QdrantGraphQueryConfig = {},
): Layer.Layer<QdrantGraphEmbeddingsQueryService> =>
Layer.succeed(
QdrantGraphEmbeddingsQueryService,
QdrantGraphEmbeddingsQueryService.of(makeQdrantGraphEmbeddingsQueryService(config)),
);

View file

@ -13,61 +13,95 @@ import {
ProducerSpec,
type ProcessorConfig,
type FlowContext,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type TriplesQueryRequest,
type TriplesQueryResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { FalkorDBTriplesQuery } from "./falkordb.js";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import {
FalkorDBTriplesQueryLive,
FalkorDBTriplesQueryService,
makeFalkorDBTriplesQueryService,
type FalkorDBQueryConfig,
} from "./falkordb.js";
export class TriplesQueryService extends FlowProcessor {
private query: FalkorDBTriplesQuery;
const onTriplesQueryMessage = Effect.fn("TriplesQueryService.onMessage")(function* (
msg: TriplesQueryRequest,
properties: Record<string, string>,
flowCtx: FlowContext<FalkorDBTriplesQueryService>,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<TriplesQueryResponse>("triples-response");
const query = yield* FalkorDBTriplesQueryService;
const triples = yield* query.queryTriples(
msg.s,
msg.p,
msg.o,
msg.limit ?? 100,
).pipe(
Effect.catch((error) =>
Effect.logError("[TriplesQuery] Query failed", {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
producer.send(requestId, {
triples: [],
error: { type: "query-error", message: error.message },
})
),
Effect.as(null),
),
),
);
if (triples === null) return;
yield* producer.send(requestId, { triples: Array.from(triples) });
});
export const makeTriplesQuerySpecs = (): ReadonlyArray<Spec<FalkorDBTriplesQueryService>> => [
new ConsumerSpec<
TriplesQueryRequest,
FlowResourceNotFoundError | MessagingDeliveryError,
FalkorDBTriplesQueryService
>("triples-request", onTriplesQueryMessage),
new ProducerSpec<TriplesQueryResponse>("triples-response"),
];
export class TriplesQueryService extends FlowProcessor<FalkorDBTriplesQueryService> {
private readonly query = makeFalkorDBTriplesQueryService();
constructor(config: ProcessorConfig) {
super(config);
this.query = new FalkorDBTriplesQuery();
this.registerSpecification(
ConsumerSpec.fromPromise<TriplesQueryRequest>("triples-request", this.onMessage.bind(this)),
);
this.registerSpecification(new ProducerSpec<TriplesQueryResponse>("triples-response"));
for (const spec of makeTriplesQuerySpecs()) {
this.registerSpecification(spec);
}
console.log("[TriplesQuery] Service initialized");
}
private async onMessage(
msg: TriplesQueryRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = flowCtx.flow.producer<TriplesQueryResponse>("triples-response");
try {
const triples = await this.query.queryTriples(
msg.s,
msg.p,
msg.o,
msg.limit ?? 100,
);
await producer.send(requestId, { triples });
} catch (err) {
console.error("[TriplesQuery] Query failed:", err);
await producer.send(requestId, {
triples: [],
error: { type: "query-error", message: String(err) },
});
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(
FalkorDBTriplesQueryService,
FalkorDBTriplesQueryService.of(this.query),
),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig & FalkorDBQueryConfig, never, FalkorDBTriplesQueryService>({
id: "triples-query",
make: (config) => new TriplesQueryService(config),
specs: () => makeTriplesQuerySpecs(),
layer: (config) => FalkorDBTriplesQueryLive(config),
});
export async function run(): Promise<void> {
await TriplesQueryService.launch("triples-query");
await Effect.runPromise(program);
}

View file

@ -7,7 +7,9 @@
*/
import { createClient, Graph } from "falkordb";
import type { Term, Triple } from "@trustgraph/base";
import { errorMessage, type Term, type Triple } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface FalkorDBQueryConfig {
url?: string;
@ -264,3 +266,61 @@ export class FalkorDBTriplesQuery {
}
}
}
export class FalkorDBTriplesQueryError extends S.TaggedErrorClass<FalkorDBTriplesQueryError>()(
"FalkorDBTriplesQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface FalkorDBTriplesQueryServiceShape {
readonly queryTriples: (
s: Term | undefined,
p: Term | undefined,
o: Term | undefined,
limit: number,
) => Effect.Effect<ReadonlyArray<Triple>, FalkorDBTriplesQueryError>;
}
export class FalkorDBTriplesQueryService extends Context.Service<
FalkorDBTriplesQueryService,
FalkorDBTriplesQueryServiceShape
>()(
"@trustgraph/flow/query/triples/falkordb/FalkorDBTriplesQueryService",
) {}
const falkorDBTriplesQueryError = (operation: string, cause: unknown) =>
new FalkorDBTriplesQueryError({
operation,
message: errorMessage(cause),
cause,
});
export const makeFalkorDBTriplesQueryService = (
config: FalkorDBQueryConfig = {},
): FalkorDBTriplesQueryServiceShape => {
const query = new FalkorDBTriplesQuery(config);
return {
queryTriples: Effect.fn("FalkorDBTriplesQuery.queryTriples")((
s: Term | undefined,
p: Term | undefined,
o: Term | undefined,
limit: number,
) =>
Effect.tryPromise({
try: () => query.queryTriples(s, p, o, limit),
catch: (cause) => falkorDBTriplesQueryError("query-triples", cause),
})),
};
};
export const FalkorDBTriplesQueryLive = (
config: FalkorDBQueryConfig = {},
): Layer.Layer<FalkorDBTriplesQueryService> =>
Layer.succeed(
FalkorDBTriplesQueryService,
FalkorDBTriplesQueryService.of(makeFalkorDBTriplesQueryService(config)),
);

View file

@ -1,118 +1,166 @@
/**
* Document RAG service FlowProcessor wrapper around the DocumentRag class.
* Document RAG service.
*
* Consumes DocumentRagRequest messages, runs the document retrieval pipeline
* (embed query find similar chunks synthesize answer), emits DocumentRagResponse.
*
* Each request gets its own DocumentRag instance for security isolation.
* Consumes DocumentRagRequest messages, runs the document retrieval pipeline,
* and emits DocumentRagResponse.
*
* Python reference: trustgraph-flow/trustgraph/retrieval/document_rag/
*/
import {
FlowProcessor,
ConsumerSpec,
FlowProcessor,
ProducerSpec,
RequestResponseSpec,
type ProcessorConfig,
type FlowContext,
type DocumentRagRequest,
type DocumentRagResponse,
type TextCompletionRequest,
type TextCompletionResponse,
type EmbeddingsRequest,
type EmbeddingsResponse,
makeFlowProcessorProgram,
type DocumentEmbeddingsRequest,
type DocumentEmbeddingsResponse,
type DocumentRagRequest,
type DocumentRagResponse,
type EffectRequestOptions,
type EffectRequestResponse,
type EmbeddingsRequest,
type EmbeddingsResponse,
type FlowContext,
type FlowRequestOptions,
type FlowRequestor,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type ProcessorConfig,
type PromptRequest,
type PromptResponse,
type Spec,
type TextCompletionRequest,
type TextCompletionResponse,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { DocumentRag } from "./document-rag.js";
import { Effect } from "effect";
import {
DocumentRagEngine,
DocumentRagEngineError,
DocumentRagLive,
makeDocumentRagEngine,
type DocumentRagClients,
} from "./document-rag.js";
export class DocumentRagService extends FlowProcessor {
const toEffectRequestOptions = <TRes>(
options: FlowRequestOptions<TRes> | undefined,
): EffectRequestOptions<TRes> | undefined => {
if (options === undefined) return undefined;
return {
...(options.timeoutMs === undefined ? {} : { timeoutMs: options.timeoutMs }),
...(options.recipient === undefined
? {}
: {
recipient: (response: TRes) =>
Effect.promise(() => options.recipient?.(response) ?? Promise.resolve(true)),
}),
};
};
const toPromiseRequestor = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
): FlowRequestor<TReq, TRes> => ({
request: (request, options) =>
Effect.runPromise(requestor.request(request, toEffectRequestOptions(options))),
stop: () => Effect.runPromise(requestor.stop),
});
const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function* (
msg: DocumentRagRequest,
properties: Record<string, string>,
flowCtx: FlowContext<DocumentRagEngine>,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<DocumentRagResponse>("document-rag-response");
const engine = yield* DocumentRagEngine;
const clients: DocumentRagClients = {
llm: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm")),
embeddings: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings")),
docEmbeddings: toPromiseRequestor(
yield* flowCtx.flow.requestorEffect<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>("doc-embeddings"),
),
prompt: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt")),
};
const response = yield* engine.query(
clients,
msg.query,
{
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
},
).pipe(
Effect.catch((error: DocumentRagEngineError) =>
Effect.logError("[DocumentRag] Query failed", {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
producer.send(requestId, {
response: "",
error: { type: "rag-error", message: error.message },
}),
),
Effect.as(undefined),
),
),
);
if (response === undefined) return;
yield* producer.send(requestId, { response, endOfStream: true });
});
export const makeDocumentRagSpecs = (): ReadonlyArray<Spec<DocumentRagEngine>> => [
new ConsumerSpec<DocumentRagRequest, FlowResourceNotFoundError | MessagingDeliveryError, DocumentRagEngine>(
"document-rag-request",
onDocumentRagRequest,
),
new ProducerSpec<DocumentRagResponse>("document-rag-response"),
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
new RequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
"doc-embeddings",
"document-embeddings-request",
"document-embeddings-response",
),
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
];
export class DocumentRagService extends FlowProcessor<DocumentRagEngine> {
constructor(config: ProcessorConfig) {
super(config);
// Consumer: document RAG requests
this.registerSpecification(
ConsumerSpec.fromPromise<DocumentRagRequest>("document-rag-request", this.onRequest.bind(this)),
);
// Producer: document RAG responses
this.registerSpecification(new ProducerSpec<DocumentRagResponse>("document-rag-response"));
// Request-response clients
this.registerSpecification(
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
);
this.registerSpecification(
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
);
this.registerSpecification(
new RequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
"doc-embeddings",
"document-embeddings-request",
"document-embeddings-response",
),
);
this.registerSpecification(
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
);
console.log("[DocumentRag] Service initialized");
for (const spec of makeDocumentRagSpecs()) {
this.registerSpecification(spec);
}
}
private async onRequest(
msg: DocumentRagRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = flowCtx.flow.producer<DocumentRagResponse>("document-rag-response");
try {
const documentRag = new DocumentRag({
llm: flowCtx.flow.requestor<TextCompletionRequest, TextCompletionResponse>("llm"),
embeddings: flowCtx.flow.requestor<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
docEmbeddings: flowCtx.flow.requestor<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>("doc-embeddings"),
prompt: flowCtx.flow.requestor<PromptRequest, PromptResponse>("prompt"),
});
const response = await documentRag.query(msg.query, {
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
});
await producer.send(requestId, { response, endOfStream: true });
} catch (err) {
console.error("[DocumentRag] Query failed:", err);
await producer.send(requestId, {
response: "",
error: { type: "rag-error", message: String(err) },
});
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(DocumentRagEngine, DocumentRagEngine.of(makeDocumentRagEngine())),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram({
id: "document-rag",
make: (config) => new DocumentRagService(config),
specs: makeDocumentRagSpecs,
layer: () => DocumentRagLive,
});
export async function run(): Promise<void> {
await DocumentRagService.launch("document-rag");
await Effect.runPromise(program);
}

View file

@ -1,23 +1,23 @@
/**
* 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 {
FlowRequestor,
TextCompletionRequest,
TextCompletionResponse,
EmbeddingsRequest,
EmbeddingsResponse,
DocumentEmbeddingsRequest,
DocumentEmbeddingsResponse,
EmbeddingsRequest,
EmbeddingsResponse,
FlowRequestor,
PromptRequest,
PromptResponse,
TextCompletionRequest,
TextCompletionResponse,
} from "@trustgraph/base";
import { errorMessage } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface DocumentRagClients {
llm: FlowRequestor<TextCompletionRequest, TextCompletionResponse>;
@ -28,55 +28,110 @@ export interface DocumentRagClients {
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
export interface DocumentRagQueryOptions {
readonly collection?: string;
readonly streaming?: boolean;
readonly chunkCallback?: ChunkCallback;
}
export class DocumentRagEngineError extends S.TaggedErrorClass<DocumentRagEngineError>()(
"DocumentRagEngineError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface DocumentRagEngineShape {
readonly query: (
clients: DocumentRagClients,
queryText: string,
options?: DocumentRagQueryOptions,
) => Effect.Effect<string, DocumentRagEngineError>;
}
export class DocumentRagEngine extends Context.Service<DocumentRagEngine, DocumentRagEngineShape>()(
"@trustgraph/flow/retrieval/document-rag/DocumentRagEngine",
) {}
const documentRagError = (operation: string, cause: unknown) =>
new DocumentRagEngineError({
operation,
cause,
message: errorMessage(cause),
});
export function makeDocumentRagEngine(): DocumentRagEngineShape {
return {
query: Effect.fn("DocumentRagEngine.query")((
clients: DocumentRagClients,
queryText: string,
options?: DocumentRagQueryOptions,
) =>
Effect.tryPromise({
try: () => queryDocumentRag(clients, queryText, options),
catch: (cause) => documentRagError("query", cause),
}),
),
};
}
export const DocumentRagLive: Layer.Layer<DocumentRagEngine> = Layer.succeed(
DocumentRagEngine,
DocumentRagEngine.of(makeDocumentRagEngine()),
);
export class DocumentRag {
private readonly engine = makeDocumentRagEngine();
private readonly clients: DocumentRagClients;
constructor(clients: DocumentRagClients) {
this.clients = clients;
}
async query(
query(
queryText: string,
options?: {
collection?: string;
streaming?: boolean;
chunkCallback?: ChunkCallback;
},
options?: DocumentRagQueryOptions,
): Promise<string> {
const collection = options?.collection ?? "default";
// 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,
collection,
user: "default",
});
const chunks = (docResp as DocumentEmbeddingsResponse).chunks ?? [];
console.log(`[DocumentRag] Found ${chunks.length} matching chunks`);
// Step 3: Build context from chunks
const context = chunks
.flatMap((c) =>
c.content !== undefined && c.content.length > 0 ? [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;
return Effect.runPromise(this.engine.query(this.clients, queryText, options));
}
}
async function queryDocumentRag(
clients: DocumentRagClients,
queryText: string,
options?: DocumentRagQueryOptions,
): Promise<string> {
const collection = options?.collection ?? "default";
const embResp = await clients.embeddings.request({ text: [queryText] });
const vectors = embResp.vectors;
const docResp = await clients.docEmbeddings.request({
vectors,
limit: 10,
collection,
user: "default",
});
const chunks = docResp.chunks ?? [];
console.log(`[DocumentRag] Found ${chunks.length} matching chunks`);
const context = chunks
.flatMap((chunk) =>
chunk.content !== undefined && chunk.content.length > 0 ? [chunk.content] : [],
)
.join("\n\n---\n\n");
const promptResp = await clients.prompt.request({
name: "document-rag-synthesize",
variables: { query: queryText, context },
});
const resp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
});
return resp.response;
}

View file

@ -1,158 +1,197 @@
/**
* Graph RAG service FlowProcessor wrapper around the GraphRag class.
* Graph RAG service.
*
* Consumes GraphRagRequest messages from the agent/gateway, runs the full
* Graph RAG pipeline (concept extraction entity lookup graph traversal
* edge scoring answer synthesis), and emits GraphRagResponse.
*
* Each request gets its own GraphRag instance to prevent data leakage
* across requests (security requirement from the Python implementation).
* Graph RAG pipeline, and emits GraphRagResponse.
*
* Python reference: trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
*/
import {
FlowProcessor,
ConsumerSpec,
FlowProcessor,
ProducerSpec,
RequestResponseSpec,
type ProcessorConfig,
makeFlowProcessorProgram,
type EffectRequestOptions,
type EffectRequestResponse,
type FlowContext,
type GraphRagRequest,
type GraphRagResponse,
type TextCompletionRequest,
type TextCompletionResponse,
type EmbeddingsRequest,
type EmbeddingsResponse,
type FlowRequestOptions,
type FlowRequestor,
type FlowResourceNotFoundError,
type GraphEmbeddingsRequest,
type GraphEmbeddingsResponse,
type TriplesQueryRequest,
type TriplesQueryResponse,
type GraphRagRequest,
type GraphRagResponse,
type EmbeddingsRequest,
type EmbeddingsResponse,
type MessagingDeliveryError,
type ProcessorConfig,
type PromptRequest,
type PromptResponse,
type Spec,
type TextCompletionRequest,
type TextCompletionResponse,
type TriplesQueryRequest,
type TriplesQueryResponse,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { GraphRag } from "./graph-rag.js";
import { Effect } from "effect";
import {
GraphRagEngine,
GraphRagEngineError,
GraphRagLive,
makeGraphRagEngine,
type GraphRagClients,
type GraphRagConfig,
} from "./graph-rag.js";
export class GraphRagService extends FlowProcessor {
constructor(config: ProcessorConfig) {
super(config);
const toEffectRequestOptions = <TRes>(
options: FlowRequestOptions<TRes> | undefined,
): EffectRequestOptions<TRes> | undefined => {
if (options === undefined) return undefined;
return {
...(options.timeoutMs === undefined ? {} : { timeoutMs: options.timeoutMs }),
...(options.recipient === undefined
? {}
: {
recipient: (response: TRes) =>
Effect.promise(() => options.recipient?.(response) ?? Promise.resolve(true)),
}),
};
};
// Consumer: graph RAG requests
this.registerSpecification(
ConsumerSpec.fromPromise<GraphRagRequest>("graph-rag-request", this.onRequest.bind(this)),
);
const toPromiseRequestor = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
): FlowRequestor<TReq, TRes> => ({
request: (request, options) =>
Effect.runPromise(requestor.request(request, toEffectRequestOptions(options))),
stop: () => Effect.runPromise(requestor.stop),
});
// Producer: graph RAG responses
this.registerSpecification(new ProducerSpec<GraphRagResponse>("graph-rag-response"));
const graphRagConfigFromRequest = (msg: GraphRagRequest): GraphRagConfig => ({
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
...(msg.maxSubgraphSize !== undefined ? { maxSubgraphSize: msg.maxSubgraphSize } : {}),
...(msg.maxPathLength !== undefined ? { maxPathLength: msg.maxPathLength } : {}),
});
// Request-response clients for the pipeline
this.registerSpecification(
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
);
this.registerSpecification(
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
);
this.registerSpecification(
new RequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
"graph-embeddings",
"graph-embeddings-request",
"graph-embeddings-response",
),
);
this.registerSpecification(
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
),
);
this.registerSpecification(
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
);
const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* (
msg: GraphRagRequest,
properties: Record<string, string>,
flowCtx: FlowContext<GraphRagEngine>,
) {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
console.log("[GraphRag] Service initialized");
const producer = yield* flowCtx.flow.producerEffect<GraphRagResponse>("graph-rag-response");
const engine = yield* GraphRagEngine;
yield* Effect.log(`[GraphRagService] Received request ${requestId}: "${msg.query?.slice(0, 60)}..." collection=${msg.collection}`);
const clients: GraphRagClients = {
llm: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm")),
embeddings: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings")),
graphEmbeddings: toPromiseRequestor(
yield* flowCtx.flow.requestorEffect<GraphEmbeddingsRequest, GraphEmbeddingsResponse>("graph-embeddings"),
),
triples: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TriplesQueryRequest, TriplesQueryResponse>("triples")),
prompt: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt")),
};
const result = yield* engine.query(
clients,
msg.query,
{
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
},
graphRagConfigFromRequest(msg),
).pipe(
Effect.catch((error: GraphRagEngineError) =>
Effect.logError("[GraphRag] Query failed", {
error: error.message,
operation: error.operation,
}).pipe(
Effect.flatMap(() =>
producer.send(requestId, {
response: "",
error: { type: "rag-error", message: error.message },
}),
),
Effect.as(undefined),
),
),
);
if (result === undefined) return;
const response: GraphRagResponse = {
response: result.answer,
endOfStream: true,
};
if (result.subgraph.length > 0) {
(response as Record<string, unknown>).message_type = "explain";
(response as Record<string, unknown>).explain_id = `explain-${requestId}`;
(response as Record<string, unknown>).explain_triples = result.subgraph;
}
private async onRequest(
msg: GraphRagRequest,
properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
yield* producer.send(requestId, response);
});
const producer = flowCtx.flow.producer<GraphRagResponse>("graph-rag-response");
console.log(`[GraphRagService] Received request ${requestId}: "${msg.query?.slice(0, 60)}..." collection=${msg.collection}`);
export const makeGraphRagSpecs = (): ReadonlyArray<Spec<GraphRagEngine>> => [
new ConsumerSpec<GraphRagRequest, FlowResourceNotFoundError | MessagingDeliveryError, GraphRagEngine>(
"graph-rag-request",
onGraphRagRequest,
),
new ProducerSpec<GraphRagResponse>("graph-rag-response"),
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
new RequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
"graph-embeddings",
"graph-embeddings-request",
"graph-embeddings-response",
),
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
),
new RequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
];
try {
// Create a per-request GraphRag instance with flow clients
const graphRag = new GraphRag(
{
llm: flowCtx.flow.requestor<TextCompletionRequest, TextCompletionResponse>("llm"),
embeddings: flowCtx.flow.requestor<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
graphEmbeddings: flowCtx.flow.requestor<GraphEmbeddingsRequest, GraphEmbeddingsResponse>("graph-embeddings"),
triples: flowCtx.flow.requestor<TriplesQueryRequest, TriplesQueryResponse>("triples"),
prompt: flowCtx.flow.requestor<PromptRequest, PromptResponse>("prompt"),
},
{
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
...(msg.maxSubgraphSize !== undefined
? { maxSubgraphSize: msg.maxSubgraphSize }
: {}),
...(msg.maxPathLength !== undefined ? { maxPathLength: msg.maxPathLength } : {}),
},
);
const result = await graphRag.query(msg.query, {
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
});
// Send answer with explain data embedded in a SINGLE message.
// Non-streaming callers (agent's RequestResponse) return the first
// response — so the answer must be in that first (and only) message.
// Streaming callers (gateway) extract explain data + answer from
// the same message.
const response: GraphRagResponse = {
response: result.answer,
endOfStream: true,
};
if (result.subgraph.length > 0) {
(response as Record<string, unknown>).message_type = "explain";
(response as Record<string, unknown>).explain_id = `explain-${requestId}`;
(response as Record<string, unknown>).explain_triples = result.subgraph;
}
await producer.send(requestId, response);
} catch (err) {
console.error("[GraphRag] Query failed:", err);
await producer.send(requestId, {
response: "",
error: { type: "rag-error", message: String(err) },
});
export class GraphRagService extends FlowProcessor<GraphRagEngine> {
constructor(config: ProcessorConfig) {
super(config);
for (const spec of makeGraphRagSpecs()) {
this.registerSpecification(spec);
}
}
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(GraphRagEngine, GraphRagEngine.of(makeGraphRagEngine())),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram({
id: "graph-rag",
make: (config) => new GraphRagService(config),
specs: makeGraphRagSpecs,
layer: () => GraphRagLive,
});
export async function run(): Promise<void> {
await GraphRagService.launch("graph-rag");
await Effect.runPromise(program);
}

View file

@ -1,22 +1,15 @@
/**
* 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 {
EmbeddingsRequest,
EmbeddingsResponse,
FlowRequestor,
GraphEmbeddingsRequest,
GraphEmbeddingsResponse,
FlowRequestor,
PromptRequest,
PromptResponse,
Term,
@ -26,6 +19,10 @@ import type {
TriplesQueryRequest,
TriplesQueryResponse,
} from "@trustgraph/base";
import { errorMessage } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export interface GraphRagConfig {
entityLimit?: number;
@ -46,321 +43,373 @@ export interface GraphRagClients {
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
export interface GraphRagQueryOptions {
readonly collection?: string;
readonly streaming?: boolean;
readonly chunkCallback?: ChunkCallback;
}
export interface GraphRagResult {
answer: string;
subgraph: Triple[];
}
interface NormalizedGraphRagConfig {
entityLimit: number;
tripleLimit: number;
maxSubgraphSize: number;
maxPathLength: number;
edgeScoreLimit: number;
edgeLimit: number;
}
export class GraphRagEngineError extends S.TaggedErrorClass<GraphRagEngineError>()(
"GraphRagEngineError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface GraphRagEngineShape {
readonly query: (
clients: GraphRagClients,
queryText: string,
options?: GraphRagQueryOptions,
config?: GraphRagConfig,
) => Effect.Effect<GraphRagResult, GraphRagEngineError>;
}
export class GraphRagEngine extends Context.Service<GraphRagEngine, GraphRagEngineShape>()(
"@trustgraph/flow/retrieval/graph-rag/GraphRagEngine",
) {}
const graphRagError = (operation: string, cause: unknown) =>
new GraphRagEngineError({
operation,
cause,
message: errorMessage(cause),
});
export function normalizeGraphRagConfig(config: GraphRagConfig = {}): NormalizedGraphRagConfig {
return {
entityLimit: config.entityLimit ?? 50,
tripleLimit: config.tripleLimit ?? 30,
maxSubgraphSize: config.maxSubgraphSize ?? 1000,
maxPathLength: config.maxPathLength ?? 2,
edgeScoreLimit: config.edgeScoreLimit ?? 50,
edgeLimit: config.edgeLimit ?? 25,
};
}
export function makeGraphRagEngine(): GraphRagEngineShape {
return {
query: Effect.fn("GraphRagEngine.query")((
clients: GraphRagClients,
queryText: string,
options?: GraphRagQueryOptions,
config?: GraphRagConfig,
) =>
Effect.tryPromise({
try: () => queryGraphRag(clients, queryText, options, config),
catch: (cause) => graphRagError("query", cause),
}),
),
};
}
export const GraphRagLive: Layer.Layer<GraphRagEngine> = Layer.succeed(
GraphRagEngine,
GraphRagEngine.of(makeGraphRagEngine()),
);
export class GraphRag {
private readonly engine = makeGraphRagEngine();
private readonly clients: GraphRagClients;
private config: Required<GraphRagConfig>;
private readonly config: GraphRagConfig;
constructor(
clients: GraphRagClients,
config: GraphRagConfig = {},
) {
this.clients = clients;
this.config = {
entityLimit: config.entityLimit ?? 50,
tripleLimit: config.tripleLimit ?? 30,
maxSubgraphSize: config.maxSubgraphSize ?? 1000,
maxPathLength: config.maxPathLength ?? 2,
edgeScoreLimit: config.edgeScoreLimit ?? 50,
edgeLimit: config.edgeLimit ?? 25,
};
this.config = config;
}
async query(
query(
queryText: string,
options?: {
collection?: string;
streaming?: boolean;
chunkCallback?: ChunkCallback;
},
options?: GraphRagQueryOptions,
): Promise<GraphRagResult> {
console.log(`[GraphRag] Query: "${queryText.slice(0, 80)}..."`);
// Step 1: Extract concepts from the query via prompt + LLM
const concepts = await this.extractConcepts(queryText);
console.log(`[GraphRag] Step 1: extracted ${concepts.length} concepts: ${concepts.slice(0, 5).join(", ")}`);
// Step 2: Embed concepts concurrently
const vectors = await this.getVectors(concepts);
console.log(`[GraphRag] Step 2: got ${vectors.length} vectors (dim=${vectors[0]?.length ?? 0})`);
// Step 3: Find matching entities via graph embeddings
const entities = await this.getEntities(vectors, options?.collection);
console.log(`[GraphRag] Step 3: found ${entities.length} matching entities`);
// Step 4: Traverse the knowledge graph from entities
const subgraph = await this.followEdges(entities, options?.collection);
console.log(`[GraphRag] Step 4: traversed graph, ${subgraph.length} triples in subgraph`);
// Step 5: Score and filter edges via LLM
const scoredEdges = await this.scoreEdges(queryText, subgraph);
console.log(`[GraphRag] Step 5: scored down to ${scoredEdges.length} edges`);
// Step 6: Synthesize answer
console.log(`[GraphRag] Step 6: synthesizing answer from ${scoredEdges.length} edges...`);
const answer = await this.synthesize(
queryText,
scoredEdges,
options?.chunkCallback,
return Effect.runPromise(
this.engine.query(this.clients, queryText, options, this.config),
);
console.log(`[GraphRag] Step 6: done (${answer.length} chars)`);
return { answer, subgraph: scoredEdges };
}
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((c) => c.length > 0);
}
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[][], collection?: string): Promise<Term[]> {
const resp = await this.clients.graphEmbeddings.request({
vectors,
user: "default",
collection: collection ?? "default",
limit: this.config.entityLimit,
});
return (resp as GraphEmbeddingsResponse).entities;
}
private async followEdges(entities: Term[], collection?: string): Promise<Triple[]> {
// BFS multi-hop traversal up to maxPathLength
const visited = new Set<string>();
const subgraph: Triple[] = [];
// Current frontier: the set of entities to expand at this depth level
let currentLevel = new Set<string>(
entities.map((e) => termToString(e)),
);
for (let depth = 0; depth < this.config.maxPathLength; depth++) {
if (currentLevel.size === 0 || subgraph.length >= this.config.maxSubgraphSize) {
break;
}
// Filter out already-visited entities
const unvisited = [...currentLevel].filter((e) => !visited.has(e));
if (unvisited.length === 0) break;
// Batch triple queries for all unvisited entities at this depth
// Query each entity as subject to get outgoing edges
const queries = unvisited.map((entityStr) => {
const term = stringToTerm(entityStr);
const request: TriplesQueryRequest = {
s: term,
limit: this.config.tripleLimit,
...(collection !== undefined ? { collection } : {}),
};
return this.clients.triples.request(request);
});
const results = await Promise.all(queries);
const nextLevel = new Set<string>();
for (const result of results) {
const triples = (result as TriplesQueryResponse).triples;
for (const triple of triples) {
subgraph.push(triple);
// Collect objects as next-level entities for further expansion
// (only if we have more depth levels remaining)
if (depth < this.config.maxPathLength - 1) {
const objStr = termToString(triple.o);
if (!visited.has(objStr)) {
nextLevel.add(objStr);
}
}
if (subgraph.length >= this.config.maxSubgraphSize) {
return subgraph;
}
}
}
// Mark current level as visited and move to next
for (const e of currentLevel) {
visited.add(e);
}
currentLevel = nextLevel;
}
return subgraph.slice(0, this.config.maxSubgraphSize);
}
private async scoreEdges(query: string, triples: Triple[]): Promise<Triple[]> {
if (triples.length === 0) return [];
// If the subgraph is small enough, skip LLM scoring entirely
// 500 triples is well within LLM context limits and avoids lossy scoring
if (triples.length <= 500) {
console.log(`[GraphRag] Skipping edge scoring — ${triples.length} triples fits in context directly`);
return triples;
}
// Build a numbered list of edges for the LLM to score
const edgeDescriptions = triples.map((t, i) => ({
id: String(i),
s: termToString(t.s),
p: termToString(t.p),
o: termToString(t.o),
}));
// Limit how many edges we send for scoring to avoid overflowing context
const toScore = edgeDescriptions.slice(0, this.config.edgeScoreLimit);
const knowledgeJson = JSON.stringify(toScore, null, 2);
// Ask the LLM to score each edge for relevance to the query
const promptResp = await this.clients.prompt.request({
name: "kg-edge-scoring",
variables: {
query,
knowledge: knowledgeJson,
},
});
const llmResp = await this.clients.llm.request({
system: (promptResp as PromptResponse).system,
prompt: (promptResp as PromptResponse).prompt,
});
const responseText = (llmResp as TextCompletionResponse).response;
console.log(`[GraphRag] Edge scoring LLM response (first 500 chars): ${responseText.slice(0, 500)}`);
// Parse scores from LLM response
// Expected format: JSON array of { id: string, score: number }
// or newline-separated JSON objects
const scored: Array<{ id: string; score: number }> = [];
try {
// Try parsing as a JSON array first
const parsed = JSON.parse(responseText) as Array<{ id: string; score: number }>;
if (Array.isArray(parsed)) {
for (const item of parsed) {
if (
typeof item === "object" &&
item !== null &&
typeof item.id === "string" &&
typeof item.score === "number"
) {
scored.push({ id: item.id, score: item.score });
}
}
}
} catch {
// Fall back to parsing line-by-line JSON objects
for (const line of responseText.split("\n")) {
const trimmed = line.trim();
if (trimmed.length === 0) continue;
try {
const obj = JSON.parse(trimmed) as { id?: string; score?: number };
if (
typeof obj === "object" &&
obj !== null &&
typeof obj.id === "string" &&
typeof obj.score === "number"
) {
scored.push({ id: obj.id, score: obj.score });
}
} catch {
// Skip unparseable lines
}
}
}
// Sort by score descending and keep top N
scored.sort((a, b) => b.score - a.score);
const topN = scored.slice(0, this.config.edgeLimit);
// Map back to triples
const result: Triple[] = [];
for (const entry of topN) {
const idx = parseInt(entry.id, 10);
if (!isNaN(idx) && idx >= 0 && idx < triples.length) {
result.push(triples[idx]);
}
}
console.log(`[GraphRag] Edge scoring: LLM returned ${scored.length} scores, keeping top ${topN.length}, mapped ${result.length} triples`);
// If scoring failed entirely, fall back to returning the first edgeLimit triples
if (result.length === 0) {
return triples.slice(0, this.config.edgeLimit);
}
return result;
}
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 !== undefined) {
// 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.length > 0) {
fullText += r.response;
await chunkCallback(r.response, r.endOfStream === true);
}
return r.endOfStream === true;
},
},
);
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 {
async function queryGraphRag(
clients: GraphRagClients,
queryText: string,
options?: GraphRagQueryOptions,
rawConfig?: GraphRagConfig,
): Promise<GraphRagResult> {
const config = normalizeGraphRagConfig(rawConfig);
console.log(`[GraphRag] Query: "${queryText.slice(0, 80)}..."`);
const concepts = await extractConcepts(clients, queryText);
console.log(`[GraphRag] Step 1: extracted ${concepts.length} concepts: ${concepts.slice(0, 5).join(", ")}`);
const vectors = await getVectors(clients, concepts);
console.log(`[GraphRag] Step 2: got ${vectors.length} vectors (dim=${vectors[0]?.length ?? 0})`);
const entities = await getEntities(clients, config, vectors, options?.collection);
console.log(`[GraphRag] Step 3: found ${entities.length} matching entities`);
const subgraph = await followEdges(clients, config, entities, options?.collection);
console.log(`[GraphRag] Step 4: traversed graph, ${subgraph.length} triples in subgraph`);
const scoredEdges = await scoreEdges(clients, config, queryText, subgraph);
console.log(`[GraphRag] Step 5: scored down to ${scoredEdges.length} edges`);
console.log(`[GraphRag] Step 6: synthesizing answer from ${scoredEdges.length} edges...`);
const answer = await synthesize(
clients,
queryText,
scoredEdges,
options?.chunkCallback,
);
console.log(`[GraphRag] Step 6: done (${answer.length} chars)`);
return { answer, subgraph: scoredEdges };
}
async function extractConcepts(clients: GraphRagClients, query: string): Promise<string[]> {
const promptResp = await clients.prompt.request({
name: "extract-concepts",
variables: { query },
});
const llmResp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
});
return llmResp.response
.split("\n")
.map((concept) => concept.trim())
.filter((concept) => concept.length > 0);
}
async function getVectors(clients: GraphRagClients, concepts: string[]): Promise<number[][]> {
const resp = await clients.embeddings.request({ text: concepts });
return resp.vectors;
}
async function getEntities(
clients: GraphRagClients,
config: NormalizedGraphRagConfig,
vectors: number[][],
collection?: string,
): Promise<Term[]> {
const resp = await clients.graphEmbeddings.request({
vectors,
user: "default",
collection: collection ?? "default",
limit: config.entityLimit,
});
return resp.entities;
}
async function followEdges(
clients: GraphRagClients,
config: NormalizedGraphRagConfig,
entities: Term[],
collection?: string,
): Promise<Triple[]> {
const visited = new Set<string>();
const subgraph: Triple[] = [];
let currentLevel = new Set<string>(
entities.map((entity) => termToString(entity)),
);
for (let depth = 0; depth < config.maxPathLength; depth++) {
if (currentLevel.size === 0 || subgraph.length >= config.maxSubgraphSize) {
break;
}
const unvisited = [...currentLevel].filter((entity) => !visited.has(entity));
if (unvisited.length === 0) break;
const queries = unvisited.map((entityStr) => {
const term = stringToTerm(entityStr);
const request: TriplesQueryRequest = {
s: term,
limit: config.tripleLimit,
...(collection !== undefined ? { collection } : {}),
};
return clients.triples.request(request);
});
const results = await Promise.all(queries);
const nextLevel = new Set<string>();
for (const result of results) {
for (const triple of result.triples) {
subgraph.push(triple);
if (depth < config.maxPathLength - 1) {
const objStr = termToString(triple.o);
if (!visited.has(objStr)) {
nextLevel.add(objStr);
}
}
if (subgraph.length >= config.maxSubgraphSize) {
return subgraph;
}
}
}
for (const entity of currentLevel) {
visited.add(entity);
}
currentLevel = nextLevel;
}
return subgraph.slice(0, config.maxSubgraphSize);
}
async function scoreEdges(
clients: GraphRagClients,
config: NormalizedGraphRagConfig,
query: string,
triples: Triple[],
): Promise<Triple[]> {
if (triples.length === 0) return [];
if (triples.length <= 500) {
console.log(`[GraphRag] Skipping edge scoring - ${triples.length} triples fits in context directly`);
return triples;
}
const edgeDescriptions = triples.map((triple, index) => ({
id: String(index),
s: termToString(triple.s),
p: termToString(triple.p),
o: termToString(triple.o),
}));
const toScore = edgeDescriptions.slice(0, config.edgeScoreLimit);
const knowledgeJson = JSON.stringify(toScore, null, 2);
const promptResp = await clients.prompt.request({
name: "kg-edge-scoring",
variables: {
query,
knowledge: knowledgeJson,
},
});
const llmResp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
});
console.log(`[GraphRag] Edge scoring LLM response (first 500 chars): ${llmResp.response.slice(0, 500)}`);
const scored = parseScoredEdges(llmResp.response);
scored.sort((a, b) => b.score - a.score);
const topN = scored.slice(0, config.edgeLimit);
const result: Triple[] = [];
for (const entry of topN) {
const idx = Number.parseInt(entry.id, 10);
if (!Number.isNaN(idx) && idx >= 0 && idx < triples.length) {
result.push(triples[idx]);
}
}
console.log(`[GraphRag] Edge scoring: LLM returned ${scored.length} scores, keeping top ${topN.length}, mapped ${result.length} triples`);
if (result.length === 0) {
return triples.slice(0, config.edgeLimit);
}
return result;
}
async function synthesize(
clients: GraphRagClients,
query: string,
edges: Triple[],
chunkCallback?: ChunkCallback,
): Promise<string> {
const context = edges
.map((triple) => `${termToString(triple.s)} -> ${termToString(triple.p)} -> ${termToString(triple.o)}`)
.join("\n");
const promptResp = await clients.prompt.request({
name: "graph-rag-synthesize",
variables: { query, context },
});
if (chunkCallback !== undefined) {
let fullText = "";
await clients.llm.request(
{
system: promptResp.system,
prompt: promptResp.prompt,
streaming: true,
},
{
recipient: async (resp) => {
if (resp.response.length > 0) {
fullText += resp.response;
await chunkCallback(resp.response, resp.endOfStream === true);
}
return resp.endOfStream === true;
},
},
);
return fullText;
}
const resp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
});
return resp.response;
}
const ScoredEdge = S.Struct({
id: S.String,
score: S.Number,
});
const ScoredEdgesFromJson = S.Array(ScoredEdge).pipe(S.fromJsonString);
const ScoredEdgeFromJson = ScoredEdge.pipe(S.fromJsonString);
const decodeScoredEdges = S.decodeUnknownOption(ScoredEdgesFromJson);
const decodeScoredEdge = S.decodeUnknownOption(ScoredEdgeFromJson);
function parseScoredEdges(responseText: string): Array<typeof ScoredEdge.Type> {
const parsedArray = decodeScoredEdges(responseText);
if (O.isSome(parsedArray)) {
return Array.from(parsedArray.value);
}
const scored: Array<typeof ScoredEdge.Type> = [];
for (const line of responseText.split("\n")) {
const trimmed = line.trim();
if (trimmed.length === 0) continue;
const parsedLine = decodeScoredEdge(trimmed);
if (O.isSome(parsedLine)) {
scored.push(parsedLine.value);
}
}
return scored;
}
export function termToString(term: Term): string {
switch (term.type) {
case "IRI":
return term.iri;
@ -373,7 +422,7 @@ function termToString(term: Term): string {
}
}
function stringToTerm(value: string): Term {
export function stringToTerm(value: string): Term {
if (value.startsWith("http://") || value.startsWith("https://")) {
return { type: "IRI", iri: value };
}

View file

@ -15,83 +15,112 @@ import {
RequestResponseSpec,
type ProcessorConfig,
type FlowContext,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type MessagingTimeoutError,
type EntityContexts,
type EmbeddingsRequest,
type EmbeddingsResponse,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { QdrantGraphEmbeddingsStore } from "./qdrant-graph.js";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import {
QdrantGraphEmbeddingsStoreLive,
QdrantGraphEmbeddingsStoreService,
makeQdrantGraphEmbeddingsStoreService,
type QdrantGraphEmbeddingsConfig,
type QdrantGraphEmbeddingsStoreError,
} from "./qdrant-graph.js";
export class GraphEmbeddingsStoreService extends FlowProcessor {
private store: QdrantGraphEmbeddingsStore;
type GraphEmbeddingsStoreRequirements = QdrantGraphEmbeddingsStoreService;
type GraphEmbeddingsStoreError =
| FlowResourceNotFoundError
| MessagingDeliveryError
| MessagingTimeoutError
| QdrantGraphEmbeddingsStoreError;
const onGraphEmbeddingsStoreMessage = Effect.fn("GraphEmbeddingsStoreService.onMessage")(function* (
msg: EntityContexts,
_properties: Record<string, string>,
flowCtx: FlowContext<GraphEmbeddingsStoreRequirements>,
): Effect.fn.Return<void, GraphEmbeddingsStoreError, GraphEmbeddingsStoreRequirements> {
if (msg.entities.length === 0) return;
const embeddingsClient =
yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings-client");
const user = msg.metadata?.user ?? "default";
const collection = msg.metadata?.collection ?? "default";
const texts = msg.entities.map((entity) => entity.context);
const embResponse = yield* embeddingsClient.request({ text: texts });
if (embResponse.error !== undefined) {
yield* Effect.logError("[GraphEmbeddingsStore] Embeddings error", {
error: embResponse.error.message,
});
return;
}
const entities = msg.entities.map((entity, index) => ({
entity: entity.entity,
vector: embResponse.vectors[index],
chunkId: entity.chunkId,
}));
const store = yield* QdrantGraphEmbeddingsStoreService;
yield* store.store({ user, collection, entities });
yield* Effect.log(
`[GraphEmbeddingsStore] Stored ${entities.length} embeddings for ${user}/${collection}`,
);
});
export const makeGraphEmbeddingsStoreSpecs = (): ReadonlyArray<Spec<GraphEmbeddingsStoreRequirements>> => [
new ConsumerSpec<EntityContexts, GraphEmbeddingsStoreError, GraphEmbeddingsStoreRequirements>(
"store-graph-embeddings-input",
onGraphEmbeddingsStoreMessage,
),
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings-client",
"embeddings-request",
"embeddings-response",
),
];
export class GraphEmbeddingsStoreService extends FlowProcessor<GraphEmbeddingsStoreRequirements> {
private readonly store = makeQdrantGraphEmbeddingsStoreService();
constructor(config: ProcessorConfig) {
super(config);
this.store = new QdrantGraphEmbeddingsStore();
this.registerSpecification(
ConsumerSpec.fromPromise<EntityContexts>(
"store-graph-embeddings-input",
this.onMessage.bind(this),
),
);
this.registerSpecification(
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings-client",
"embeddings-request",
"embeddings-response",
),
);
for (const spec of makeGraphEmbeddingsStoreSpecs()) {
this.registerSpecification(spec);
}
console.log("[GraphEmbeddingsStore] Service initialized");
}
private async onMessage(
msg: EntityContexts,
_properties: Record<string, string>,
flowCtx: FlowContext,
): Promise<void> {
if (msg.entities.length === 0) return;
const embeddingsClient =
flowCtx.flow.requestor<EmbeddingsRequest, EmbeddingsResponse>("embeddings-client");
const user = msg.metadata?.user ?? "default";
const collection = msg.metadata?.collection ?? "default";
// Get text contexts for vectorization
const texts = msg.entities.map((e) => e.context);
// Call embeddings service
const embResponse = await embeddingsClient.request({ text: texts });
if (embResponse.error !== undefined) {
console.error(
"[GraphEmbeddingsStore] Embeddings error:",
embResponse.error.message,
);
return;
}
// Store entity+vector pairs
const entities = msg.entities.map((e, i) => ({
entity: e.entity,
vector: embResponse.vectors[i],
chunkId: e.chunkId,
}));
await this.store.store({ user, collection, entities });
console.log(
`[GraphEmbeddingsStore] Stored ${entities.length} embeddings for ${user}/${collection}`,
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(
QdrantGraphEmbeddingsStoreService,
QdrantGraphEmbeddingsStoreService.of(this.store),
),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<
ProcessorConfig & QdrantGraphEmbeddingsConfig,
never,
GraphEmbeddingsStoreRequirements
>({
id: "graph-embeddings-store",
make: (config) => new GraphEmbeddingsStoreService(config),
specs: () => makeGraphEmbeddingsStoreSpecs(),
layer: (config) => QdrantGraphEmbeddingsStoreLive(config),
});
export async function run(): Promise<void> {
await GraphEmbeddingsStoreService.launch("graph-embeddings-store");
await Effect.runPromise(program);
}

View file

@ -9,7 +9,9 @@
*/
import { QdrantClient } from "@qdrant/js-client-rest";
import type { Term } from "@trustgraph/base";
import { errorMessage, type Term } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface QdrantGraphEmbeddingsConfig {
url?: string;
@ -127,3 +129,67 @@ export class QdrantGraphEmbeddingsStore {
);
}
}
export class QdrantGraphEmbeddingsStoreError extends S.TaggedErrorClass<QdrantGraphEmbeddingsStoreError>()(
"QdrantGraphEmbeddingsStoreError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantGraphEmbeddingsStoreServiceShape {
readonly store: (
message: GraphEmbeddingsMessage,
) => Effect.Effect<void, QdrantGraphEmbeddingsStoreError>;
readonly deleteCollection: (
user: string,
collection: string,
) => Effect.Effect<void, QdrantGraphEmbeddingsStoreError>;
}
export class QdrantGraphEmbeddingsStoreService extends Context.Service<
QdrantGraphEmbeddingsStoreService,
QdrantGraphEmbeddingsStoreServiceShape
>()(
"@trustgraph/flow/storage/embeddings/qdrant-graph/QdrantGraphEmbeddingsStoreService",
) {}
const qdrantGraphEmbeddingsStoreError = (operation: string, cause: unknown) =>
new QdrantGraphEmbeddingsStoreError({
operation,
message: errorMessage(cause),
cause,
});
export const makeQdrantGraphEmbeddingsStoreService = (
config: QdrantGraphEmbeddingsConfig = {},
): QdrantGraphEmbeddingsStoreServiceShape => {
const store = new QdrantGraphEmbeddingsStore(config);
return {
store: Effect.fn("QdrantGraphEmbeddingsStore.store")(function* (message) {
return yield* Effect.tryPromise({
try: () => store.store(message),
catch: (cause) => qdrantGraphEmbeddingsStoreError("store", cause),
});
}),
deleteCollection: Effect.fn("QdrantGraphEmbeddingsStore.deleteCollection")(function* (
user,
collection,
) {
return yield* Effect.tryPromise({
try: () => store.deleteCollection(user, collection),
catch: (cause) => qdrantGraphEmbeddingsStoreError("delete-collection", cause),
});
}),
};
};
export const QdrantGraphEmbeddingsStoreLive = (
config: QdrantGraphEmbeddingsConfig = {},
): Layer.Layer<QdrantGraphEmbeddingsStoreService> =>
Layer.succeed(
QdrantGraphEmbeddingsStoreService,
QdrantGraphEmbeddingsStoreService.of(makeQdrantGraphEmbeddingsStoreService(config)),
);

View file

@ -14,47 +14,72 @@ import {
type ProcessorConfig,
type FlowContext,
type Triples,
type Spec,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import { FalkorDBTriplesStore } from "./falkordb.js";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import {
FalkorDBTriplesStoreLive,
FalkorDBTriplesStoreService,
makeFalkorDBTriplesStoreService,
type FalkorDBConfig,
type FalkorDBTriplesStoreError,
} from "./falkordb.js";
export class TriplesStoreService extends FlowProcessor {
private store: FalkorDBTriplesStore;
const onStoreTriplesMessage = Effect.fn("TriplesStoreService.onMessage")(function* (
msg: Triples,
_properties: Record<string, string>,
_flowCtx: FlowContext<FalkorDBTriplesStoreService>,
): Effect.fn.Return<void, FalkorDBTriplesStoreError, FalkorDBTriplesStoreService> {
if (msg.triples.length === 0) return;
const user = msg.metadata?.user ?? "default";
const collection = msg.metadata?.collection ?? "default";
const store = yield* FalkorDBTriplesStoreService;
yield* store.storeTriples(msg.triples, user, collection);
yield* Effect.log(
`[TriplesStore] Stored ${msg.triples.length} triples for ${user}/${collection}`,
);
});
export const makeTriplesStoreSpecs = (): ReadonlyArray<Spec<FalkorDBTriplesStoreService>> => [
new ConsumerSpec<Triples, FalkorDBTriplesStoreError, FalkorDBTriplesStoreService>(
"store-triples-input",
onStoreTriplesMessage,
),
];
export class TriplesStoreService extends FlowProcessor<FalkorDBTriplesStoreService> {
private readonly store = makeFalkorDBTriplesStoreService();
constructor(config: ProcessorConfig) {
super(config);
this.store = new FalkorDBTriplesStore();
this.registerSpecification(
ConsumerSpec.fromPromise<Triples>("store-triples-input", this.onMessage.bind(this)),
);
for (const spec of makeTriplesStoreSpecs()) {
this.registerSpecification(spec);
}
console.log("[TriplesStore] Service initialized");
}
private async onMessage(
msg: Triples,
_properties: Record<string, string>,
_flowCtx: FlowContext,
): Promise<void> {
if (msg.triples.length === 0) return;
const user = msg.metadata?.user ?? "default";
const collection = msg.metadata?.collection ?? "default";
await this.store.storeTriples(msg.triples, user, collection);
console.log(
`[TriplesStore] Stored ${msg.triples.length} triples for ${user}/${collection}`,
override startEffect() {
return super.startEffect().pipe(
Effect.provideService(
FalkorDBTriplesStoreService,
FalkorDBTriplesStoreService.of(this.store),
),
);
}
}
export const program = makeProcessorProgram({
export const program = makeFlowProcessorProgram<ProcessorConfig & FalkorDBConfig, never, FalkorDBTriplesStoreService>({
id: "triples-store",
make: (config) => new TriplesStoreService(config),
specs: () => makeTriplesStoreSpecs(),
layer: (config) => FalkorDBTriplesStoreLive(config),
});
export async function run(): Promise<void> {
await TriplesStoreService.launch("triples-store");
await Effect.runPromise(program);
}

View file

@ -8,7 +8,9 @@
*/
import { createClient, Graph } from "falkordb";
import type { Term, Triple } from "@trustgraph/base";
import { errorMessage, type Term, type Triple } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface FalkorDBConfig {
url?: string;
@ -130,3 +132,71 @@ export class FalkorDBTriplesStore {
);
}
}
export class FalkorDBTriplesStoreError extends S.TaggedErrorClass<FalkorDBTriplesStoreError>()(
"FalkorDBTriplesStoreError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface FalkorDBTriplesStoreServiceShape {
readonly storeTriples: (
triples: ReadonlyArray<Triple>,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly deleteCollection: (
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
}
export class FalkorDBTriplesStoreService extends Context.Service<
FalkorDBTriplesStoreService,
FalkorDBTriplesStoreServiceShape
>()(
"@trustgraph/flow/storage/triples/falkordb/FalkorDBTriplesStoreService",
) {}
const falkorDBTriplesStoreError = (operation: string, cause: unknown) =>
new FalkorDBTriplesStoreError({
operation,
message: errorMessage(cause),
cause,
});
export const makeFalkorDBTriplesStoreService = (
config: FalkorDBConfig = {},
): FalkorDBTriplesStoreServiceShape => {
const store = new FalkorDBTriplesStore(config);
return {
storeTriples: Effect.fn("FalkorDBTriplesStore.storeTriples")((
triples: ReadonlyArray<Triple>,
user: string,
collection: string,
) =>
Effect.tryPromise({
try: () => store.storeTriples(Array.from(triples), user, collection),
catch: (cause) => falkorDBTriplesStoreError("store-triples", cause),
})),
deleteCollection: Effect.fn("FalkorDBTriplesStore.deleteCollection")((
user: string,
collection: string,
) =>
Effect.tryPromise({
try: () => store.deleteCollection(user, collection),
catch: (cause) => falkorDBTriplesStoreError("delete-collection", cause),
})),
};
};
export const FalkorDBTriplesStoreLive = (
config: FalkorDBConfig = {},
): Layer.Layer<FalkorDBTriplesStoreService> =>
Layer.succeed(
FalkorDBTriplesStoreService,
FalkorDBTriplesStoreService.of(makeFalkorDBTriplesStoreService(config)),
);