Enforce strict Effect tsgo migrations

This commit is contained in:
elpresidank 2026-06-01 23:19:54 -05:00
parent 64fb23e7d0
commit f6878d4dd7
49 changed files with 5547 additions and 3250 deletions

View file

@ -0,0 +1,113 @@
import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { describe, expect, it } from "vitest";
import {
ConfigServiceError,
makeConfigService,
} from "../config/service.js";
import type {
BackendConsumer,
BackendProducer,
ConfigRequest,
CreateConsumerOptions,
CreateProducerOptions,
PubSubBackend,
} from "@trustgraph/base";
class NoopPubSub implements PubSubBackend {
async createProducer<T>(_options: CreateProducerOptions): Promise<BackendProducer<T>> {
return {
send: async () => undefined,
flush: async () => undefined,
close: async () => undefined,
};
}
async createConsumer<T>(_options: CreateConsumerOptions): Promise<BackendConsumer<T>> {
return {
receive: async () => null,
acknowledge: async () => undefined,
negativeAcknowledge: async () => undefined,
unsubscribe: async () => undefined,
close: async () => undefined,
};
}
async close(): Promise<void> {}
}
const makeService = (persistPath?: string) =>
makeConfigService({
id: "config-test",
manageProcessSignals: false,
pubsub: new NoopPubSub(),
...(persistPath === undefined ? {} : { persistPath }),
});
describe("ConfigService operations", () => {
it("uses tagged errors for invalid mutations", async () => {
const service = makeService();
const putError = await service.handlePut({ operation: "put" } as ConfigRequest)
.catch((caught: unknown) => caught);
const deleteError = await service.handleDelete({ operation: "delete" } as ConfigRequest)
.catch((caught: unknown) => caught);
expect(putError).toBeInstanceOf(ConfigServiceError);
expect(putError).toMatchObject({ _tag: "ConfigServiceError", operation: "put" });
expect(deleteError).toBeInstanceOf(ConfigServiceError);
expect(deleteError).toMatchObject({ _tag: "ConfigServiceError", operation: "delete" });
});
it("persists the workspace-aware config shape through Effect.tryPromise", async () => {
const dir = await mkdtemp(join(tmpdir(), "trustgraph-config-service-"));
const persistPath = join(dir, "config.json");
const service = makeService(persistPath);
await service.handlePut({
operation: "put",
values: [
{ workspace: "alpha", type: "prompt", key: "system", value: "hello" },
],
} as ConfigRequest);
const persisted = await Bun.file(persistPath).json();
await rm(dir, { recursive: true, force: true });
expect(persisted).toEqual({
version: 1,
workspaces: {
alpha: {
prompt: {
system: "hello",
},
},
},
});
});
it("loads the legacy persisted data shape without try/catch", async () => {
const dir = await mkdtemp(join(tmpdir(), "trustgraph-config-service-"));
const persistPath = join(dir, "config.json");
await Bun.write(
persistPath,
`{"version":7,"data":{"prompt":{"system":"legacy"}}}`,
);
const service = makeService(persistPath);
await service.loadFromDisk();
const response = service.handleGet({
operation: "get",
keys: ["prompt", "system"],
} as ConfigRequest);
await rm(dir, { recursive: true, force: true });
expect(response).toEqual({
version: 7,
values: {
system: "legacy",
},
});
});
});

View file

@ -75,7 +75,7 @@ const mcpToolError = (
cause: unknown,
tool?: string,
): McpToolError =>
new McpToolError({
McpToolError.make({
operation,
message: errorMessage(cause),
...(tool === undefined ? {} : { tool }),
@ -166,12 +166,9 @@ const invokeConfiguredTool = Effect.fn("McpToolRuntime.invokeTool")(function* (
const result = yield* Effect.acquireUseRelease(
Effect.tryPromise({
try: async () => {
await client.connect(transport as unknown as Parameters<Client["connect"]>[0]);
return client;
},
try: () => client.connect(transport as unknown as Parameters<Client["connect"]>[0]),
catch: (cause) => mcpToolError("connect", cause, name),
}),
}).pipe(Effect.as(client)),
(connectedClient) =>
Effect.tryPromise({
try: () =>
@ -318,6 +315,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, McpToolR
layer: () => McpToolRuntimeLive,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -132,79 +132,80 @@ const toPromiseRequestor = <TReq, TRes>(
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>;
): Effect.Effect<AgentTool | null> =>
Effect.gen(function* () {
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 () => "",
};
if (name.length === 0) {
yield* Effect.logWarning(`[AgentService] Skipping tool with no name: ${toolId}`);
return null;
}
default:
console.warn(`[AgentService] Unknown tool type "${implType}" for ${name}`);
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: () => Promise.resolve(""),
};
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: () => Promise.resolve(""),
};
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: () => Promise.resolve(""),
};
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: () => Promise.resolve(""),
};
}
default:
yield* Effect.logWarning(`[AgentService] Unknown tool type "${implType}" for ${name}`);
return null;
}
});
const loadConfiguredTools = Effect.fn("AgentRuntime.loadConfiguredTools")(function* (
config: Record<string, unknown>,
@ -231,7 +232,7 @@ const loadConfiguredTools = Effect.fn("AgentRuntime.loadConfiguredTools")(functi
continue;
}
const tool = buildConfiguredTool(toolId, decoded.value);
const tool = yield* buildConfiguredTool(toolId, decoded.value);
if (tool === null) continue;
tools.push(tool);
@ -348,7 +349,7 @@ const executeTool = (
): Effect.Effect<string> =>
Effect.tryPromise({
try: () => tool.execute(input),
catch: (cause) => new AgentToolExecutionError({ message: errorMessage(cause) }),
catch: (cause) => AgentToolExecutionError.make({ message: errorMessage(cause) }),
}).pipe(
Effect.catch((error: AgentToolExecutionError) =>
Effect.succeed(`Error executing tool: ${error.message}`),
@ -473,12 +474,12 @@ const onAgentRequest = Effect.fn("AgentService.onRequest")(function* (
}).pipe(
Effect.catch((error: unknown) =>
Effect.logError(`[AgentService] Error processing request ${requestId}`, {
error: error instanceof Error ? error.message : String(error),
error: errorMessage(error),
}).pipe(
Effect.flatMap(() =>
responseProducer.send(requestId, {
chunk_type: "error",
content: `Agent error: ${error instanceof Error ? error.message : String(error)}`,
content: `Agent error: ${errorMessage(error)}`,
end_of_message: true,
end_of_dialog: true,
}),
@ -538,7 +539,7 @@ export function makeAgentService(config: ProcessorConfig): AgentService {
Effect.provideService(AgentRuntime, runtime),
)),
);
console.log("[AgentService] Service initialized");
Effect.runSync(Effect.log("[AgentService] Service initialized"));
return service;
}
@ -629,6 +630,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, AgentRun
layer: () => AgentRuntimeLive,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -18,9 +18,14 @@ import type {
Term,
Triple,
} from "@trustgraph/base";
import { Effect } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
import type { AgentTool, ToolArg } from "./types.js";
const decodeJsonUnknown = S.decodeUnknownOption(S.UnknownFromJsonString);
/**
* Format a Term to a human-readable string.
*/
@ -41,17 +46,15 @@ function termToString(term: Term): string {
* Parse tool input -- accepts either raw JSON or a plain string question.
*/
function parseQuestion(input: string): string {
try {
const parsed = JSON.parse(input) as Record<string, unknown>;
if (typeof parsed === "object" && parsed !== null && "question" in parsed) {
return String(parsed.question);
}
// If it's a string JSON value, use it directly
if (typeof parsed === "string") {
return parsed;
}
} catch {
// Not valid JSON -- treat as plain text
const decoded = decodeJsonUnknown(input);
if (O.isNone(decoded)) return input;
const parsed = decoded.value;
if (typeof parsed === "object" && parsed !== null && "question" in parsed) {
return String(parsed.question);
}
if (typeof parsed === "string") {
return parsed;
}
return input;
}
@ -83,15 +86,15 @@ export function createKnowledgeQueryTool(
description: "The question to ask the knowledge graph",
},
],
async execute(input: string): Promise<string> {
execute: (input: string): Promise<string> => Effect.runPromise(Effect.gen(function* () {
const question = parseQuestion(input);
console.log(`[KnowledgeQuery] Executing: "${question.slice(0, 60)}..." collection=${collection}`);
yield* Effect.log(`[KnowledgeQuery] Executing: "${question.slice(0, 60)}..." collection=${collection}`);
const request: GraphRagRequest = {
query: question,
...(collection !== undefined ? { collection } : {}),
};
const res = await client.request(request);
console.log(`[KnowledgeQuery] Response (${res.response?.length ?? 0} chars): ${res.error !== undefined ? `ERROR: ${res.error.message}` : `${res.response?.slice(0, 300)}...`}`);
const res = yield* Effect.tryPromise(() => client.request(request));
yield* Effect.log(`[KnowledgeQuery] Response (${res.response?.length ?? 0} chars): ${res.error !== undefined ? `ERROR: ${res.error.message}` : `${res.response?.slice(0, 300)}...`}`);
// Extract explain data if embedded in the response
const rawRes = res as Record<string, unknown>;
@ -100,15 +103,15 @@ export function createKnowledgeQueryTool(
rawRes.explain_triples !== undefined &&
onExplain !== undefined
) {
onExplain({
yield* Effect.sync(() => onExplain({
explainId: (rawRes.explain_id as string) ?? "",
triples: rawRes.explain_triples as Triple[],
});
}));
}
if (res.error !== undefined) return `Error: ${res.error.message}`;
return res.response;
},
})),
};
}
@ -130,16 +133,16 @@ export function createDocumentQueryTool(
description: "The question to search documents for",
},
],
async execute(input: string): Promise<string> {
execute: (input: string): Promise<string> => Effect.runPromise(Effect.gen(function* () {
const question = parseQuestion(input);
const request: DocumentRagRequest = {
query: question,
...(collection !== undefined ? { collection } : {}),
};
const res = await client.request(request);
const res = yield* Effect.tryPromise(() => client.request(request));
if (res.error !== undefined) return `Error: ${res.error.message}`;
return res.response;
},
})),
};
}
@ -152,39 +155,42 @@ function parseTriplesInput(input: string): {
o?: Term;
limit?: number;
} {
try {
const parsed = JSON.parse(input) as Record<string, unknown>;
const toTerm = (val: unknown): Term | undefined => {
if (typeof val === "string") {
return { type: "LITERAL", value: val };
}
if (typeof val === "object" && val !== null && "type" in val) {
return val as Term;
}
return undefined;
};
const result: {
s?: Term;
p?: Term;
o?: Term;
limit?: number;
} = {};
const s = toTerm(parsed.subject ?? parsed.s);
const p = toTerm(parsed.predicate ?? parsed.p);
const o = toTerm(parsed.object ?? parsed.o);
if (s !== undefined) result.s = s;
if (p !== undefined) result.p = p;
if (o !== undefined) result.o = o;
if (typeof parsed.limit === "number") result.limit = parsed.limit;
return result;
} catch {
// If not valid JSON, treat as a subject search
const decoded = decodeJsonUnknown(input);
if (
O.isNone(decoded) ||
typeof decoded.value !== "object" ||
decoded.value === null
) {
return {
s: { type: "LITERAL", value: input },
};
}
const parsed = decoded.value as Record<string, unknown>;
const toTerm = (val: unknown): Term | undefined => {
if (typeof val === "string") {
return { type: "LITERAL", value: val };
}
if (typeof val === "object" && val !== null && "type" in val) {
return val as Term;
}
return undefined;
};
const result: {
s?: Term;
p?: Term;
o?: Term;
limit?: number;
} = {};
const s = toTerm(parsed.subject ?? parsed.s);
const p = toTerm(parsed.predicate ?? parsed.p);
const o = toTerm(parsed.object ?? parsed.o);
if (s !== undefined) result.s = s;
if (p !== undefined) result.p = p;
if (o !== undefined) result.o = o;
if (typeof parsed.limit === "number") result.limit = parsed.limit;
return result;
}
/**
@ -216,7 +222,7 @@ export function createTriplesQueryTool(
description: "The object entity to search for (optional)",
},
],
async execute(input: string): Promise<string> {
execute: (input: string): Promise<string> => Effect.runPromise(Effect.gen(function* () {
const { s, p, o, limit } = parseTriplesInput(input);
const request: TriplesQueryRequest = {
limit: limit ?? 20,
@ -225,7 +231,7 @@ export function createTriplesQueryTool(
...(o !== undefined ? { o } : {}),
...(collection !== undefined ? { collection } : {}),
};
const res = await client.request(request);
const res = yield* Effect.tryPromise(() => client.request(request));
if (res.error !== undefined) return `Error: ${res.error.message}`;
@ -238,7 +244,7 @@ export function createTriplesQueryTool(
`(${termToString(t.s)}) -[${termToString(t.p)}]-> (${termToString(t.o)})`,
);
return lines.join("\n");
},
})),
};
}
@ -258,12 +264,12 @@ export function createMcpTool(
name: toolName,
description,
args,
async execute(input: string): Promise<string> {
const res = await client.request({ name: toolName, parameters: input });
execute: (input: string): Promise<string> => Effect.runPromise(Effect.gen(function* () {
const res = yield* Effect.tryPromise(() => client.request({ name: toolName, parameters: input }));
if (res.error !== undefined) return `Error: ${res.error.message}`;
if (res.text !== undefined) return res.text;
if (res.object !== undefined) return res.object;
return "No content";
},
})),
};
}

View file

@ -91,7 +91,7 @@ export function makeChunkingService(config: ProcessorConfig): ChunkingService {
const service = makeFlowProcessor(config, {
specifications: makeChunkingSpecs(),
});
console.log("[ChunkingService] Service initialized");
Effect.runSync(Effect.log("[ChunkingService] Service initialized"));
return service;
}
@ -102,6 +102,6 @@ export const program = makeFlowProcessorProgram({
specs: () => makeChunkingSpecs(),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -11,7 +11,7 @@
* Python reference: trustgraph-flow/trustgraph/config/service/service.py
*/
import { Effect } from "effect";
import { Duration, Effect } from "effect";
import * as S from "effect/Schema";
import {
makeAsyncProcessor,
@ -45,6 +45,20 @@ const ConfigPushSchema = S.Struct({
config: S.Record(S.String, S.Unknown),
});
export class ConfigServiceError extends S.TaggedErrorClass<ConfigServiceError>()(
"ConfigServiceError",
{
message: S.String,
operation: S.String,
},
) {}
const configServiceError = (operation: string, cause: unknown): ConfigServiceError =>
ConfigServiceError.make({
operation,
message: errorMessage(cause),
});
const DEFAULT_WORKSPACE = "default";
interface ConfigKeyLike {
@ -62,6 +76,14 @@ interface ConfigValueLike {
type NamespaceStore = Map<string, unknown>;
type WorkspaceStore = Map<string, NamespaceStore>;
const PersistedConfigSchema = S.Struct({
version: S.optionalKey(S.Number),
data: S.optionalKey(S.Record(S.String, S.Record(S.String, S.Unknown))),
workspaces: S.optionalKey(S.Record(S.String, S.Record(S.String, S.Record(S.String, S.Unknown)))),
});
const PersistedConfigJsonSchema = PersistedConfigSchema.pipe(S.fromJsonString);
type PersistedConfig = typeof PersistedConfigSchema.Type;
function isRecord(value: unknown): value is Record<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
@ -70,13 +92,63 @@ function optionalString(value: unknown): string | undefined {
return typeof value === "string" && value.length > 0 ? value : undefined;
}
function toPersistedWorkspaces(
store: Map<string, WorkspaceStore>,
): Record<string, Record<string, Record<string, unknown>>> {
const workspaces: Record<string, Record<string, Record<string, unknown>>> = {};
for (const [workspace, ws] of 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;
}
workspaces[workspace] = workspaceData;
}
return workspaces;
}
function hydratePersistedConfig(
store: Map<string, WorkspaceStore>,
parsed: PersistedConfig,
): void {
store.clear();
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);
}
store.set(workspace, ws);
}
return;
}
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);
}
store.set(DEFAULT_WORKSPACE, ws);
}
export type ConfigService = AsyncProcessorRuntime & Record<string, any>;
export function makeConfigService(config: ConfigServiceConfig): ConfigService {
const service = makeAsyncProcessor(config, {
run: async () => {
await service.run();
},
run: () => service.run(),
}) as ConfigService;
const baseStop = service.stop;
service.store = new Map<string, WorkspaceStore>();
@ -88,115 +160,183 @@ export function makeConfigService(config: ConfigServiceConfig): ConfigService {
Object.assign(service, {
run: async function(this: ConfigService): Promise<void> {
// Optionally load persisted state
if (this.persistPath !== null) {
await this.loadFromDisk();
}
run: function(this: ConfigService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
// Optionally load persisted state
if (service.persistPath !== null) {
yield* Effect.tryPromise({
try: () => service.loadFromDisk(),
catch: (cause) => configServiceError("load", cause),
});
}
// Create producers
this.responseProducer = await this.pubsub.createProducer<ConfigResponse>({
topic: topics.configResponse,
schema: ConfigResponseSchema,
});
this.pushProducer = await this.pubsub.createProducer<ConfigPush>({
topic: topics.configPush,
schema: ConfigPushSchema,
});
// Create producers
service.responseProducer = yield* Effect.tryPromise({
try: () =>
service.pubsub.createProducer<ConfigResponse>({
topic: topics.configResponse,
schema: ConfigResponseSchema,
}),
catch: (cause) => configServiceError("response-producer", cause),
});
service.pushProducer = yield* Effect.tryPromise({
try: () =>
service.pubsub.createProducer<ConfigPush>({
topic: topics.configPush,
schema: ConfigPushSchema,
}),
catch: (cause) => configServiceError("push-producer", cause),
});
// Create consumer for config requests
this.consumer = await this.pubsub.createConsumer<ConfigRequest>({
topic: topics.configRequest,
subscription: `${this.config.id}-config-request`,
schema: ConfigRequestSchema,
});
// Create consumer for config requests
service.consumer = yield* Effect.tryPromise({
try: () =>
service.pubsub.createConsumer<ConfigRequest>({
topic: topics.configRequest,
subscription: `${service.config.id}-config-request`,
schema: ConfigRequestSchema,
}),
catch: (cause) => configServiceError("consumer", cause),
});
// Push initial config
await this.pushConfig();
// Push initial config
yield* Effect.tryPromise({
try: () => service.pushConfig(),
catch: (cause) => configServiceError("push-initial-config", cause),
});
console.log(`[ConfigService] Listening on ${topics.configRequest}`);
yield* Effect.log(`[ConfigService] Listening on ${topics.configRequest}`);
// Main consume loop
while (this.running) {
try {
const consumer = this.consumer;
if (consumer === null) throw new Error("Config consumer not started");
// Main consume loop
while (service.running) {
const shouldContinue = yield* Effect.gen(function* () {
const consumer = service.consumer;
if (consumer === null) {
return yield* configServiceError("consume", "Config consumer not started");
}
const msg = await consumer.receive(2000);
if (msg === null) continue;
const msg = yield* Effect.tryPromise({
try: () => consumer.receive(2000),
catch: (cause) => configServiceError("consume-receive", cause),
});
if (msg === null) return true;
await this.handleMessage(msg);
await consumer.acknowledge(msg);
} catch (err) {
if (!this.running) break;
console.error("[ConfigService] Error in consume loop:", err);
await sleep(1000);
}
}
yield* Effect.tryPromise({
try: () => service.handleMessage(msg),
catch: (cause) => configServiceError("consume-handle", cause),
});
yield* Effect.tryPromise({
try: () => consumer.acknowledge(msg),
catch: (cause) => configServiceError("consume-acknowledge", cause),
});
return true;
}).pipe(
Effect.catch((err) => {
if (!service.running) return Effect.succeed(false);
return Effect.logError("[ConfigService] Error in consume loop", { error: err.message }).pipe(
Effect.flatMap(() => Effect.sleep(Duration.millis(1000))),
Effect.as(true),
);
}),
);
if (!shouldContinue) break;
}
}),
);
},
handleMessage: async function(this: ConfigService, msg: Message<ConfigRequest>): Promise<void> {
const request = await Effect.runPromise(S.decodeUnknownEffect(ConfigRequestSchema)(msg.value()));
const props = msg.properties();
const requestId = props.id;
handleMessage: function(this: ConfigService, msg: Message<ConfigRequest>): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const request = yield* S.decodeUnknownEffect(ConfigRequestSchema)(msg.value()).pipe(
Effect.mapError((cause) => configServiceError("decode", cause)),
);
const props = msg.properties();
const requestId = props.id;
if (requestId === undefined || requestId.length === 0) {
console.warn("[ConfigService] Received request without id, ignoring");
return;
}
if (requestId === undefined || requestId.length === 0) {
yield* Effect.logWarning("[ConfigService] Received request without id, ignoring");
return;
}
try {
const response = await this.handleOperation(request);
const responseProducer = this.responseProducer;
if (responseProducer === null) throw new Error("Config response producer not started");
await responseProducer.send(response, { id: requestId });
} catch (err) {
const message = errorMessage(err);
const responseProducer = this.responseProducer;
if (responseProducer === null) throw new Error("Config response producer not started");
await responseProducer.send(
{
error: { type: "config-error", message },
},
{ id: requestId },
);
}
const sendResponse = (response: ConfigResponse): Effect.Effect<void, ConfigServiceError> =>
Effect.gen(function* () {
const responseProducer = service.responseProducer;
if (responseProducer === null) {
return yield* configServiceError("respond", "Config response producer not started");
}
yield* Effect.tryPromise({
try: () => responseProducer.send(response, { id: requestId }),
catch: (cause) => configServiceError("respond", cause),
});
});
yield* Effect.gen(function* () {
const response = yield* Effect.tryPromise<ConfigResponse, ConfigServiceError>({
try: () => service.handleOperation(request),
catch: (cause) => configServiceError("operation", cause),
});
yield* sendResponse(response);
}).pipe(
Effect.catch((err) =>
sendResponse({
error: { type: "config-error", message: err.message },
}),
),
);
}),
);
},
handleOperation: async function(this: ConfigService, request: ConfigRequest): Promise<ConfigResponse> {
const op: ConfigOperation = request.operation;
handleOperation: function(this: ConfigService, request: ConfigRequest): Promise<ConfigResponse> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const op: ConfigOperation = request.operation;
switch (op) {
case "get":
return this.handleGet(request);
switch (op) {
case "get":
return service.handleGet(request);
case "put":
return await this.handlePut(request);
case "put":
return yield* Effect.tryPromise<ConfigResponse, ConfigServiceError>({
try: () => service.handlePut(request),
catch: (cause) => configServiceError("put", cause),
});
case "delete":
return await this.handleDelete(request);
case "delete":
return yield* Effect.tryPromise<ConfigResponse, ConfigServiceError>({
try: () => service.handleDelete(request),
catch: (cause) => configServiceError("delete", cause),
});
case "list":
return this.handleList(request);
case "list":
return service.handleList(request);
case "config":
return this.handleConfigDump(request);
case "config":
return service.handleConfigDump(request);
case "getvalues":
return this.handleGetValues(request);
case "getvalues":
return service.handleGetValues(request);
case "getvalues-all-ws":
return this.handleGetValuesAllWorkspaces(request);
case "getvalues-all-ws":
return service.handleGetValuesAllWorkspaces(request);
default:
throw new Error(`Unknown config operation: ${op as string}`);
}
default:
return yield* configServiceError("operation", `Unknown config operation: ${op as string}`);
}
}),
);
},
@ -364,76 +504,104 @@ export function makeConfigService(config: ConfigServiceConfig): ConfigService {
handlePut: async function(this: ConfigService, request: ConfigRequest): Promise<ConfigResponse> {
const values = this.configValues(request);
if (values.length === 0) throw new Error("Put requires config values");
handlePut: function(this: ConfigService, request: ConfigRequest): Promise<ConfigResponse> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const values = service.configValues(request);
if (values.length === 0) return yield* configServiceError("put", "Put requires config values");
for (const item of values) {
this.namespaceStore(item.workspace ?? DEFAULT_WORKSPACE, item.type, true)?.set(item.key, item.value);
}
for (const item of values) {
service.namespaceStore(item.workspace ?? DEFAULT_WORKSPACE, item.type, true)?.set(item.key, item.value);
}
this.version++;
await this.persist();
await this.pushConfig();
service.version++;
yield* Effect.tryPromise({
try: () => service.persist(),
catch: (cause) => configServiceError("persist", cause),
});
yield* Effect.tryPromise({
try: () => service.pushConfig(),
catch: (cause) => configServiceError("push-config", cause),
});
return { version: this.version };
return { version: service.version };
}),
);
},
handleDelete: async function(this: ConfigService, 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);
handleDelete: function(this: ConfigService, request: ConfigRequest): Promise<ConfigResponse> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const workspace = service.workspaceFor(request);
const objectKeys = service.objectKeys(request);
if (objectKeys.length > 0) {
for (const key of objectKeys) {
const ws = service.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);
}
}
service.version++;
yield* Effect.tryPromise({
try: () => service.persist(),
catch: (cause) => configServiceError("persist", cause),
});
yield* Effect.tryPromise({
try: () => service.pushConfig(),
catch: (cause) => configServiceError("push-config", cause),
});
return { version: service.version };
}
}
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
ws.delete(namespace);
} else {
// Delete specific keys within namespace
const subMap = ws.get(namespace);
if (subMap !== undefined) {
for (let i = 1; i < keys.length; i++) {
subMap.delete(keys[i]);
const keys = service.stringKeys(request);
if (keys.length === 0) {
return yield* configServiceError("delete", "Delete requires at least one key");
}
if (subMap.size === 0) {
const namespace = keys[0];
const ws = service.workspaceStore(workspace, false);
if (ws === undefined) return { version: service.version };
if (keys.length === 1) {
// Delete entire namespace
ws.delete(namespace);
} else {
// Delete specific keys within 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) {
ws.delete(namespace);
}
}
}
}
}
this.version++;
await this.persist();
await this.pushConfig();
service.version++;
yield* Effect.tryPromise({
try: () => service.persist(),
catch: (cause) => configServiceError("persist", cause),
});
yield* Effect.tryPromise({
try: () => service.pushConfig(),
catch: (cause) => configServiceError("push-config", cause),
});
return { version: this.version };
return { version: service.version };
}),
);
},
@ -528,130 +696,140 @@ export function makeConfigService(config: ConfigServiceConfig): ConfigService {
pushConfig: async function(this: ConfigService): Promise<void> {
const pushProducer = this.pushProducer;
if (pushProducer === null) return;
pushConfig: function(this: ConfigService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const pushProducer = service.pushProducer;
if (pushProducer === null) return;
const config: Record<string, unknown> = {};
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;
}
config[namespace] = obj;
}
await pushProducer.send({
version: this.version,
config,
});
console.log(`[ConfigService] Pushed configuration version ${this.version}`);
},
persist: async function(this: ConfigService): Promise<void> {
const persistPath = this.persistPath;
if (persistPath === null) return;
try {
const workspaces: Record<string, Record<string, Record<string, unknown>>> = {};
for (const [workspace, ws] of this.store) {
const workspaceData: Record<string, Record<string, unknown>> = {};
for (const [namespace, subMap] of ws) {
const config: Record<string, unknown> = {};
const ws = service.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;
}
workspaceData[namespace] = obj;
config[namespace] = obj;
}
workspaces[workspace] = workspaceData;
}
const json = JSON.stringify(
{ version: this.version, workspaces },
null,
2,
);
yield* Effect.tryPromise({
try: () =>
pushProducer.send({
version: service.version,
config,
}),
catch: (cause) => configServiceError("push-config", cause),
});
await writeTextFile(persistPath, json);
} catch (err) {
await Effect.runPromise(Effect.logError("[ConfigService] Failed to persist config", { error: errorMessage(err) }));
}
yield* Effect.log(`[ConfigService] Pushed configuration version ${service.version}`);
}),
);
},
loadFromDisk: async function(this: ConfigService): Promise<void> {
const persistPath = this.persistPath;
if (persistPath === null) return;
persist: function(this: ConfigService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const persistPath = service.persistPath;
if (persistPath === null) return;
const payload = {
version: service.version,
workspaces: toPersistedWorkspaces(service.store),
};
try {
const raw = await readTextFile(persistPath);
const parsed = JSON.parse(raw) as {
version: number;
data?: Record<string, Record<string, unknown>>;
workspaces?: Record<string, Record<string, Record<string, unknown>>>;
};
const json = yield* S.encodeUnknownEffect(S.UnknownFromJsonString)(payload).pipe(
Effect.mapError((cause) => configServiceError("persist-encode", cause)),
);
this.version = parsed.version ?? 0;
this.store.clear();
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);
}
} 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}, workspaces=${this.store.size})`,
);
} catch {
// File doesn't exist yet or is invalid — start fresh
await Effect.runPromise(Effect.log("[ConfigService] No persisted config found, starting fresh"));
}
yield* Effect.tryPromise({
try: () => writeTextFile(persistPath, json),
catch: (cause) => configServiceError("persist-write", cause),
});
}).pipe(
Effect.catch((err) =>
Effect.logError("[ConfigService] Failed to persist config", { error: err.message }),
),
),
);
},
stop: async function(this: ConfigService): Promise<void> {
if (this.consumer !== null) {
await this.consumer.close();
this.consumer = null;
}
if (this.responseProducer !== null) {
await this.responseProducer.close();
this.responseProducer = null;
}
if (this.pushProducer !== null) {
await this.pushProducer.close();
this.pushProducer = null;
}
await baseStop();
loadFromDisk: function(this: ConfigService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const persistPath = service.persistPath;
if (persistPath === null) return;
const parsed = yield* Effect.gen(function* () {
const raw = yield* Effect.tryPromise({
try: () => readTextFile(persistPath),
catch: (cause) => configServiceError("persist-read", cause),
});
return yield* S.decodeUnknownEffect(PersistedConfigJsonSchema)(raw).pipe(
Effect.mapError((cause) => configServiceError("persist-decode", cause)),
);
}).pipe(
Effect.catch(() =>
Effect.log("[ConfigService] No persisted config found, starting fresh").pipe(
Effect.as(null as PersistedConfig | null),
),
),
);
if (parsed === null) return;
service.version = parsed.version ?? 0;
hydratePersistedConfig(service.store, parsed);
yield* Effect.log(`[ConfigService] Loaded persisted config (version=${service.version}, workspaces=${service.store.size})`);
}),
);
},
stop: function(this: ConfigService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const consumer = service.consumer;
if (consumer !== null) {
yield* Effect.tryPromise({
try: () => consumer.close(),
catch: (cause) => configServiceError("close-consumer", cause),
});
service.consumer = null;
}
const responseProducer = service.responseProducer;
if (responseProducer !== null) {
yield* Effect.tryPromise({
try: () => responseProducer.close(),
catch: (cause) => configServiceError("close-response-producer", cause),
});
service.responseProducer = null;
}
const pushProducer = service.pushProducer;
if (pushProducer !== null) {
yield* Effect.tryPromise({
try: () => pushProducer.close(),
catch: (cause) => configServiceError("close-push-producer", cause),
});
service.pushProducer = null;
}
yield* Effect.tryPromise({
try: () => baseStop(),
catch: (cause) => configServiceError("stop", cause),
});
}),
);
}
});
@ -660,10 +838,6 @@ export function makeConfigService(config: ConfigServiceConfig): ConfigService {
export const ConfigService = makeConfigService;
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
export const loadConfigServiceRuntimeConfig = Effect.fn("loadConfigServiceRuntimeConfig")(function* () {
const processorConfig = yield* loadProcessorRuntimeConfig("config-svc", {
manageProcessSignals: false,
@ -681,6 +855,6 @@ export const program = makeProcessorProgram({
make: (config) => makeConfigService(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -12,6 +12,7 @@
import {
makeAsyncProcessor,
makeProcessorProgram,
type ProcessorConfig,
type AsyncProcessorRuntime,
topics,
@ -19,10 +20,11 @@ import {
type KnowledgeResponse,
type Triple,
type Term,
errorMessage,
} from "@trustgraph/base";
import { makeProcessorProgram } from "@trustgraph/base";
import type { Message } from "@trustgraph/base";
import { Effect } from "effect";
import { Config, Duration, Effect } from "effect";
import * as S from "effect/Schema";
import { ensureDirectory, joinPath, readTextFile, writeTextFile } from "../runtime/effect-files.js";
export interface KnowledgeCoreServiceConfig extends ProcessorConfig {
@ -42,18 +44,88 @@ interface DocumentEmbeddingsCore {
export type KnowledgeCoreService = AsyncProcessorRuntime & Record<string, any>;
export class KnowledgeCoreServiceError extends S.TaggedErrorClass<KnowledgeCoreServiceError>()(
"KnowledgeCoreServiceError",
{
message: S.String,
operation: S.String,
},
) {}
interface KnowledgeResponseProducer {
send(response: KnowledgeResponse, properties: { id: string }): Promise<void>;
close(): Promise<void>;
}
interface CloseableResource {
close(): Promise<void>;
}
const knowledgeCoreServiceError = (operation: string, cause: unknown): KnowledgeCoreServiceError =>
KnowledgeCoreServiceError.make({
operation,
message: errorMessage(cause),
});
const tryPromise = <A>(
operation: string,
evaluate: () => Promise<A>,
): Effect.Effect<A, KnowledgeCoreServiceError> =>
Effect.tryPromise({
try: evaluate,
catch: (cause) => knowledgeCoreServiceError(operation, cause),
});
const trySync = <A>(
operation: string,
evaluate: () => A,
): Effect.Effect<A, KnowledgeCoreServiceError> =>
Effect.try({
try: evaluate,
catch: (cause) => knowledgeCoreServiceError(operation, cause),
});
const failPromise = (operation: string, cause: unknown): Promise<never> =>
Effect.runPromise(Effect.fail(knowledgeCoreServiceError(operation, cause)));
const sendResponse = (
service: KnowledgeCoreService,
response: KnowledgeResponse,
requestId: string,
operation = "respond",
): Effect.Effect<void, KnowledgeCoreServiceError> =>
Effect.gen(function* () {
const responseProducer = service.responseProducer as KnowledgeResponseProducer | null | undefined;
if (responseProducer === null || responseProducer === undefined) {
return yield* knowledgeCoreServiceError(operation, "Knowledge response producer not started");
}
yield* tryPromise(operation, () => responseProducer.send(response, { id: requestId }));
});
const closeResource = (
resource: CloseableResource,
operation: string,
): Effect.Effect<void> =>
tryPromise(operation, () => resource.close()).pipe(
Effect.catch((error) =>
Effect.logError("[KnowledgeCoreService] Failed to close resource", {
error: error.message,
operation: error.operation,
}),
),
);
export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): KnowledgeCoreService {
const service = makeAsyncProcessor(config, {
run: async () => {
await service.run();
},
run: () => service.run(),
}) as KnowledgeCoreService;
const baseStop = service.stop;
service.cores = new Map<string, KnowledgeCore>();
service.deCores = new Map<string, DocumentEmbeddingsCore[]>();
service.consumer = null;
service.responseProducer = null;
const dataDir = config.dataDir ?? process.env.KNOWLEDGE_DATA_DIR ?? "./data/knowledge";
const dataDir = config.dataDir ?? "./data/knowledge";
service.dataDir = dataDir;
service.persistPath = joinPath(dataDir, "knowledge-state.json");
Object.assign(service, {
@ -66,68 +138,106 @@ export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): Kn
run: async function(this: KnowledgeCoreService): Promise<void> {
await ensureDirectory(this.dataDir);
// Load persisted state
await this.loadFromDisk();
run: function(this: KnowledgeCoreService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
if (config.dataDir === undefined) {
const configuredDataDir = yield* Config.string("KNOWLEDGE_DATA_DIR").pipe(
Config.withDefault("./data/knowledge"),
);
service.dataDir = configuredDataDir;
service.persistPath = joinPath(configuredDataDir, "knowledge-state.json");
}
// Create producer
this.responseProducer = await this.pubsub.createProducer<KnowledgeResponse>({
topic: topics.knowledgeResponse,
});
yield* tryPromise("ensure-directory", () => ensureDirectory(service.dataDir));
// Load persisted state
yield* tryPromise("load", () => service.loadFromDisk());
// Create consumer
this.consumer = await this.pubsub.createConsumer<KnowledgeRequest>({
topic: topics.knowledgeRequest,
subscription: `${this.config.id}-knowledge-request`,
});
// Create producer
service.responseProducer = yield* tryPromise("response-producer", () =>
service.pubsub.createProducer<KnowledgeResponse>({
topic: topics.knowledgeResponse,
}),
);
console.log(`[KnowledgeCoreService] Listening on ${topics.knowledgeRequest}`);
// Create consumer
service.consumer = yield* tryPromise("consumer", () =>
service.pubsub.createConsumer<KnowledgeRequest>({
topic: topics.knowledgeRequest,
subscription: `${service.config.id}-knowledge-request`,
}),
);
// Main consume loop
while (this.running) {
try {
const msg = await this.consumer.receive(2000);
if (msg === null) continue;
yield* Effect.log(`[KnowledgeCoreService] Listening on ${topics.knowledgeRequest}`);
await this.handleMessage(msg);
await this.consumer.acknowledge(msg);
} catch (err) {
if (!this.running) break;
console.error("[KnowledgeCoreService] Error in consume loop:", err);
await sleep(1000);
}
}
// Main consume loop
while (service.running) {
const shouldContinue = yield* Effect.gen(function* () {
const consumer = service.consumer;
if (consumer === null || consumer === undefined) {
return yield* knowledgeCoreServiceError("consume", "Knowledge request consumer not started");
}
const msg = yield* tryPromise("consume-receive", () => consumer.receive(2000));
if (msg === null) return true;
yield* tryPromise("consume-handle", () => service.handleMessage(msg));
yield* tryPromise("consume-acknowledge", () => consumer.acknowledge(msg));
return true;
}).pipe(
Effect.catch((error) => {
if (!service.running) return Effect.succeed(false);
return Effect.logError("[KnowledgeCoreService] Error in consume loop", {
error: error.message,
}).pipe(
Effect.flatMap(() => Effect.sleep(Duration.millis(1000))),
Effect.as(true),
);
}),
);
if (!shouldContinue) break;
}
}),
);
},
handleMessage: async function(this: KnowledgeCoreService, msg: Message<KnowledgeRequest>): Promise<void> {
const request = msg.value();
const props = msg.properties();
const requestId = props.id;
handleMessage: function(this: KnowledgeCoreService, msg: Message<KnowledgeRequest>): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const request = msg.value();
const props = msg.properties();
const requestId = props.id;
if (requestId === undefined || requestId.length === 0) {
console.warn("[KnowledgeCoreService] Received request without id, ignoring");
return;
}
if (requestId === undefined || requestId.length === 0) {
yield* Effect.logWarning("[KnowledgeCoreService] Received request without id, ignoring");
return;
}
try {
await this.handleOperation(request, requestId);
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
await this.responseProducer!.send(
{ error: { type: "knowledge-error", message } },
{ id: requestId },
);
}
yield* tryPromise("operation", () => service.handleOperation(request, requestId)).pipe(
Effect.catch((error) =>
sendResponse(
service,
{ error: { type: "knowledge-error", message: error.message } },
requestId,
"respond-error",
),
),
);
}),
);
},
handleOperation: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
handleOperation: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
switch (request.operation) {
case "list-kg-cores":
return this.listKgCores(request, requestId);
@ -152,7 +262,7 @@ export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): Kn
case "load-de-core":
return this.loadDeCore(request, requestId);
default:
throw new Error(`Unknown knowledge operation: ${request.operation as string}`);
return failPromise("operation", `Unknown knowledge operation: ${request.operation as string}`);
}
},
@ -185,231 +295,297 @@ export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): Kn
listKgCores: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const prefix = user.length > 0 ? `${user}:` : "";
listKgCores: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const prefix = user.length > 0 ? `${user}:` : "";
const ids: string[] = [];
for (const key of (this.cores as Map<string, KnowledgeCore>).keys()) {
if (prefix.length === 0 || key.startsWith(prefix)) {
// Extract the ID portion after the user prefix
const id = key.slice(prefix.length);
ids.push(id);
}
}
const ids: string[] = [];
for (const key of (service.cores as Map<string, KnowledgeCore>).keys()) {
if (prefix.length === 0 || key.startsWith(prefix)) {
// Extract the ID portion after the user prefix
const id = key.slice(prefix.length);
ids.push(id);
}
}
await this.responseProducer!.send({ ids }, { id: requestId });
},
getKgCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
const core = this.cores.get(key);
if (core === undefined) {
throw new Error(`Knowledge core not found: ${key}`);
}
// Send triples and embeddings in batches
const BATCH_SIZE = 100;
// Send triples in batches
for (let i = 0; i < core.triples.length; i += BATCH_SIZE) {
const batch = core.triples.slice(i, i + BATCH_SIZE);
const isLast = i + BATCH_SIZE >= core.triples.length && core.graphEmbeddings.length === 0;
await this.responseProducer!.send(
{ triples: batch, eos: isLast },
{ id: requestId },
);
}
// Send graph embeddings in batches
for (let i = 0; i < core.graphEmbeddings.length; i += BATCH_SIZE) {
const batch = core.graphEmbeddings.slice(i, i + BATCH_SIZE);
const isLast = i + BATCH_SIZE >= core.graphEmbeddings.length;
await this.responseProducer!.send(
{ graphEmbeddings: batch, "graph-embeddings": batch, eos: isLast } as KnowledgeResponse,
{ id: requestId },
);
}
// If core was empty, send a final eos
if (core.triples.length === 0 && core.graphEmbeddings.length === 0) {
await this.responseProducer!.send({ eos: true }, { id: requestId });
}
},
deleteKgCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
this.cores.delete(key);
await this.persist();
console.log(`[KnowledgeCoreService] Deleted core: ${key}`);
await this.responseProducer!.send({}, { id: requestId });
},
putKgCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
let core = this.cores.get(key);
if (core === undefined) {
core = { triples: [], graphEmbeddings: [] };
this.cores.set(key, core);
}
// Append triples if provided
if (request.triples !== undefined && request.triples.length > 0) {
core.triples.push(...request.triples);
}
// Append graph embeddings if provided
const graphEmbeddings = this.graphEmbeddings(request);
if (graphEmbeddings.length > 0) {
core.graphEmbeddings.push(...graphEmbeddings);
}
await this.persist();
console.log(
`[KnowledgeCoreService] Updated core ${key}: triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length}`,
yield* sendResponse(service, { ids }, requestId);
}),
);
await this.responseProducer!.send({}, { id: requestId });
},
loadKgCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
getKgCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
const core = this.cores.get(key);
if (core === undefined) {
throw new Error(`Knowledge core not found: ${key}`);
}
const core = service.cores.get(key);
if (core === undefined) {
return yield* knowledgeCoreServiceError("get-kg-core", `Knowledge core not found: ${key}`);
}
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();
}
}
// Send triples and embeddings in batches
const BATCH_SIZE = 100;
console.log(
`[KnowledgeCoreService] Loaded core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length})`,
// Send triples in batches
for (let i = 0; i < core.triples.length; i += BATCH_SIZE) {
const batch = core.triples.slice(i, i + BATCH_SIZE);
const isLast = i + BATCH_SIZE >= core.triples.length && core.graphEmbeddings.length === 0;
yield* sendResponse(
service,
{ triples: batch, eos: isLast },
requestId,
"respond-kg-triples",
);
}
// Send graph embeddings in batches
for (let i = 0; i < core.graphEmbeddings.length; i += BATCH_SIZE) {
const batch = core.graphEmbeddings.slice(i, i + BATCH_SIZE);
const isLast = i + BATCH_SIZE >= core.graphEmbeddings.length;
yield* sendResponse(
service,
{ graphEmbeddings: batch, "graph-embeddings": batch, eos: isLast } as KnowledgeResponse,
requestId,
"respond-kg-embeddings",
);
}
// If core was empty, send a final eos
if (core.triples.length === 0 && core.graphEmbeddings.length === 0) {
yield* sendResponse(service, { eos: true }, requestId, "respond-kg-empty");
}
}),
);
await this.responseProducer!.send({}, { id: requestId });
},
unloadKgCore: async function(this: KnowledgeCoreService, _request: KnowledgeRequest, requestId: string): Promise<void> {
await this.responseProducer!.send({}, { id: requestId });
deleteKgCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
service.cores.delete(key);
yield* tryPromise("persist-delete-kg-core", () => service.persist());
yield* Effect.log(`[KnowledgeCoreService] Deleted core: ${key}`);
yield* sendResponse(service, {}, requestId);
}),
);
},
listDeCores: async function(this: KnowledgeCoreService, 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 });
putKgCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
let core = service.cores.get(key);
if (core === undefined) {
core = { triples: [], graphEmbeddings: [] };
service.cores.set(key, core);
}
// Append triples if provided
if (request.triples !== undefined && request.triples.length > 0) {
core.triples.push(...request.triples);
}
// Append graph embeddings if provided
const graphEmbeddings = service.graphEmbeddings(request);
if (graphEmbeddings.length > 0) {
core.graphEmbeddings.push(...graphEmbeddings);
}
yield* tryPromise("persist-put-kg-core", () => service.persist());
yield* Effect.log(
`[KnowledgeCoreService] Updated core ${key}: triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length}`,
);
yield* sendResponse(service, {}, requestId);
}),
);
},
getDeCore: async function(this: KnowledgeCoreService, 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}`);
loadKgCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
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 });
}
const core = service.cores.get(key);
if (core === undefined) {
return yield* knowledgeCoreServiceError("load-kg-core", `Knowledge core not found: ${key}`);
}
if (core.triples.length > 0) {
yield* Effect.acquireUseRelease(
tryPromise("triples-producer", () =>
service.pubsub.createProducer<unknown>({ topic: "tg.flow.triples" }),
),
(producer) =>
tryPromise("send-triples", () =>
producer.send({
metadata: {
id: coreId,
root: coreId,
user,
collection: request.collection ?? "default",
},
triples: core.triples,
}),
),
(producer) => closeResource(producer, "close-triples-producer"),
);
}
yield* Effect.log(
`[KnowledgeCoreService] Loaded core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length})`,
);
yield* sendResponse(service, {}, requestId);
}),
);
},
deleteDeCore: async function(this: KnowledgeCoreService, 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 });
unloadKgCore: function(this: KnowledgeCoreService, _request: KnowledgeRequest, requestId: string): Promise<void> {
return Effect.runPromise(sendResponse(this, {}, requestId));
},
putDeCore: async function(this: KnowledgeCoreService, 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 });
listDeCores: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const prefix = user.length > 0 ? `${user}:` : "";
const ids = [...service.deCores.keys()]
.filter((key) => prefix.length === 0 || key.startsWith(prefix))
.map((key) => key.slice(prefix.length));
yield* sendResponse(service, { ids }, requestId);
}),
);
},
loadDeCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = this.coreKey(user, coreId);
if (!(this.deCores as Map<string, DocumentEmbeddingsCore[]>).has(key)) throw new Error(`Document embeddings core not found: ${key}`);
await this.responseProducer!.send({}, { id: requestId });
getDeCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
const core = service.deCores.get(key);
if (core === undefined) {
return yield* knowledgeCoreServiceError("get-de-core", `Document embeddings core not found: ${key}`);
}
for (let i = 0; i < core.length; i++) {
const isLast = i === core.length - 1;
yield* sendResponse(
service,
{
documentEmbeddings: core[i],
"document-embeddings": core[i],
eos: isLast,
} as KnowledgeResponse,
requestId,
"respond-de-core",
);
}
if (core.length === 0) {
yield* sendResponse(service, { eos: true }, requestId, "respond-de-empty");
}
}),
);
},
deleteDeCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
service.deCores.delete(service.coreKey(user, coreId));
yield* tryPromise("persist-delete-de-core", () => service.persist());
yield* sendResponse(service, {}, requestId);
}),
);
},
putDeCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
const item = service.documentEmbeddings(request);
if (item === undefined) {
return yield* knowledgeCoreServiceError("put-de-core", "put-de-core requires document-embeddings");
}
const core = service.deCores.get(key) ?? [];
core.push(item);
service.deCores.set(key, core);
yield* tryPromise("persist-put-de-core", () => service.persist());
yield* sendResponse(service, {}, requestId);
}),
);
},
loadDeCore: function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const user = request.user ?? "";
const coreId = request.id ?? "";
const key = service.coreKey(user, coreId);
if (!(service.deCores as Map<string, DocumentEmbeddingsCore[]>).has(key)) {
return yield* knowledgeCoreServiceError("load-de-core", `Document embeddings core not found: ${key}`);
}
yield* sendResponse(service, {}, requestId);
}),
);
},
@ -417,69 +593,88 @@ export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): Kn
// ---------- Persistence ----------
persist: async function(this: KnowledgeCoreService): Promise<void> {
try {
// Serialize Map to object
const data: {
kg: Record<string, KnowledgeCore>;
de: Record<string, DocumentEmbeddingsCore[]>;
} = { kg: {}, de: {} };
for (const [key, core] of this.cores) {
data.kg[key] = core;
}
for (const [key, core] of this.deCores) {
data.de[key] = core;
}
const json = JSON.stringify(data, null, 2);
await writeTextFile(this.persistPath, json);
} catch (err) {
console.error("[KnowledgeCoreService] Failed to persist state:", err);
}
},
loadFromDisk: async function(this: KnowledgeCoreService): Promise<void> {
try {
const raw = await readTextFile(this.persistPath);
const parsed = JSON.parse(raw) as Record<string, KnowledgeCore> | {
kg?: Record<string, KnowledgeCore>;
de?: Record<string, DocumentEmbeddingsCore[]>;
};
this.cores.clear();
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);
persist: function(this: KnowledgeCoreService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
// Serialize Map to object
const data: {
kg: Record<string, KnowledgeCore>;
de: Record<string, DocumentEmbeddingsCore[]>;
} = { kg: {}, de: {} };
for (const [key, core] of service.cores) {
data.kg[key] = core;
}
for (const [key, core] of service.deCores) {
data.de[key] = core;
}
}
console.log(`[KnowledgeCoreService] Loaded persisted state (kg=${this.cores.size}, de=${this.deCores.size})`);
} catch {
console.log("[KnowledgeCoreService] No persisted state found, starting fresh");
}
const json = yield* trySync("persist-serialize", () => JSON.stringify(data, null, 2));
yield* tryPromise("persist-write", () => writeTextFile(service.persistPath, json));
}).pipe(
Effect.catch((error) =>
Effect.logError("[KnowledgeCoreService] Failed to persist state", {
error: error.message,
}),
),
),
);
},
stop: async function(this: KnowledgeCoreService): Promise<void> {
if (this.consumer !== null) {
await this.consumer.close();
this.consumer = null;
}
if (this.responseProducer !== null) {
await this.responseProducer.close();
this.responseProducer = null;
}
await baseStop();
loadFromDisk: function(this: KnowledgeCoreService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
const raw = yield* tryPromise("load-read", () => readTextFile(service.persistPath));
const parsed = yield* trySync("load-parse", () =>
JSON.parse(raw) as Record<string, KnowledgeCore> | {
kg?: Record<string, KnowledgeCore>;
de?: Record<string, DocumentEmbeddingsCore[]>;
},
);
service.cores.clear();
service.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)) {
service.cores.set(key, core);
}
if ("de" in parsed && parsed.de !== undefined) {
for (const [key, core] of Object.entries(parsed.de)) {
service.deCores.set(key, core);
}
}
yield* Effect.log(`[KnowledgeCoreService] Loaded persisted state (kg=${service.cores.size}, de=${service.deCores.size})`);
}).pipe(
Effect.catch(() =>
Effect.log("[KnowledgeCoreService] No persisted state found, starting fresh"),
),
),
);
},
stop: function(this: KnowledgeCoreService): Promise<void> {
const service = this;
return Effect.runPromise(
Effect.gen(function* () {
if (service.consumer !== null) {
yield* tryPromise("close-consumer", () => service.consumer.close());
service.consumer = null;
}
if (service.responseProducer !== null) {
yield* tryPromise("close-response-producer", () => service.responseProducer.close());
service.responseProducer = null;
}
yield* tryPromise("base-stop", () => baseStop());
}),
);
}
});
@ -488,15 +683,11 @@ export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): Kn
export const KnowledgeCoreService = makeKnowledgeCoreService;
function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
export const program = makeProcessorProgram({
id: "knowledge-svc",
make: (config) => makeKnowledgeCoreService(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -37,7 +37,7 @@ import {
errorMessage,
} from "@trustgraph/base";
import { makeFlowProcessorProgram } from "@trustgraph/base";
import { Effect } from "effect";
import { Clock, Effect } from "effect";
import * as S from "effect/Schema";
export class PdfDecoderError extends S.TaggedErrorClass<PdfDecoderError>()(
@ -63,7 +63,7 @@ const pdfDecoderError = (
documentId: string,
cause: unknown,
) =>
new PdfDecoderError({
PdfDecoderError.make({
operation,
documentId,
message: errorMessage(cause),
@ -76,18 +76,24 @@ const loadPdf = (documentId: string, pdfBuffer: Buffer) =>
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 loadPageText = Effect.fn("loadPageText")(function*(
documentId: string,
pageNumber: number,
pdf: PdfDocument,
) {
const page = yield* Effect.tryPromise({
try: () => pdf.getPage(pageNumber),
catch: (cause) => pdfDecoderError("load-page", documentId, cause),
});
const textContent = yield* Effect.tryPromise({
try: () => page.getTextContent(),
catch: (cause) => pdfDecoderError("load-page-text", documentId, cause),
});
return textContent.items
.filter((item): item is TextItem => "str" in item)
.map((item) => item.str)
.join(" ");
});
const onPdfDecodeMessage = Effect.fn("PdfDecoderService.onMessage")(function* (
msg: Document,
@ -156,6 +162,7 @@ const onPdfDecodeMessage = Effect.fn("PdfDecoderService.onMessage")(function* (
continue;
}
const timestamp = yield* Clock.currentTimeMillis;
const childResp = yield* librarian.request({
operation: "add-child-document",
documentMetadata: {
@ -165,7 +172,7 @@ const onPdfDecodeMessage = Effect.fn("PdfDecoderService.onMessage")(function* (
title: `Page ${i}`,
parentId: documentId,
documentType: "page",
time: Date.now(),
time: timestamp,
comments: "",
tags: [],
},
@ -226,7 +233,7 @@ export function makePdfDecoderService(config: ProcessorConfig): PdfDecoderServic
const service = makeFlowProcessor(config, {
specifications: makePdfDecoderSpecs(),
});
console.log("[PdfDecoder] Service initialized");
Effect.runSync(Effect.log("[PdfDecoder] Service initialized"));
return service;
}
@ -245,6 +252,6 @@ export const program = makeFlowProcessorProgram({
specs: () => makePdfDecoderSpecs(),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -4,11 +4,13 @@
* Python reference: trustgraph-flow/trustgraph/embeddings/ollama/processor.py
*/
import { Effect, Layer } from "effect";
import { Config, Effect, Layer } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
import {
Embeddings,
embeddingsError,
EmbeddingsError,
errorMessage,
makeEmbeddingsService,
makeEmbeddingsSpecs,
type EmbeddingsServiceShape,
@ -22,18 +24,58 @@ export interface OllamaEmbeddingsConfig extends ProcessorConfig {
fetch?: typeof fetch;
}
interface OllamaEmbedResponse {
embeddings: number[][];
const EmbeddingVector = S.Array(S.Number);
const OllamaEmbedResponse = S.Struct({
embeddings: S.Array(EmbeddingVector),
});
type OllamaEmbedResponse = typeof OllamaEmbedResponse.Type;
interface ResolvedOllamaEmbeddingsConfig {
readonly defaultModel: string;
readonly ollamaHost: string;
readonly fetchImpl: typeof fetch;
}
const ollamaEmbeddingsError = (operation: string, cause: unknown): EmbeddingsError =>
EmbeddingsError.make({
operation,
message: errorMessage(cause),
provider: "ollama",
});
const ollamaEmbeddingsMessageError = (operation: string, message: string): EmbeddingsError =>
EmbeddingsError.make({
operation,
message,
provider: "ollama",
});
const optionalStringConfig = Effect.fn("OllamaEmbeddings.optionalStringConfig")(function*(name: string) {
return O.getOrUndefined(yield* Config.string(name).pipe(Config.option));
});
const loadOllamaEmbeddingsConfig = Effect.fn("OllamaEmbeddings.loadConfig")(function*(
config: OllamaEmbeddingsConfig,
) {
return {
defaultModel: config.model ?? "mxbai-embed-large",
ollamaHost:
config.ollamaHost ??
(yield* optionalStringConfig("OLLAMA_URL")) ??
(yield* optionalStringConfig("OLLAMA_HOST")) ??
"http://localhost:11434",
fetchImpl: config.fetch ?? globalThis.fetch,
} satisfies ResolvedOllamaEmbeddingsConfig;
});
export function makeOllamaEmbeddings(config: OllamaEmbeddingsConfig): EmbeddingsServiceShape {
const defaultModel = config.model ?? "mxbai-embed-large";
const ollamaHost =
config.ollamaHost ??
process.env.OLLAMA_URL ??
process.env.OLLAMA_HOST ??
"http://localhost:11434";
const fetchImpl = config.fetch ?? globalThis.fetch;
const {
defaultModel,
ollamaHost,
fetchImpl,
} = Effect.runSync(loadOllamaEmbeddingsConfig(config)) satisfies ResolvedOllamaEmbeddingsConfig;
return {
embed: Effect.fn("OllamaEmbeddings.embed")((texts: ReadonlyArray<string>, model?: string) => {
@ -49,29 +91,38 @@ export function makeOllamaEmbeddings(config: OllamaEmbeddingsConfig): Embeddings
model: useModel,
input: Array.from(texts),
}).pipe(
Effect.mapError((error) => embeddingsError("ollama.encode-request", error, "ollama")),
Effect.mapError((error) => ollamaEmbeddingsError("ollama.encode-request", error))
);
return yield* Effect.tryPromise({
try: async () => {
const response = await fetchImpl(url, {
const response = yield* Effect.tryPromise({
try: () =>
fetchImpl(url, {
method: "POST",
headers: { "Content-Type": "application/json" },
body,
});
if (!response.ok) {
const errorBody = await response.text();
throw new Error(
`Ollama embeddings request failed (${response.status}): ${errorBody}`,
);
}
const data = (await response.json()) as OllamaEmbedResponse;
return data.embeddings;
},
catch: (error) => embeddingsError("ollama.embed", error, "ollama"),
}),
catch: (error) => ollamaEmbeddingsError("ollama.fetch", error),
});
if (!response.ok) {
const errorBody = yield* Effect.tryPromise({
try: () => response.text(),
catch: (error) => ollamaEmbeddingsError("ollama.error-body", error),
});
return yield* ollamaEmbeddingsMessageError(
"ollama.embed",
`Ollama embeddings request failed (${response.status}): ${errorBody}`,
);
}
const data = yield* Effect.tryPromise({
try: () => response.json() as Promise<unknown>,
catch: (error) => ollamaEmbeddingsError("ollama.response-json", error),
});
const decoded = yield* S.decodeUnknownEffect(OllamaEmbedResponse)(data).pipe(
Effect.mapError((error) => ollamaEmbeddingsError("ollama.decode-response", error))
);
return Array.from(decoded.embeddings, (vector) => Array.from(vector));
});
}),
};
@ -88,9 +139,10 @@ export type OllamaEmbeddingsProcessor = ReturnType<typeof makeOllamaEmbeddingsPr
export function makeOllamaEmbeddingsProcessor(config: OllamaEmbeddingsConfig) {
const embeddings = makeOllamaEmbeddings(config);
console.log(
`[OllamaEmbeddings] Initialized (host=${config.ollamaHost ?? process.env.OLLAMA_URL ?? process.env.OLLAMA_HOST ?? "http://localhost:11434"}, model=${config.model ?? "mxbai-embed-large"})`,
);
const resolved = Effect.runSync(loadOllamaEmbeddingsConfig(config)) satisfies ResolvedOllamaEmbeddingsConfig;
Effect.runSync(Effect.log(
`[OllamaEmbeddings] Initialized (host=${resolved.ollamaHost}, model=${resolved.defaultModel})`,
));
return makeEmbeddingsService(config, embeddings);
}
@ -102,6 +154,6 @@ export const program = makeFlowProcessorProgram<OllamaEmbeddingsConfig, never, E
layer: (config) => OllamaEmbeddingsLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -289,7 +289,7 @@ export function makeKnowledgeExtractService(config: ProcessorConfig): KnowledgeE
const service = makeFlowProcessor(config, {
specifications: makeKnowledgeExtractSpecs(),
});
console.log("[KnowledgeExtract] Service initialized");
Effect.runSync(Effect.log("[KnowledgeExtract] Service initialized"));
return service;
}
@ -324,7 +324,9 @@ export function parseJsonResponse<T>(raw: string): T | null {
if (O.isSome(decoded)) return decoded.value as T;
}
console.warn("[KnowledgeExtract] Failed to parse JSON from LLM response:", raw.slice(0, 300));
Effect.runSync(Effect.logWarning("[KnowledgeExtract] Failed to parse JSON from LLM response", {
response: raw.slice(0, 300),
}));
return null;
}
@ -333,7 +335,9 @@ function parseRelationshipsResponse(raw: string): ReadonlyArray<ExtractedRelatio
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));
Effect.runSync(Effect.logWarning("[KnowledgeExtract] Failed to parse relationships from LLM response", {
response: raw.slice(0, 300),
}));
return null;
}
@ -342,7 +346,9 @@ function parseDefinitionsResponse(raw: string): ReadonlyArray<ExtractedDefinitio
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));
Effect.runSync(Effect.logWarning("[KnowledgeExtract] Failed to parse definitions from LLM response", {
response: raw.slice(0, 300),
}));
return null;
}
@ -380,6 +386,6 @@ export const program = makeFlowProcessorProgram({
specs: () => makeKnowledgeExtractSpecs(),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

File diff suppressed because it is too large Load diff

View file

@ -8,7 +8,7 @@
* Python reference: trustgraph-flow/trustgraph/gateway/dispatch/manager.py
*/
import { Effect, Exit, Scope, SynchronizedRef } from "effect";
import { Clock, Effect, Exit, Random, Scope, SynchronizedRef } from "effect";
import {
loadMessagingRuntimeConfig,
makeNatsBackend,
@ -146,13 +146,13 @@ export function makeDispatcherManager(config: GatewayConfig): DispatcherManager
const pubsub: PubSubBackend = config.pubsub ?? makeNatsBackend(config.natsUrl ?? "nats://localhost:4222");
let runtime: DispatcherRuntime | null = null;
const start = async (): Promise<void> => {
if (runtime !== null) return;
const start = (): Promise<void> => {
if (runtime !== null) return Promise.resolve();
runtime = await Effect.runPromise(
return Effect.runPromise(
Effect.gen(function* () {
const scope = yield* Scope.make();
return yield* Effect.gen(function* () {
const nextRuntime = yield* Effect.gen(function* () {
const messagingConfig = yield* loadMessagingRuntimeConfig();
const requestors = yield* SynchronizedRef.make<RequestorMap>(new Map());
return {
@ -163,62 +163,79 @@ export function makeDispatcherManager(config: GatewayConfig): DispatcherManager
}).pipe(
Effect.onError((cause) => Scope.close(scope, Exit.failCause(cause))),
);
runtime = nextRuntime;
}),
);
};
const stop = async (): Promise<void> => {
const current = runtime;
runtime = null;
const stop = (): Promise<void> =>
Effect.runPromise(
Effect.gen(function* () {
const current = runtime;
runtime = null;
if (current !== null) {
await Effect.runPromise(Scope.close(current.scope, Exit.void));
}
if (current !== null) {
yield* Scope.close(current.scope, Exit.void);
}
await pubsub.close();
};
yield* Effect.tryPromise({
try: () => pubsub.close(),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "close-pubsub", cause),
});
}),
);
// ---------- Internal helpers ----------
const ensureRuntime = async (): Promise<DispatcherRuntime> => {
if (runtime === null) {
await start();
}
if (runtime === null) {
throw messagingLifecycleError("gateway-dispatcher", "start", "Dispatcher manager failed to start");
}
return runtime;
};
const ensureRuntime = (): Promise<DispatcherRuntime> =>
Effect.runPromise(
Effect.gen(function* () {
if (runtime === null) {
yield* Effect.tryPromise({
try: () => start(),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "start", cause),
});
}
if (runtime === null) {
return yield* messagingLifecycleError("gateway-dispatcher", "start", "Dispatcher manager failed to start");
}
return runtime;
}),
);
const getRequestor = async (
const getRequestor = (
requestTopic: string,
responseTopic: string,
key: string,
): Promise<EffectRequestResponse<unknown, unknown>> => {
const current = await ensureRuntime();
): Promise<EffectRequestResponse<unknown, unknown>> =>
Effect.runPromise(
Effect.gen(function* () {
const current = yield* Effect.tryPromise({
try: () => ensureRuntime(),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "ensure-runtime", cause),
});
return await Effect.runPromise(
SynchronizedRef.modifyEffect(current.requestors, (requestors) => {
const cached = requestors.get(key);
if (cached !== undefined) {
return Effect.succeed([cached, requestors] as const);
}
return yield* SynchronizedRef.modifyEffect(current.requestors, (requestors) => {
const cached = requestors.get(key);
if (cached !== undefined) {
return Effect.succeed([cached, requestors] as const);
}
return current.factory.make<unknown, unknown>({
requestTopic,
responseTopic,
subscription: `gateway-${key}`,
}).pipe(
Scope.provide(current.scope),
Effect.map((requestor) => {
const next = new Map(requestors);
next.set(key, requestor);
return [requestor, next] as const;
}),
);
return current.factory.make<unknown, unknown>({
requestTopic,
responseTopic,
subscription: `gateway-${key}`,
}).pipe(
Scope.provide(current.scope),
Effect.map((requestor) => {
const next = new Map(requestors);
next.set(key, requestor);
return [requestor, next] as const;
}),
);
});
}),
);
};
const resolveGlobalTopics = (
kind: string,
@ -256,93 +273,107 @@ export function makeDispatcherManager(config: GatewayConfig): DispatcherManager
// ---------- Global service dispatch ----------
const dispatchGlobalService = async (
const dispatchGlobalService = (
kind: string,
request: Record<string, unknown>,
): Promise<unknown> => {
const { requestTopic, responseTopic } = resolveGlobalTopics(kind);
const rr = await getRequestor(requestTopic, responseTopic, `global:${kind}`);
): Promise<unknown> =>
Effect.runPromise(
Effect.gen(function* () {
const { requestTopic, responseTopic } = resolveGlobalTopics(kind);
const rr = yield* Effect.tryPromise({
try: () => getRequestor(requestTopic, responseTopic, `global:${kind}`),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "get-requestor", cause),
});
const translated = translateRequest(kind, request);
const response = await Effect.runPromise(rr.request(translated));
return translateResponse(kind, response);
};
const translated = translateRequest(kind, request);
const response = yield* rr.request(translated);
return translateResponse(kind, response);
}),
);
const dispatchGlobalServiceStreaming = async (
const dispatchGlobalServiceStreaming = (
kind: string,
request: Record<string, unknown>,
responder: Responder,
): Promise<void> => {
const { requestTopic, responseTopic } = resolveGlobalTopics(kind);
const rr = await getRequestor(requestTopic, responseTopic, `global:${kind}`);
const translated = translateRequest(kind, request);
): Promise<void> =>
Effect.runPromise(
Effect.gen(function* () {
const { requestTopic, responseTopic } = resolveGlobalTopics(kind);
const rr = yield* Effect.tryPromise({
try: () => getRequestor(requestTopic, responseTopic, `global:${kind}`),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "get-requestor", cause),
});
const translated = translateRequest(kind, request);
await Effect.runPromise(
rr.request(translated, {
recipient: (response) => {
const translatedRes = translateResponse(kind, response);
const complete = dispatcherManagerIsCompleteResponse(translatedRes);
return Effect.tryPromise({
try: async () => {
await responder(translatedRes, complete);
return complete;
},
catch: (error) => messagingDeliveryError(responseTopic, "stream-responder", error),
});
},
yield* rr.request(translated, {
recipient: (response) => {
const translatedRes = translateResponse(kind, response);
const complete = dispatcherManagerIsCompleteResponse(translatedRes);
return Effect.tryPromise({
try: () => responder(translatedRes, complete).then(() => complete),
catch: (error) => messagingDeliveryError(responseTopic, "stream-responder", error),
});
},
});
}),
);
};
// ---------- Flow-scoped service dispatch ----------
const dispatchFlowService = async (
const dispatchFlowService = (
flow: string,
kind: string,
request: Record<string, unknown>,
): Promise<unknown> => {
const { requestTopic, responseTopic } = resolveFlowTopics(kind);
const rr = await getRequestor(
requestTopic,
responseTopic,
`flow:${flow}:${kind}`,
): Promise<unknown> =>
Effect.runPromise(
Effect.gen(function* () {
const { requestTopic, responseTopic } = resolveFlowTopics(kind);
const rr = yield* Effect.tryPromise({
try: () => getRequestor(
requestTopic,
responseTopic,
`flow:${flow}:${kind}`,
),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "get-requestor", cause),
});
const translated = translateRequest(kind, request);
const response = yield* rr.request(translated);
return translateResponse(kind, response);
}),
);
const translated = translateRequest(kind, request);
const response = await Effect.runPromise(rr.request(translated));
return translateResponse(kind, response);
};
const dispatchFlowServiceStreaming = async (
const dispatchFlowServiceStreaming = (
flow: string,
kind: string,
request: Record<string, unknown>,
responder: Responder,
): Promise<void> => {
const { requestTopic, responseTopic } = resolveFlowTopics(kind);
const rr = await getRequestor(
requestTopic,
responseTopic,
`flow:${flow}:${kind}`,
);
const translated = translateRequest(kind, request);
): Promise<void> =>
Effect.runPromise(
Effect.gen(function* () {
const { requestTopic, responseTopic } = resolveFlowTopics(kind);
const rr = yield* Effect.tryPromise({
try: () => getRequestor(
requestTopic,
responseTopic,
`flow:${flow}:${kind}`,
),
catch: (cause) => messagingLifecycleError("gateway-dispatcher", "get-requestor", cause),
});
const translated = translateRequest(kind, request);
await Effect.runPromise(
rr.request(translated, {
recipient: (response) => {
const translatedRes = translateResponse(kind, response);
const complete = dispatcherManagerIsCompleteResponse(translatedRes);
return Effect.tryPromise({
try: async () => {
await responder(translatedRes, complete);
return complete;
},
catch: (error) => messagingDeliveryError(responseTopic, "stream-responder", error),
});
},
yield* rr.request(translated, {
recipient: (response) => {
const translatedRes = translateResponse(kind, response);
const complete = dispatcherManagerIsCompleteResponse(translatedRes);
return Effect.tryPromise({
try: () => responder(translatedRes, complete).then(() => complete),
catch: (error) => messagingDeliveryError(responseTopic, "stream-responder", error),
});
},
});
}),
);
};
// ---------- Fire-and-forget publish ----------
@ -350,12 +381,27 @@ export function makeDispatcherManager(config: GatewayConfig): DispatcherManager
* Publish a single message to an arbitrary topic (no request/response).
* Used for injecting documents into the processing pipeline.
*/
const publishToTopic = async (topic: string, message: unknown, id?: string): Promise<void> => {
const producer = await pubsub.createProducer<unknown>({ topic });
const messageId = id ?? `pub-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
await producer.send(message, { id: messageId });
await producer.close();
};
const publishToTopic = (topic: string, message: unknown, id?: string): Promise<void> =>
Effect.runPromise(
Effect.gen(function* () {
const producer = yield* Effect.tryPromise({
try: () => pubsub.createProducer<unknown>({ topic }),
catch: (cause) => messagingDeliveryError(topic, "create-producer", cause),
});
const timestamp = yield* Clock.currentTimeMillis;
const suffix = yield* Random.nextIntBetween(0, 36 ** 6, { halfOpen: true });
const messageId = id ?? `pub-${timestamp}-${suffix.toString(36).padStart(6, "0")}`;
yield* Effect.tryPromise({
try: () => producer.send(message, { id: messageId }),
catch: (cause) => messagingDeliveryError(topic, "send", cause),
});
yield* Effect.tryPromise({
try: () => producer.close(),
catch: (cause) => messagingDeliveryError(topic, "close-producer", cause),
});
}),
);
return {
start,

View file

@ -19,6 +19,7 @@
*/
import type { Term, Triple } from "@trustgraph/base";
import * as S from "effect/Schema";
// ---------- Client wire format type definitions ----------
@ -46,6 +47,14 @@ interface ClientTripleTerm {
type ClientTerm = ClientIriTerm | ClientBlankTerm | ClientLiteralTerm | ClientTripleTerm;
export class DispatchSerializationError extends S.TaggedErrorClass<DispatchSerializationError>()(
"DispatchSerializationError",
{
message: S.String,
operation: S.String,
},
) {}
interface ClientTriple {
s: ClientTerm;
p: ClientTerm;
@ -70,7 +79,10 @@ export function clientTermToInternal(wire: ClientTerm): Term {
};
case "t": {
if (wire.tr === undefined) {
throw new Error("Client triple term is missing tr");
throw DispatchSerializationError.make({
operation: "client-term-to-internal",
message: "Client triple term is missing tr",
});
}
return {
type: "TRIPLE",

View file

@ -34,37 +34,44 @@ export const makeSocketRpcProtocol = Effect.gen(function* () {
});
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))!),
);
}
};
const encodeDefect = (cause: unknown) =>
Effect.sync(() => parser.encode(RpcMessage.ResponseDefectEncoded(cause))!);
const write = (response: RpcMessage.FromServerEncoded) =>
Effect.sync(() => parser.encode(response)).pipe(
Effect.flatMap((encoded) =>
encoded === undefined ? Effect.void : Effect.orDie(writeRaw(encoded)),
),
Effect.catchDefect((cause: unknown) =>
encodeDefect(cause).pipe(
Effect.flatMap((encoded) => Effect.orDie(writeRaw(encoded))),
Effect.orDie,
),
),
);
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(
yield* socket.runRaw((data) =>
Effect.sync(() => parser.decode(data) as ReadonlyArray<RpcMessage.FromClientEncoded>).pipe(
Effect.flatMap((decoded) =>
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 }),
),
Effect.catchDefect((cause: unknown) =>
encodeDefect(cause).pipe(
Effect.flatMap((encoded) => writeRaw(encoded)),
),
),
)
).pipe(
Effect.catchReason("SocketError", "SocketCloseError", () => Effect.void),
Effect.orDie,
);

View file

@ -42,7 +42,7 @@ const makeGatewayRpcHandlers = (dispatcher: DispatcherManager) =>
Dispatch: (payload) =>
Effect.tryPromise({
try: () => dispatchOne(dispatcher, payload),
catch: (cause) => new DispatchError({ message: errorMessage(cause) }),
catch: (cause) => DispatchError.make({ message: errorMessage(cause) }),
}),
DispatchStream: Effect.fn("GatewayRpc.DispatchStream")(function* (payload) {
const context = yield* Effect.context<never>();
@ -52,11 +52,10 @@ const makeGatewayRpcHandlers = (dispatcher: DispatcherManager) =>
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) }),
dispatchStream(dispatcher, payload, (response, complete) =>
runPromise(Queue.offer(queue, DispatchStreamChunk.make({ response, complete }))).then(() => complete),
),
catch: (cause) => DispatchError.make({ message: errorMessage(cause) }),
}).pipe(
Effect.flatMap(() => Queue.end(queue)),
Effect.catch((error) => Queue.fail(queue, error)),
@ -82,26 +81,24 @@ function dispatchOne(
return dispatcher.dispatchGlobalService(payload.service, payload.request);
}
async function dispatchStream(
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);
};
const send = (response: unknown, complete: boolean): Promise<void> =>
responder(response, complete).then(() => undefined);
if (payload.scope === "flow") {
await dispatcher.dispatchFlowServiceStreaming(
return dispatcher.dispatchFlowServiceStreaming(
payload.flow ?? "default",
payload.service,
payload.request,
send,
);
return;
}
await dispatcher.dispatchGlobalServiceStreaming(
return dispatcher.dispatchGlobalServiceStreaming(
payload.service,
payload.request,
send,

View file

@ -7,13 +7,13 @@
* Python reference: trustgraph-flow/trustgraph/gateway/service.py
*/
import Fastify from "fastify";
import Fastify, { type FastifyReply } from "fastify";
import websocketPlugin from "@fastify/websocket";
import { Config, Effect, Exit, Scope } from "effect";
import { Clock, Config, Effect, Exit, Random, Scope } from "effect";
import * as O from "effect/Option";
import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
import * as EffectSocket from "effect/unstable/socket/Socket";
import { optionalStringConfig, registry, toTgError, type PubSubBackend } from "@trustgraph/base";
import { messagingLifecycleError, optionalStringConfig, registry, toTgError, type PubSubBackend } from "@trustgraph/base";
import { makeDispatcherManager } from "./dispatch/manager.js";
import { makeGatewayRpcServer } from "./rpc-server.js";
@ -25,187 +25,227 @@ export interface GatewayConfig {
pubsub?: PubSubBackend;
}
export async function createGateway(config: GatewayConfig) {
export function createGateway(config: GatewayConfig) {
const app = Fastify({ logger: true });
await app.register(websocketPlugin);
const dispatcher = makeDispatcherManager(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/rpc")) return; // RPC socket auth via query param
if (config.secret !== undefined && config.secret.length > 0) {
const auth = request.headers.authorization;
if (auth === undefined || auth !== `Bearer ${config.secret}`) {
reply.code(401).send({ error: "Unauthorized" });
}
const sendDispatchResult = (reply: FastifyReply, result: unknown): unknown => {
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);
}
});
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;
const body = request.body as Record<string, unknown>;
try {
const result = await dispatcher.dispatchGlobalService(kind, body) as Record<string, unknown>;
const err = result?.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/flow/:flow/service/:kind (flow-scoped services)
app.post<{ Params: { flow: string; kind: string } }>(
"/api/v1/flow/:flow/service/:kind",
async (request, reply) => {
const { flow, kind } = request.params;
const body = request.body as Record<string, unknown>;
try {
const result = await dispatcher.dispatchFlowService(flow, kind, body) as Record<string, unknown>;
const err = result?.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/flow/:flow/load (trigger document processing)
app.post<{ Params: { flow: string } }>(
"/api/v1/flow/:flow/load",
async (request, reply) => {
const { flow } = request.params;
const body = request.body as {
documentId?: string;
user?: string;
collection?: string;
};
if (body.documentId === undefined || body.documentId.length === 0) {
return reply.code(400).send({
error: { type: "bad-request", message: "documentId is required" },
});
}
try {
const user = body.user ?? "default";
const collection = body.collection ?? "default";
const documentId = body.documentId;
// Publish Document message to the decode-input topic
const topic = "tg.flow.document";
const metadata = {
id: `load-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`,
root: documentId,
user,
collection,
};
await dispatcher.publishToTopic(topic, { metadata, documentId });
return { status: "processing", documentId, flow };
} catch (err) {
reply.code(500).send({
error: toTgError(err),
});
}
},
);
// 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) {
socket.close(4001, "Unauthorized");
return;
}
const program = Effect.scoped(
Effect.gen(function* () {
const effectSocket = yield* EffectSocket.fromWebSocket(
Effect.succeed(socket as unknown as globalThis.WebSocket),
{ closeCodeIsError: (code) => code !== 1000 },
);
yield* rpcServer.onSocket(effectSocket, headersFrom(request.headers));
}),
);
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");
}
});
});
// Metrics endpoint — returns Prometheus metrics from prom-client
app.get("/api/v1/metrics", async (_, reply) => {
reply.header("content-type", registry.contentType);
return registry.metrics();
});
return {
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();
},
return result;
};
const sendDispatchError = (reply: FastifyReply, error: unknown): unknown =>
reply.code(500).send({ error: toTgError(error) });
return Effect.runPromise(
Effect.gen(function* () {
yield* Effect.tryPromise({
try: () => app.register(websocketPlugin),
catch: (cause) => messagingLifecycleError("gateway", "register-websocket", cause),
});
yield* Effect.tryPromise({
try: () => dispatcher.start(),
catch: (cause) => messagingLifecycleError("gateway", "dispatcher-start", cause),
});
const rpcScope = yield* Scope.make();
const rpcServer = yield* makeGatewayRpcServer(dispatcher).pipe(
Effect.provideService(RpcSerialization.RpcSerialization, RpcSerialization.ndjson),
Scope.provide(rpcScope),
);
return { rpcScope, rpcServer };
}),
).then(({ rpcScope, rpcServer }) => {
// Authentication middleware
app.addHook("onRequest", (request, reply) => {
if (request.url === "/api/v1/metrics") return;
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;
if (auth === undefined || auth !== `Bearer ${config.secret}`) {
reply.code(401).send({ error: "Unauthorized" });
}
}
});
app.post<{
Body: {
scope?: string;
service?: string;
flow?: string;
request?: Record<string, unknown>;
};
}>("/api/v1/workbench/dispatch", (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" },
});
}
return Effect.runPromise(
Effect.tryPromise({
try: () =>
body.scope === "flow"
? dispatcher.dispatchFlowService(body.flow ?? "default", service, payload)
: dispatcher.dispatchGlobalService(service, payload),
catch: (cause) => messagingLifecycleError("gateway", "workbench-dispatch", cause),
}).pipe(
Effect.map((result) => sendDispatchResult(reply, result)),
Effect.catch((error) => Effect.succeed(sendDispatchError(reply, error))),
),
);
});
// REST endpoint: POST /api/v1/:kind (global services)
app.post<{ Params: { kind: string } }>("/api/v1/:kind", (request, reply) => {
const { kind } = request.params;
const body = request.body as Record<string, unknown>;
return Effect.runPromise(
Effect.tryPromise({
try: () => dispatcher.dispatchGlobalService(kind, body),
catch: (cause) => messagingLifecycleError("gateway", "global-dispatch", cause),
}).pipe(
Effect.map((result) => sendDispatchResult(reply, result)),
Effect.catch((error) => Effect.succeed(sendDispatchError(reply, error))),
),
);
});
// REST endpoint: POST /api/v1/flow/:flow/service/:kind (flow-scoped services)
app.post<{ Params: { flow: string; kind: string } }>(
"/api/v1/flow/:flow/service/:kind",
(request, reply) => {
const { flow, kind } = request.params;
const body = request.body as Record<string, unknown>;
return Effect.runPromise(
Effect.tryPromise({
try: () => dispatcher.dispatchFlowService(flow, kind, body),
catch: (cause) => messagingLifecycleError("gateway", "flow-dispatch", cause),
}).pipe(
Effect.map((result) => sendDispatchResult(reply, result)),
Effect.catch((error) => Effect.succeed(sendDispatchError(reply, error))),
),
);
},
);
// REST endpoint: POST /api/v1/flow/:flow/load (trigger document processing)
app.post<{ Params: { flow: string } }>(
"/api/v1/flow/:flow/load",
(request, reply) => {
const { flow } = request.params;
const body = request.body as {
documentId?: string;
user?: string;
collection?: string;
};
if (body.documentId === undefined || body.documentId.length === 0) {
return reply.code(400).send({
error: { type: "bad-request", message: "documentId is required" },
});
}
return Effect.runPromise(
Effect.gen(function* () {
const user = body.user ?? "default";
const collection = body.collection ?? "default";
const documentId = body.documentId;
const timestamp = yield* Clock.currentTimeMillis;
const suffix = yield* Random.nextIntBetween(0, 36 ** 6, { halfOpen: true });
// Publish Document message to the decode-input topic
const topic = "tg.flow.document";
const metadata = {
id: `load-${timestamp}-${suffix.toString(36).padStart(6, "0")}`,
root: documentId,
user,
collection,
};
yield* Effect.tryPromise({
try: () => dispatcher.publishToTopic(topic, { metadata, documentId }),
catch: (cause) => messagingLifecycleError("gateway", "publish-load", cause),
});
return { status: "processing", documentId, flow };
}).pipe(
Effect.catch((error) => Effect.succeed(sendDispatchError(reply, error))),
),
);
},
);
// 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) {
socket.close(4001, "Unauthorized");
return;
}
const program = Effect.scoped(
Effect.gen(function* () {
const effectSocket = yield* EffectSocket.fromWebSocket(
Effect.succeed(socket as unknown as globalThis.WebSocket),
{ closeCodeIsError: (code) => code !== 1000 },
);
yield* rpcServer.onSocket(effectSocket, headersFrom(request.headers));
}),
);
void Effect.runPromise(program.pipe(Scope.provide(rpcScope))).catch((error) => {
void Effect.runPromise(
Effect.logError("[Gateway] RPC WebSocket error", { error: toTgError(error).message }).pipe(
Effect.flatMap(() =>
Effect.sync(() => {
if (socket.readyState === 1) {
socket.close(1011, "Internal server error");
}
}),
),
),
);
});
});
// Metrics endpoint — returns Prometheus metrics from prom-client
app.get("/api/v1/metrics", (_, reply) => {
reply.header("content-type", registry.contentType);
return registry.metrics();
});
return {
start: () => app.listen({ port: config.port, host: "0.0.0.0" }),
stop: () =>
Effect.runPromise(
Effect.gen(function* () {
yield* Effect.tryPromise({
try: () => app.close(),
catch: (cause) => messagingLifecycleError("gateway", "app-close", cause),
});
yield* Scope.close(rpcScope, Exit.void);
yield* Effect.tryPromise({
try: () => dispatcher.stop(),
catch: (cause) => messagingLifecycleError("gateway", "dispatcher-stop", cause),
});
}),
),
};
});
}
function headersFrom(headers: Record<string, string | string[] | number | undefined>): ReadonlyArray<[string, string]> {
@ -217,8 +257,8 @@ function headersFrom(headers: Record<string, string | string[] | number | undefi
});
}
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}
export const loadGatewayConfig = Effect.fn("loadGatewayConfig")(function* () {

File diff suppressed because it is too large Load diff

View file

@ -19,9 +19,15 @@ import {
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerStatusError,
requiredString,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type AzureOpenAIProcessorConfig = ProcessorConfig & {
model?: string;
@ -32,32 +38,65 @@ export type AzureOpenAIProcessorConfig = ProcessorConfig & {
maxOutput?: number;
};
export function makeAzureOpenAIProvider(config: AzureOpenAIProcessorConfig): LlmProvider {
const defaultModel = config.model ?? process.env.AZURE_MODEL ?? "gpt-4o";
const defaultTemperature = config.temperature ?? 0.0;
const maxOutput = config.maxOutput ?? 4096;
const apiKey = config.apiKey ?? process.env.AZURE_TOKEN;
if (apiKey === undefined || apiKey.length === 0) {
throw new Error("Azure OpenAI API key not specified");
}
const endpoint = config.endpoint ?? process.env.AZURE_ENDPOINT;
if (endpoint === undefined || endpoint.length === 0) {
throw new Error("Azure OpenAI endpoint not specified");
}
type ResolvedAzureOpenAIConfig = {
readonly defaultModel: string;
readonly defaultTemperature: number;
readonly maxOutput: number;
readonly apiKey: string;
readonly endpoint: string;
readonly apiVersion: string;
};
const loadAzureOpenAIConfig = Effect.fn("loadAzureOpenAIConfig")(function* (
config: AzureOpenAIProcessorConfig,
) {
const defaultModel =
config.model ?? (yield* optionalStringConfig("AzureOpenAI", "AZURE_MODEL")) ?? "gpt-4o";
const apiKey = yield* requiredString(
config.apiKey ?? (yield* optionalStringConfig("AzureOpenAI", "AZURE_TOKEN")),
"AzureOpenAI",
"AZURE_TOKEN",
"Azure OpenAI API key not specified",
);
const endpoint = yield* requiredString(
config.endpoint ?? (yield* optionalStringConfig("AzureOpenAI", "AZURE_ENDPOINT")),
"AzureOpenAI",
"AZURE_ENDPOINT",
"Azure OpenAI endpoint not specified",
);
const apiVersion =
config.apiVersion ??
process.env.AZURE_API_VERSION ??
(yield* optionalStringConfig("AzureOpenAI", "AZURE_API_VERSION")) ??
"2024-12-01-preview";
return {
defaultModel,
defaultTemperature: config.temperature ?? 0.0,
maxOutput: config.maxOutput ?? 4096,
apiKey,
endpoint,
apiVersion,
};
});
const mapAzureOpenAIError = (error: unknown): TextCompletionRuntimeError =>
providerStatusError("AzureOpenAI", error);
export function makeAzureOpenAIProvider(config: AzureOpenAIProcessorConfig): LlmProvider {
const {
defaultModel,
defaultTemperature,
maxOutput,
apiKey,
endpoint,
apiVersion,
} = Effect.runSync(loadAzureOpenAIConfig(config)) satisfies ResolvedAzureOpenAIConfig;
const client = new AzureOpenAI({ apiKey, apiVersion, endpoint });
console.log("[AzureOpenAI] LLM service initialized");
Effect.runSync(Effect.log("[AzureOpenAI] LLM service initialized"));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -66,87 +105,106 @@ export function makeAzureOpenAIProvider(config: AzureOpenAIProcessorConfig): Llm
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const resp = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
});
return {
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
};
} catch (err) {
if ((err as any)?.status === 429) {
throw tooManyRequestsError();
}
throw err;
}
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
}),
catch: mapAzureOpenAIError,
}).pipe(
Effect.map((resp): LlmResult => ({
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
})),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const stream = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
stream: true,
stream_options: { include_usage: true },
});
const stream = Stream.fromEffect(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
stream: true,
stream_options: { include_usage: true },
}),
catch: mapAzureOpenAIError,
}),
).pipe(
Stream.flatMap((openAIStream) => {
const iterator = openAIStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
const content = chunk.choices[0]?.delta?.content;
if (content !== null && content !== undefined && content.length > 0) {
yield {
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
}
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapAzureOpenAIError,
});
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if ((err as any)?.status === 429) {
throw tooManyRequestsError();
}
throw err;
}
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
const content = chunk.choices[0]?.delta?.content;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
if (content !== null && content !== undefined && content.length > 0) {
return [{
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapAzureOpenAIError);
},
};
}
@ -171,6 +229,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -15,9 +15,15 @@ import {
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerStatusError,
requiredString,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type ClaudeProcessorConfig = ProcessorConfig & {
model?: string;
@ -26,21 +32,46 @@ export type ClaudeProcessorConfig = ProcessorConfig & {
maxOutput?: number;
};
type ResolvedClaudeConfig = {
readonly defaultModel: string;
readonly defaultTemperature: number;
readonly maxOutput: number;
readonly apiKey: string;
};
const loadClaudeConfig = Effect.fn("loadClaudeConfig")(function*(config: ClaudeProcessorConfig) {
const apiKey = yield* requiredString(
config.apiKey ?? (yield* optionalStringConfig("Claude", "CLAUDE_KEY")),
"Claude",
"CLAUDE_KEY",
"Claude API key not specified",
);
return {
defaultModel: config.model ?? "claude-sonnet-4-20250514",
defaultTemperature: config.temperature ?? 0.0,
maxOutput: config.maxOutput ?? 8192,
apiKey,
} satisfies ResolvedClaudeConfig;
});
const mapClaudeError = (error: unknown): TextCompletionRuntimeError =>
providerStatusError("Claude", error);
export function makeClaudeProvider(config: ClaudeProcessorConfig): LlmProvider {
const defaultModel = config.model ?? "claude-sonnet-4-20250514";
const defaultTemperature = config.temperature ?? 0.0;
const maxOutput = config.maxOutput ?? 8192;
const apiKey = config.apiKey ?? process.env.CLAUDE_KEY;
if (apiKey === undefined || apiKey.length === 0) {
throw new Error("Claude API key not specified");
}
const {
defaultModel,
defaultTemperature,
maxOutput,
apiKey,
} = Effect.runSync(loadClaudeConfig(config)) satisfies ResolvedClaudeConfig;
const client = new Anthropic({ apiKey });
console.log("[Claude] LLM service initialized");
Effect.runSync(Effect.log("[Claude] LLM service initialized"));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -49,88 +80,120 @@ export function makeClaudeProvider(config: ClaudeProcessorConfig): LlmProvider {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const response = await client.messages.create({
model: modelName,
max_tokens: maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
});
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.messages.create({
model: modelName,
max_tokens: maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
}),
catch: mapClaudeError,
}).pipe(
Effect.map((response): LlmResult => {
const firstContent = response.content[0];
const text = firstContent?.type === "text"
? firstContent.text
: "";
const text = response.content[0].type === "text"
? response.content[0].text
: "";
return {
text,
inToken: response.usage.input_tokens,
outToken: response.usage.output_tokens,
model: modelName,
};
} catch (err) {
if (err instanceof Anthropic.RateLimitError) {
throw tooManyRequestsError();
}
throw err;
}
return {
text,
inToken: response.usage.input_tokens,
outToken: response.usage.output_tokens,
model: modelName,
};
}),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const stream = client.messages.stream({
model: modelName,
max_tokens: maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
});
for await (const event of stream) {
if (event.type === "content_block_delta" && event.delta.type === "text_delta") {
yield {
text: event.delta.text,
inToken: null,
outToken: null,
const stream = Stream.fromEffect(
Effect.try({
try: () =>
client.messages.stream({
model: modelName,
isFinal: false,
};
}
}
max_tokens: maxOutput,
temperature: temp,
system,
messages: [
{ role: "user", content: prompt },
],
}),
catch: mapClaudeError,
}),
).pipe(
Stream.flatMap((anthropicStream) => {
const iterator = anthropicStream[Symbol.asyncIterator]();
const finalMessage = await stream.finalMessage();
yield {
text: "",
inToken: finalMessage.usage.input_tokens,
outToken: finalMessage.usage.output_tokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if (err instanceof Anthropic.RateLimitError) {
throw tooManyRequestsError();
}
throw err;
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapClaudeError,
});
if (next.done === true) {
const finalMessage = yield* Effect.tryPromise({
try: () => anthropicStream.finalMessage(),
catch: mapClaudeError,
});
return [{
text: "",
inToken: finalMessage.usage.input_tokens,
outToken: finalMessage.usage.output_tokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const event = next.value;
if (
event.type === "content_block_delta" &&
event.delta.type === "text_delta"
) {
return [{
text: event.delta.text,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapClaudeError);
},
};
}
export type ClaudeProcessor = ReturnType<typeof makeClaudeProcessor>;
export function makeClaudeProcessor(config: ClaudeProcessorConfig): ReturnType<typeof makeLlmService> {
export function makeClaudeProcessor(
config: ClaudeProcessorConfig,
): ReturnType<typeof makeLlmService> {
return makeLlmService(config, makeClaudeProvider(config));
}
@ -146,6 +209,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -0,0 +1,101 @@
import {
TooManyRequestsError,
errorMessage,
type LlmChunk,
} from "@trustgraph/base";
import { Config, Effect } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export class TextCompletionConfigError extends S.TaggedErrorClass<TextCompletionConfigError>()(
"TextCompletionConfigError",
{
message: S.String,
provider: S.String,
key: S.String,
},
) {}
export class TextCompletionProviderError extends S.TaggedErrorClass<TextCompletionProviderError>()(
"TextCompletionProviderError",
{
message: S.String,
provider: S.String,
},
) {}
export type TextCompletionRuntimeError =
| TextCompletionProviderError
| TooManyRequestsError;
export const optionalStringConfig = Effect.fn("TextCompletion.optionalStringConfig")(function*(
provider: string,
name: string,
) {
const value = yield* Config.string(name).pipe(
Config.option,
Effect.mapError((cause) =>
TextCompletionConfigError.make({
provider,
key: name,
message: errorMessage(cause),
})
),
);
return O.getOrUndefined(value);
});
export const requiredString = (
value: string | undefined,
provider: string,
key: string,
message: string,
) =>
value !== undefined && value.length > 0
? Effect.succeed(value)
: Effect.fail(TextCompletionConfigError.make({ provider, key, message }));
export const providerRuntimeError = (
provider: string,
error: unknown,
): TextCompletionRuntimeError =>
TextCompletionProviderError.make({
provider,
message: errorMessage(error),
});
export const providerStatusError = (
provider: string,
error: unknown,
): TextCompletionRuntimeError => {
const status = typeof error === "object" && error !== null && "status" in error
? (error as { readonly status?: unknown }).status
: undefined;
const statusCode = typeof error === "object" && error !== null && "statusCode" in error
? (error as { readonly statusCode?: unknown }).statusCode
: undefined;
return status === 429 || statusCode === 429
? TooManyRequestsError.make({ message: "Rate limit exceeded" })
: providerRuntimeError(provider, error);
};
export const toAsyncGenerator = (
iterable: AsyncIterable<LlmChunk>,
mapError: (error: unknown) => TextCompletionRuntimeError,
): AsyncGenerator<LlmChunk> => {
const iterator = iterable[Symbol.asyncIterator]();
let generator: AsyncGenerator<LlmChunk>;
generator = {
next: (value?: unknown) => iterator.next(value as never),
return: (value?: unknown) =>
iterator.return === undefined
? Promise.resolve({ done: true, value: value as LlmChunk })
: iterator.return(value as never) as Promise<IteratorResult<LlmChunk>>,
throw: (error?: unknown) =>
iterator.throw === undefined
? Effect.runPromise(Effect.fail(mapError(error))) as Promise<IteratorResult<LlmChunk>>
: iterator.throw(error) as Promise<IteratorResult<LlmChunk>>,
[Symbol.asyncIterator]: () => generator,
} as AsyncGenerator<LlmChunk>;
return generator;
};

View file

@ -17,9 +17,15 @@ import {
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerStatusError,
requiredString,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type MistralProcessorConfig = ProcessorConfig & {
model?: string;
@ -28,22 +34,49 @@ export type MistralProcessorConfig = ProcessorConfig & {
maxOutput?: number;
};
type ResolvedMistralConfig = {
readonly defaultModel: string;
readonly defaultTemperature: number;
readonly maxOutput: number;
readonly apiKey: string;
};
const loadMistralConfig = Effect.fn("loadMistralConfig")(function*(config: MistralProcessorConfig) {
const apiKey = yield* requiredString(
config.apiKey ?? (yield* optionalStringConfig("Mistral", "MISTRAL_TOKEN")),
"Mistral",
"MISTRAL_TOKEN",
"Mistral API key not specified",
);
return {
defaultModel:
config.model ??
(yield* optionalStringConfig("Mistral", "MISTRAL_MODEL")) ??
"ministral-8b-latest",
defaultTemperature: config.temperature ?? 0.0,
maxOutput: config.maxOutput ?? 4096,
apiKey,
} satisfies ResolvedMistralConfig;
});
const mapMistralError = (error: unknown): TextCompletionRuntimeError =>
providerStatusError("Mistral", error);
export function makeMistralProvider(config: MistralProcessorConfig): LlmProvider {
const defaultModel =
config.model ?? process.env.MISTRAL_MODEL ?? "ministral-8b-latest";
const defaultTemperature = config.temperature ?? 0.0;
const maxOutput = config.maxOutput ?? 4096;
const apiKey = config.apiKey ?? process.env.MISTRAL_TOKEN;
if (apiKey === undefined || apiKey.length === 0) {
throw new Error("Mistral API key not specified");
}
const {
defaultModel,
defaultTemperature,
maxOutput,
apiKey,
} = Effect.runSync(loadMistralConfig(config)) satisfies ResolvedMistralConfig;
const client = new Mistral({ apiKey });
console.log("[Mistral] LLM service initialized");
Effect.runSync(Effect.log("[Mistral] LLM service initialized"));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -52,93 +85,114 @@ export function makeMistralProvider(config: MistralProcessorConfig): LlmProvider
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const resp = await client.chat.complete({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
maxTokens: maxOutput,
});
return {
text: (resp.choices?.[0]?.message?.content as string) ?? "",
inToken: resp.usage?.promptTokens ?? 0,
outToken: resp.usage?.completionTokens ?? 0,
model: modelName,
};
} catch (err) {
if ((err as any)?.statusCode === 429 || (err as any)?.status === 429) {
throw tooManyRequestsError();
}
throw err;
}
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.chat.complete({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
maxTokens: maxOutput,
}),
catch: mapMistralError,
}).pipe(
Effect.map((resp): LlmResult => ({
text: (resp.choices?.[0]?.message?.content as string) ?? "",
inToken: resp.usage?.promptTokens ?? 0,
outToken: resp.usage?.completionTokens ?? 0,
model: modelName,
})),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const stream = await client.chat.stream({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
maxTokens: maxOutput,
});
const stream = Stream.fromEffect(
Effect.tryPromise({
try: () =>
client.chat.stream({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
maxTokens: maxOutput,
}),
catch: mapMistralError,
}),
).pipe(
Stream.flatMap((mistralStream) => {
const iterator = mistralStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
const delta = chunk.data?.choices?.[0]?.delta;
const content = delta?.content;
if (typeof content === "string" && content.length > 0) {
yield {
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
if (chunk.data?.usage !== undefined) {
totalInputTokens = chunk.data.usage.promptTokens ?? 0;
totalOutputTokens = chunk.data.usage.completionTokens ?? 0;
}
}
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapMistralError,
});
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if ((err as any)?.statusCode === 429 || (err as any)?.status === 429) {
throw tooManyRequestsError();
}
throw err;
}
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
const delta = chunk.data?.choices?.[0]?.delta;
const content = delta?.content;
if (chunk.data?.usage !== undefined) {
totalInputTokens = chunk.data.usage.promptTokens ?? 0;
totalOutputTokens = chunk.data.usage.completionTokens ?? 0;
}
if (typeof content === "string" && content.length > 0) {
return [{
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapMistralError);
},
};
}
export type MistralProcessor = ReturnType<typeof makeMistralProcessor>;
export function makeMistralProcessor(config: MistralProcessorConfig): ReturnType<typeof makeLlmService> {
export function makeMistralProcessor(
config: MistralProcessorConfig,
): ReturnType<typeof makeLlmService> {
return makeLlmService(config, makeMistralProvider(config));
}
@ -154,6 +208,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -18,32 +18,51 @@ import {
type LlmResult,
type LlmChunk,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerRuntimeError,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type OllamaProcessorConfig = ProcessorConfig & {
model?: string;
ollamaUrl?: string;
};
export function makeOllamaProvider(config: OllamaProcessorConfig): LlmProvider {
const defaultModel =
config.model ??
process.env.OLLAMA_MODEL ??
"qwen2.5:0.5b";
type ResolvedOllamaConfig = {
readonly defaultModel: string;
readonly host: string;
};
const host =
config.ollamaUrl ??
process.env.OLLAMA_URL ??
"http://localhost:11434";
const loadOllamaConfig = Effect.fn("loadOllamaConfig")(function*(config: OllamaProcessorConfig) {
return {
defaultModel:
config.model ??
(yield* optionalStringConfig("Ollama", "OLLAMA_MODEL")) ??
"qwen2.5:0.5b",
host:
config.ollamaUrl ??
(yield* optionalStringConfig("Ollama", "OLLAMA_URL")) ??
"http://localhost:11434",
} satisfies ResolvedOllamaConfig;
});
const mapOllamaError = (error: unknown): TextCompletionRuntimeError =>
providerRuntimeError("Ollama", error);
export function makeOllamaProvider(config: OllamaProcessorConfig): LlmProvider {
const { defaultModel, host } = Effect.runSync(loadOllamaConfig(config)) satisfies ResolvedOllamaConfig;
const client = new Ollama({ host });
console.log(
Effect.runSync(Effect.log(
`[Ollama] LLM service initialized (host=${host}, model=${defaultModel})`,
);
));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -52,73 +71,107 @@ export function makeOllamaProvider(config: OllamaProcessorConfig): LlmProvider {
const modelName = model ?? defaultModel;
const fullPrompt = system + "\n\n" + prompt;
const resp = await client.generate({
model: modelName,
prompt: fullPrompt,
stream: false,
});
return {
text: resp.response,
inToken: resp.prompt_eval_count ?? 0,
outToken: resp.eval_count ?? 0,
model: modelName,
};
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.generate({
model: modelName,
prompt: fullPrompt,
stream: false,
}),
catch: mapOllamaError,
}).pipe(
Effect.map((resp): LlmResult => ({
text: resp.response,
inToken: resp.prompt_eval_count ?? 0,
outToken: resp.eval_count ?? 0,
model: modelName,
})),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
_temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const fullPrompt = system + "\n\n" + prompt;
const stream = await client.generate({
model: modelName,
prompt: fullPrompt,
stream: true,
});
const stream = Stream.fromEffect(
Effect.tryPromise({
try: () =>
client.generate({
model: modelName,
prompt: fullPrompt,
stream: true,
}),
catch: mapOllamaError,
}),
).pipe(
Stream.flatMap((ollamaStream) => {
const iterator = ollamaStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
// Token counts accumulate across chunks; keep the latest values
if (chunk.prompt_eval_count !== undefined) {
totalInputTokens = chunk.prompt_eval_count;
}
if (chunk.eval_count !== undefined) {
totalOutputTokens = chunk.eval_count;
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
if (chunk.response.length > 0) {
yield {
text: chunk.response,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
}
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapOllamaError,
});
// Final chunk with accumulated token counts
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
if (chunk.prompt_eval_count !== undefined) {
totalInputTokens = chunk.prompt_eval_count;
}
if (chunk.eval_count !== undefined) {
totalOutputTokens = chunk.eval_count;
}
if (chunk.response.length > 0) {
return [{
text: chunk.response,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapOllamaError);
},
};
}
export type OllamaProcessor = ReturnType<typeof makeOllamaProcessor>;
export function makeOllamaProcessor(config: OllamaProcessorConfig): ReturnType<typeof makeLlmService> {
export function makeOllamaProcessor(
config: OllamaProcessorConfig,
): ReturnType<typeof makeLlmService> {
return makeLlmService(config, makeOllamaProvider(config));
}
@ -134,6 +187,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -21,7 +21,14 @@ import {
type LlmResult,
type LlmChunk,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerStatusError,
requiredString,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type OpenAICompatibleProcessorConfig = ProcessorConfig & {
model?: string;
@ -31,30 +38,57 @@ export type OpenAICompatibleProcessorConfig = ProcessorConfig & {
maxOutput?: number;
};
type ResolvedOpenAICompatibleConfig = {
readonly defaultModel: string;
readonly defaultTemperature: number;
readonly maxOutput: number;
readonly apiKey: string;
readonly baseURL: string;
};
const loadOpenAICompatibleConfig = Effect.fn("loadOpenAICompatibleConfig")(function*(
config: OpenAICompatibleProcessorConfig,
) {
const defaultModel =
config.model ?? (yield* optionalStringConfig("OpenAI-Compatible", "OPENAI_COMPAT_MODEL")) ?? "default";
const baseURL = yield* requiredString(
config.baseUrl ?? (yield* optionalStringConfig("OpenAI-Compatible", "OPENAI_COMPAT_URL")),
"OpenAI-Compatible",
"OPENAI_COMPAT_URL",
"OpenAI-compatible server URL not specified (set OPENAI_COMPAT_URL)",
);
const apiKey =
config.apiKey ?? (yield* optionalStringConfig("OpenAI-Compatible", "OPENAI_COMPAT_KEY")) ?? "sk-no-key-required";
return {
defaultModel,
defaultTemperature: config.temperature ?? 0.0,
maxOutput: config.maxOutput ?? 4096,
apiKey,
baseURL,
} satisfies ResolvedOpenAICompatibleConfig;
});
const mapOpenAICompatibleError = (error: unknown): TextCompletionRuntimeError =>
providerStatusError("OpenAI-Compatible", error);
export function makeOpenAICompatibleProvider(
config: OpenAICompatibleProcessorConfig,
): LlmProvider {
const defaultModel =
config.model ?? process.env.OPENAI_COMPAT_MODEL ?? "default";
const defaultTemperature = config.temperature ?? 0.0;
const maxOutput = config.maxOutput ?? 4096;
const baseURL = config.baseUrl ?? process.env.OPENAI_COMPAT_URL;
if (baseURL === undefined || baseURL.length === 0) {
throw new Error(
"OpenAI-compatible server URL not specified (set OPENAI_COMPAT_URL)",
);
}
const apiKey =
config.apiKey ?? process.env.OPENAI_COMPAT_KEY ?? "sk-no-key-required";
const {
defaultModel,
defaultTemperature,
maxOutput,
apiKey,
baseURL,
} = Effect.runSync(loadOpenAICompatibleConfig(config)) satisfies ResolvedOpenAICompatibleConfig;
const client = new OpenAI({ baseURL, apiKey });
console.log("[OpenAI-Compatible] LLM service initialized");
Effect.runSync(Effect.log("[OpenAI-Compatible] LLM service initialized"));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -63,72 +97,105 @@ export function makeOpenAICompatibleProvider(
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
const resp = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_tokens: maxOutput,
});
return {
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
};
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_tokens: maxOutput,
}),
catch: mapOpenAICompatibleError,
}).pipe(
Effect.map((resp): LlmResult => ({
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
})),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
const stream = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_tokens: maxOutput,
stream: true,
});
const stream = Stream.fromEffect(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_tokens: maxOutput,
stream: true,
}),
catch: mapOpenAICompatibleError,
}),
).pipe(
Stream.flatMap((openAIStream) => {
const iterator = openAIStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
const content = chunk.choices[0]?.delta?.content;
if (content !== null && content !== undefined && content.length > 0) {
yield {
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
}
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapOpenAICompatibleError,
});
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
const content = chunk.choices[0]?.delta?.content;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
if (content !== null && content !== undefined && content.length > 0) {
return [{
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapOpenAICompatibleError);
},
};
}
@ -153,6 +220,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -15,9 +15,15 @@ import {
type ProcessorConfig,
type LlmResult,
type LlmChunk,
tooManyRequestsError,
} from "@trustgraph/base";
import { Effect, Layer } from "effect";
import { Effect, Layer, Stream } from "effect";
import {
optionalStringConfig,
providerStatusError,
requiredString,
toAsyncGenerator,
type TextCompletionRuntimeError,
} from "./common.ts";
export type OpenAIProcessorConfig = ProcessorConfig & {
model?: string;
@ -27,24 +33,52 @@ export type OpenAIProcessorConfig = ProcessorConfig & {
maxOutput?: number;
};
type ResolvedOpenAIConfig = {
readonly defaultModel: string;
readonly defaultTemperature: number;
readonly maxOutput: number;
readonly apiKey: string;
readonly baseURL: string | undefined;
};
const loadOpenAIConfig = Effect.fn("loadOpenAIConfig")(function*(config: OpenAIProcessorConfig) {
const apiKey = yield* requiredString(
config.apiKey ?? (yield* optionalStringConfig("OpenAI", "OPENAI_TOKEN")),
"OpenAI",
"OPENAI_TOKEN",
"OpenAI API key not specified",
);
return {
defaultModel: config.model ?? "gpt-4o",
defaultTemperature: config.temperature ?? 0.0,
maxOutput: config.maxOutput ?? 4096,
apiKey,
baseURL: config.baseUrl ?? (yield* optionalStringConfig("OpenAI", "OPENAI_BASE_URL")),
} satisfies ResolvedOpenAIConfig;
});
const mapOpenAIError = (error: unknown): TextCompletionRuntimeError =>
providerStatusError("OpenAI", error);
export function makeOpenAIProvider(config: OpenAIProcessorConfig): LlmProvider {
const defaultModel = config.model ?? "gpt-4o";
const defaultTemperature = config.temperature ?? 0.0;
const maxOutput = config.maxOutput ?? 4096;
const apiKey = config.apiKey ?? process.env.OPENAI_TOKEN;
if (apiKey === undefined || apiKey.length === 0) {
throw new Error("OpenAI API key not specified");
}
const {
defaultModel,
defaultTemperature,
maxOutput,
apiKey,
baseURL,
} = Effect.runSync(loadOpenAIConfig(config)) satisfies ResolvedOpenAIConfig;
const client = new OpenAI({
apiKey,
baseURL: config.baseUrl ?? process.env.OPENAI_BASE_URL,
});
apiKey,
baseURL,
});
console.log("[OpenAI] LLM service initialized");
Effect.runSync(Effect.log("[OpenAI] LLM service initialized"));
return {
generateContent: async (
generateContent: (
system: string,
prompt: string,
model?: string,
@ -53,94 +87,115 @@ export function makeOpenAIProvider(config: OpenAIProcessorConfig): LlmProvider {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const resp = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
});
return {
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
};
} catch (err) {
if (err instanceof OpenAI.RateLimitError) {
throw tooManyRequestsError();
}
throw err;
}
return Effect.runPromise(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
}),
catch: mapOpenAIError,
}).pipe(
Effect.map((resp): LlmResult => ({
text: resp.choices[0].message.content ?? "",
inToken: resp.usage?.prompt_tokens ?? 0,
outToken: resp.usage?.completion_tokens ?? 0,
model: modelName,
})),
),
);
},
supportsStreaming: () => true,
generateContentStream: async function* (
generateContentStream: (
system: string,
prompt: string,
model?: string,
temperature?: number,
): AsyncGenerator<LlmChunk> {
): AsyncGenerator<LlmChunk> => {
const modelName = model ?? defaultModel;
const temp = temperature ?? defaultTemperature;
try {
const stream = await client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
stream: true,
stream_options: { include_usage: true },
});
const stream = Stream.fromEffect(
Effect.tryPromise({
try: () =>
client.chat.completions.create({
model: modelName,
messages: [
{ role: "system", content: system },
{ role: "user", content: prompt },
],
temperature: temp,
max_completion_tokens: maxOutput,
stream: true,
stream_options: { include_usage: true },
}),
catch: mapOpenAIError,
}),
).pipe(
Stream.flatMap((openAIStream) => {
const iterator = openAIStream[Symbol.asyncIterator]();
let totalInputTokens = 0;
let totalOutputTokens = 0;
for await (const chunk of stream) {
const content = chunk.choices[0]?.delta?.content;
if (content !== null && content !== undefined && content.length > 0) {
yield {
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
};
}
return Stream.unfold<"pulling" | "done", LlmChunk, TextCompletionRuntimeError, never>(
"pulling",
(state) => {
if (state === "done") return Effect.void as Effect.Effect<undefined>;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
}
return Effect.gen(function* () {
while (true) {
const next = yield* Effect.tryPromise({
try: () => iterator.next(),
catch: mapOpenAIError,
});
yield {
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
};
} catch (err) {
if (err instanceof OpenAI.RateLimitError) {
throw tooManyRequestsError();
}
throw err;
}
if (next.done === true) {
return [{
text: "",
inToken: totalInputTokens,
outToken: totalOutputTokens,
model: modelName,
isFinal: true,
}, "done"] as const;
}
const chunk = next.value;
const content = chunk.choices[0]?.delta?.content;
if (chunk.usage !== null && chunk.usage !== undefined) {
totalInputTokens = chunk.usage.prompt_tokens;
totalOutputTokens = chunk.usage.completion_tokens;
}
if (content !== null && content !== undefined && content.length > 0) {
return [{
text: content,
inToken: null,
outToken: null,
model: modelName,
isFinal: false,
}, "pulling"] as const;
}
}
});
},
);
}),
);
return toAsyncGenerator(Stream.toAsyncIterable(stream), mapOpenAIError);
},
};
}
export type OpenAIProcessor = ReturnType<typeof makeOpenAIProcessor>;
export function makeOpenAIProcessor(config: OpenAIProcessorConfig): ReturnType<typeof makeLlmService> {
export function makeOpenAIProcessor(
config: OpenAIProcessorConfig,
): ReturnType<typeof makeLlmService> {
return makeLlmService(config, makeOpenAIProvider(config));
}
@ -156,6 +211,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig, never, Llm>({
),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -167,7 +167,7 @@ export function makePromptTemplateService(config: PromptTemplateConfig): PromptT
Effect.runPromise(handler(pushedConfig, version)),
);
}
console.log("[PromptTemplate] Service initialized");
Effect.runSync(Effect.log("[PromptTemplate] Service initialized"));
return service;
}
@ -195,6 +195,6 @@ export const program = makeFlowProcessorProgram({
configHandlers: (config: PromptTemplateConfig) => promptTemplateRuntime(config).configHandlers,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -101,7 +101,7 @@ export function makeDocEmbeddingsQueryService(config: ProcessorConfig): DocEmbed
),
),
});
console.log("[DocEmbeddingsQuery] Service initialized");
Effect.runSync(Effect.log("[DocEmbeddingsQuery] Service initialized"));
return service;
}
@ -113,6 +113,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig & QdrantDocQuery
layer: (config) => QdrantDocEmbeddingsQueryLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -9,7 +9,8 @@
import { QdrantClient } from "@qdrant/js-client-rest";
import { errorMessage } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import { Config, Context, Effect, Layer } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export interface QdrantDocQueryConfig {
@ -30,24 +31,58 @@ export interface DocEmbeddingsQueryRequest {
limit: number;
}
export class QdrantDocEmbeddingsQueryError extends S.TaggedErrorClass<QdrantDocEmbeddingsQueryError>()(
"QdrantDocEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
const qdrantDocEmbeddingsQueryError = (operation: string, cause: unknown) =>
QdrantDocEmbeddingsQueryError.make({
operation,
message: errorMessage(cause),
cause,
});
interface ResolvedQdrantDocQueryConfig {
readonly url: string;
readonly apiKey?: string;
}
const loadQdrantDocQueryConfig = Effect.fn("QdrantDocEmbeddingsQuery.loadConfig")(function* (
config: QdrantDocQueryConfig,
) {
const envApiKey = O.getOrUndefined(yield* Config.string("QDRANT_API_KEY").pipe(Config.option));
const apiKey = config.apiKey ?? envApiKey;
return {
url: config.url ?? (yield* Config.string("QDRANT_URL").pipe(Config.withDefault("http://localhost:6333"))),
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
} satisfies ResolvedQdrantDocQueryConfig;
});
export interface QdrantDocEmbeddingsQuery {
readonly query: (request: DocEmbeddingsQueryRequest) => Promise<ChunkMatch[]>;
readonly queryEffect: (
request: DocEmbeddingsQueryRequest,
) => Effect.Effect<ReadonlyArray<ChunkMatch>, QdrantDocEmbeddingsQueryError>;
}
export function makeQdrantDocEmbeddingsQuery(
config: QdrantDocQueryConfig = {},
): QdrantDocEmbeddingsQuery {
const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333";
const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY;
const resolved = Effect.runSync(loadQdrantDocQueryConfig(config));
const client = new QdrantClient({
url,
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
url: resolved.url,
...(resolved.apiKey !== undefined ? { apiKey: resolved.apiKey } : {}),
});
console.log("[QdrantDocQuery] Query service initialized");
Effect.runSync(Effect.log("[QdrantDocQuery] Query service initialized"));
const query = async (request: DocEmbeddingsQueryRequest): Promise<ChunkMatch[]> => {
const queryEffect = Effect.fn("QdrantDocEmbeddingsQuery.query")(function* (request: DocEmbeddingsQueryRequest) {
const { vector, user, collection, limit } = request;
if (vector.length === 0) {
@ -58,18 +93,25 @@ export function makeQdrantDocEmbeddingsQuery(
const collectionName = `d_${user}_${collection}_${dim}`;
// Check if collection exists -- return empty if not
const exists = await client.collectionExists(collectionName);
const exists = yield* Effect.tryPromise({
try: () => client.collectionExists(collectionName),
catch: (cause) => qdrantDocEmbeddingsQueryError("collection-exists", cause),
});
if (!exists.exists) {
console.log(
yield* Effect.log(
`[QdrantDocQuery] Collection ${collectionName} does not exist, returning empty results`,
);
return [];
}
const searchResult = await client.search(collectionName, {
vector,
limit,
with_payload: true,
const searchResult = yield* Effect.tryPromise({
try: () =>
client.search(collectionName, {
vector,
limit,
with_payload: true,
}),
catch: (cause) => qdrantDocEmbeddingsQueryError("search", cause),
});
const chunks: ChunkMatch[] = [];
@ -86,20 +128,14 @@ export function makeQdrantDocEmbeddingsQuery(
}
return chunks;
});
return {
query: (request) => Effect.runPromise(queryEffect(request)),
queryEffect,
};
return { query };
}
export class QdrantDocEmbeddingsQueryError extends S.TaggedErrorClass<QdrantDocEmbeddingsQueryError>()(
"QdrantDocEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantDocEmbeddingsQueryServiceShape {
readonly query: (
request: DocEmbeddingsQueryRequest,
@ -113,24 +149,12 @@ export class QdrantDocEmbeddingsQueryService extends Context.Service<
"@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 = makeQdrantDocEmbeddingsQuery(config);
return {
query: Effect.fn("QdrantDocEmbeddingsQuery.query")(function* (request) {
return yield* Effect.tryPromise({
try: () => query.query(request),
catch: (cause) => qdrantDocEmbeddingsQueryError("query", cause),
});
}),
query: query.queryEffect,
};
};

View file

@ -102,7 +102,7 @@ export function makeGraphEmbeddingsQueryService(config: ProcessorConfig): GraphE
),
),
});
console.log("[GraphEmbeddingsQuery] Service initialized");
Effect.runSync(Effect.log("[GraphEmbeddingsQuery] Service initialized"));
return service;
}
@ -114,6 +114,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig & QdrantGraphQue
layer: (config) => QdrantGraphEmbeddingsQueryLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -12,7 +12,8 @@
import { QdrantClient } from "@qdrant/js-client-rest";
import { errorMessage, type Term } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import { Config, Context, Effect, Layer } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export interface QdrantGraphQueryConfig {
@ -32,6 +33,38 @@ export interface GraphEmbeddingsQueryRequest {
limit: number;
}
export class QdrantGraphEmbeddingsQueryError extends S.TaggedErrorClass<QdrantGraphEmbeddingsQueryError>()(
"QdrantGraphEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
const qdrantGraphEmbeddingsQueryError = (operation: string, cause: unknown) =>
QdrantGraphEmbeddingsQueryError.make({
operation,
message: errorMessage(cause),
cause,
});
interface ResolvedQdrantGraphQueryConfig {
readonly url: string;
readonly apiKey?: string;
}
const loadQdrantGraphQueryConfig = Effect.fn("QdrantGraphEmbeddingsQuery.loadConfig")(function* (
config: QdrantGraphQueryConfig,
) {
const envApiKey = O.getOrUndefined(yield* Config.string("QDRANT_API_KEY").pipe(Config.option));
const apiKey = config.apiKey ?? envApiKey;
return {
url: config.url ?? (yield* Config.string("QDRANT_URL").pipe(Config.withDefault("http://localhost:6333"))),
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
} satisfies ResolvedQdrantGraphQueryConfig;
});
function createTerm(value: string): Term {
if (value.startsWith("http://") || value.startsWith("https://")) {
return { type: "IRI", iri: value };
@ -41,22 +74,26 @@ function createTerm(value: string): Term {
export interface QdrantGraphEmbeddingsQuery {
readonly query: (request: GraphEmbeddingsQueryRequest) => Promise<EntityMatch[]>;
readonly queryEffect: (
request: GraphEmbeddingsQueryRequest,
) => Effect.Effect<ReadonlyArray<EntityMatch>, QdrantGraphEmbeddingsQueryError>;
}
export function makeQdrantGraphEmbeddingsQuery(
config: QdrantGraphQueryConfig = {},
): QdrantGraphEmbeddingsQuery {
const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333";
const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY;
const resolved = Effect.runSync(loadQdrantGraphQueryConfig(config));
const client = new QdrantClient({
url,
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
url: resolved.url,
...(resolved.apiKey !== undefined ? { apiKey: resolved.apiKey } : {}),
});
console.log("[QdrantGraphQuery] Query service initialized");
Effect.runSync(Effect.log("[QdrantGraphQuery] Query service initialized"));
const query = async (request: GraphEmbeddingsQueryRequest): Promise<EntityMatch[]> => {
const queryEffect = Effect.fn("QdrantGraphEmbeddingsQuery.query")(function* (
request: GraphEmbeddingsQueryRequest,
) {
const { vector, user, collection, limit } = request;
if (vector.length === 0) {
@ -67,9 +104,12 @@ export function makeQdrantGraphEmbeddingsQuery(
const collectionName = `t_${user}_${collection}_${dim}`;
// Check if collection exists -- return empty if not
const exists = await client.collectionExists(collectionName);
const exists = yield* Effect.tryPromise({
try: () => client.collectionExists(collectionName),
catch: (cause) => qdrantGraphEmbeddingsQueryError("collection-exists", cause),
});
if (!exists.exists) {
console.log(
yield* Effect.log(
`[QdrantGraphQuery] Collection ${collectionName} does not exist, returning empty results`,
);
return [];
@ -77,10 +117,14 @@ export function makeQdrantGraphEmbeddingsQuery(
// Query 2x the limit so we have a better chance of getting `limit`
// unique entities after deduplication (same heuristic as Python impl)
const searchResult = await client.search(collectionName, {
vector,
limit: limit * 2,
with_payload: true,
const searchResult = yield* Effect.tryPromise({
try: () =>
client.search(collectionName, {
vector,
limit: limit * 2,
with_payload: true,
}),
catch: (cause) => qdrantGraphEmbeddingsQueryError("search", cause),
});
const entitySet = new Set<string>();
@ -106,20 +150,14 @@ export function makeQdrantGraphEmbeddingsQuery(
}
return entities;
});
return {
query: (request) => Effect.runPromise(queryEffect(request)),
queryEffect,
};
return { query };
}
export class QdrantGraphEmbeddingsQueryError extends S.TaggedErrorClass<QdrantGraphEmbeddingsQueryError>()(
"QdrantGraphEmbeddingsQueryError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantGraphEmbeddingsQueryServiceShape {
readonly query: (
request: GraphEmbeddingsQueryRequest,
@ -133,24 +171,12 @@ export class QdrantGraphEmbeddingsQueryService extends Context.Service<
"@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 = makeQdrantGraphEmbeddingsQuery(config);
return {
query: Effect.fn("QdrantGraphEmbeddingsQuery.query")(function* (request) {
return yield* Effect.tryPromise({
try: () => query.query(request),
catch: (cause) => qdrantGraphEmbeddingsQueryError("query", cause),
});
}),
query: query.queryEffect,
};
};

View file

@ -88,7 +88,7 @@ export function makeTriplesQueryService(config: ProcessorConfig): TriplesQuerySe
),
),
});
console.log("[TriplesQuery] Service initialized");
void Effect.runPromise(Effect.log("[TriplesQuery] Service initialized"));
return service;
}
@ -100,6 +100,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig & FalkorDBQueryC
layer: (config) => FalkorDBTriplesQueryLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -1,5 +1,5 @@
/**
* FalkorDB triples query service queries RDF triples from FalkorDB.
* FalkorDB triples query service - queries RDF triples from FalkorDB.
*
* Implements all SPO query patterns (S, P, O, SP, SO, PO, SPO, *).
*
@ -8,7 +8,7 @@
import { createClient, Graph } from "falkordb";
import { errorMessage, type Term, type Triple } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import { Config, Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface FalkorDBQueryConfig {
@ -19,10 +19,14 @@ export interface FalkorDBQueryConfig {
function termToValue(term: Term | undefined): string | null {
if (term === undefined) return null;
switch (term.type) {
case "IRI": return term.iri;
case "LITERAL": return term.value;
case "BLANK": return term.id;
default: return null;
case "IRI":
return term.iri;
case "LITERAL":
return term.value;
case "BLANK":
return term.id;
default:
return null;
}
}
@ -36,7 +40,6 @@ function createTerm(value: string): Term {
return { type: "LITERAL", value };
}
/** Extract a string field from a FalkorDB result row (returns object with named keys). */
function field(row: unknown, key: string): string {
return (row as Record<string, unknown>)?.[key] as string ?? "";
}
@ -50,231 +53,6 @@ export interface FalkorDBTriplesQuery {
) => Promise<Triple[]>;
}
export function makeFalkorDBTriplesQuery(
config: FalkorDBQueryConfig = {},
): FalkorDBTriplesQuery {
const url = config.url ?? process.env.FALKORDB_URL ?? "redis://localhost:6379";
const database = config.database ?? "falkordb";
const client = createClient({ url });
const graph = new Graph(client, database);
const connectPromise = client.connect().then(() => {
console.log(`[FalkorDBTriplesQuery] Connected to ${url}, graph: ${database}`);
}).catch((err) => {
console.error(`[FalkorDBTriplesQuery] Connection failed:`, err);
throw err;
});
const ensureConnected = async (): Promise<void> => {
await connectPromise;
};
const matchPattern = async (
out: [string, string, string][],
sv: string, pv: string, ov: string, limit: number,
): Promise<void> => {
for (const destType of ["Literal", "Node"] as const) {
const destKey = destType === "Literal" ? "value" : "uri";
const result = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri LIMIT ${limit}`,
{ params: { src: sv, rel: pv, dest: ov } },
);
for (const _rec of (result.data ?? [])) {
out.push([sv, pv, ov]);
}
}
};
const matchSP = async (
out: [string, string, string][],
sv: string, pv: string, limit: number,
): Promise<void> => {
// Literals
const litResult = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:Literal) ` +
`RETURN dest.value as dest LIMIT ${limit}`,
{ params: { src: sv, rel: pv } },
);
for (const rec of (litResult.data ?? [])) {
out.push([sv, pv, field(rec, "dest")]);
}
// Nodes
const nodeResult = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:Node) ` +
`RETURN dest.uri as dest LIMIT ${limit}`,
{ params: { src: sv, rel: pv } },
);
for (const rec of (nodeResult.data ?? [])) {
out.push([sv, pv, field(rec, "dest")]);
}
};
const matchSO = async (
out: [string, string, string][],
sv: string, ov: string, limit: number,
): Promise<void> => {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const result = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN rel.uri as rel LIMIT ${limit}`,
{ params: { src: sv, dest: ov } },
);
for (const rec of (result.data ?? [])) {
out.push([sv, field(rec, "rel"), ov]);
}
}
};
const matchPO = async (
out: [string, string, string][],
pv: string, ov: string, limit: number,
): Promise<void> => {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const result = await graph.query(
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri as src LIMIT ${limit}`,
{ params: { rel: pv, dest: ov } },
);
for (const rec of (result.data ?? [])) {
out.push([field(rec, "src"), pv, ov]);
}
}
};
const matchS = async (
out: [string, string, string][],
sv: string, limit: number,
): Promise<void> => {
const litResult = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:Literal) ` +
`RETURN rel.uri as rel, dest.value as dest LIMIT ${limit}`,
{ params: { src: sv } },
);
for (const rec of (litResult.data ?? [])) {
out.push([sv, field(rec, "rel"), field(rec, "dest")]);
}
const nodeResult = await graph.query(
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:Node) ` +
`RETURN rel.uri as rel, dest.uri as dest LIMIT ${limit}`,
{ params: { src: sv } },
);
for (const rec of (nodeResult.data ?? [])) {
out.push([sv, field(rec, "rel"), field(rec, "dest")]);
}
};
const matchP = async (
out: [string, string, string][],
pv: string, limit: number,
): Promise<void> => {
const litResult = await graph.query(
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:Literal) ` +
`RETURN src.uri as src, dest.value as dest LIMIT ${limit}`,
{ params: { rel: pv } },
);
for (const rec of (litResult.data ?? [])) {
out.push([field(rec, "src"), pv, field(rec, "dest")]);
}
const nodeResult = await graph.query(
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:Node) ` +
`RETURN src.uri as src, dest.uri as dest LIMIT ${limit}`,
{ params: { rel: pv } },
);
for (const rec of (nodeResult.data ?? [])) {
out.push([field(rec, "src"), pv, field(rec, "dest")]);
}
};
const matchO = async (
out: [string, string, string][],
ov: string, limit: number,
): Promise<void> => {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const result = await graph.query(
`MATCH (src:Node)-[rel:Rel]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri as src, rel.uri as rel LIMIT ${limit}`,
{ params: { dest: ov } },
);
for (const rec of (result.data ?? [])) {
out.push([field(rec, "src"), field(rec, "rel"), ov]);
}
}
};
const matchAll = async (
out: [string, string, string][],
limit: number,
): Promise<void> => {
const litResult = await graph.query(
`MATCH (src:Node)-[rel:Rel]->(dest:Literal) ` +
`RETURN src.uri as src, rel.uri as rel, dest.value as dest LIMIT ${limit}`,
);
for (const rec of (litResult.data ?? [])) {
out.push([field(rec, "src"), field(rec, "rel"), field(rec, "dest")]);
}
const nodeResult = await graph.query(
`MATCH (src:Node)-[rel:Rel]->(dest:Node) ` +
`RETURN src.uri as src, rel.uri as rel, dest.uri as dest LIMIT ${limit}`,
);
for (const rec of (nodeResult.data ?? [])) {
out.push([field(rec, "src"), field(rec, "rel"), field(rec, "dest")]);
}
};
const queryTriples = async (
s?: Term,
p?: Term,
o?: Term,
limit = 100,
): Promise<Triple[]> => {
await ensureConnected();
const sv = termToValue(s);
const pv = termToValue(p);
const ov = termToValue(o);
const rawTriples: [string, string, string][] = [];
// Query both Node and Literal targets for each pattern
if (sv !== null && pv !== null && ov !== null) {
// SPO — exact match
await matchPattern(rawTriples, sv, pv, ov, limit);
} else if (sv !== null && pv !== null) {
// SP — known subject + predicate
await matchSP(rawTriples, sv, pv, limit);
} else if (sv !== null && ov !== null) {
// SO — known subject + object
await matchSO(rawTriples, sv, ov, limit);
} else if (pv !== null && ov !== null) {
// PO — known predicate + object
await matchPO(rawTriples, pv, ov, limit);
} else if (sv !== null) {
// S only
await matchS(rawTriples, sv, limit);
} else if (pv !== null) {
// P only
await matchP(rawTriples, pv, limit);
} else if (ov !== null) {
// O only
await matchO(rawTriples, ov, limit);
} else {
// Wildcard — all triples
await matchAll(rawTriples, limit);
}
return rawTriples
.filter(([s, p, o]) => s !== null && p !== null && o !== null)
.slice(0, limit)
.map(([s, p, o]) => ({
s: createTerm(s),
p: createTerm(p),
o: createTerm(o),
}));
};
return { queryTriples };
}
export class FalkorDBTriplesQueryError extends S.TaggedErrorClass<FalkorDBTriplesQueryError>()(
"FalkorDBTriplesQueryError",
{
@ -300,31 +78,358 @@ export class FalkorDBTriplesQueryService extends Context.Service<
"@trustgraph/flow/query/triples/falkordb/FalkorDBTriplesQueryService",
) {}
const falkorDBTriplesQueryError = (operation: string, cause: unknown) =>
new FalkorDBTriplesQueryError({
const falkorDBTriplesQueryError = (operation: string, cause: unknown): FalkorDBTriplesQueryError =>
FalkorDBTriplesQueryError.make({
operation,
message: errorMessage(cause),
cause,
});
export const makeFalkorDBTriplesQueryService = (
interface FalkorDBQueryConnection {
readonly graph: Graph;
}
type FalkorDBQueryOptions = Parameters<Graph["query"]>[1];
const resolveFalkorDBQueryConfig = Effect.fn("FalkorDBTriplesQuery.resolveConfig")(function* (
config: FalkorDBQueryConfig,
) {
const url = config.url ?? (yield* Config.string("FALKORDB_URL").pipe(
Config.withDefault("redis://localhost:6379"),
Effect.mapError((cause) => falkorDBTriplesQueryError("config", cause)),
));
return {
url,
database: config.database ?? "falkordb",
};
});
const connectFalkorDBTriplesQuery = (
config: FalkorDBQueryConfig,
): Effect.Effect<FalkorDBQueryConnection, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const { url, database } = yield* resolveFalkorDBQueryConfig(config);
const { client, graph } = yield* Effect.try({
try: () => {
const client = createClient({ url });
return { client, graph: new Graph(client, database) };
},
catch: (cause) => falkorDBTriplesQueryError("create-client", cause),
});
yield* Effect.tryPromise({
try: () => client.connect(),
catch: (cause) => falkorDBTriplesQueryError("connect", cause),
}).pipe(
Effect.tapError((error) =>
Effect.logError("[FalkorDBTriplesQuery] Connection failed", {
error: error.message,
operation: error.operation,
})
),
);
yield* Effect.log(`[FalkorDBTriplesQuery] Connected to ${url}, graph: ${database}`);
return { graph };
});
const queryRows = (
graph: Graph,
operation: string,
query: string,
options?: FalkorDBQueryOptions,
): Effect.Effect<ReadonlyArray<unknown>, FalkorDBTriplesQueryError> =>
Effect.tryPromise({
try: () => graph.query<unknown>(query, options),
catch: (cause) => falkorDBTriplesQueryError(operation, cause),
}).pipe(
Effect.map((result) => result.data ?? []),
);
const matchPattern = (
graph: Graph,
out: [string, string, string][],
sv: string,
pv: string,
ov: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
for (const destType of ["Literal", "Node"] as const) {
const destKey = destType === "Literal" ? "value" : "uri";
const rows = yield* queryRows(
graph,
"match-spo",
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri LIMIT ${limit}`,
{ params: { src: sv, rel: pv, dest: ov } },
);
for (const _rec of rows) {
out.push([sv, pv, ov]);
}
}
});
const matchSP = (
graph: Graph,
out: [string, string, string][],
sv: string,
pv: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const litRows = yield* queryRows(
graph,
"match-sp-literal",
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:Literal) ` +
`RETURN dest.value as dest LIMIT ${limit}`,
{ params: { src: sv, rel: pv } },
);
for (const rec of litRows) {
out.push([sv, pv, field(rec, "dest")]);
}
const nodeRows = yield* queryRows(
graph,
"match-sp-node",
`MATCH (src:Node {uri: $src})-[rel:Rel {uri: $rel}]->(dest:Node) ` +
`RETURN dest.uri as dest LIMIT ${limit}`,
{ params: { src: sv, rel: pv } },
);
for (const rec of nodeRows) {
out.push([sv, pv, field(rec, "dest")]);
}
});
const matchSO = (
graph: Graph,
out: [string, string, string][],
sv: string,
ov: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const rows = yield* queryRows(
graph,
"match-so",
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN rel.uri as rel LIMIT ${limit}`,
{ params: { src: sv, dest: ov } },
);
for (const rec of rows) {
out.push([sv, field(rec, "rel"), ov]);
}
}
});
const matchPO = (
graph: Graph,
out: [string, string, string][],
pv: string,
ov: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const rows = yield* queryRows(
graph,
"match-po",
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri as src LIMIT ${limit}`,
{ params: { rel: pv, dest: ov } },
);
for (const rec of rows) {
out.push([field(rec, "src"), pv, ov]);
}
}
});
const matchS = (
graph: Graph,
out: [string, string, string][],
sv: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const litRows = yield* queryRows(
graph,
"match-s-literal",
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:Literal) ` +
`RETURN rel.uri as rel, dest.value as dest LIMIT ${limit}`,
{ params: { src: sv } },
);
for (const rec of litRows) {
out.push([sv, field(rec, "rel"), field(rec, "dest")]);
}
const nodeRows = yield* queryRows(
graph,
"match-s-node",
`MATCH (src:Node {uri: $src})-[rel:Rel]->(dest:Node) ` +
`RETURN rel.uri as rel, dest.uri as dest LIMIT ${limit}`,
{ params: { src: sv } },
);
for (const rec of nodeRows) {
out.push([sv, field(rec, "rel"), field(rec, "dest")]);
}
});
const matchP = (
graph: Graph,
out: [string, string, string][],
pv: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const litRows = yield* queryRows(
graph,
"match-p-literal",
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:Literal) ` +
`RETURN src.uri as src, dest.value as dest LIMIT ${limit}`,
{ params: { rel: pv } },
);
for (const rec of litRows) {
out.push([field(rec, "src"), pv, field(rec, "dest")]);
}
const nodeRows = yield* queryRows(
graph,
"match-p-node",
`MATCH (src:Node)-[rel:Rel {uri: $rel}]->(dest:Node) ` +
`RETURN src.uri as src, dest.uri as dest LIMIT ${limit}`,
{ params: { rel: pv } },
);
for (const rec of nodeRows) {
out.push([field(rec, "src"), pv, field(rec, "dest")]);
}
});
const matchO = (
graph: Graph,
out: [string, string, string][],
ov: string,
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) {
const rows = yield* queryRows(
graph,
"match-o",
`MATCH (src:Node)-[rel:Rel]->(dest:${destType} {${destKey}: $dest}) ` +
`RETURN src.uri as src, rel.uri as rel LIMIT ${limit}`,
{ params: { dest: ov } },
);
for (const rec of rows) {
out.push([field(rec, "src"), field(rec, "rel"), ov]);
}
}
});
const matchAll = (
graph: Graph,
out: [string, string, string][],
limit: number,
): Effect.Effect<void, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const litRows = yield* queryRows(
graph,
"match-all-literal",
`MATCH (src:Node)-[rel:Rel]->(dest:Literal) ` +
`RETURN src.uri as src, rel.uri as rel, dest.value as dest LIMIT ${limit}`,
);
for (const rec of litRows) {
out.push([field(rec, "src"), field(rec, "rel"), field(rec, "dest")]);
}
const nodeRows = yield* queryRows(
graph,
"match-all-node",
`MATCH (src:Node)-[rel:Rel]->(dest:Node) ` +
`RETURN src.uri as src, rel.uri as rel, dest.uri as dest LIMIT ${limit}`,
);
for (const rec of nodeRows) {
out.push([field(rec, "src"), field(rec, "rel"), field(rec, "dest")]);
}
});
const queryTriplesEffect = (
getConnection: () => Effect.Effect<FalkorDBQueryConnection, FalkorDBTriplesQueryError>,
s: Term | undefined,
p: Term | undefined,
o: Term | undefined,
limit: number,
): Effect.Effect<ReadonlyArray<Triple>, FalkorDBTriplesQueryError> =>
Effect.gen(function* () {
const { graph } = yield* getConnection();
const sv = termToValue(s);
const pv = termToValue(p);
const ov = termToValue(o);
const rawTriples: [string, string, string][] = [];
if (sv !== null && pv !== null && ov !== null) {
yield* matchPattern(graph, rawTriples, sv, pv, ov, limit);
} else if (sv !== null && pv !== null) {
yield* matchSP(graph, rawTriples, sv, pv, limit);
} else if (sv !== null && ov !== null) {
yield* matchSO(graph, rawTriples, sv, ov, limit);
} else if (pv !== null && ov !== null) {
yield* matchPO(graph, rawTriples, pv, ov, limit);
} else if (sv !== null) {
yield* matchS(graph, rawTriples, sv, limit);
} else if (pv !== null) {
yield* matchP(graph, rawTriples, pv, limit);
} else if (ov !== null) {
yield* matchO(graph, rawTriples, ov, limit);
} else {
yield* matchAll(graph, rawTriples, limit);
}
return rawTriples
.slice(0, limit)
.map(([subject, predicate, object]) => ({
s: createTerm(subject),
p: createTerm(predicate),
o: createTerm(object),
}));
});
const makeFalkorDBTriplesQueryEffect = (
config: FalkorDBQueryConfig = {},
): FalkorDBTriplesQueryServiceShape => {
const query = makeFalkorDBTriplesQuery(config);
let cachedConnection: Effect.Effect<FalkorDBQueryConnection, FalkorDBTriplesQueryError> | undefined;
const getConnection = Effect.fn("FalkorDBTriplesQuery.connection")(function* () {
if (cachedConnection === undefined) {
cachedConnection = yield* Effect.cached(connectFalkorDBTriplesQuery(config));
}
return yield* cachedConnection;
});
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),
})),
) => queryTriplesEffect(getConnection, s, p, o, limit)),
};
};
export function makeFalkorDBTriplesQuery(
config: FalkorDBQueryConfig = {},
): FalkorDBTriplesQuery {
const query = makeFalkorDBTriplesQueryEffect(config);
return {
queryTriples: (s, p, o, limit = 100) =>
Effect.runPromise(query.queryTriples(s, p, o, limit)).then((triples) => Array.from(triples)),
};
}
export const makeFalkorDBTriplesQueryService = (
config: FalkorDBQueryConfig = {},
): FalkorDBTriplesQueryServiceShape => makeFalkorDBTriplesQueryEffect(config);
export const FalkorDBTriplesQueryLive = (
config: FalkorDBQueryConfig = {},
): Layer.Layer<FalkorDBTriplesQueryService> =>

View file

@ -161,6 +161,6 @@ export const program = makeFlowProcessorProgram({
layer: () => DocumentRagLive,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -56,7 +56,7 @@ export class DocumentRagEngine extends Context.Service<DocumentRagEngine, Docume
) {}
const documentRagError = (operation: string, cause: unknown) =>
new DocumentRagEngineError({
DocumentRagEngineError.make({
operation,
cause,
message: errorMessage(cause),
@ -68,11 +68,7 @@ export function makeDocumentRagEngine(): DocumentRagEngineShape {
clients: DocumentRagClients,
queryText: string,
options?: DocumentRagQueryOptions,
) =>
Effect.tryPromise({
try: () => queryDocumentRag(clients, queryText, options),
catch: (cause) => documentRagError("query", cause),
}),
) => queryDocumentRag(clients, queryText, options),
),
};
}
@ -97,40 +93,54 @@ export function makeDocumentRag(clients: DocumentRagClients): DocumentRag {
};
}
async function queryDocumentRag(
function queryDocumentRag(
clients: DocumentRagClients,
queryText: string,
options?: DocumentRagQueryOptions,
): Promise<string> {
const collection = options?.collection ?? "default";
): Effect.Effect<string, DocumentRagEngineError> {
return Effect.gen(function* () {
const collection = options?.collection ?? "default";
const embResp = await clients.embeddings.request({ text: [queryText] });
const vectors = embResp.vectors;
const embResp = yield* Effect.tryPromise({
try: () => clients.embeddings.request({ text: [queryText] }),
catch: (cause) => documentRagError("embeddings", cause),
});
const vectors = embResp.vectors;
const docResp = await clients.docEmbeddings.request({
vectors,
limit: 10,
collection,
user: "default",
const docResp = yield* Effect.tryPromise({
try: () => clients.docEmbeddings.request({
vectors,
limit: 10,
collection,
user: "default",
}),
catch: (cause) => documentRagError("document-embeddings", cause),
});
const chunks = docResp.chunks ?? [];
yield* Effect.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 = yield* Effect.tryPromise({
try: () => clients.prompt.request({
name: "document-rag-synthesize",
variables: { query: queryText, context },
}),
catch: (cause) => documentRagError("prompt", cause),
});
const resp = yield* Effect.tryPromise({
try: () => clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => documentRagError("llm", cause),
});
return resp.response;
});
const chunks = docResp.chunks ?? [];
console.log(`[DocumentRag] Found ${chunks.length} matching chunks`);
const context = chunks
.flatMap((chunk) =>
chunk.content !== undefined && chunk.content.length > 0 ? [chunk.content] : [],
)
.join("\n\n---\n\n");
const promptResp = await clients.prompt.request({
name: "document-rag-synthesize",
variables: { query: queryText, context },
});
const resp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
});
return resp.response;
}

View file

@ -192,6 +192,6 @@ export const program = makeFlowProcessorProgram({
layer: () => GraphRagLive,
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -86,7 +86,7 @@ export class GraphRagEngine extends Context.Service<GraphRagEngine, GraphRagEngi
) {}
const graphRagError = (operation: string, cause: unknown) =>
new GraphRagEngineError({
GraphRagEngineError.make({
operation,
cause,
message: errorMessage(cause),
@ -110,11 +110,7 @@ export function makeGraphRagEngine(): GraphRagEngineShape {
queryText: string,
options?: GraphRagQueryOptions,
config?: GraphRagConfig,
) =>
Effect.tryPromise({
try: () => queryGraphRag(clients, queryText, options, config),
catch: (cause) => graphRagError("query", cause),
}),
) => queryGraphRag(clients, queryText, options, config),
),
};
}
@ -142,239 +138,283 @@ export function makeGraphRag(
};
}
async function queryGraphRag(
function queryGraphRag(
clients: GraphRagClients,
queryText: string,
options?: GraphRagQueryOptions,
rawConfig?: GraphRagConfig,
): Promise<GraphRagResult> {
const config = normalizeGraphRagConfig(rawConfig);
console.log(`[GraphRag] Query: "${queryText.slice(0, 80)}..."`);
): Effect.Effect<GraphRagResult, GraphRagEngineError> {
return Effect.gen(function* () {
const config = normalizeGraphRagConfig(rawConfig);
yield* Effect.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 concepts = yield* extractConcepts(clients, queryText);
yield* Effect.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 vectors = yield* getVectors(clients, concepts);
yield* Effect.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 entities = yield* getEntities(clients, config, vectors, options?.collection);
yield* Effect.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 subgraph = yield* followEdges(clients, config, entities, options?.collection);
yield* Effect.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`);
const scoredEdges = yield* scoreEdges(clients, config, queryText, subgraph);
yield* Effect.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)`);
yield* Effect.log(`[GraphRag] Step 6: synthesizing answer from ${scoredEdges.length} edges...`);
const answer = yield* synthesize(
clients,
queryText,
scoredEdges,
options?.chunkCallback,
);
yield* Effect.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 },
return { answer, subgraph: scoredEdges };
});
}
const llmResp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
function extractConcepts(clients: GraphRagClients, query: string): Effect.Effect<string[], GraphRagEngineError> {
return Effect.gen(function* () {
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
name: "extract-concepts",
variables: { query },
}),
catch: (cause) => graphRagError("extract-concepts-prompt", cause),
});
const llmResp = yield* Effect.tryPromise({
try: () => clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => graphRagError("extract-concepts-llm", cause),
});
return llmResp.response
.split("\n")
.map((concept) => concept.trim())
.filter((concept) => concept.length > 0);
});
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;
function getVectors(clients: GraphRagClients, concepts: string[]): Effect.Effect<number[][], GraphRagEngineError> {
return Effect.gen(function* () {
const resp = yield* Effect.tryPromise({
try: () => clients.embeddings.request({ text: concepts }),
catch: (cause) => graphRagError("get-vectors", cause),
});
return resp.vectors;
});
}
async function getEntities(
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,
): Effect.Effect<Term[], GraphRagEngineError> {
return Effect.gen(function* () {
const resp = yield* Effect.tryPromise({
try: () => clients.graphEmbeddings.request({
vectors,
user: "default",
collection: collection ?? "default",
limit: config.entityLimit,
}),
catch: (cause) => graphRagError("get-entities", cause),
});
return resp.entities;
});
return resp.entities;
}
async function followEdges(
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)),
);
): Effect.Effect<Triple[], GraphRagEngineError> {
return Effect.gen(function* () {
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;
}
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 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 queries = unvisited.map((entityStr) => {
const term = stringToTerm(entityStr);
const request: TriplesQueryRequest = {
s: term,
limit: config.tripleLimit,
...(collection !== undefined ? { collection } : {}),
};
return Effect.tryPromise({
try: () => clients.triples.request(request),
catch: (cause) => graphRagError("follow-edges-query", cause),
});
});
const results = await Promise.all(queries);
const nextLevel = new Set<string>();
const results = yield* Effect.all(queries);
const nextLevel = new Set<string>();
for (const result of results) {
for (const triple of result.triples) {
subgraph.push(triple);
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 (depth < config.maxPathLength - 1) {
const objStr = termToString(triple.o);
if (!visited.has(objStr)) {
nextLevel.add(objStr);
}
}
if (subgraph.length >= config.maxSubgraphSize) {
return subgraph;
}
}
if (subgraph.length >= config.maxSubgraphSize) {
return subgraph;
}
}
for (const entity of currentLevel) {
visited.add(entity);
}
currentLevel = nextLevel;
}
for (const entity of currentLevel) {
visited.add(entity);
}
currentLevel = nextLevel;
}
return subgraph.slice(0, config.maxSubgraphSize);
return subgraph.slice(0, config.maxSubgraphSize);
});
}
async function scoreEdges(
function scoreEdges(
clients: GraphRagClients,
config: NormalizedGraphRagConfig,
query: string,
triples: Triple[],
): Promise<Triple[]> {
if (triples.length === 0) return [];
): Effect.Effect<Triple[], GraphRagEngineError> {
return Effect.gen(function* () {
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]);
if (triples.length <= 500) {
yield* Effect.log(`[GraphRag] Skipping edge scoring - ${triples.length} triples fits in context directly`);
return triples;
}
}
console.log(`[GraphRag] Edge scoring: LLM returned ${scored.length} scores, keeping top ${topN.length}, mapped ${result.length} triples`);
const edgeDescriptions = triples.map((triple, index) => ({
id: String(index),
s: termToString(triple.s),
p: termToString(triple.p),
o: termToString(triple.o),
}));
if (result.length === 0) {
return triples.slice(0, config.edgeLimit);
}
const toScore = edgeDescriptions.slice(0, config.edgeScoreLimit);
const knowledgeJson = yield* S.encodeUnknownEffect(S.UnknownFromJsonString)(toScore).pipe(
Effect.mapError((cause) => graphRagError("edge-score-encode", cause)),
);
return result;
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
name: "kg-edge-scoring",
variables: {
query,
knowledge: knowledgeJson,
},
}),
catch: (cause) => graphRagError("edge-score-prompt", cause),
});
const llmResp = yield* Effect.tryPromise({
try: () => clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => graphRagError("edge-score-llm", cause),
});
yield* Effect.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]);
}
}
yield* Effect.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(
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");
): Effect.Effect<string, GraphRagEngineError> {
return Effect.gen(function* () {
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 },
});
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
name: "graph-rag-synthesize",
variables: { query, context },
}),
catch: (cause) => graphRagError("synthesize-prompt", cause),
});
if (chunkCallback !== undefined) {
let fullText = "";
await clients.llm.request(
{
if (chunkCallback !== undefined) {
let fullText = "";
yield* Effect.tryPromise({
try: () => clients.llm.request(
{
system: promptResp.system,
prompt: promptResp.prompt,
streaming: true,
},
{
recipient: (resp) => {
if (resp.response.length === 0) return Promise.resolve(resp.endOfStream === true);
fullText += resp.response;
return chunkCallback(resp.response, resp.endOfStream === true).then(() => resp.endOfStream === true);
},
},
),
catch: (cause) => graphRagError("synthesize-stream", cause),
});
return fullText;
}
const resp = yield* Effect.tryPromise({
try: () => 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;
}
}),
catch: (cause) => graphRagError("synthesize-llm", cause),
});
const resp = await clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
return resp.response;
});
return resp.response;
}
const ScoredEdge = S.Struct({

View file

@ -23,8 +23,8 @@ export function readTextFile(path: string): Promise<string> {
return Bun.file(path).text();
}
export async function readBinaryFile(path: string): Promise<Uint8Array> {
return new Uint8Array(await Bun.file(path).arrayBuffer());
export function readBinaryFile(path: string): Promise<Uint8Array> {
return Bun.file(path).arrayBuffer().then((buffer) => new Uint8Array(buffer));
}
export function writeTextFile(path: string, data: string): Promise<void> {

View file

@ -103,7 +103,7 @@ export function makeGraphEmbeddingsStoreService(config: ProcessorConfig): GraphE
),
),
});
console.log("[GraphEmbeddingsStore] Service initialized");
Effect.runSync(Effect.log("[GraphEmbeddingsStore] Service initialized"));
return service;
}
@ -119,6 +119,6 @@ export const program = makeFlowProcessorProgram<
layer: (config) => QdrantGraphEmbeddingsStoreLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -9,6 +9,10 @@
*/
import { QdrantClient } from "@qdrant/js-client-rest";
import { errorMessage } from "@trustgraph/base";
import { Config, Effect, Random } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export interface QdrantDocEmbeddingsConfig {
url?: string;
@ -27,43 +31,110 @@ export interface DocEmbeddingsMessage {
chunks: DocEmbeddingChunk[];
}
export class QdrantDocEmbeddingsStoreError extends S.TaggedErrorClass<QdrantDocEmbeddingsStoreError>()(
"QdrantDocEmbeddingsStoreError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
const qdrantDocEmbeddingsStoreError = (operation: string, cause: unknown) =>
QdrantDocEmbeddingsStoreError.make({
operation,
message: errorMessage(cause),
cause,
});
interface ResolvedQdrantDocEmbeddingsConfig {
readonly url: string;
readonly apiKey?: string;
}
const loadQdrantDocEmbeddingsConfig = Effect.fn("QdrantDocEmbeddings.loadConfig")(function* (
config: QdrantDocEmbeddingsConfig,
) {
const envApiKey = O.getOrUndefined(yield* Config.string("QDRANT_API_KEY").pipe(Config.option));
const apiKey = config.apiKey ?? envApiKey;
return {
url: config.url ?? (yield* Config.string("QDRANT_URL").pipe(Config.withDefault("http://localhost:6333"))),
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
} satisfies ResolvedQdrantDocEmbeddingsConfig;
});
const randomHex = Effect.fn("QdrantDocEmbeddings.randomHex")(function* (digits: number) {
let result = "";
for (let index = 0; index < digits; index++) {
const value = yield* Random.nextIntBetween(0, 16);
result += value.toString(16);
}
return result;
});
const randomPointId = Effect.fn("QdrantDocEmbeddings.randomPointId")(function* () {
const part1 = yield* randomHex(8);
const part2 = yield* randomHex(4);
const versionRest = yield* randomHex(3);
const variant = yield* Random.nextIntBetween(8, 12);
const variantRest = yield* randomHex(3);
const part5 = yield* randomHex(12);
return `${part1}-${part2}-4${versionRest}-${variant.toString(16)}${variantRest}-${part5}`;
});
export interface QdrantDocEmbeddingsStore {
readonly store: (message: DocEmbeddingsMessage) => Promise<void>;
readonly deleteCollection: (user: string, collection: string) => Promise<void>;
readonly storeEffect: (
message: DocEmbeddingsMessage,
) => Effect.Effect<void, QdrantDocEmbeddingsStoreError>;
readonly deleteCollectionEffect: (
user: string,
collection: string,
) => Effect.Effect<void, QdrantDocEmbeddingsStoreError>;
}
export function makeQdrantDocEmbeddingsStore(
config: QdrantDocEmbeddingsConfig = {},
): QdrantDocEmbeddingsStore {
const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333";
const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY;
const resolved = Effect.runSync(loadQdrantDocEmbeddingsConfig(config));
const client = new QdrantClient({
url,
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
url: resolved.url,
...(resolved.apiKey !== undefined ? { apiKey: resolved.apiKey } : {}),
});
const knownCollections = new Set<string>();
console.log("[QdrantDocEmbeddings] Store initialized");
Effect.runSync(Effect.log("[QdrantDocEmbeddings] Store initialized"));
const collectionName = (user: string, collection: string, dim: number): string =>
`d_${user}_${collection}_${dim}`;
const ensureCollection = async (name: string, dim: number): Promise<void> => {
const ensureCollectionEffect = Effect.fn("QdrantDocEmbeddings.ensureCollection")(function* (
name: string,
dim: number,
) {
if (knownCollections.has(name)) return;
const exists = await client.collectionExists(name);
const exists = yield* Effect.tryPromise({
try: () => client.collectionExists(name),
catch: (cause) => qdrantDocEmbeddingsStoreError("collection-exists", cause),
});
if (!exists.exists) {
console.log(`[QdrantDocEmbeddings] Creating collection ${name} (dim=${dim})`);
await client.createCollection(name, {
vectors: { size: dim, distance: "Cosine" },
yield* Effect.log(`[QdrantDocEmbeddings] Creating collection ${name} (dim=${dim})`);
yield* Effect.tryPromise({
try: () =>
client.createCollection(name, {
vectors: { size: dim, distance: "Cosine" },
}),
catch: (cause) => qdrantDocEmbeddingsStoreError("create-collection", cause),
});
}
knownCollections.add(name);
};
});
const store = async (message: DocEmbeddingsMessage): Promise<void> => {
const storeEffect = Effect.fn("QdrantDocEmbeddings.store")(function* (message: DocEmbeddingsMessage) {
for (const chunk of message.chunks) {
if (chunk.chunkId.length === 0) continue;
if (chunk.vector.length === 0) continue;
@ -71,48 +142,68 @@ export function makeQdrantDocEmbeddingsStore(
const dim = chunk.vector.length;
const name = collectionName(message.user, message.collection, dim);
await ensureCollection(name, dim);
yield* ensureCollectionEffect(name, dim);
await client.upsert(name, {
points: [
{
id: crypto.randomUUID(),
vector: chunk.vector,
payload: {
chunk_id: chunk.chunkId,
...(chunk.content !== undefined && chunk.content.length > 0
? { content: chunk.content }
: {}),
},
},
],
const id = yield* randomPointId();
yield* Effect.tryPromise({
try: () =>
client.upsert(name, {
points: [
{
id,
vector: chunk.vector,
payload: {
chunk_id: chunk.chunkId,
...(chunk.content !== undefined && chunk.content.length > 0
? { content: chunk.content }
: {}),
},
},
],
}),
catch: (cause) => qdrantDocEmbeddingsStoreError("upsert", cause),
});
}
};
});
const deleteCollection = async (user: string, collection: string): Promise<void> => {
const deleteCollectionEffect = Effect.fn("QdrantDocEmbeddings.deleteCollection")(function* (
user: string,
collection: string,
) {
const prefix = `d_${user}_${collection}_`;
const allCollections = await client.getCollections();
const allCollections = yield* Effect.tryPromise({
try: () => client.getCollections(),
catch: (cause) => qdrantDocEmbeddingsStoreError("get-collections", cause),
});
const matching = allCollections.collections.filter((c) =>
c.name.startsWith(prefix),
);
if (matching.length === 0) {
console.log(`[QdrantDocEmbeddings] No collections matching prefix ${prefix}`);
yield* Effect.log(`[QdrantDocEmbeddings] No collections matching prefix ${prefix}`);
return;
}
for (const coll of matching) {
await client.deleteCollection(coll.name);
yield* Effect.tryPromise({
try: () => client.deleteCollection(coll.name),
catch: (cause) => qdrantDocEmbeddingsStoreError("delete-collection", cause),
});
knownCollections.delete(coll.name);
console.log(`[QdrantDocEmbeddings] Deleted collection: ${coll.name}`);
yield* Effect.log(`[QdrantDocEmbeddings] Deleted collection: ${coll.name}`);
}
console.log(
yield* Effect.log(
`[QdrantDocEmbeddings] Deleted ${matching.length} collection(s) for ${user}/${collection}`,
);
};
});
return { store, deleteCollection };
return {
store: (message) => Effect.runPromise(storeEffect(message)),
deleteCollection: (user, collection) =>
Effect.runPromise(deleteCollectionEffect(user, collection)),
storeEffect,
deleteCollectionEffect,
};
}

View file

@ -10,7 +10,8 @@
import { QdrantClient } from "@qdrant/js-client-rest";
import { errorMessage, type Term } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import { Config, Context, Effect, Layer, Random } from "effect";
import * as O from "effect/Option";
import * as S from "effect/Schema";
export interface QdrantGraphEmbeddingsConfig {
@ -30,6 +31,57 @@ export interface GraphEmbeddingsMessage {
entities: GraphEmbeddingEntity[];
}
export class QdrantGraphEmbeddingsStoreError extends S.TaggedErrorClass<QdrantGraphEmbeddingsStoreError>()(
"QdrantGraphEmbeddingsStoreError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
const qdrantGraphEmbeddingsStoreError = (operation: string, cause: unknown) =>
QdrantGraphEmbeddingsStoreError.make({
operation,
message: errorMessage(cause),
cause,
});
interface ResolvedQdrantGraphEmbeddingsConfig {
readonly url: string;
readonly apiKey?: string;
}
const loadQdrantGraphEmbeddingsConfig = Effect.fn("QdrantGraphEmbeddings.loadConfig")(function* (
config: QdrantGraphEmbeddingsConfig,
) {
const envApiKey = O.getOrUndefined(yield* Config.string("QDRANT_API_KEY").pipe(Config.option));
const apiKey = config.apiKey ?? envApiKey;
return {
url: config.url ?? (yield* Config.string("QDRANT_URL").pipe(Config.withDefault("http://localhost:6333"))),
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
} satisfies ResolvedQdrantGraphEmbeddingsConfig;
});
const randomHex = Effect.fn("QdrantGraphEmbeddings.randomHex")(function* (digits: number) {
let result = "";
for (let index = 0; index < digits; index++) {
const value = yield* Random.nextIntBetween(0, 16);
result += value.toString(16);
}
return result;
});
const randomPointId = Effect.fn("QdrantGraphEmbeddings.randomPointId")(function* () {
const part1 = yield* randomHex(8);
const part2 = yield* randomHex(4);
const versionRest = yield* randomHex(3);
const variant = yield* Random.nextIntBetween(8, 12);
const variantRest = yield* randomHex(3);
const part5 = yield* randomHex(12);
return `${part1}-${part2}-4${versionRest}-${variant.toString(16)}${variantRest}-${part5}`;
});
function getTermValue(term: Term): string | null {
switch (term.type) {
case "IRI":
@ -46,40 +98,56 @@ function getTermValue(term: Term): string | null {
export interface QdrantGraphEmbeddingsStore {
readonly store: (message: GraphEmbeddingsMessage) => Promise<void>;
readonly deleteCollection: (user: string, collection: string) => Promise<void>;
readonly storeEffect: (
message: GraphEmbeddingsMessage,
) => Effect.Effect<void, QdrantGraphEmbeddingsStoreError>;
readonly deleteCollectionEffect: (
user: string,
collection: string,
) => Effect.Effect<void, QdrantGraphEmbeddingsStoreError>;
}
export function makeQdrantGraphEmbeddingsStore(
config: QdrantGraphEmbeddingsConfig = {},
): QdrantGraphEmbeddingsStore {
const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333";
const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY;
const resolved = Effect.runSync(loadQdrantGraphEmbeddingsConfig(config));
const client = new QdrantClient({
url,
...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}),
url: resolved.url,
...(resolved.apiKey !== undefined ? { apiKey: resolved.apiKey } : {}),
});
const knownCollections = new Set<string>();
console.log("[QdrantGraphEmbeddings] Store initialized");
Effect.runSync(Effect.log("[QdrantGraphEmbeddings] Store initialized"));
const collectionName = (user: string, collection: string, dim: number): string =>
`t_${user}_${collection}_${dim}`;
const ensureCollection = async (name: string, dim: number): Promise<void> => {
const ensureCollectionEffect = Effect.fn("QdrantGraphEmbeddings.ensureCollection")(function* (
name: string,
dim: number,
) {
if (knownCollections.has(name)) return;
const exists = await client.collectionExists(name);
const exists = yield* Effect.tryPromise({
try: () => client.collectionExists(name),
catch: (cause) => qdrantGraphEmbeddingsStoreError("collection-exists", cause),
});
if (!exists.exists) {
console.log(`[QdrantGraphEmbeddings] Creating collection ${name} (dim=${dim})`);
await client.createCollection(name, {
vectors: { size: dim, distance: "Cosine" },
yield* Effect.log(`[QdrantGraphEmbeddings] Creating collection ${name} (dim=${dim})`);
yield* Effect.tryPromise({
try: () =>
client.createCollection(name, {
vectors: { size: dim, distance: "Cosine" },
}),
catch: (cause) => qdrantGraphEmbeddingsStoreError("create-collection", cause),
});
}
knownCollections.add(name);
};
});
const store = async (message: GraphEmbeddingsMessage): Promise<void> => {
const storeEffect = Effect.fn("QdrantGraphEmbeddings.store")(function* (message: GraphEmbeddingsMessage) {
for (const entry of message.entities) {
const entityValue = getTermValue(entry.entity);
if (entityValue === null || entityValue.length === 0) continue;
@ -88,61 +156,72 @@ export function makeQdrantGraphEmbeddingsStore(
const dim = entry.vector.length;
const name = collectionName(message.user, message.collection, dim);
await ensureCollection(name, dim);
yield* ensureCollectionEffect(name, dim);
const payload: Record<string, unknown> = { entity: entityValue };
if (entry.chunkId !== undefined && entry.chunkId.length > 0) {
payload.chunk_id = entry.chunkId;
}
await client.upsert(name, {
points: [
{
id: crypto.randomUUID(),
vector: entry.vector,
payload,
},
],
const id = yield* randomPointId();
yield* Effect.tryPromise({
try: () =>
client.upsert(name, {
points: [
{
id,
vector: entry.vector,
payload,
},
],
}),
catch: (cause) => qdrantGraphEmbeddingsStoreError("upsert", cause),
});
}
};
});
const deleteCollection = async (user: string, collection: string): Promise<void> => {
const deleteCollectionEffect = Effect.fn("QdrantGraphEmbeddings.deleteCollection")(function* (
user: string,
collection: string,
) {
const prefix = `t_${user}_${collection}_`;
const allCollections = await client.getCollections();
const allCollections = yield* Effect.tryPromise({
try: () => client.getCollections(),
catch: (cause) => qdrantGraphEmbeddingsStoreError("get-collections", cause),
});
const matching = allCollections.collections.filter((c) =>
c.name.startsWith(prefix),
);
if (matching.length === 0) {
console.log(`[QdrantGraphEmbeddings] No collections matching prefix ${prefix}`);
yield* Effect.log(`[QdrantGraphEmbeddings] No collections matching prefix ${prefix}`);
return;
}
for (const coll of matching) {
await client.deleteCollection(coll.name);
yield* Effect.tryPromise({
try: () => client.deleteCollection(coll.name),
catch: (cause) => qdrantGraphEmbeddingsStoreError("delete-collection", cause),
});
knownCollections.delete(coll.name);
console.log(`[QdrantGraphEmbeddings] Deleted collection: ${coll.name}`);
yield* Effect.log(`[QdrantGraphEmbeddings] Deleted collection: ${coll.name}`);
}
console.log(
yield* Effect.log(
`[QdrantGraphEmbeddings] Deleted ${matching.length} collection(s) for ${user}/${collection}`,
);
});
return {
store: (message) => Effect.runPromise(storeEffect(message)),
deleteCollection: (user, collection) =>
Effect.runPromise(deleteCollectionEffect(user, collection)),
storeEffect,
deleteCollectionEffect,
};
return { store, deleteCollection };
}
export class QdrantGraphEmbeddingsStoreError extends S.TaggedErrorClass<QdrantGraphEmbeddingsStoreError>()(
"QdrantGraphEmbeddingsStoreError",
{
message: S.String,
operation: S.String,
cause: S.DefectWithStack,
},
) {}
export interface QdrantGraphEmbeddingsStoreServiceShape {
readonly store: (
message: GraphEmbeddingsMessage,
@ -160,33 +239,13 @@ export class QdrantGraphEmbeddingsStoreService extends Context.Service<
"@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 = makeQdrantGraphEmbeddingsStore(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),
});
}),
store: store.storeEffect,
deleteCollection: store.deleteCollectionEffect,
};
};

View file

@ -66,7 +66,7 @@ export function makeTriplesStoreService(config: ProcessorConfig): TriplesStoreSe
),
),
});
console.log("[TriplesStore] Service initialized");
void Effect.runPromise(Effect.log("[TriplesStore] Service initialized"));
return service;
}
@ -78,6 +78,6 @@ export const program = makeFlowProcessorProgram<ProcessorConfig & FalkorDBConfig
layer: (config) => FalkorDBTriplesStoreLive(config),
});
export async function run(): Promise<void> {
await Effect.runPromise(program);
export function run(): Promise<void> {
return Effect.runPromise(program);
}

View file

@ -1,5 +1,5 @@
/**
* FalkorDB triples store writes RDF triples to a FalkorDB graph.
* FalkorDB triples store - writes RDF triples to a FalkorDB graph.
*
* FalkorDB is Redis-based and uses Cypher queries, same as the Python impl.
* Pairs well with Graphiti which also uses FalkorDB as its backend.
@ -9,7 +9,7 @@
import { createClient, Graph } from "falkordb";
import { errorMessage, type Term, type Triple } from "@trustgraph/base";
import { Context, Effect, Layer } from "effect";
import { Config, Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface FalkorDBConfig {
@ -26,7 +26,7 @@ function getTermValue(term: Term): string {
case "BLANK":
return term.id;
case "TRIPLE":
return getTermValue(term.triple.s); // fallback
return getTermValue(term.triple.s);
}
}
@ -55,113 +55,6 @@ export interface FalkorDBTriplesStore {
readonly deleteCollection: (user: string, collection: string) => Promise<void>;
}
export function makeFalkorDBTriplesStore(config: FalkorDBConfig = {}): FalkorDBTriplesStore {
const url = config.url ?? process.env.FALKORDB_URL ?? "redis://localhost:6379";
const database = config.database ?? "falkordb";
const client = createClient({ url });
const graph = new Graph(client, database);
const connectPromise = client.connect().then(() => {
console.log(`[FalkorDBTriplesStore] Connected to ${url}, graph: ${database}`);
}).catch((err) => {
console.error(`[FalkorDBTriplesStore] Connection failed:`, err);
throw err;
});
const ensureConnected = async (): Promise<void> => {
await connectPromise;
};
const createNode = async (uri: string, user: string, collection: string): Promise<void> => {
await ensureConnected();
await graph.query(
"MERGE (n:Node {uri: $uri, user: $user, collection: $collection})",
{ params: { uri, user, collection } },
);
};
const createLiteral = async (value: string, user: string, collection: string): Promise<void> => {
await ensureConnected();
await graph.query(
"MERGE (n:Literal {value: $value, user: $user, collection: $collection})",
{ params: { value, user, collection } },
);
};
const relateNode = async (
src: string, uri: string, dest: string,
user: string, collection: string,
): Promise<void> => {
await ensureConnected();
await graph.query(
"MATCH (src:Node {uri: $src, user: $user, collection: $collection}) " +
"MATCH (dest:Node {uri: $dest, user: $user, collection: $collection}) " +
"MERGE (src)-[:Rel {uri: $uri, user: $user, collection: $collection}]->(dest)",
{ params: { src, dest, uri, user, collection } },
);
};
const relateLiteral = async (
src: string, uri: string, dest: string,
user: string, collection: string,
): Promise<void> => {
await ensureConnected();
await graph.query(
"MATCH (src:Node {uri: $src, user: $user, collection: $collection}) " +
"MATCH (dest:Literal {value: $dest, user: $user, collection: $collection}) " +
"MERGE (src)-[:Rel {uri: $uri, user: $user, collection: $collection}]->(dest)",
{ params: { src, dest, uri, user, collection } },
);
};
const storeTriples = async (
triples: Triple[],
user = "default",
collection = "default",
): Promise<void> => {
for (const t of triples) {
const s = getTermValue(t.s);
const p = getTermValue(t.p);
const o = getTermValue(t.o);
await createNode(s, user, collection);
if (t.o.type === "IRI") {
await createNode(o, user, collection);
await relateNode(s, p, o, user, collection);
} else {
await createLiteral(o, user, collection);
await relateLiteral(s, p, o, user, collection);
}
}
};
const deleteCollection = async (user: string, collection: string): Promise<void> => {
await ensureConnected();
await graph.query(
"MATCH (n:Node {user: $user, collection: $collection}) DETACH DELETE n",
{ params: { user, collection } },
);
await graph.query(
"MATCH (n:Literal {user: $user, collection: $collection}) DETACH DELETE n",
{ params: { user, collection } },
);
await graph.query(
"MATCH (c:CollectionMetadata {user: $user, collection: $collection}) DELETE c",
{ params: { user, collection } },
);
};
return {
createNode,
createLiteral,
relateNode,
relateLiteral,
storeTriples,
deleteCollection,
};
}
export class FalkorDBTriplesStoreError extends S.TaggedErrorClass<FalkorDBTriplesStoreError>()(
"FalkorDBTriplesStoreError",
{
@ -190,35 +83,268 @@ export class FalkorDBTriplesStoreService extends Context.Service<
"@trustgraph/flow/storage/triples/falkordb/FalkorDBTriplesStoreService",
) {}
const falkorDBTriplesStoreError = (operation: string, cause: unknown) =>
new FalkorDBTriplesStoreError({
const falkorDBTriplesStoreError = (operation: string, cause: unknown): FalkorDBTriplesStoreError =>
FalkorDBTriplesStoreError.make({
operation,
message: errorMessage(cause),
cause,
});
interface FalkorDBStoreConnection {
readonly graph: Graph;
}
type FalkorDBQueryOptions = Parameters<Graph["query"]>[1];
interface FalkorDBTriplesStoreEffectShape {
readonly createNode: (
uri: string,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly createLiteral: (
value: string,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly relateNode: (
src: string,
uri: string,
dest: string,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly relateLiteral: (
src: string,
uri: string,
dest: string,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly storeTriples: (
triples: ReadonlyArray<Triple>,
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
readonly deleteCollection: (
user: string,
collection: string,
) => Effect.Effect<void, FalkorDBTriplesStoreError>;
}
const resolveFalkorDBStoreConfig = Effect.fn("FalkorDBTriplesStore.resolveConfig")(function* (
config: FalkorDBConfig,
) {
const url = config.url ?? (yield* Config.string("FALKORDB_URL").pipe(
Config.withDefault("redis://localhost:6379"),
Effect.mapError((cause) => falkorDBTriplesStoreError("config", cause)),
));
return {
url,
database: config.database ?? "falkordb",
};
});
const connectFalkorDBTriplesStore = (
config: FalkorDBConfig,
): Effect.Effect<FalkorDBStoreConnection, FalkorDBTriplesStoreError> =>
Effect.gen(function* () {
const { url, database } = yield* resolveFalkorDBStoreConfig(config);
const { client, graph } = yield* Effect.try({
try: () => {
const client = createClient({ url });
return { client, graph: new Graph(client, database) };
},
catch: (cause) => falkorDBTriplesStoreError("create-client", cause),
});
yield* Effect.tryPromise({
try: () => client.connect(),
catch: (cause) => falkorDBTriplesStoreError("connect", cause),
}).pipe(
Effect.tapError((error) =>
Effect.logError("[FalkorDBTriplesStore] Connection failed", {
error: error.message,
operation: error.operation,
})
),
);
yield* Effect.log(`[FalkorDBTriplesStore] Connected to ${url}, graph: ${database}`);
return { graph };
});
const runGraphQuery = (
graph: Graph,
operation: string,
query: string,
options?: FalkorDBQueryOptions,
): Effect.Effect<void, FalkorDBTriplesStoreError> =>
Effect.tryPromise({
try: () => graph.query(query, options),
catch: (cause) => falkorDBTriplesStoreError(operation, cause),
}).pipe(
Effect.asVoid,
);
const makeFalkorDBTriplesStoreEffect = (
config: FalkorDBConfig = {},
): FalkorDBTriplesStoreEffectShape => {
let cachedConnection: Effect.Effect<FalkorDBStoreConnection, FalkorDBTriplesStoreError> | undefined;
const getConnection = Effect.fn("FalkorDBTriplesStore.connection")(function* () {
if (cachedConnection === undefined) {
cachedConnection = yield* Effect.cached(connectFalkorDBTriplesStore(config));
}
return yield* cachedConnection;
});
const createNode = Effect.fn("FalkorDBTriplesStore.createNode")(function* (
uri: string,
user: string,
collection: string,
) {
const { graph } = yield* getConnection();
yield* runGraphQuery(
graph,
"create-node",
"MERGE (n:Node {uri: $uri, user: $user, collection: $collection})",
{ params: { uri, user, collection } },
);
});
const createLiteral = Effect.fn("FalkorDBTriplesStore.createLiteral")(function* (
value: string,
user: string,
collection: string,
) {
const { graph } = yield* getConnection();
yield* runGraphQuery(
graph,
"create-literal",
"MERGE (n:Literal {value: $value, user: $user, collection: $collection})",
{ params: { value, user, collection } },
);
});
const relateNode = Effect.fn("FalkorDBTriplesStore.relateNode")(function* (
src: string,
uri: string,
dest: string,
user: string,
collection: string,
) {
const { graph } = yield* getConnection();
yield* runGraphQuery(
graph,
"relate-node",
"MATCH (src:Node {uri: $src, user: $user, collection: $collection}) " +
"MATCH (dest:Node {uri: $dest, user: $user, collection: $collection}) " +
"MERGE (src)-[:Rel {uri: $uri, user: $user, collection: $collection}]->(dest)",
{ params: { src, dest, uri, user, collection } },
);
});
const relateLiteral = Effect.fn("FalkorDBTriplesStore.relateLiteral")(function* (
src: string,
uri: string,
dest: string,
user: string,
collection: string,
) {
const { graph } = yield* getConnection();
yield* runGraphQuery(
graph,
"relate-literal",
"MATCH (src:Node {uri: $src, user: $user, collection: $collection}) " +
"MATCH (dest:Literal {value: $dest, user: $user, collection: $collection}) " +
"MERGE (src)-[:Rel {uri: $uri, user: $user, collection: $collection}]->(dest)",
{ params: { src, dest, uri, user, collection } },
);
});
const storeTriples = Effect.fn("FalkorDBTriplesStore.storeTriples")(function* (
triples: ReadonlyArray<Triple>,
user: string,
collection: string,
) {
for (const triple of triples) {
const s = getTermValue(triple.s);
const p = getTermValue(triple.p);
const o = getTermValue(triple.o);
yield* createNode(s, user, collection);
if (triple.o.type === "IRI") {
yield* createNode(o, user, collection);
yield* relateNode(s, p, o, user, collection);
} else {
yield* createLiteral(o, user, collection);
yield* relateLiteral(s, p, o, user, collection);
}
}
});
const deleteCollection = Effect.fn("FalkorDBTriplesStore.deleteCollection")(function* (
user: string,
collection: string,
) {
const { graph } = yield* getConnection();
yield* runGraphQuery(
graph,
"delete-collection-nodes",
"MATCH (n:Node {user: $user, collection: $collection}) DETACH DELETE n",
{ params: { user, collection } },
);
yield* runGraphQuery(
graph,
"delete-collection-literals",
"MATCH (n:Literal {user: $user, collection: $collection}) DETACH DELETE n",
{ params: { user, collection } },
);
yield* runGraphQuery(
graph,
"delete-collection-metadata",
"MATCH (c:CollectionMetadata {user: $user, collection: $collection}) DELETE c",
{ params: { user, collection } },
);
});
return {
createNode,
createLiteral,
relateNode,
relateLiteral,
storeTriples,
deleteCollection,
};
};
export function makeFalkorDBTriplesStore(config: FalkorDBConfig = {}): FalkorDBTriplesStore {
const store = makeFalkorDBTriplesStoreEffect(config);
return {
createNode: (uri, user, collection) =>
Effect.runPromise(store.createNode(uri, user, collection)),
createLiteral: (value, user, collection) =>
Effect.runPromise(store.createLiteral(value, user, collection)),
relateNode: (src, uri, dest, user, collection) =>
Effect.runPromise(store.relateNode(src, uri, dest, user, collection)),
relateLiteral: (src, uri, dest, user, collection) =>
Effect.runPromise(store.relateLiteral(src, uri, dest, user, collection)),
storeTriples: (triples, user = "default", collection = "default") =>
Effect.runPromise(store.storeTriples(triples, user, collection)),
deleteCollection: (user, collection) =>
Effect.runPromise(store.deleteCollection(user, collection)),
};
}
export const makeFalkorDBTriplesStoreService = (
config: FalkorDBConfig = {},
): FalkorDBTriplesStoreServiceShape => {
const store = makeFalkorDBTriplesStore(config);
const store = makeFalkorDBTriplesStoreEffect(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),
})),
storeTriples: store.storeTriples,
deleteCollection: store.deleteCollection,
};
};