mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-01 09:29:38 +02:00
Advance TS port Effect workbench
This commit is contained in:
parent
92dae8c374
commit
3515106670
116 changed files with 12286 additions and 9584 deletions
|
|
@ -12,7 +12,6 @@
|
|||
},
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.39.0",
|
||||
"@effect/platform-bun": "4.0.0-beta.65",
|
||||
"@fastify/websocket": "^11.0.0",
|
||||
"@qdrant/js-client-rest": "^1.13.0",
|
||||
"@trustgraph/base": "workspace:*",
|
||||
|
|
@ -20,13 +19,28 @@
|
|||
"fastify": "^5.2.0",
|
||||
"ollama": "^0.6.3",
|
||||
"@mistralai/mistralai": "^1.0.0",
|
||||
"@effect/platform-node": "4.0.0-beta.74",
|
||||
"@effect/platform-node-shared": "4.0.0-beta.74",
|
||||
"@effect/ai-anthropic": "4.0.0-beta.74",
|
||||
"@effect/ai-openai": "4.0.0-beta.74",
|
||||
"@effect/ai-openrouter": "4.0.0-beta.74",
|
||||
"@effect/atom-react": "4.0.0-beta.74",
|
||||
"@effect/openapi-generator": "4.0.0-beta.74",
|
||||
"@effect/opentelemetry": "4.0.0-beta.74",
|
||||
"@effect/platform-browser": "4.0.0-beta.74",
|
||||
"@effect/platform-bun": "4.0.0-beta.74",
|
||||
"@effect/tsgo": "0.13.0",
|
||||
"@effect/sql-pg": "4.0.0-beta.74",
|
||||
"@effect/sql-sqlite-bun": "4.0.0-beta.74",
|
||||
"@effect/sql-sqlite-node": "4.0.0-beta.74",
|
||||
"@effect/vitest": "4.0.0-beta.74",
|
||||
"@modelcontextprotocol/sdk": "^1.12.0",
|
||||
"effect": "4.0.0-beta.65",
|
||||
"effect": "4.0.0-beta.74",
|
||||
"openai": "^4.85.0",
|
||||
"pdfjs-dist": "^5.6.205"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@effect/vitest": "4.0.0-beta.65",
|
||||
"@effect/vitest": "4.0.0-beta.74",
|
||||
"@types/node": "^22.0.0",
|
||||
"typescript": "^5.8.0",
|
||||
"vitest": "^4.1.6"
|
||||
|
|
|
|||
131
ts/packages/flow/src/__tests__/retrieval-rag.test.ts
Normal file
131
ts/packages/flow/src/__tests__/retrieval-rag.test.ts
Normal 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");
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
|
@ -1 +1 @@
|
|||
export { McpToolService } from "./service.js";
|
||||
export { McpToolService, run } from "./service.js";
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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));
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
35
ts/packages/flow/src/gateway/rpc-contract.ts
Normal file
35
ts/packages/flow/src/gateway/rpc-contract.ts
Normal 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);
|
||||
92
ts/packages/flow/src/gateway/rpc-protocol.ts
Normal file
92
ts/packages/flow/src/gateway/rpc-protocol.ts
Normal 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;
|
||||
});
|
||||
109
ts/packages/flow/src/gateway/rpc-server.ts
Normal file
109
ts/packages/flow/src/gateway/rpc-server.ts
Normal 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,
|
||||
);
|
||||
}
|
||||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue