From dca278682856a9215f0260e7180c52693d3a71b6 Mon Sep 17 00:00:00 2001 From: elpresidank Date: Mon, 1 Jun 2026 20:26:47 -0500 Subject: [PATCH] Remove native classes from TS runtime --- ts/CLASS_EFFECT_GOAL.md | 151 + ts/bun.lock | 3 + ts/package.json | 1 + .../base/src/__tests__/consumer.test.ts | 31 +- .../__tests__/flow-processor-runtime.test.ts | 6 +- .../src/__tests__/flow-spec-runtime.test.ts | 19 +- .../src/__tests__/messaging-runtime.test.ts | 4 +- .../src/__tests__/runtime-services.test.ts | 94 +- ts/packages/base/src/backend/index.ts | 2 +- ts/packages/base/src/backend/nats.ts | 364 +- ts/packages/base/src/backend/pubsub.ts | 6 +- ts/packages/base/src/errors.ts | 2 +- ts/packages/base/src/messaging/consumer.ts | 122 +- ts/packages/base/src/messaging/index.ts | 8 +- ts/packages/base/src/messaging/producer.ts | 78 +- .../base/src/messaging/request-response.ts | 135 +- ts/packages/base/src/messaging/subscriber.ts | 215 +- ts/packages/base/src/metrics/index.ts | 8 +- ts/packages/base/src/metrics/prometheus.ts | 72 +- .../base/src/processor/async-processor.ts | 283 +- .../base/src/processor/flow-processor.ts | 185 +- ts/packages/base/src/processor/flow.ts | 343 +- ts/packages/base/src/processor/index.ts | 8 + ts/packages/base/src/processor/program.ts | 30 +- ts/packages/base/src/schema/messages.ts | 17 +- .../base/src/services/embeddings-service.ts | 41 +- ts/packages/base/src/services/index.ts | 2 + ts/packages/base/src/services/llm-service.ts | 66 +- ts/packages/base/src/spec/consumer-spec.ts | 108 +- ts/packages/base/src/spec/index.ts | 8 +- ts/packages/base/src/spec/parameter-spec.ts | 29 +- ts/packages/base/src/spec/producer-spec.ts | 35 +- .../base/src/spec/request-response-spec.ts | 54 +- .../client/src/__tests__/rpc-timeout.test.ts | 66 +- .../client/src/socket/effect-rpc-client.ts | 228 +- .../client/src/socket/trustgraph-socket.ts | 3809 +++++++++-------- .../flow/src/agent/mcp-tool/service.ts | 47 +- ts/packages/flow/src/agent/react/index.ts | 2 +- ts/packages/flow/src/agent/react/parser.ts | 161 +- ts/packages/flow/src/agent/react/service.ts | 62 +- ts/packages/flow/src/chunking/service.ts | 37 +- ts/packages/flow/src/config/service.ts | 997 +++-- ts/packages/flow/src/cores/service.ts | 800 ++-- ts/packages/flow/src/decoding/pdf-decoder.ts | 35 +- ts/packages/flow/src/embeddings/ollama.ts | 27 +- .../flow/src/extract/knowledge-extract.ts | 37 +- ts/packages/flow/src/flow-manager/service.ts | 878 ++-- .../flow/src/gateway/dispatch/manager.ts | 165 +- ts/packages/flow/src/gateway/index.ts | 8 +- ts/packages/flow/src/gateway/server.ts | 4 +- ts/packages/flow/src/index.ts | 24 +- .../flow/src/librarian/collection-manager.ts | 108 +- ts/packages/flow/src/librarian/service.ts | 1608 +++---- .../src/model/text-completion/azure-openai.ts | 216 +- .../flow/src/model/text-completion/claude.ts | 205 +- .../flow/src/model/text-completion/mistral.ts | 211 +- .../flow/src/model/text-completion/ollama.ts | 160 +- .../text-completion/openai-compatible.ts | 194 +- .../flow/src/model/text-completion/openai.ts | 211 +- ts/packages/flow/src/prompt/template.ts | 42 +- .../query/embeddings/qdrant-doc-service.ts | 46 +- .../flow/src/query/embeddings/qdrant-doc.ts | 36 +- .../query/embeddings/qdrant-graph-service.ts | 46 +- .../flow/src/query/embeddings/qdrant-graph.ts | 36 +- .../src/query/triples/falkordb-service.ts | 46 +- .../flow/src/query/triples/falkordb.ts | 374 +- .../src/retrieval/document-rag-service.ts | 44 +- .../flow/src/retrieval/document-rag.ts | 23 +- .../flow/src/retrieval/graph-rag-service.ts | 46 +- ts/packages/flow/src/retrieval/graph-rag.ts | 33 +- .../embeddings/graph-embeddings-service.ts | 46 +- .../flow/src/storage/embeddings/qdrant-doc.ts | 66 +- .../src/storage/embeddings/qdrant-graph.ts | 68 +- .../src/storage/triples/falkordb-service.ts | 42 +- .../flow/src/storage/triples/falkordb.ts | 131 +- ts/packages/workbench/package.json | 29 +- .../src/components/error-boundary.tsx | 90 +- ts/packages/workbench/src/qa/mock-api.ts | 81 +- ts/scripts/inventory-native-classes.ts | 170 + 79 files changed, 7622 insertions(+), 6703 deletions(-) create mode 100644 ts/CLASS_EFFECT_GOAL.md create mode 100644 ts/scripts/inventory-native-classes.ts diff --git a/ts/CLASS_EFFECT_GOAL.md b/ts/CLASS_EFFECT_GOAL.md new file mode 100644 index 00000000..1505725c --- /dev/null +++ b/ts/CLASS_EFFECT_GOAL.md @@ -0,0 +1,151 @@ +# TrustGraph TS Native Class To Functional Effect Services Goal + +Use the `grill-me` skill from: + +`/home/elpresidank/YeeBois/projects/beep-effect/.agents/skills/grill-me/SKILL.md` + +## Objective + +Refactor the TrustGraph TypeScript port so production runtime code no longer +uses native TypeScript or JavaScript class syntax except where an Effect module +API truly requires class syntax and no practical functional form exists. + +The preferred style is functional and pure: + +- Functions, factories, closures, and object-literal service implementations. +- Functions returning functions or service objects instead of constructors. +- `Context.Service` tags plus `Layer` builders for dependency provision. +- `Effect.gen`, `Effect.fn`, `Effect.acquireRelease`, `Effect.addFinalizer`, + and scoped Layers for lifecycle. +- Top-level `Effect.runPromise` only at CLI/bootstrap boundaries. + +Terminal condition: repeated inventory agents find zero blocking native classes +in production runtime source like these current smells: + +- `ts/packages/base/src/processor/async-processor.ts#L38-162` +- `ts/packages/base/src/processor/flow-processor.ts#L284-358` + +## Class Policy + +Blocking scope is production runtime source under `ts/packages/**/src`, +excluding `__tests__`, `*.test.ts`, and `*.spec.ts`. + +Tests and scripts must still be inventoried separately, but they do not block +completion unless they are required to verify migrated production behavior. + +Allowed class syntax is intentionally narrow and proof-based. Do not preserve +class syntax just because it is Effect-adjacent. For every candidate exemption, +prove that the specific Effect API requires class syntax and that a functional +alternative is not better. + +Candidate exemptions to investigate, not blindly preserve: + +- `S.Class`, `S.TaggedClass`, `S.TaggedErrorClass`, `S.ErrorClass` +- `Data.TaggedError` +- `Context.Service` +- Effect RPC / HTTP API class forms such as `Rpc.make` and `HttpApi.make` + +If a functional Effect equivalent exists, use the functional form. The target +style is functions returning functions or service objects, not class-shaped +Effect code. + +Blocking class forms include: + +- Abstract base classes and inheritance-based processors. +- Service classes extending `AsyncProcessor`, `FlowProcessor`, or `LlmService`. +- Client API wrapper classes. +- Backend adapter classes. +- Spec/resource classes. +- Query/store/engine classes. +- Metrics/parser/lifecycle classes. +- React class components such as error boundaries. + +## Required Workflow + +1. Read and follow `grill-me`. +2. Before asking the user questions, use parallel read-only sub-agents to + inventory class declarations. +3. Use AST-based inspection where possible; regex alone is not enough. +4. Split inventory agents by subsystem: + - Base processor/messaging/runtime/specs/backend. + - Flow service processors and LLM/embedding/RAG services. + - Query/store engines and adapters. + - Client socket/API wrappers. + - Workbench production source. +5. Consolidate findings into: + - Blocking production native classes. + - Effect-native class forms with proof. + - Functional replacements available for prior exemptions. + - Non-blocking test/script classes. +6. Use `grill-me` to resolve any remaining API or rollout tradeoffs. +7. Produce a decision-complete implementation plan. +8. After approval, implement in phases and rerun inventory until blocking count + is zero. + +## Refactor Direction + +Replace inheritance and mutable class instances with: + +- Plain TypeScript service interfaces. +- `Context.Service` tags only where needed for Effect dependency injection. +- `Layer.succeed`, `Layer.effect`, and `Layer.scoped` for construction. +- Closure-held state only where necessary, preferably scoped and finalized. +- Object-literal services instead of class instances. +- Factory functions like `makeXService`, `makeXLayer`, and `runX`. +- `Effect.tryPromise` for external Promise APIs. +- `Effect.fn` for reusable effectful service methods. + +Preserve behavior and in-repo call semantics, but do not preserve native class +constructors as public API. Replace exported classes with interfaces, service +tags, factory functions, or layer builders, and update all in-repo callers. + +## Verification Requirements + +Required gates: + +- `cd ts && bun run check:tsgo` +- `cd ts && bun run build` +- `cd ts && bun run test` +- `cd ts && bun run workbench:qa` +- Add or provide an inventory command/script that fails on blocking production + classes. +- Run the inventory command after every migration phase. +- Final inventory must report zero blocking production native classes. + +Live smoke if service surfaces changed: + +- `cd ts/deploy && docker compose -f docker-compose.yml -f docker-compose.dev.yml up -d --build` +- `cd ts && bun run seed` +- `cd ts && bun run seed:demo` +- `cd ts && bun run seed:flows` +- `cd ts && SKIP_LLM=1 bun run test:pipeline` + +## New Session Prompt + +Paste this into a new Codex session: + +```text +Use the grill-me skill from /home/elpresidank/YeeBois/projects/beep-effect/.agents/skills/grill-me/SKILL.md. + +The repo is /home/elpresidank/YeeBois/dev/trustgraph and the TypeScript port is in ./ts. I want to remove native TypeScript/JavaScript class syntax from production runtime code and convert the port to functional, pure Effect services. + +Before asking me questions, use parallel read-only sub-agents to inventory every class declaration under ts/packages/**/src. Production runtime source is the blocking scope; tests and scripts should be inventoried separately but are non-blocking. Use AST-based inspection where possible, not regex alone. + +Allowed class syntax is extremely narrow and proof-based. Even Effect-native-looking classes such as S.Class, S.TaggedClass, S.TaggedErrorClass, S.ErrorClass, Data.TaggedError, Context.Service, Rpc.make, and HttpApi.make must be treated as candidate exemptions, not automatic exemptions. If a functional Effect equivalent exists, use the functional form. I prefer functions, factories, closures, object-literal service implementations, Context.Service tags, and Layers over all class-shaped code. + +Use these files as representative smells: +- ts/packages/base/src/processor/async-processor.ts#L38-162 +- ts/packages/base/src/processor/flow-processor.ts#L284-358 + +Inventory by subsystem: base processor/messaging/runtime/specs/backend; flow service processors and LLM/embedding/RAG services; query/store engines and adapters; client socket/API wrappers; workbench production source. + +After inventory, use grill-me to resolve any remaining API or rollout tradeoffs, then produce a decision-complete implementation plan. After I approve the plan, implement in phases and repeat inventory until zero blocking production native classes remain. + +Required verification: cd ts && bun run check:tsgo; cd ts && bun run build; cd ts && bun run test; cd ts && bun run workbench:qa; add or provide an inventory command/script that fails on blocking production classes. If service surfaces changed, also run the live docker/seed/pipeline smoke from this file. +``` + +## Goal Command + +```text +/goal Use ts/CLASS_EFFECT_GOAL.md to run the TrustGraph TS native-class-to-functional-Effect-services migration. First use grill-me and parallel read-only sub-agents to inventory every native class declaration. Then produce and execute a decision-complete refactor plan that replaces classes with pure functions, factories, closures, object-literal services, Context.Service tags, and Layers. Terminal condition: repeated inventory agents and the repo inventory command report zero blocking native classes in production runtime source under ts/packages/**/src. Even Effect-native class forms such as S.Class, TaggedError classes, Context.Service, Rpc.make, and HttpApi.make must be treated as narrow proof-based exemptions; if a functional Effect equivalent exists, use it. Verify with check:tsgo, build, test, workbench:qa, and live smoke when service surfaces change. Preserve unrelated local files. +``` diff --git a/ts/bun.lock b/ts/bun.lock index 4a78932a..e67f8654 100644 --- a/ts/bun.lock +++ b/ts/bun.lock @@ -179,6 +179,7 @@ "lucide-react": "^0.513.0", "react": "^19.1.0", "react-dom": "^19.1.0", + "react-error-boundary": "^6.1.2", "react-force-graph-2d": "^1.29.1", "react-markdown": "^10.1.0", "react-router": "^7.6.0", @@ -1139,6 +1140,8 @@ "react-dom": ["react-dom@19.2.4", "", { "dependencies": { "scheduler": "0.27.0" }, "peerDependencies": { "react": "19.2.4" } }, "sha512-AXJdLo8kgMbimY95O2aKQqsz2iWi9jMgKJhRBAxECE4IFxfcazB2LmzloIoibJI3C12IlY20+KFaLv+71bUJeQ=="], + "react-error-boundary": ["react-error-boundary@6.1.2", "", { "peerDependencies": { "react": "^18.0.0 || ^19.0.0" } }, "sha512-3DpCr5HVdZ0caUjYE/kIHBEJN0mNP3ZCgf16c48uJ5TbWjorKVp+YG8W3XqlJ7vJAVNw6wNIImyPXmFydwmyng=="], + "react-force-graph-2d": ["react-force-graph-2d@1.29.1", "", { "dependencies": { "force-graph": "1.51.2", "prop-types": "15.8.1", "react-kapsule": "2.5.7" }, "peerDependencies": { "react": "19.2.4" } }, "sha512-1Rl/1Z3xy2iTHKj6a0jRXGyiI86xUti81K+jBQZ+Oe46csaMikp47L5AjrzA9hY9fNGD63X8ffrqnvaORukCuQ=="], "react-is": ["react-is@16.13.1", "", {}, "sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ=="], diff --git a/ts/package.json b/ts/package.json index d84d8552..2e5d5c6b 100644 --- a/ts/package.json +++ b/ts/package.json @@ -7,6 +7,7 @@ "lint": "bunx --bun turbo lint", "test": "bunx --bun turbo test", "check:tsgo": "tsgo -b tsconfig.json", + "inventory:classes": "bun scripts/inventory-native-classes.ts", "workbench:qa": "bun run --cwd packages/workbench qa:browser", "prepare": "effect-tsgo patch", "clean": "turbo clean", diff --git a/ts/packages/base/src/__tests__/consumer.test.ts b/ts/packages/base/src/__tests__/consumer.test.ts index 6c650ba3..e8270d41 100644 --- a/ts/packages/base/src/__tests__/consumer.test.ts +++ b/ts/packages/base/src/__tests__/consumer.test.ts @@ -1,5 +1,5 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { Consumer, type ConsumerOptions, type FlowContext } from "../messaging/consumer.js"; +import { makeConsumer, type ConsumerOptions, type FlowContext } from "../messaging/consumer.js"; import type { PubSubBackend, BackendConsumer, @@ -75,20 +75,21 @@ describe("Consumer", () => { // ── Constructor ────────────────────────────────────────────────── it("stores options and applies defaults", () => { const handler = vi.fn(); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "my-topic", subscription: "my-sub", handler, }); - // Access private fields via any-cast to verify defaults - expect((consumer as any).concurrency).toBe(1); - expect((consumer as any).rateLimitRetryMs).toBe(10_000); + expect(consumer).toMatchObject({ + start: expect.any(Function), + stop: expect.any(Function), + }); }); it("accepts custom concurrency and rateLimitRetryMs", () => { - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "t", subscription: "s", @@ -97,8 +98,10 @@ describe("Consumer", () => { rateLimitRetryMs: 5_000, }); - expect((consumer as any).concurrency).toBe(4); - expect((consumer as any).rateLimitRetryMs).toBe(5_000); + expect(consumer).toMatchObject({ + start: expect.any(Function), + stop: expect.any(Function), + }); }); // ── start() creates consumer and calls handler ───────────────── @@ -116,7 +119,7 @@ describe("Consumer", () => { return null; }); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "topic-a", subscription: "sub-a", @@ -147,7 +150,7 @@ describe("Consumer", () => { return null; }); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "t", subscription: "s", @@ -174,7 +177,7 @@ describe("Consumer", () => { return null; }); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "t", subscription: "s", @@ -217,7 +220,7 @@ describe("Consumer", () => { return null; }); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "t", subscription: "s", @@ -249,7 +252,7 @@ describe("Consumer", () => { return null; }); - const consumer = new Consumer({ + const consumer = makeConsumer({ pubsub, topic: "t", subscription: "s", @@ -268,6 +271,6 @@ describe("Consumer", () => { await startPromise; expect(backendConsumer.close).toHaveBeenCalled(); - expect((consumer as any).running).toBe(false); + await expect(consumer.stop()).resolves.toBeUndefined(); }); }); diff --git a/ts/packages/base/src/__tests__/flow-processor-runtime.test.ts b/ts/packages/base/src/__tests__/flow-processor-runtime.test.ts index e8ee0426..6c4d2f39 100644 --- a/ts/packages/base/src/__tests__/flow-processor-runtime.test.ts +++ b/ts/packages/base/src/__tests__/flow-processor-runtime.test.ts @@ -3,7 +3,7 @@ import { ConfigProvider, Effect, Fiber } from "effect"; import { FlowProcessor, MessagingRuntimeLive, - ProducerSpec, + makeProducerSpec, PubSub, runFlowProcessorDefinitionScoped, runProcessorScoped, @@ -146,7 +146,7 @@ class TestFlowProcessor extends FlowProcessor { private readonly events: Array, ) { super(config); - this.registerSpecification(new ProducerSpec("output")); + this.registerSpecification(makeProducerSpec("output")); this.registerConfigHandler(async (_config, version) => { this.events.push(`handler:${version}`); }); @@ -225,7 +225,7 @@ describe("Effect-native FlowProcessor runtime", () => { const fiber = yield* runFlowProcessorDefinitionScoped({ id: "functional-flow-processor-test", pubsub: backend, - specifications: [new ProducerSpec("output")], + specifications: [makeProducerSpec("output")], configHandlers: [ (_config, version) => Effect.sync(() => { events.push(`handler:${version}`); diff --git a/ts/packages/base/src/__tests__/flow-spec-runtime.test.ts b/ts/packages/base/src/__tests__/flow-spec-runtime.test.ts index 5df40a30..abe3723f 100644 --- a/ts/packages/base/src/__tests__/flow-spec-runtime.test.ts +++ b/ts/packages/base/src/__tests__/flow-spec-runtime.test.ts @@ -2,13 +2,14 @@ import { describe, expect, it } from "@effect/vitest"; import { ConfigProvider, Duration, Effect, Fiber } from "effect"; import * as TestClock from "effect/testing/TestClock"; import { - ConsumerSpec, + makeConsumerSpec, + makeConsumerSpecFromPromise, Flow, MessagingRuntimeLive, - ParameterSpec, - ProducerSpec, + makeParameterSpec, + makeProducerSpec, PubSub, - RequestResponseSpec, + makeRequestResponseSpec, type BackendConsumer, type BackendProducer, type CreateConsumerOptions, @@ -156,7 +157,7 @@ describe("Effect-native flow specifications", () => { "processor", backend, { topics: { output: "actual-output" } }, - [new ProducerSpec("output")], + [makeProducerSpec("output")], ); yield* Effect.scoped( @@ -179,7 +180,7 @@ describe("Effect-native flow specifications", () => { ); it.effect( - "runs Promise handlers through the explicit ConsumerSpec compatibility helper", + "runs Promise handlers through the explicit makeConsumerSpec compatibility helper", Effect.fnUntraced(function* () { const message = createMessage("payload", { id: "request-1" }); const consumer = new ScriptedConsumer([message]); @@ -191,7 +192,7 @@ describe("Effect-native flow specifications", () => { backend, {}, [ - ConsumerSpec.fromPromise( + makeConsumerSpecFromPromise( "input", async (value, properties, flowContext: FlowContext) => { handled.push(`${flowContext.name}:${properties.id}:${value}`); @@ -237,7 +238,7 @@ describe("Effect-native flow specifications", () => { response: "actual-response", }, }, - [new RequestResponseSpec("rr", "request", "response")], + [makeRequestResponseSpec("rr", "request", "response")], ); const response = yield* Effect.scoped( @@ -270,7 +271,7 @@ describe("Effect-native flow specifications", () => { "processor", backend, { parameters: { present: 42 } }, - [new ParameterSpec("present")], + [makeParameterSpec("present")], ); const errors = yield* Effect.scoped( diff --git a/ts/packages/base/src/__tests__/messaging-runtime.test.ts b/ts/packages/base/src/__tests__/messaging-runtime.test.ts index 81570fd8..baf02f85 100644 --- a/ts/packages/base/src/__tests__/messaging-runtime.test.ts +++ b/ts/packages/base/src/__tests__/messaging-runtime.test.ts @@ -6,7 +6,7 @@ import { defaultMessagingRuntimeConfig, makeEffectRequestResponseFromPubSub, MessagingRuntimeLive, - ProducerSpec, + makeProducerSpec, runEffectConsumerScoped, runEffectProducerScoped, runFlowScoped, @@ -260,7 +260,7 @@ describe("Effect-native messaging runtime", () => { "processor", backend, {}, - [new ProducerSpec("flow-output")], + [makeProducerSpec("flow-output")], ); yield* Effect.scoped( diff --git a/ts/packages/base/src/__tests__/runtime-services.test.ts b/ts/packages/base/src/__tests__/runtime-services.test.ts index 86aadba8..49e6aaf0 100644 --- a/ts/packages/base/src/__tests__/runtime-services.test.ts +++ b/ts/packages/base/src/__tests__/runtime-services.test.ts @@ -1,8 +1,8 @@ import { describe, expect, it } from "@effect/vitest"; import { Effect } from "effect"; import { - AsyncProcessor, PubSub, + makeAsyncProcessor, runProcessorScoped, type BackendConsumer, type BackendProducer, @@ -79,53 +79,49 @@ class FailingProducerBackend extends FakePubSubBackend { } } -class RecordingProcessor extends AsyncProcessor { - constructor( - config: ProcessorConfig, - private readonly events: Array, - ) { - super(config); - } +const makeRecordingProcessor = ( + config: ProcessorConfig, + events: Array, +) => { + const processor = makeAsyncProcessor(config, { + run: async (runtime) => { + events.push(`run:${runtime.config.manageProcessSignals === false ? "effect-signals" : "class-signals"}`); + }, + }); + const stop = processor.stop; + processor.stop = async () => { + events.push("stop"); + await stop(); + }; + return processor; +}; - protected async run(): Promise { - this.events.push(`run:${this.config.manageProcessSignals === false ? "effect-signals" : "class-signals"}`); - } +const makeFailingProcessor = (config: ProcessorConfig) => + makeAsyncProcessor(config, { + run: async () => { + throw new Error("processor failed"); + }, + }); - override async stop(): Promise { - this.events.push("stop"); - await super.stop(); - } -} - -class FailingProcessor extends AsyncProcessor { - protected async run(): Promise { - throw new Error("processor failed"); - } -} - -class NativeRecordingProcessor extends AsyncProcessor { - constructor( - config: ProcessorConfig, - private readonly events: Array, - ) { - super(config); - } - - protected override runEffect() { - const events = this.events; - const config = this.config; - return Effect.gen(function* () { - const pubsub = yield* PubSub; - events.push(`native:${config.manageProcessSignals === false ? "effect-signals" : "class-signals"}`); - events.push(`pubsub:${pubsub.backend.constructor.name}`); - }); - } - - override stopEffect() { - this.events.push("native-stop"); - return super.stopEffect(); - } -} +const makeNativeRecordingProcessor = ( + config: ProcessorConfig, + events: Array, +) => { + const processor = makeAsyncProcessor(config, { + runEffect: (runtime) => + Effect.gen(function* () { + const pubsub = yield* PubSub; + events.push(`native:${runtime.config.manageProcessSignals === false ? "effect-signals" : "class-signals"}`); + events.push(`pubsub:${pubsub.backend.constructor.name}`); + }), + }); + const stopEffect = processor.stopEffect; + processor.stopEffect = () => { + events.push("native-stop"); + return stopEffect(); + }; + return processor; +}; describe("Effect runtime services", () => { it.effect( @@ -180,7 +176,7 @@ describe("Effect runtime services", () => { metricsPort: 8000, manageProcessSignals: true, }, - (config) => new RecordingProcessor(config, events), + (config) => makeRecordingProcessor(config, events), ).pipe(Effect.provide(PubSub.layer(backend))), ); @@ -203,7 +199,7 @@ describe("Effect runtime services", () => { metricsPort: 8000, manageProcessSignals: true, }, - (config) => new NativeRecordingProcessor(config, events), + (config) => makeNativeRecordingProcessor(config, events), ).pipe(Effect.provide(PubSub.layer(backend))), ); @@ -224,7 +220,7 @@ describe("Effect runtime services", () => { metricsPort: 8000, manageProcessSignals: true, }, - (config) => new FailingProcessor(config), + makeFailingProcessor, ).pipe( Effect.provide(PubSub.layer(backend)), Effect.flip, diff --git a/ts/packages/base/src/backend/index.ts b/ts/packages/base/src/backend/index.ts index 8ce6264b..0e664f77 100644 --- a/ts/packages/base/src/backend/index.ts +++ b/ts/packages/base/src/backend/index.ts @@ -9,7 +9,7 @@ export type { InitialPosition, } from "./types.js"; -export { NatsBackend } from "./nats.js"; +export { makeNatsBackend } from "./nats.js"; export { PubSub, NatsPubSubLive, diff --git a/ts/packages/base/src/backend/nats.ts b/ts/packages/base/src/backend/nats.ts index e1d060e0..ef85fec3 100644 --- a/ts/packages/base/src/backend/nats.ts +++ b/ts/packages/base/src/backend/nats.ts @@ -32,239 +32,207 @@ import type { const sc = StringCodec(); -class NatsMessage implements Message { +interface NatsMessage extends Message { /** Exposed so acknowledge/negativeAcknowledge can access the raw JsMsg */ readonly _jsMsg: JsMsg; - private readonly decoded: T; +} - constructor(msg: JsMsg, decoded: T) { - this._jsMsg = msg; - this.decoded = decoded; - } - - value(): T { - return this.decoded; - } - - properties(): Record { - const headers = this._jsMsg.headers; - const props: Record = {}; - if (headers !== undefined) { - for (const [key, values] of headers) { - const value = values[0]; - if (value !== undefined) { - props[key] = value; +function makeNatsMessage(msg: JsMsg, decoded: T): NatsMessage { + return { + _jsMsg: msg, + value: () => decoded, + properties: () => { + const headers = msg.headers; + const props: Record = {}; + if (headers !== undefined) { + for (const [key, values] of headers) { + const value = values[0]; + if (value !== undefined) { + props[key] = value; + } } } - } - return props; - } + return props; + }, + }; } -class NatsProducer implements BackendProducer { - private readonly js: JetStreamClient; - private readonly subject: string; - private readonly schema: S.Top | undefined; +function makeNatsProducer( + js: JetStreamClient, + subject: string, + schema?: S.Top, +): BackendProducer { + return { + send: async (message, properties) => { + const encoded = schema !== undefined + ? S.encodeUnknownSync(schema as S.Codec)(message) + : message; + const data = sc.encode(JSON.stringify(encoded)); + const opts: Record = {}; - constructor(js: JetStreamClient, subject: string, schema?: S.Top) { - this.js = js; - this.subject = subject; - this.schema = schema; - } - - async send(message: T, properties?: Record): Promise { - const encoded = this.schema !== undefined - ? S.encodeUnknownSync(this.schema as S.Codec)(message) - : message; - const data = sc.encode(JSON.stringify(encoded)); - const opts: Record = {}; - - if (properties !== undefined && Object.keys(properties).length > 0) { - const { headers } = await import("nats"); - const hdrs = headers(); - for (const [key, val] of Object.entries(properties)) { - hdrs.append(key, val); + if (properties !== undefined && Object.keys(properties).length > 0) { + const { headers } = await import("nats"); + const hdrs = headers(); + for (const [key, val] of Object.entries(properties)) { + hdrs.append(key, val); + } + opts.headers = hdrs; } - opts.headers = hdrs; - } - await this.js.publish(this.subject, data, opts); - } - - async flush(): Promise { - // NATS publishes are flushed on the connection level - } - - async close(): Promise { - // No per-producer cleanup needed for NATS - } + await js.publish(subject, data, opts); + }, + flush: async () => { + // NATS publishes are flushed on the connection level. + }, + close: async () => { + // No per-producer cleanup needed for NATS. + }, + }; } -class NatsConsumer implements BackendConsumer { - private consumer: NatsJsConsumer | null = null; - private readonly js: JetStreamClient; - private readonly jsm: JetStreamManager; - private readonly subject: string; - private readonly subscription: string; - private readonly initialPosition: "latest" | "earliest"; - private readonly streamName: string; - private readonly schema: S.Top | undefined; - - constructor( - js: JetStreamClient, - jsm: JetStreamManager, - subject: string, - subscription: string, - initialPosition: "latest" | "earliest", - streamName: string, - schema?: S.Top, - ) { - this.js = js; - this.jsm = jsm; - this.subject = subject; - this.subscription = subscription; - this.initialPosition = initialPosition; - this.streamName = streamName; - this.schema = schema; - } - - async init(): Promise { - // Stream is already ensured by NatsBackend.ensureStream(). - // Create or bind to durable consumer. - try { - this.consumer = await this.js.consumers.get(this.streamName, this.subscription); - } catch { - const deliverPolicy = - this.initialPosition === "earliest" - ? DeliverPolicy.All - : DeliverPolicy.New; - - await this.jsm.consumers.add(this.streamName, { - durable_name: this.subscription, - ack_policy: AckPolicy.Explicit, - deliver_policy: deliverPolicy, - filter_subject: this.subject, - }); - - this.consumer = await this.js.consumers.get(this.streamName, this.subscription); - } - } - - async receive(timeoutMs = 2000): Promise | null> { - if (this.consumer === null) throw new Error("Consumer not initialized"); - - // Pull a single message with a timeout using the pull-based API. - // consumer.next() returns a JsMsg or null when the timeout expires. - const msg = await this.consumer.next({ expires: timeoutMs }); - if (msg === null) return null; - - const parsed = JSON.parse(sc.decode(msg.data)); - const decoded = this.schema !== undefined - ? S.decodeUnknownSync(this.schema as S.Codec)(parsed) as T - : parsed as T; - return new NatsMessage(msg, decoded); - } - - async acknowledge(message: Message): Promise { - const natsMsg = message as NatsMessage; - natsMsg._jsMsg.ack(); - } - - async negativeAcknowledge(message: Message): Promise { - const natsMsg = message as NatsMessage; - natsMsg._jsMsg.nak(); - } - - async unsubscribe(): Promise { - // The pull-based consumer does not have a persistent subscription to drain. - // Clearing the reference is sufficient; the durable consumer persists server-side. - this.consumer = null; - } - - async close(): Promise { - this.consumer = null; - } +interface InitializableBackendConsumer extends BackendConsumer { + readonly init: () => Promise; } -export class NatsBackend implements PubSubBackend { - private connection: NatsConnection | null = null; - private js: JetStreamClient | null = null; - private jsm: JetStreamManager | null = null; - private initializedStreams = new Set(); - private readonly url: string; +function makeNatsConsumer( + js: JetStreamClient, + jsm: JetStreamManager, + subject: string, + subscription: string, + initialPosition: "latest" | "earliest", + streamName: string, + schema?: S.Top, +): InitializableBackendConsumer { + let consumer: NatsJsConsumer | null = null; - constructor(url = "nats://localhost:4222") { - this.url = url; - } + return { + init: async () => { + // Stream is already ensured by makeNatsBackend(). Create or bind to a durable consumer. + try { + consumer = await js.consumers.get(streamName, subscription); + } catch { + const deliverPolicy = + initialPosition === "earliest" + ? DeliverPolicy.All + : DeliverPolicy.New; - private async ensureConnected(): Promise { - if (this.connection === null) { - this.connection = await connect({ servers: this.url }); - this.js = this.connection.jetstream(); - this.jsm = await this.connection.jetstreamManager(); + await jsm.consumers.add(streamName, { + durable_name: subscription, + ack_policy: AckPolicy.Explicit, + deliver_policy: deliverPolicy, + filter_subject: subject, + }); + + consumer = await js.consumers.get(streamName, subscription); + } + }, + receive: async (timeoutMs = 2000) => { + if (consumer === null) throw new Error("Consumer not initialized"); + + // Pull a single message with a timeout using the pull-based API. + // consumer.next() returns a JsMsg or null when the timeout expires. + const msg = await consumer.next({ expires: timeoutMs }); + if (msg === null) return null; + + const parsed = JSON.parse(sc.decode(msg.data)); + const decoded = schema !== undefined + ? S.decodeUnknownSync(schema as S.Codec)(parsed) as T + : parsed as T; + return makeNatsMessage(msg, decoded); + }, + acknowledge: async (message) => { + const natsMsg = message as NatsMessage; + natsMsg._jsMsg.ack(); + }, + negativeAcknowledge: async (message) => { + const natsMsg = message as NatsMessage; + natsMsg._jsMsg.nak(); + }, + unsubscribe: async () => { + // The pull-based consumer does not have a persistent subscription to drain. + // Clearing the reference is sufficient; the durable consumer persists server-side. + consumer = null; + }, + close: async () => { + consumer = null; + }, + }; +} + +export function makeNatsBackend(url = "nats://localhost:4222"): PubSubBackend { + let connection: NatsConnection | null = null; + let js: JetStreamClient | null = null; + let jsm: JetStreamManager | null = null; + const initializedStreams = new Set(); + + const ensureConnected = async (): Promise => { + if (connection === null) { + connection = await connect({ servers: url }); + js = connection.jetstream(); + jsm = await connection.jetstreamManager(); } - } + }; /** * Ensure the stream for a given subject exists with a wildcard filter. * E.g. subject "tg.flow.config-request" → stream "tg_flow" with subjects ["tg.flow.>"] */ - private async ensureStream(subject: string): Promise { + const ensureStream = async (subject: string): Promise => { const parts = subject.split("."); const streamName = parts.slice(0, 2).join("_"); - if (this.initializedStreams.has(streamName)) return streamName; + if (initializedStreams.has(streamName)) return streamName; const wildcardSubject = `${parts.slice(0, 2).join(".")}.>`; - const jsm = this.jsm; - if (jsm === null) throw new Error("NATS backend not connected"); + const manager = jsm; + if (manager === null) throw new Error("NATS backend not connected"); try { - await jsm.streams.info(streamName); + await manager.streams.info(streamName); } catch { - await jsm.streams.add({ + await manager.streams.add({ name: streamName, subjects: [wildcardSubject], }); } - this.initializedStreams.add(streamName); + initializedStreams.add(streamName); return streamName; - } + }; - async createProducer(options: CreateProducerOptions): Promise> { - await this.ensureConnected(); - await this.ensureStream(options.topic); - const js = this.js; - if (js === null) throw new Error("NATS backend not connected"); - return new NatsProducer(js, options.topic, options.schema); - } - - async createConsumer(options: CreateConsumerOptions): Promise> { - await this.ensureConnected(); - const streamName = await this.ensureStream(options.topic); - const js = this.js; - const jsm = this.jsm; - if (js === null || jsm === null) throw new Error("NATS backend not connected"); - const consumer = new NatsConsumer( - js, - jsm, - options.topic, - options.subscription, - options.initialPosition ?? "latest", - streamName, - options.schema, - ); - await consumer.init(); - return consumer; - } - - async close(): Promise { - if (this.connection !== null) { - await this.connection.drain(); - this.connection = null; - this.js = null; - this.jsm = null; - } - } + return { + createProducer: async (options: CreateProducerOptions) => { + await ensureConnected(); + await ensureStream(options.topic); + const client = js; + if (client === null) throw new Error("NATS backend not connected"); + return makeNatsProducer(client, options.topic, options.schema); + }, + createConsumer: async (options: CreateConsumerOptions) => { + await ensureConnected(); + const streamName = await ensureStream(options.topic); + const client = js; + const manager = jsm; + if (client === null || manager === null) throw new Error("NATS backend not connected"); + const consumer = makeNatsConsumer( + client, + manager, + options.topic, + options.subscription, + options.initialPosition ?? "latest", + streamName, + options.schema, + ); + await consumer.init(); + return consumer; + }, + close: async () => { + if (connection !== null) { + await connection.drain(); + connection = null; + js = null; + jsm = null; + } + }, + }; } diff --git a/ts/packages/base/src/backend/pubsub.ts b/ts/packages/base/src/backend/pubsub.ts index 5be30c60..d24f9de4 100644 --- a/ts/packages/base/src/backend/pubsub.ts +++ b/ts/packages/base/src/backend/pubsub.ts @@ -14,7 +14,7 @@ import type { CreateProducerOptions, PubSubBackend, } from "./types.js"; -import { NatsBackend } from "./nats.js"; +import { makeNatsBackend } from "./nats.js"; import { pubSubError } from "../errors.js"; export interface PubSubService { @@ -78,14 +78,14 @@ export function pubSubLayer(backend: PubSubBackend): Layer.Layer { } export function makeNatsPubSubLayer(url = "nats://localhost:4222"): Layer.Layer { - return pubSubLayer(new NatsBackend(url)); + return pubSubLayer(makeNatsBackend(url)); } export const NatsPubSubLive = Layer.effect(PubSub)( Effect.gen(function* () { const natsUrl = O.getOrUndefined(yield* Config.string("NATS_URL").pipe(Config.option)); const pulsarHost = O.getOrUndefined(yield* Config.string("PULSAR_HOST").pipe(Config.option)); - const service = makePubSubService(new NatsBackend(natsUrl ?? pulsarHost ?? "nats://localhost:4222")); + const service = makePubSubService(makeNatsBackend(natsUrl ?? pulsarHost ?? "nats://localhost:4222")); yield* Effect.addFinalizer(() => service.close.pipe( Effect.catch((error) => diff --git a/ts/packages/base/src/errors.ts b/ts/packages/base/src/errors.ts index 024a544e..658be3de 100644 --- a/ts/packages/base/src/errors.ts +++ b/ts/packages/base/src/errors.ts @@ -5,7 +5,7 @@ */ import * as S from "effect/Schema"; -import type { TgError } from "./schema/primitives.js"; +import type { TgError } from "./schema/index.ts"; export class TooManyRequestsError extends S.TaggedErrorClass()( "TooManyRequestsError", diff --git a/ts/packages/base/src/messaging/consumer.ts b/ts/packages/base/src/messaging/consumer.ts index 3d2c4dfb..07e1e369 100644 --- a/ts/packages/base/src/messaging/consumer.ts +++ b/ts/packages/base/src/messaging/consumer.ts @@ -33,67 +33,55 @@ export interface ConsumerOptions { rateLimitTimeoutMs?: number; } -export class Consumer { - private backend: BackendConsumer | null = null; - private running = false; - private abortController = new AbortController(); - private readonly options: ConsumerOptions; +declare const ConsumerMessageType: unique symbol; - private readonly concurrency: number; - private readonly rateLimitRetryMs: number; +export interface Consumer { + readonly [ConsumerMessageType]?: (_: T) => T; + readonly start: (flow: FlowContext) => Promise; + readonly stop: () => Promise; +} - constructor(options: ConsumerOptions) { - this.options = options; - this.concurrency = options.concurrency ?? 1; - this.rateLimitRetryMs = options.rateLimitRetryMs ?? 10_000; - } +export function makeConsumer(options: ConsumerOptions): Consumer { + let backend: BackendConsumer | null = null; + let running = false; + let abortController = new AbortController(); + const concurrency = options.concurrency ?? 1; + const rateLimitRetryMs = options.rateLimitRetryMs ?? 10_000; - async start(flow: FlowContext): Promise { - this.backend = await this.options.pubsub.createConsumer({ - topic: this.options.topic, - subscription: this.options.subscription, - initialPosition: this.options.initialPosition ?? "latest", - }); - - this.running = true; - - // Spawn concurrent consumer tasks - const tasks = Array.from({ length: this.concurrency }, () => - this.consumeLoop(flow), - ); - // Run all concurrently — first rejection stops all - await Promise.all(tasks); - } - - async stop(): Promise { - this.running = false; - this.abortController.abort(); - if (this.backend !== null) { - await this.backend.close(); - this.backend = null; + const handleWithRetry = async (msg: Message, flow: FlowContext): Promise => { + try { + await options.handler(msg.value(), msg.properties(), flow); + } catch (err) { + if (S.is(TooManyRequestsError)(err)) { + console.warn(`[Consumer] Rate limited, retrying in ${rateLimitRetryMs}ms`); + await sleep(rateLimitRetryMs); + await options.handler(msg.value(), msg.properties(), flow); + } else { + throw err; + } } - } + }; - private async consumeLoop(flow: FlowContext): Promise { - while (this.running) { + const consumeLoop = async (flow: FlowContext): Promise => { + while (running) { let msg: Message | null = null; try { - const backend = this.backend; - if (backend === null) throw new Error("Consumer backend not started"); + const currentBackend = backend; + if (currentBackend === null) throw new Error("Consumer backend not started"); - msg = await backend.receive(2000); + msg = await currentBackend.receive(2000); if (msg === null) continue; - await this.handleWithRetry(msg, flow); - await backend.acknowledge(msg); + await handleWithRetry(msg, flow); + await currentBackend.acknowledge(msg); } catch (err) { - if (!this.running) break; + if (!running) break; console.error("[Consumer] Error in consume loop:", err); if (msg !== null) { try { - const backend = this.backend; - if (backend !== null) { - await backend.negativeAcknowledge(msg); + const currentBackend = backend; + if (currentBackend !== null) { + await currentBackend.negativeAcknowledge(msg); } } catch (nakErr) { console.error("[Consumer] Failed to nak message:", nakErr); @@ -102,21 +90,35 @@ export class Consumer { await sleep(1000); } } - } + }; - private async handleWithRetry(msg: Message, flow: FlowContext): Promise { - try { - await this.options.handler(msg.value(), msg.properties(), flow); - } catch (err) { - if (S.is(TooManyRequestsError)(err)) { - console.warn(`[Consumer] Rate limited, retrying in ${this.rateLimitRetryMs}ms`); - await sleep(this.rateLimitRetryMs); - await this.options.handler(msg.value(), msg.properties(), flow); - } else { - throw err; + return { + start: async (flow) => { + backend = await options.pubsub.createConsumer({ + topic: options.topic, + subscription: options.subscription, + initialPosition: options.initialPosition ?? "latest", + }); + + running = true; + + // Spawn concurrent consumer tasks. + const tasks = Array.from({ length: concurrency }, () => + consumeLoop(flow), + ); + // Run all concurrently: first rejection stops all. + await Promise.all(tasks); + }, + stop: async () => { + running = false; + abortController.abort(); + abortController = new AbortController(); + if (backend !== null) { + await backend.close(); + backend = null; } - } - } + }, + }; } function sleep(ms: number): Promise { diff --git a/ts/packages/base/src/messaging/index.ts b/ts/packages/base/src/messaging/index.ts index 81b47775..8cbef12e 100644 --- a/ts/packages/base/src/messaging/index.ts +++ b/ts/packages/base/src/messaging/index.ts @@ -1,7 +1,7 @@ -export { Producer } from "./producer.js"; -export { Consumer, type MessageHandler, type FlowContext, type ConsumerOptions } from "./consumer.js"; -export { Subscriber, AsyncQueue } from "./subscriber.js"; -export { RequestResponse, type RequestResponseOptions } from "./request-response.js"; +export { makeProducer, type Producer } from "./producer.js"; +export { makeConsumer, type Consumer, type MessageHandler, type FlowContext, type ConsumerOptions } from "./consumer.js"; +export { makeAsyncQueue, makeSubscriber, type Subscriber, type AsyncQueue } from "./subscriber.js"; +export { makeRequestResponse, type RequestResponse, type RequestResponseOptions } from "./request-response.js"; export { ConsumerFactory, ConsumerFactoryLive, diff --git a/ts/packages/base/src/messaging/producer.ts b/ts/packages/base/src/messaging/producer.ts index f12125f7..2c7da1a7 100644 --- a/ts/packages/base/src/messaging/producer.ts +++ b/ts/packages/base/src/messaging/producer.ts @@ -4,47 +4,47 @@ * Python reference: trustgraph-base/trustgraph/base/producer.py */ -import type { PubSubBackend, BackendProducer } from "../backend/types.js"; +import type { PubSubBackend } from "../backend/types.js"; import type { ProducerMetrics } from "../metrics/prometheus.js"; import { Effect } from "effect"; import { makeEffectProducerHandle, type EffectProducer } from "./runtime.js"; -export class Producer { - private backend: BackendProducer | null = null; - private effectProducer: EffectProducer | null = null; - private readonly pubsub: PubSubBackend; - private readonly topic: string; - private readonly metrics: ProducerMetrics | undefined; - - constructor(pubsub: PubSubBackend, topic: string, metrics?: ProducerMetrics) { - this.pubsub = pubsub; - this.topic = topic; - this.metrics = metrics; - } - - async start(): Promise { - this.backend = await this.pubsub.createProducer({ topic: this.topic }); - this.effectProducer = makeEffectProducerHandle(this.backend, { - topic: this.topic, - ...(this.metrics === undefined ? {} : { metrics: this.metrics }), - }); - } - - async send(id: string, message: T): Promise { - if (this.effectProducer === null) throw new Error("Producer not started"); - - await Effect.runPromise(this.effectProducer.send(id, message)); - } - - async stop(): Promise { - if (this.effectProducer !== null) { - await Effect.runPromise( - this.effectProducer.flush.pipe( - Effect.flatMap(() => this.effectProducer === null ? Effect.void : this.effectProducer.close), - ), - ); - this.effectProducer = null; - this.backend = null; - } - } +export interface Producer { + readonly start: () => Promise; + readonly send: (id: string, message: T) => Promise; + readonly stop: () => Promise; +} + +export function makeProducer( + pubsub: PubSubBackend, + topic: string, + metrics?: ProducerMetrics, +): Producer { + let effectProducer: EffectProducer | null = null; + + return { + start: async () => { + const backend = await pubsub.createProducer({ topic }); + effectProducer = makeEffectProducerHandle(backend, { + topic, + ...(metrics === undefined ? {} : { metrics }), + }); + }, + send: async (id, message) => { + if (effectProducer === null) throw new Error("Producer not started"); + + await Effect.runPromise(effectProducer.send(id, message)); + }, + stop: async () => { + if (effectProducer !== null) { + const producer = effectProducer; + await Effect.runPromise( + producer.flush.pipe( + Effect.flatMap(() => producer.close), + ), + ); + effectProducer = null; + } + }, + }; } diff --git a/ts/packages/base/src/messaging/request-response.ts b/ts/packages/base/src/messaging/request-response.ts index 410e337c..6f227f31 100644 --- a/ts/packages/base/src/messaging/request-response.ts +++ b/ts/packages/base/src/messaging/request-response.ts @@ -8,8 +8,8 @@ */ import { randomUUID } from "node:crypto"; -import { Producer } from "./producer.js"; -import { Subscriber } from "./subscriber.js"; +import { makeProducer, type Producer } from "./producer.js"; +import { makeSubscriber, type Subscriber } from "./subscriber.js"; import type { PubSubBackend } from "../backend/types.js"; export interface RequestResponseOptions { @@ -19,73 +19,76 @@ export interface RequestResponseOptions { subscription: string; } -export class RequestResponse { - private producer: Producer; - private subscriber: Subscriber; - - constructor(options: RequestResponseOptions) { - this.producer = new Producer(options.pubsub, options.requestTopic); - this.subscriber = new Subscriber( - options.pubsub, - options.responseTopic, - options.subscription, - ); - } - - async start(): Promise { - await this.producer.start(); - await this.subscriber.start(); - } - - async stop(): Promise { - await this.producer.stop(); - await this.subscriber.stop(); - } - - /** - * Send a request and wait for responses. - * - * @param request - The request payload - * @param options.timeoutMs - Total timeout in milliseconds (default: 300s) - * @param options.recipient - Optional callback for streaming responses. - * Return `true` to indicate the final response has been received. - * If omitted, returns the first response. - */ - async request( +export interface RequestResponse { + readonly start: () => Promise; + readonly stop: () => Promise; + readonly request: ( request: TReq, options?: { timeoutMs?: number; recipient?: (response: TRes) => Promise; }, - ): Promise { - const id = randomUUID(); - const timeoutMs = options?.timeoutMs ?? 300_000; - const recipient = options?.recipient; - - const queue = this.subscriber.subscribe(id); - - try { - await this.producer.send(id, request); - - const deadline = Date.now() + timeoutMs; - - while (true) { - const remaining = deadline - Date.now(); - if (remaining <= 0) { - throw new Error(`Request timed out after ${timeoutMs}ms`); - } - - const response = await queue.pop(remaining); - - if (recipient !== undefined) { - const isFinal = await recipient(response); - if (isFinal) return response; - } else { - return response; - } - } - } finally { - this.subscriber.unsubscribe(id); - } - } + ) => Promise; +} + +export function makeRequestResponse( + options: RequestResponseOptions, +): RequestResponse { + const producer: Producer = makeProducer(options.pubsub, options.requestTopic); + const subscriber: Subscriber = makeSubscriber( + options.pubsub, + options.responseTopic, + options.subscription, + ); + + return { + start: async () => { + await producer.start(); + await subscriber.start(); + }, + stop: async () => { + await producer.stop(); + await subscriber.stop(); + }, + /** + * Send a request and wait for responses. + * + * @param request - The request payload + * @param options.timeoutMs - Total timeout in milliseconds (default: 300s) + * @param options.recipient - Optional callback for streaming responses. + * Return `true` to indicate the final response has been received. + * If omitted, returns the first response. + */ + request: async (request, requestOptions) => { + const id = randomUUID(); + const timeoutMs = requestOptions?.timeoutMs ?? 300_000; + const recipient = requestOptions?.recipient; + + const queue = subscriber.subscribe(id); + + try { + await producer.send(id, request); + + const deadline = Date.now() + timeoutMs; + + while (true) { + const remaining = deadline - Date.now(); + if (remaining <= 0) { + throw new Error(`Request timed out after ${timeoutMs}ms`); + } + + const response = await queue.pop(remaining); + + if (recipient !== undefined) { + const isFinal = await recipient(response); + if (isFinal) return response; + } else { + return response; + } + } + } finally { + subscriber.unsubscribe(id); + } + }, + }; } diff --git a/ts/packages/base/src/messaging/subscriber.ts b/ts/packages/base/src/messaging/subscriber.ts index 801781e8..f15c4a5e 100644 --- a/ts/packages/base/src/messaging/subscriber.ts +++ b/ts/packages/base/src/messaging/subscriber.ts @@ -13,114 +13,84 @@ type Resolver = { /** * Simple async queue for inter-task communication (replaces asyncio.Queue). */ -export class AsyncQueue { - private buffer: T[] = []; - private waiters: Array<(value: T) => void> = []; - - push(item: T): void { - const waiter = this.waiters.shift(); - if (waiter !== undefined) { - waiter(item); - } else { - this.buffer.push(item); - } - } - - async pop(timeoutMs?: number): Promise { - const buffered = this.buffer.shift(); - if (buffered !== undefined) return buffered; - - return new Promise((resolve, reject) => { - let timer: ReturnType | undefined; - - const waiter = (value: T) => { - if (timer !== undefined) clearTimeout(timer); - resolve(value); - }; - - this.waiters.push(waiter); - - if (timeoutMs !== undefined) { - timer = setTimeout(() => { - const idx = this.waiters.indexOf(waiter); - if (idx !== -1) this.waiters.splice(idx, 1); - reject(new Error(`Queue.pop timed out after ${timeoutMs}ms`)); - }, timeoutMs); - } - }); - } - - get length(): number { - return this.buffer.length; - } +export interface AsyncQueue { + readonly push: (item: T) => void; + readonly pop: (timeoutMs?: number) => Promise; + readonly length: number; } -export class Subscriber { - private backend: BackendConsumer | null = null; - private running = false; - private readonly pubsub: PubSubBackend; - private readonly topic: string; - private readonly subscription: string; +export function makeAsyncQueue(): AsyncQueue { + const buffer: T[] = []; + const waiters: Array<(value: T) => void> = []; + + return { + push: (item) => { + const waiter = waiters.shift(); + if (waiter !== undefined) { + waiter(item); + } else { + buffer.push(item); + } + }, + pop: async (timeoutMs) => { + const buffered = buffer.shift(); + if (buffered !== undefined) return buffered; + + return new Promise((resolve, reject) => { + let timer: ReturnType | undefined; + + const waiter = (value: T) => { + if (timer !== undefined) clearTimeout(timer); + resolve(value); + }; + + waiters.push(waiter); + + if (timeoutMs !== undefined) { + timer = setTimeout(() => { + const idx = waiters.indexOf(waiter); + if (idx !== -1) waiters.splice(idx, 1); + reject(new Error(`Queue.pop timed out after ${timeoutMs}ms`)); + }, timeoutMs); + } + }); + }, + get length() { + return buffer.length; + }, + }; +} + +export interface Subscriber { + readonly start: () => Promise; + readonly stop: () => Promise; + readonly subscribe: (id: string) => AsyncQueue; + readonly subscribeAll: (id: string) => AsyncQueue; + readonly unsubscribe: (id: string) => void; + readonly unsubscribeAll: (id: string) => void; +} + +export function makeSubscriber( + pubsub: PubSubBackend, + topic: string, + subscription: string, +): Subscriber { + let backend: BackendConsumer | null = null; + let running = false; // ID-specific subscriptions (request/response correlation) - private idSubscribers = new Map>(); + const idSubscribers = new Map>(); // Wildcard subscribers (receive all messages) - private allSubscribers = new Map>(); + const allSubscribers = new Map>(); - constructor(pubsub: PubSubBackend, topic: string, subscription: string) { - this.pubsub = pubsub; - this.topic = topic; - this.subscription = subscription; - } - - async start(): Promise { - this.backend = await this.pubsub.createConsumer({ - topic: this.topic, - subscription: this.subscription, - }); - this.running = true; - // Start the dispatch loop (fire and forget — runs until stop) - this.dispatchLoop().catch((err) => { - if (this.running === true) console.error("[Subscriber] dispatch loop error:", err); - }); - } - - async stop(): Promise { - this.running = false; - if (this.backend !== null) { - await this.backend.close(); - this.backend = null; - } - } - - subscribe(id: string): AsyncQueue { - const queue = new AsyncQueue(); - this.idSubscribers.set(id, { queue }); - return queue; - } - - subscribeAll(id: string): AsyncQueue { - const queue = new AsyncQueue(); - this.allSubscribers.set(id, { queue }); - return queue; - } - - unsubscribe(id: string): void { - this.idSubscribers.delete(id); - } - - unsubscribeAll(id: string): void { - this.allSubscribers.delete(id); - } - - private async dispatchLoop(): Promise { + const dispatchLoop = async (): Promise => { let consecutiveErrors = 0; - while (this.running) { + while (running) { try { - const backend = this.backend; - if (backend === null) throw new Error("Subscriber backend not started"); + const currentBackend = backend; + if (currentBackend === null) throw new Error("Subscriber backend not started"); - const msg = await backend.receive(2000); + const msg = await currentBackend.receive(2000); if (msg === null) continue; consecutiveErrors = 0; @@ -131,20 +101,20 @@ export class Subscriber { // Route to ID-specific subscriber if (id !== undefined && id.length > 0) { - const sub = this.idSubscribers.get(id); + const sub = idSubscribers.get(id); if (sub !== undefined) { sub.queue.push(value); } } // Broadcast to all-subscribers - for (const sub of this.allSubscribers.values()) { + for (const sub of allSubscribers.values()) { sub.queue.push(value); } - await backend.acknowledge(msg); + await currentBackend.acknowledge(msg); } catch (err) { - if (!this.running) break; + if (!running) break; consecutiveErrors++; if (consecutiveErrors <= 3) { console.error("[Subscriber] Error:", err); @@ -156,5 +126,42 @@ export class Subscriber { await new Promise((r) => setTimeout(r, delay)); } } - } + }; + + return { + start: async () => { + backend = await pubsub.createConsumer({ + topic, + subscription, + }); + running = true; + // Start the dispatch loop (fire and forget; runs until stop). + dispatchLoop().catch((err) => { + if (running === true) console.error("[Subscriber] dispatch loop error:", err); + }); + }, + stop: async () => { + running = false; + if (backend !== null) { + await backend.close(); + backend = null; + } + }, + subscribe: (id) => { + const queue = makeAsyncQueue(); + idSubscribers.set(id, { queue }); + return queue; + }, + subscribeAll: (id) => { + const queue = makeAsyncQueue(); + allSubscribers.set(id, { queue }); + return queue; + }, + unsubscribe: (id) => { + idSubscribers.delete(id); + }, + unsubscribeAll: (id) => { + allSubscribers.delete(id); + }, + }; } diff --git a/ts/packages/base/src/metrics/index.ts b/ts/packages/base/src/metrics/index.ts index ccff4d2b..27011ef5 100644 --- a/ts/packages/base/src/metrics/index.ts +++ b/ts/packages/base/src/metrics/index.ts @@ -1 +1,7 @@ -export { ConsumerMetrics, ProducerMetrics, registry } from "./prometheus.js"; +export { + makeConsumerMetrics, + makeProducerMetrics, + registry, + type ConsumerMetrics, + type ProducerMetrics, +} from "./prometheus.js"; diff --git a/ts/packages/base/src/metrics/prometheus.ts b/ts/packages/base/src/metrics/prometheus.ts index f88b3ab8..7e00133f 100644 --- a/ts/packages/base/src/metrics/prometheus.ts +++ b/ts/packages/base/src/metrics/prometheus.ts @@ -9,64 +9,64 @@ import { Counter, Histogram, Registry, collectDefaultMetrics } from "prom-client export const registry = new Registry(); collectDefaultMetrics({ register: registry }); -export class ConsumerMetrics { - private requestHistogram: Histogram; - private processingCounter: Counter; - private rateLimitCounter: Counter; - private readonly labels: { processor: string; flow: string; name: string }; +export interface ConsumerMetrics { + readonly recordTime: (seconds: number) => void; + readonly process: (status: "success" | "error") => void; + readonly rateLimit: () => void; +} - constructor(processor: string, flow: string, name: string) { - this.labels = { processor, flow, name }; - this.requestHistogram = new Histogram({ +export function makeConsumerMetrics( + processor: string, + flow: string, + name: string, +): ConsumerMetrics { + const labels = { processor, flow, name }; + const requestHistogram = new Histogram({ name: "tg_consumer_request_duration_seconds", help: "Consumer request processing time", labelNames: ["processor", "flow", "name"], registers: [registry], - }); + }); - this.processingCounter = new Counter({ + const processingCounter = new Counter({ name: "tg_consumer_processing_total", help: "Consumer processing outcomes", labelNames: ["processor", "flow", "name", "status"], registers: [registry], - }); + }); - this.rateLimitCounter = new Counter({ + const rateLimitCounter = new Counter({ name: "tg_consumer_rate_limit_total", help: "Consumer rate limit events", labelNames: ["processor", "flow", "name"], registers: [registry], - }); - } + }); - recordTime(seconds: number): void { - this.requestHistogram.observe(this.labels, seconds); - } - - process(status: "success" | "error"): void { - this.processingCounter.inc({ ...this.labels, status }); - } - - rateLimit(): void { - this.rateLimitCounter.inc(this.labels); - } + return { + recordTime: (seconds) => requestHistogram.observe(labels, seconds), + process: (status) => processingCounter.inc({ ...labels, status }), + rateLimit: () => rateLimitCounter.inc(labels), + }; } -export class ProducerMetrics { - private counter: Counter; - private readonly labels: { processor: string; flow: string; name: string }; +export interface ProducerMetrics { + readonly inc: () => void; +} - constructor(processor: string, flow: string, name: string) { - this.labels = { processor, flow, name }; - this.counter = new Counter({ +export function makeProducerMetrics( + processor: string, + flow: string, + name: string, +): ProducerMetrics { + const labels = { processor, flow, name }; + const counter = new Counter({ name: "tg_producer_items_total", help: "Producer items sent", labelNames: ["processor", "flow", "name"], registers: [registry], - }); - } + }); - inc(): void { - this.counter.inc(this.labels); - } + return { + inc: () => counter.inc(labels), + }; } diff --git a/ts/packages/base/src/processor/async-processor.ts b/ts/packages/base/src/processor/async-processor.ts index 1d6ab205..0b7adfb1 100644 --- a/ts/packages/base/src/processor/async-processor.ts +++ b/ts/packages/base/src/processor/async-processor.ts @@ -7,7 +7,7 @@ */ import type { PubSubBackend } from "../backend/types.js"; -import { NatsBackend } from "../backend/nats.js"; +import { makeNatsBackend } from "../backend/nats.js"; import { Effect } from "effect"; import { processorLifecycleError, type ProcessorLifecycleError } from "../errors.js"; import { loadProcessorRuntimeConfig } from "../runtime/config.js"; @@ -30,105 +30,72 @@ export type EffectConfigHandler = ( version: number, ) => Effect.Effect; +declare const processorRunErrorType: unique symbol; +declare const processorRunRequirementsType: unique symbol; + +export interface ProcessorRuntime { + readonly [processorRunErrorType]?: RunError; + readonly [processorRunRequirementsType]?: RunRequirements; + readonly start: () => Promise; + readonly stop: () => Promise; + startEffect(): unknown; + stopEffect(): unknown; +} + +export interface AsyncProcessorRuntime< + RunError = ProcessorLifecycleError, + RunRequirements = never, +> extends ProcessorRuntime { + readonly config: ProcessorConfig; + readonly pubsub: PubSubBackend; + readonly configHandlers: ConfigHandler[]; + readonly running: boolean; + readonly isRunning: () => boolean; + readonly registerConfigHandler: (handler: ConfigHandler) => void; + readonly onShutdown: (callback: () => Promise) => void; + readonly run: () => Promise; + runEffect(): unknown; +} + +export interface AsyncProcessorRuntimeOptions< + RunError = ProcessorLifecycleError, + RunRequirements = never, +> { + readonly run?: ( + processor: AsyncProcessorRuntime, + ) => Promise; + readonly runEffect?: ( + processor: AsyncProcessorRuntime, + ) => Effect.Effect; +} + interface RegisteredSignalHandler { readonly signal: NodeJS.Signals; readonly handler: () => void; } -export abstract class AsyncProcessor { - protected pubsub: PubSubBackend; - protected running = false; - protected configHandlers: ConfigHandler[] = []; - private shutdownCallbacks: Array<() => Promise> = []; - private signalHandlers: RegisteredSignalHandler[] = []; - private readonly ownsPubSub: boolean; - protected readonly config: ProcessorConfig; +export function makeAsyncProcessor< + RunError = ProcessorLifecycleError, + RunRequirements = never, +>( + config: ProcessorConfig, + options: AsyncProcessorRuntimeOptions = {}, +): AsyncProcessorRuntime { + const pubsub = config.pubsub ?? makeNatsBackend(config.pubsubUrl ?? "nats://localhost:4222"); + const ownsPubSub = config.pubsub === undefined; + const configHandlers: ConfigHandler[] = []; + const shutdownCallbacks: Array<() => Promise> = []; + let running = false; + let signalHandlers: RegisteredSignalHandler[] = []; - constructor(config: ProcessorConfig) { - this.config = config; - this.pubsub = config.pubsub ?? new NatsBackend(config.pubsubUrl ?? "nats://localhost:4222"); - this.ownsPubSub = config.pubsub === undefined; - } - - registerConfigHandler(handler: ConfigHandler): void { - this.configHandlers.push(handler); - } - - async start(): Promise { - await Effect.runPromise( - this.startEffect() as Effect.Effect, - ); - } - - async stop(): Promise { - await Effect.runPromise(this.stopEffect()); - } - - protected onShutdown(callback: () => Promise): void { - this.shutdownCallbacks.push(callback); - } - - startEffect(): Effect.Effect { - const processor = this; - return Effect.gen(function* () { - yield* Effect.sync(() => { - processor.running = true; - processor.registerProcessSignalHandlers(); - }); - - yield* processor.runEffect(); - }).pipe( - Effect.withSpan("trustgraph.processor.start", { - attributes: { - "trustgraph.processor.id": processor.config.id, - }, - }), - ); - } - - stopEffect(): Effect.Effect { - const processor = this; - return Effect.gen(function* () { - yield* Effect.sync(() => { - processor.running = false; - processor.unregisterProcessSignalHandlers(); - }); - - for (const cb of processor.shutdownCallbacks) { - yield* Effect.tryPromise({ - try: () => cb(), - catch: (error) => processorLifecycleError(processor.config.id, "shutdown-callback", error), - }); - } - - if (processor.ownsPubSub) { - yield* Effect.tryPromise({ - try: () => processor.pubsub.close(), - catch: (error) => processorLifecycleError(processor.config.id, "close-pubsub", error), - }); - } - }); - } - - protected run(): Promise { - return Effect.runPromise(this.runEffect() as unknown as Effect.Effect); - } - - protected runEffect(): Effect.Effect { - return Effect.tryPromise({ - try: () => this.run(), - catch: (error) => processorLifecycleError(this.config.id, "start", error), - }) as unknown as Effect.Effect; - } - - private registerProcessSignalHandlers(): void { - if (this.config.manageProcessSignals === false || this.signalHandlers.length > 0) { + const registerProcessSignalHandlers = (): void => { + if (config.manageProcessSignals === false || signalHandlers.length > 0) { return; } const shutdown = () => { - console.log(`[${this.config.id}] Shutting down...`); - void this.stop().then(() => process.exit(0)); + console.log(`[${config.id}] Shutting down...`); + void processor.stop().then(() => process.exit(0)); }; const handlers: RegisteredSignalHandler[] = [ { signal: "SIGINT", handler: shutdown }, @@ -137,26 +104,128 @@ export abstract class AsyncProcessor { + for (const { signal, handler } of signalHandlers) { process.off(signal, handler); } - this.signalHandlers = []; - } + signalHandlers = []; + }; - /** - * Static launch helper — parses env/args and starts the processor. - * Subclasses call: `MyProcessor.launch("my-service")` - */ - static async launch>( + const processor: AsyncProcessorRuntime = { + config, + pubsub, + configHandlers, + get running() { + return running; + }, + isRunning: () => running, + registerConfigHandler: (handler) => { + configHandlers.push(handler); + }, + start: async () => { + await Effect.runPromise( + processor.startEffect() as Effect.Effect, + ); + }, + stop: async () => { + await Effect.runPromise( + processor.stopEffect() as Effect.Effect, + ); + }, + onShutdown: (callback) => { + shutdownCallbacks.push(callback); + }, + startEffect() { + const startProcessor = Effect.fn("trustgraph.processor.start")(function* () { + yield* Effect.sync(() => { + running = true; + registerProcessSignalHandlers(); + }); + + yield* ( + processor.runEffect() as Effect.Effect + ); + }); + return startProcessor().pipe( + Effect.withSpan("trustgraph.processor.start", { + attributes: { + "trustgraph.processor.id": config.id, + }, + }), + ); + }, + stopEffect() { + const stopProcessor = Effect.fn("trustgraph.processor.stop")(function* () { + yield* Effect.sync(() => { + running = false; + unregisterProcessSignalHandlers(); + }); + + for (const cb of shutdownCallbacks) { + yield* Effect.tryPromise({ + try: () => cb(), + catch: (error) => processorLifecycleError(config.id, "shutdown-callback", error), + }); + } + + if (ownsPubSub) { + yield* Effect.tryPromise({ + try: () => pubsub.close(), + catch: (error) => processorLifecycleError(config.id, "close-pubsub", error), + }); + } + }); + return stopProcessor(); + }, + run: () => + Effect.runPromise( + processor.runEffect() as unknown as Effect.Effect, + ), + runEffect: () => { + if (options.runEffect !== undefined) { + return options.runEffect(processor); + } + return Effect.tryPromise({ + try: () => options.run?.(processor) ?? Promise.resolve(), + catch: (error) => processorLifecycleError(config.id, "start", error), + }) as unknown as Effect.Effect; + }, + }; + + return processor; +} + +export type AsyncProcessor< + RunError = ProcessorLifecycleError, + RunRequirements = never, +> = AsyncProcessorRuntime; + +export const AsyncProcessor = Object.assign( + function AsyncProcessor(config: ProcessorConfig) { + return makeAsyncProcessor(config); + }, + { + async launch>( + this: new (config: ProcessorConfig) => T, + id: string, + ): Promise { + const config = await Effect.runPromise(loadProcessorRuntimeConfig(id)); + const processor = new this(config); + await processor.start(); + }, + }, +) as unknown as { + new ( + config: ProcessorConfig, + ): AsyncProcessor; + ( + config: ProcessorConfig, + ): AsyncProcessor; + launch>( this: new (config: ProcessorConfig) => T, id: string, - ): Promise { - const config = await Effect.runPromise(loadProcessorRuntimeConfig(id)); - const processor = new this(config); - await processor.start(); - } -} + ): Promise; +}; diff --git a/ts/packages/base/src/processor/flow-processor.ts b/ts/packages/base/src/processor/flow-processor.ts index 44a85e09..1a28103a 100644 --- a/ts/packages/base/src/processor/flow-processor.ts +++ b/ts/packages/base/src/processor/flow-processor.ts @@ -8,8 +8,11 @@ */ import { - AsyncProcessor, + makeAsyncProcessor, + type AsyncProcessorRuntime, + type ConfigHandler, type EffectConfigHandler, + type ProcessorRuntime, type ProcessorConfig, } from "./async-processor.js"; import type { Spec } from "../spec/types.js"; @@ -60,6 +63,44 @@ export interface FlowProcessorRuntimeOptions< readonly isRunning?: () => boolean; } +type FlowProcessorRuntimeRequirements = + | PubSub + | FlowRuntime + | ProducerFactory + | ConsumerFactory + | RequestResponseFactory + | Scope.Scope + | FlowRequirements; + +export type FlowProcessorStartEffect = Effect.Effect< + void, + PubSubError | FlowRuntimeError | ProcessorLifecycleError, + FlowProcessorRuntimeRequirements +>; + +export interface FlowProcessorRuntime + extends ProcessorRuntime< + PubSubError | FlowRuntimeError | ProcessorLifecycleError, + FlowProcessorRuntimeRequirements + > { + readonly config: ProcessorConfig; + readonly pubsub: PubSubBackend; + readonly configHandlers: ConfigHandler[]; + readonly isRunning: () => boolean; + readonly registerConfigHandler: (handler: ConfigHandler) => void; + readonly registerSpecification: ( + spec: Spec, + ) => void; + readonly specifications: ReadonlyArray>; +} + +export interface MakeFlowProcessorOptions { + readonly specifications?: ReadonlyArray>; + readonly provide?: ( + effect: FlowProcessorStartEffect, + ) => FlowProcessorStartEffect; +} + const ConfigPushSchema = S.Struct({ version: S.Number, config: S.Record(S.String, S.Unknown), @@ -281,78 +322,76 @@ export function runFlowProcessorDefinitionScoped< }); } -export abstract class FlowProcessor extends AsyncProcessor< - PubSubError | FlowRuntimeError | ProcessorLifecycleError, - | PubSub - | FlowRuntime - | ProducerFactory - | ConsumerFactory - | RequestResponseFactory - | Scope.Scope - | FlowRequirements -> { - private specifications: Array> = []; - - protected constructor(config: ProcessorConfig) { - super(config); - } - - registerSpecification( - spec: Spec, - ): void { - this.specifications.push(spec as Spec); - } - - override async start(): Promise { - const pubsub = makePubSubService(this.pubsub); - const messagingConfig = await Effect.runPromise(loadMessagingRuntimeConfig()); - const start = this.startEffect().pipe( - Effect.provideService(PubSub, pubsub), - Effect.provideService(ProducerFactory, ProducerFactory.of(makeProducerFactoryService(pubsub))), - Effect.provideService(ConsumerFactory, ConsumerFactory.of(makeConsumerFactoryService(pubsub, messagingConfig))), - Effect.provideService( - RequestResponseFactory, - RequestResponseFactory.of(makeRequestResponseFactoryService(pubsub, messagingConfig)), - ), - Effect.provideService(FlowRuntime, FlowRuntime.of({ run: runFlowRuntimeScoped })), - ) as Effect.Effect; - await Effect.runPromise( - Effect.scoped( - start, - ), - ); - } - - protected override runEffect(): Effect.Effect< - void, +export function makeFlowProcessor( + config: ProcessorConfig, + options: MakeFlowProcessorOptions = {}, +): FlowProcessorRuntime { + const specifications: Array> = [ + ...(options.specifications ?? []), + ]; + let processor: FlowProcessorRuntime; + const base: AsyncProcessorRuntime< PubSubError | FlowRuntimeError | ProcessorLifecycleError, - | PubSub - | FlowRuntime - | ProducerFactory - | ConsumerFactory - | RequestResponseFactory - | Scope.Scope - | FlowRequirements - > { - const processor = this; - const configHandlers = processor.configHandlers.map( - (handler): EffectConfigHandler => - (config, version) => - Effect.tryPromise({ - try: () => handler(config, version), - catch: (error) => pubSubError("config-handler", error), - }), - ); - return runFlowProcessorDefinitionScoped({ - id: processor.config.id, - pubsub: processor.pubsub, - specifications: processor.specifications, - configHandlers, - isRunning: () => processor.running, - }); - } + FlowProcessorRuntimeRequirements + > = makeAsyncProcessor(config, { + runEffect: (runtime) => { + const configHandlers = runtime.configHandlers.map( + (handler): EffectConfigHandler => + (pushedConfig, version) => + Effect.tryPromise({ + try: () => handler(pushedConfig, version), + catch: (error) => pubSubError("config-handler", error), + }), + ); + return runFlowProcessorDefinitionScoped({ + id: runtime.config.id, + pubsub: runtime.pubsub, + specifications, + configHandlers, + isRunning: runtime.isRunning, + }); + }, + }); - override stopEffect(): Effect.Effect { - return super.stopEffect(); - } + const startEffect = (): FlowProcessorStartEffect => { + const effect = base.startEffect() as FlowProcessorStartEffect; + return options.provide?.(effect) ?? effect; + }; + + processor = { + ...base, + specifications, + registerSpecification: (spec) => { + specifications.push(spec as Spec); + }, + startEffect, + start: async () => { + const pubsub = makePubSubService(base.pubsub); + const messagingConfig = await Effect.runPromise(loadMessagingRuntimeConfig()); + const start = startEffect().pipe( + Effect.provideService(PubSub, pubsub), + Effect.provideService(ProducerFactory, ProducerFactory.of(makeProducerFactoryService(pubsub))), + Effect.provideService(ConsumerFactory, ConsumerFactory.of(makeConsumerFactoryService(pubsub, messagingConfig))), + Effect.provideService( + RequestResponseFactory, + RequestResponseFactory.of(makeRequestResponseFactoryService(pubsub, messagingConfig)), + ), + Effect.provideService(FlowRuntime, FlowRuntime.of({ run: runFlowRuntimeScoped })), + ) as Effect.Effect; + await Effect.runPromise(Effect.scoped(start)); + }, + }; + + return processor; } + +export type FlowProcessor = FlowProcessorRuntime; + +export const FlowProcessor = makeFlowProcessor as unknown as { + new ( + config: ProcessorConfig, + ): FlowProcessor; + ( + config: ProcessorConfig, + ): FlowProcessor; +}; diff --git a/ts/packages/base/src/processor/flow.ts b/ts/packages/base/src/processor/flow.ts index b218a013..db11c93f 100644 --- a/ts/packages/base/src/processor/flow.ts +++ b/ts/packages/base/src/processor/flow.ts @@ -57,183 +57,30 @@ export interface FlowRequestor { readonly stop: () => Promise; } -export class Flow { - private producers = new Map>(); - private consumers = new Map(); - private requestors = new Map>(); - private parameters = new Map(); - private compatibilityScope: Scope.Closeable | null = null; - public readonly name: string; - public readonly processorId: string; - private readonly pubsub: PubSubBackend; - private readonly definition: FlowDefinition; - private readonly specifications: ReadonlyArray>; +export function makeFlow( + name: string, + processorId: string, + pubsub: PubSubBackend, + definition: FlowDefinition, + specifications: ReadonlyArray>, +) { + const producers = new Map>(); + const consumers = new Map(); + const requestors = new Map>(); + const parameters = new Map(); + let compatibilityScope: Scope.Closeable | null = null; - constructor( - name: string, - processorId: string, - pubsub: PubSubBackend, - definition: FlowDefinition, - specifications: ReadonlyArray>, - ) { - this.name = name; - this.processorId = processorId; - this.pubsub = pubsub; - this.definition = definition; - this.specifications = specifications; - } - - startEffect(): Effect.Effect { - const flow = this; - return Effect.gen(function* () { - for (const spec of flow.specifications) { - yield* spec.addEffect(flow, flow.definition); - } - }); - } - - async start(): Promise { - if (this.compatibilityScope !== null) { - await this.stop(); + const ensureCompatibilityScope = async (): Promise => { + if (compatibilityScope !== null) { + return compatibilityScope; } - await this.runInCompatibilityScope( - this.startEffect() as Effect.Effect, - this.pubsub, - ); - } + compatibilityScope = await Effect.runPromise(Scope.make()); + return compatibilityScope; + }; - async stop(): Promise { - const scope = this.compatibilityScope; - this.compatibilityScope = null; - if (scope !== null) { - await Effect.runPromise(Scope.close(scope, Exit.void)); - } - this.clearResources(); - } - - async runInCompatibilityScope( - effect: Effect.Effect, - pubsub: PubSubBackend, - ): Promise { - const scope = await this.ensureCompatibilityScope(); - const pubsubService = makePubSubService(pubsub); - const messagingConfig = await Effect.runPromise(loadMessagingRuntimeConfig()); - return await Effect.runPromise( - effect.pipe( - Effect.provideService(ProducerFactory, ProducerFactory.of(makeProducerFactoryService(pubsubService))), - Effect.provideService(ConsumerFactory, ConsumerFactory.of(makeConsumerFactoryService(pubsubService, messagingConfig))), - Effect.provideService( - RequestResponseFactory, - RequestResponseFactory.of(makeRequestResponseFactoryService(pubsubService, messagingConfig)), - ), - Scope.provide(scope), - ), - ); - } - - clearResources(): void { - this.producers.clear(); - this.consumers.clear(); - this.requestors.clear(); - this.parameters.clear(); - } - - registerProducer(name: string, producer: EffectProducer): void { - this.producers.set(name, producer); - } - - registerConsumer(name: string, consumer: EffectConsumer): void { - this.consumers.set(name, consumer); - } - - registerRequestor(name: string, rr: EffectRequestResponse): void { - this.requestors.set(name, rr); - } - - setParameter(name: string, value: unknown): void { - this.parameters.set(name, value); - } - - producerEffect(name: string): Effect.Effect, FlowResourceNotFoundError> { - const p = this.producers.get(name); - return p === undefined - ? Effect.fail(flowResourceNotFoundError(this.name, "producer", name)) - : Effect.succeed(p as EffectProducer); - } - - consumerEffect(name: string): Effect.Effect { - const c = this.consumers.get(name); - return c === undefined - ? Effect.fail(flowResourceNotFoundError(this.name, "consumer", name)) - : Effect.succeed(c); - } - - requestorEffect( - name: string, - ): Effect.Effect, FlowResourceNotFoundError> { - const rr = this.requestors.get(name); - return rr === undefined - ? Effect.fail(flowResourceNotFoundError(this.name, "requestor", name)) - : Effect.succeed(rr as EffectRequestResponse); - } - - parameterEffect(name: string): Effect.Effect { - const v = this.parameters.get(name); - return v === undefined - ? Effect.fail(flowResourceNotFoundError(this.name, "parameter", name)) - : Effect.succeed(v as T); - } - - producer(name: string): FlowProducer { - const p = this.producers.get(name); - if (p === undefined) throw flowResourceNotFoundError(this.name, "producer", name); - return { - send: (id, message) => Effect.runPromise((p as EffectProducer).send(id, message)), - flush: () => Effect.runPromise(p.flush), - stop: () => Effect.runPromise(p.flush.pipe(Effect.flatMap(() => p.close))), - }; - } - - consumer(name: string): FlowConsumer { - const c = this.consumers.get(name); - if (c === undefined) throw flowResourceNotFoundError(this.name, "consumer", name); - return { - stop: () => Effect.runPromise(c.stop), - }; - } - - requestor(name: string): FlowRequestor { - const rr = this.requestors.get(name); - if (rr === undefined) throw flowResourceNotFoundError(this.name, "requestor", name); - return { - request: (request, options) => - Effect.runPromise( - (rr as EffectRequestResponse).request( - request, - this.toEffectRequestOptions(options), - ), - ), - stop: () => Effect.runPromise(rr.stop), - }; - } - - parameter(name: string): T { - const v = this.parameters.get(name); - if (v === undefined) throw flowResourceNotFoundError(this.name, "parameter", name); - return v as T; - } - - private async ensureCompatibilityScope(): Promise { - if (this.compatibilityScope !== null) { - return this.compatibilityScope; - } - this.compatibilityScope = await Effect.runPromise(Scope.make()); - return this.compatibilityScope; - } - - private toEffectRequestOptions( + const toEffectRequestOptions = ( options: FlowRequestOptions | undefined, - ): EffectRequestOptions | undefined { + ): EffectRequestOptions | undefined => { if (options === undefined) { return undefined; } @@ -246,5 +93,153 @@ export class Flow { recipient: (response: TRes) => Effect.promise(() => recipient(response)), }), }; - } + }; + + const flow = { + name, + processorId, + startEffect(): Effect.Effect { + return Effect.gen(function* () { + for (const spec of specifications) { + yield* spec.addEffect(flow, definition); + } + }); + }, + async start(): Promise { + if (compatibilityScope !== null) { + await flow.stop(); + } + await flow.runInCompatibilityScope( + flow.startEffect() as Effect.Effect, + pubsub, + ); + }, + async stop(): Promise { + const scope = compatibilityScope; + compatibilityScope = null; + if (scope !== null) { + await Effect.runPromise(Scope.close(scope, Exit.void)); + } + flow.clearResources(); + }, + async runInCompatibilityScope( + effect: Effect.Effect, + runtimePubsub: PubSubBackend, + ): Promise { + const scope = await ensureCompatibilityScope(); + const pubsubService = makePubSubService(runtimePubsub); + const messagingConfig = await Effect.runPromise(loadMessagingRuntimeConfig()); + return await Effect.runPromise( + effect.pipe( + Effect.provideService(ProducerFactory, ProducerFactory.of(makeProducerFactoryService(pubsubService))), + Effect.provideService(ConsumerFactory, ConsumerFactory.of(makeConsumerFactoryService(pubsubService, messagingConfig))), + Effect.provideService( + RequestResponseFactory, + RequestResponseFactory.of(makeRequestResponseFactoryService(pubsubService, messagingConfig)), + ), + Scope.provide(scope), + ), + ); + }, + clearResources(): void { + producers.clear(); + consumers.clear(); + requestors.clear(); + parameters.clear(); + }, + registerProducer(registerName: string, producer: EffectProducer): void { + producers.set(registerName, producer); + }, + registerConsumer(registerName: string, consumer: EffectConsumer): void { + consumers.set(registerName, consumer); + }, + registerRequestor(registerName: string, rr: EffectRequestResponse): void { + requestors.set(registerName, rr); + }, + setParameter(parameterName: string, value: unknown): void { + parameters.set(parameterName, value); + }, + producerEffect(producerName: string): Effect.Effect, FlowResourceNotFoundError> { + const p = producers.get(producerName); + return p === undefined + ? Effect.fail(flowResourceNotFoundError(name, "producer", producerName)) + : Effect.succeed(p as EffectProducer); + }, + consumerEffect(consumerName: string): Effect.Effect { + const c = consumers.get(consumerName); + return c === undefined + ? Effect.fail(flowResourceNotFoundError(name, "consumer", consumerName)) + : Effect.succeed(c); + }, + requestorEffect( + requestorName: string, + ): Effect.Effect, FlowResourceNotFoundError> { + const rr = requestors.get(requestorName); + return rr === undefined + ? Effect.fail(flowResourceNotFoundError(name, "requestor", requestorName)) + : Effect.succeed(rr as EffectRequestResponse); + }, + parameterEffect(parameterName: string): Effect.Effect { + const v = parameters.get(parameterName); + return v === undefined + ? Effect.fail(flowResourceNotFoundError(name, "parameter", parameterName)) + : Effect.succeed(v as T); + }, + producer(producerName: string): FlowProducer { + const p = producers.get(producerName); + if (p === undefined) throw flowResourceNotFoundError(name, "producer", producerName); + return { + send: (id, message) => Effect.runPromise((p as EffectProducer).send(id, message)), + flush: () => Effect.runPromise(p.flush), + stop: () => Effect.runPromise(p.flush.pipe(Effect.flatMap(() => p.close))), + }; + }, + consumer(consumerName: string): FlowConsumer { + const c = consumers.get(consumerName); + if (c === undefined) throw flowResourceNotFoundError(name, "consumer", consumerName); + return { + stop: () => Effect.runPromise(c.stop), + }; + }, + requestor(requestorName: string): FlowRequestor { + const rr = requestors.get(requestorName); + if (rr === undefined) throw flowResourceNotFoundError(name, "requestor", requestorName); + return { + request: (request, options) => + Effect.runPromise( + (rr as EffectRequestResponse).request( + request, + toEffectRequestOptions(options), + ), + ), + stop: () => Effect.runPromise(rr.stop), + }; + }, + parameter(parameterName: string): T { + const v = parameters.get(parameterName); + if (v === undefined) throw flowResourceNotFoundError(name, "parameter", parameterName); + return v as T; + }, + }; + + return flow; } + +export type Flow = ReturnType>; + +export const Flow = makeFlow as unknown as { + new ( + name: string, + processorId: string, + pubsub: PubSubBackend, + definition: FlowDefinition, + specifications: ReadonlyArray>, + ): Flow; + ( + name: string, + processorId: string, + pubsub: PubSubBackend, + definition: FlowDefinition, + specifications: ReadonlyArray>, + ): Flow; +}; diff --git a/ts/packages/base/src/processor/index.ts b/ts/packages/base/src/processor/index.ts index 167030d4..ed56e73a 100644 --- a/ts/packages/base/src/processor/index.ts +++ b/ts/packages/base/src/processor/index.ts @@ -1,13 +1,21 @@ export { AsyncProcessor, + makeAsyncProcessor, type ConfigHandler, type EffectConfigHandler, + type AsyncProcessorRuntime, + type AsyncProcessorRuntimeOptions, type ProcessorConfig, + type ProcessorRuntime, } from "./async-processor.js"; export { FlowProcessor, + makeFlowProcessor, runFlowProcessorDefinitionScoped, + type FlowProcessorRuntime, type FlowProcessorRuntimeOptions, + type FlowProcessorStartEffect, + type MakeFlowProcessorOptions, } from "./flow-processor.js"; export { Flow, diff --git a/ts/packages/base/src/processor/program.ts b/ts/packages/base/src/processor/program.ts index 066cc8b1..61565cdd 100644 --- a/ts/packages/base/src/processor/program.ts +++ b/ts/packages/base/src/processor/program.ts @@ -12,7 +12,7 @@ import { type ProcessorLifecycleError, type PubSubError, } from "../errors.js"; -import { NatsBackend } from "../backend/nats.js"; +import { makeNatsBackend } from "../backend/nats.js"; import { makePubSubService, PubSub } from "../backend/pubsub.js"; import { ConsumerFactory, @@ -30,21 +30,21 @@ import { } from "../runtime/config.js"; import { loadMessagingRuntimeConfig } from "../runtime/messaging-config.js"; import type { - AsyncProcessor, EffectConfigHandler, ProcessorConfig, + ProcessorRuntime, } from "./async-processor.js"; import { runFlowProcessorDefinitionScoped } from "./flow-processor.js"; import type { Spec } from "../spec/types.js"; -type ProcessorRunError = Processor extends AsyncProcessor ? Error : never; -type ProcessorRunRequirements = Processor extends AsyncProcessor ? Requirements : never; +type ProcessorRunError = Processor extends ProcessorRuntime ? Error : never; +type ProcessorRunRequirements = Processor extends ProcessorRuntime ? Requirements : never; export interface ProcessorProgramOptions< Config extends ProcessorConfig, Error, Requirements, - Processor extends AsyncProcessor, + Processor extends ProcessorRuntime, > { readonly id: string; readonly make: (config: Config) => Processor; @@ -70,7 +70,7 @@ export interface FlowProcessorProgramOptions< export function runProcessorScoped< Config extends ProcessorConfig, - Processor extends AsyncProcessor, + Processor extends ProcessorRuntime, >( config: Config, make: (config: Config) => Processor, @@ -103,11 +103,13 @@ export function runProcessorScoped< ), ); - const typedProcessor = processor as unknown as AsyncProcessor< - ProcessorRunError, - ProcessorRunRequirements - >; - yield* typedProcessor.startEffect(); + yield* ( + processor.startEffect() as Effect.Effect< + void, + ProcessorRunError | ProcessorLifecycleError, + ProcessorRunRequirements + > + ); }); } @@ -115,7 +117,7 @@ export function makeProcessorProgram< Config extends ProcessorConfig, Error = never, Requirements = never, - Processor extends AsyncProcessor = AsyncProcessor, + Processor extends ProcessorRuntime = ProcessorRuntime, >( options: ProcessorProgramOptions, ) { @@ -133,7 +135,7 @@ export function makeProcessorProgram< manageProcessSignals: false, } as Config; - const pubsub = makePubSubService(new NatsBackend(runtimeConfig.pubsubUrl ?? "nats://localhost:4222")); + const pubsub = makePubSubService(makeNatsBackend(runtimeConfig.pubsubUrl ?? "nats://localhost:4222")); const messagingConfig = yield* loadMessagingRuntimeConfig(); yield* Effect.addFinalizer(() => pubsub.close.pipe( @@ -191,7 +193,7 @@ export function makeFlowProcessorProgram< manageProcessSignals: false, } as Config; - const pubsub = makePubSubService(new NatsBackend(runtimeConfig.pubsubUrl ?? "nats://localhost:4222")); + const pubsub = makePubSubService(makeNatsBackend(runtimeConfig.pubsubUrl ?? "nats://localhost:4222")); const messagingConfig = yield* loadMessagingRuntimeConfig(); yield* Effect.addFinalizer(() => pubsub.close.pipe( diff --git a/ts/packages/base/src/schema/messages.ts b/ts/packages/base/src/schema/messages.ts index 70dede80..355123c5 100644 --- a/ts/packages/base/src/schema/messages.ts +++ b/ts/packages/base/src/schema/messages.ts @@ -5,7 +5,7 @@ */ import * as S from "effect/Schema"; -import { Term, TgError, Triple } from "./primitives.js"; +import {Term, TgError, Triple} from "./primitives.js"; const UnknownRecord = S.Record(S.String, S.Unknown); const MutableArray = (schema: A) => schema.pipe(S.Array, S.mutable); @@ -98,13 +98,14 @@ export const AgentRequest = S.Struct({ export type AgentRequest = typeof AgentRequest.Type; export const AgentResponse = S.Struct({ - chunk_type: S.optionalKey(S.Union([ - S.Literal("thought"), - S.Literal("observation"), - S.Literal("answer"), - S.Literal("error"), - S.Literal("explain"), - ])), + chunk_type: S.optionalKey(S.Literals( + [ + "thought", + "observation", + "answer", + "error", + "explain", + ])), content: S.optionalKey(S.String), end_of_message: S.optionalKey(S.Boolean), end_of_dialog: S.optionalKey(S.Boolean), diff --git a/ts/packages/base/src/services/embeddings-service.ts b/ts/packages/base/src/services/embeddings-service.ts index a7569247..0789acd5 100644 --- a/ts/packages/base/src/services/embeddings-service.ts +++ b/ts/packages/base/src/services/embeddings-service.ts @@ -12,12 +12,12 @@ import { type MessagingDeliveryError, } from "../errors.js"; import type { FlowContext } from "../messaging/consumer.js"; -import { FlowProcessor } from "../processor/flow-processor.js"; -import type { ProcessorConfig } from "../processor/async-processor.js"; +import { makeFlowProcessor } from "../processor/index.ts"; +import type { FlowProcessorRuntime, ProcessorConfig } from "../processor/index.ts"; import type { EmbeddingsRequest, EmbeddingsResponse } from "../schema/messages.js"; -import { ConsumerSpec } from "../spec/consumer-spec.js"; -import { ParameterSpec } from "../spec/parameter-spec.js"; -import { ProducerSpec } from "../spec/producer-spec.js"; +import { makeConsumerSpec } from "../spec/index.ts"; +import { makeParameterSpec } from "../spec/index.ts"; +import { makeProducerSpec } from "../spec/index.ts"; import type { Spec } from "../spec/types.js"; export interface EmbeddingsServiceShape { @@ -66,20 +66,31 @@ const onEmbeddingsRequest = Effect.fn("EmbeddingsService.onRequest")(function* ( }); export const makeEmbeddingsSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "embeddings-request", onEmbeddingsRequest, ), - new ProducerSpec("embeddings-response"), - new ParameterSpec("model"), + makeProducerSpec("embeddings-response"), + makeParameterSpec("model"), ]; -export class EmbeddingsService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); +export type EmbeddingsService = FlowProcessorRuntime; - for (const spec of makeEmbeddingsSpecs()) { - this.registerSpecification(spec); - } - } +export function makeEmbeddingsService( + config: ProcessorConfig, + embeddings?: EmbeddingsServiceShape, +): EmbeddingsService { + return makeFlowProcessor(config, { + specifications: makeEmbeddingsSpecs(), + ...(embeddings === undefined + ? {} + : { + provide: (effect) => + effect.pipe( + Effect.provideService(Embeddings, Embeddings.of(embeddings)), + ), + }), + }); } + +export const EmbeddingsService = makeEmbeddingsService; diff --git a/ts/packages/base/src/services/index.ts b/ts/packages/base/src/services/index.ts index 0b0f5ba3..bc4054b7 100644 --- a/ts/packages/base/src/services/index.ts +++ b/ts/packages/base/src/services/index.ts @@ -2,6 +2,7 @@ export { Llm, LlmService, LlmServiceError, + makeLlmService, makeLlmServiceShape, makeLlmSpecs, type LlmProvider, @@ -10,6 +11,7 @@ export { export { Embeddings, EmbeddingsService, + makeEmbeddingsService, makeEmbeddingsSpecs, type EmbeddingsServiceShape, } from "./embeddings-service.js"; diff --git a/ts/packages/base/src/services/llm-service.ts b/ts/packages/base/src/services/llm-service.ts index 6e33f74a..9de4f820 100644 --- a/ts/packages/base/src/services/llm-service.ts +++ b/ts/packages/base/src/services/llm-service.ts @@ -12,16 +12,16 @@ import { type MessagingDeliveryError, } from "../errors.js"; import type { FlowContext } from "../messaging/consumer.js"; -import { FlowProcessor } from "../processor/flow-processor.js"; -import type { ProcessorConfig } from "../processor/async-processor.js"; +import { makeFlowProcessor } from "../processor/index.ts"; +import type { FlowProcessorRuntime, ProcessorConfig } from "../processor/index.ts"; import type { TextCompletionRequest, TextCompletionResponse, } from "../schema/messages.js"; -import type { LlmChunk, LlmResult } from "../schema/primitives.js"; -import { ConsumerSpec } from "../spec/consumer-spec.js"; -import { ParameterSpec } from "../spec/parameter-spec.js"; -import { ProducerSpec } from "../spec/producer-spec.js"; +import type { LlmChunk, LlmResult } from "../schema/index.ts"; +import { makeConsumerSpec } from "../spec/index.ts"; +import { makeParameterSpec } from "../spec/index.ts"; +import { makeProducerSpec } from "../spec/index.ts"; import type { Spec } from "../spec/types.js"; export class LlmServiceError extends S.TaggedErrorClass()( @@ -203,45 +203,29 @@ const onLlmRequest = Effect.fn("LlmService.onRequest")(function* ( }); export const makeLlmSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "text-completion-request", onLlmRequest, ), - new ProducerSpec("text-completion-response"), - new ParameterSpec("model"), - new ParameterSpec("temperature"), + makeProducerSpec("text-completion-response"), + makeParameterSpec("model"), + makeParameterSpec("temperature"), ]; -export abstract class LlmService extends FlowProcessor implements LlmProvider { - protected constructor(config: ProcessorConfig) { - super(config); +export type LlmService = FlowProcessorRuntime & LlmProvider; - for (const spec of makeLlmSpecs()) { - this.registerSpecification(spec); - } - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(Llm, Llm.of(makeLlmServiceShape(this))), - ); - } - - abstract generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise; - - abstract generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator; - - supportsStreaming(): boolean { - return false; - } +export function makeLlmService( + config: ProcessorConfig, + provider: LlmProvider, +): LlmService { + const service = makeFlowProcessor(config, { + specifications: makeLlmSpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService(Llm, Llm.of(makeLlmServiceShape(provider))), + ), + }); + return Object.assign(service, provider); } + +export const LlmService = makeLlmService; diff --git a/ts/packages/base/src/spec/consumer-spec.ts b/ts/packages/base/src/spec/consumer-spec.ts index 6cd9e1a6..2f9473ad 100644 --- a/ts/packages/base/src/spec/consumer-spec.ts +++ b/ts/packages/base/src/spec/consumer-spec.ts @@ -8,7 +8,6 @@ import { Effect } from "effect"; import * as S from "effect/Schema"; import type { Spec } from "./types.js"; import type { SpecRuntimeRequirements } from "./types.js"; -import type { PubSubBackend } from "../backend/types.js"; import type { Flow, FlowDefinition } from "../processor/flow.js"; import { type MessageHandler } from "../messaging/consumer.js"; import { @@ -24,64 +23,71 @@ import { const isTooManyRequestsError = S.is(TooManyRequestsError); -export class ConsumerSpec implements Spec { - public readonly name: string; - private readonly handler: EffectMessageHandler; - private readonly concurrency: number; +declare const ConsumerSpecType: unique symbol; - constructor( - name: string, - handler: EffectMessageHandler, - concurrency = 1, +export interface ConsumerSpec extends Spec { + readonly [ConsumerSpecType]?: { + readonly message: T; + readonly error: E; + }; + readonly addEffect: ( + flow: Flow, + definition: FlowDefinition, + ) => Effect.Effect; +} + +export function makeConsumerSpec( + name: string, + handler: EffectMessageHandler, + concurrency = 1, +): ConsumerSpec { + const addEffect = Effect.fn("ConsumerSpec.addEffect")(function* ( + flow: Flow, + definition: FlowDefinition, ) { - this.name = name; - this.handler = handler; - this.concurrency = concurrency; - } - - static fromPromise( - name: string, - handler: MessageHandler, - concurrency = 1, - ): ConsumerSpec { - return new ConsumerSpec( - name, - (message, properties, flow) => - Effect.tryPromise({ - try: () => handler(message, properties, flow), - catch: (error) => - isTooManyRequestsError(error) - ? error - : messagingHandlerError(name, `${flow.id}-${flow.name}-${name}`, error), - }), - concurrency, - ); - } - - addEffect(flow: Flow, definition: FlowDefinition) { - const spec = this; - return Effect.gen(function* () { - const topic = definition.topics?.[spec.name] ?? spec.name; + const topic = definition.topics?.[name] ?? name; const factory = yield* ConsumerFactory; const consumer = yield* factory.run( { topic, - subscription: `${flow.processorId}-${flow.name}-${spec.name}`, - handler: spec.handler, - concurrency: spec.concurrency, + subscription: `${flow.processorId}-${flow.name}-${name}`, + handler, + concurrency, }, { id: flow.processorId, name: flow.name, flow }, ); - flow.registerConsumer(spec.name, consumer); - }); - } + flow.registerConsumer(name, consumer); + }); - async add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise { - const effect = this.addEffect(flow, definition) as Effect.Effect< - void, - PubSubError, - SpecRuntimeRequirements - >; - await flow.runInCompatibilityScope(effect, pubsub); - } + return { + name, + addEffect, + add: async (flow, pubsub, definition) => { + const effect = addEffect(flow as Flow, definition) as Effect.Effect< + void, + PubSubError, + SpecRuntimeRequirements + >; + await flow.runInCompatibilityScope(effect, pubsub); + }, + }; +} + +export function makeConsumerSpecFromPromise( + name: string, + handler: MessageHandler, + concurrency = 1, +): ConsumerSpec { + return makeConsumerSpec( + name, + (message, properties, flow) => + Effect.tryPromise({ + try: () => handler(message, properties, flow), + catch: (error) => + isTooManyRequestsError(error) + ? error + : messagingHandlerError(name, `${flow.id}-${flow.name}-${name}`, error), + }), + concurrency, + ); } diff --git a/ts/packages/base/src/spec/index.ts b/ts/packages/base/src/spec/index.ts index 3395f113..2ceeae1c 100644 --- a/ts/packages/base/src/spec/index.ts +++ b/ts/packages/base/src/spec/index.ts @@ -1,5 +1,5 @@ export type { Spec, SpecRuntimeError, SpecRuntimeRequirements } from "./types.js"; -export { ConsumerSpec } from "./consumer-spec.js"; -export { ProducerSpec } from "./producer-spec.js"; -export { ParameterSpec } from "./parameter-spec.js"; -export { RequestResponseSpec } from "./request-response-spec.js"; +export { makeConsumerSpec, makeConsumerSpecFromPromise, type ConsumerSpec } from "./consumer-spec.js"; +export { makeProducerSpec, type ProducerSpec } from "./producer-spec.js"; +export { makeParameterSpec, type ParameterSpec } from "./parameter-spec.js"; +export { makeRequestResponseSpec, type RequestResponseSpec } from "./request-response-spec.js"; diff --git a/ts/packages/base/src/spec/parameter-spec.ts b/ts/packages/base/src/spec/parameter-spec.ts index 94797f47..9b20a152 100644 --- a/ts/packages/base/src/spec/parameter-spec.ts +++ b/ts/packages/base/src/spec/parameter-spec.ts @@ -6,25 +6,22 @@ import { Effect } from "effect"; import type { Spec } from "./types.js"; -import type { PubSubBackend } from "../backend/types.js"; import type { Flow, FlowDefinition } from "../processor/flow.js"; -export class ParameterSpec implements Spec { - public readonly name: string; +export interface ParameterSpec extends Spec {} - constructor(name: string) { - this.name = name; - } - - addEffect(flow: Flow, definition: FlowDefinition) { - const spec = this; - return Effect.sync(() => { - const value = definition.parameters?.[spec.name]; - flow.setParameter(spec.name, value); +export function makeParameterSpec(name: string): ParameterSpec { + const addEffect = (flow: Flow, definition: FlowDefinition) => + Effect.sync(() => { + const value = definition.parameters?.[name]; + flow.setParameter(name, value); }); - } - async add(flow: Flow, _pubsub: PubSubBackend, definition: FlowDefinition): Promise { - await Effect.runPromise(this.addEffect(flow, definition)); - } + return { + name, + addEffect, + add: async (flow, _pubsub, definition) => { + await Effect.runPromise(addEffect(flow, definition)); + }, + }; } diff --git a/ts/packages/base/src/spec/producer-spec.ts b/ts/packages/base/src/spec/producer-spec.ts index b5ebf88e..955846b4 100644 --- a/ts/packages/base/src/spec/producer-spec.ts +++ b/ts/packages/base/src/spec/producer-spec.ts @@ -6,31 +6,34 @@ import { Effect } from "effect"; import type { Spec } from "./types.js"; -import type { PubSubBackend } from "../backend/types.js"; import type { Flow, FlowDefinition } from "../processor/flow.js"; import { ProducerFactory, type EffectProducer, } from "../messaging/runtime.js"; -export class ProducerSpec implements Spec { - public readonly name: string; +declare const ProducerSpecType: unique symbol; - constructor(name: string) { - this.name = name; - } +export interface ProducerSpec extends Spec { + readonly [ProducerSpecType]?: (_: T) => T; +} - addEffect(flow: Flow, definition: FlowDefinition) { - const spec = this; - return Effect.gen(function* () { - const topic = definition.topics?.[spec.name] ?? spec.name; +export function makeProducerSpec(name: string): ProducerSpec { + const addEffect = Effect.fn("ProducerSpec.addEffect")(function* ( + flow: Flow, + definition: FlowDefinition, + ) { + const topic = definition.topics?.[name] ?? name; const factory = yield* ProducerFactory; const producer = yield* factory.make({ topic }); - flow.registerProducer(spec.name, producer as EffectProducer); - }); - } + flow.registerProducer(name, producer as EffectProducer); + }); - async add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise { - await flow.runInCompatibilityScope(this.addEffect(flow, definition), pubsub); - } + return { + name, + addEffect, + add: async (flow, pubsub, definition) => { + await flow.runInCompatibilityScope(addEffect(flow, definition), pubsub); + }, + }; } diff --git a/ts/packages/base/src/spec/request-response-spec.ts b/ts/packages/base/src/spec/request-response-spec.ts index 229213cf..b840c8be 100644 --- a/ts/packages/base/src/spec/request-response-spec.ts +++ b/ts/packages/base/src/spec/request-response-spec.ts @@ -9,44 +9,46 @@ import { Effect } from "effect"; import type { Spec } from "./types.js"; -import type { PubSubBackend } from "../backend/types.js"; import type { Flow, FlowDefinition } from "../processor/flow.js"; import { RequestResponseFactory, type EffectRequestResponse, } from "../messaging/runtime.js"; -export class RequestResponseSpec implements Spec { - public readonly name: string; - private readonly requestTopicName: string; - private readonly responseTopicName: string; +declare const RequestResponseSpecType: unique symbol; - constructor( - name: string, - requestTopicName: string, - responseTopicName: string, +export interface RequestResponseSpec extends Spec { + readonly [RequestResponseSpecType]?: { + readonly request: TReq; + readonly response: TRes; + }; +} + +export function makeRequestResponseSpec( + name: string, + requestTopicName: string, + responseTopicName: string, +): RequestResponseSpec { + const addEffect = Effect.fn("RequestResponseSpec.addEffect")(function* ( + flow: Flow, + definition: FlowDefinition, ) { - this.name = name; - this.requestTopicName = requestTopicName; - this.responseTopicName = responseTopicName; - } - - addEffect(flow: Flow, definition: FlowDefinition) { - const spec = this; - return Effect.gen(function* () { - const requestTopic = definition.topics?.[spec.requestTopicName] ?? spec.requestTopicName; - const responseTopic = definition.topics?.[spec.responseTopicName] ?? spec.responseTopicName; + const requestTopic = definition.topics?.[requestTopicName] ?? requestTopicName; + const responseTopic = definition.topics?.[responseTopicName] ?? responseTopicName; const factory = yield* RequestResponseFactory; const requestor = yield* factory.make({ requestTopic, responseTopic, - subscription: `${flow.processorId}-${flow.name}-${spec.name}`, + subscription: `${flow.processorId}-${flow.name}-${name}`, }); - flow.registerRequestor(spec.name, requestor as EffectRequestResponse); - }); - } + flow.registerRequestor(name, requestor as EffectRequestResponse); + }); - async add(flow: Flow, pubsub: PubSubBackend, definition: FlowDefinition): Promise { - await flow.runInCompatibilityScope(this.addEffect(flow, definition), pubsub); - } + return { + name, + addEffect, + add: async (flow, pubsub, definition) => { + await flow.runInCompatibilityScope(addEffect(flow, definition), pubsub); + }, + }; } diff --git a/ts/packages/client/src/__tests__/rpc-timeout.test.ts b/ts/packages/client/src/__tests__/rpc-timeout.test.ts index 186d1b51..4ed38291 100644 --- a/ts/packages/client/src/__tests__/rpc-timeout.test.ts +++ b/ts/packages/client/src/__tests__/rpc-timeout.test.ts @@ -1,8 +1,8 @@ -import { Effect, Stream } from "effect"; +import { Effect } from "effect"; import { describe, expect, it, vi } from "vitest"; -import { DispatchError, DispatchStreamChunk } from "../rpc/contract"; -import { EffectRpcClient, type DispatchInput } from "../socket/effect-rpc-client"; -import { BaseApi } from "../socket/trustgraph-socket"; +import { DispatchError } from "../rpc/contract"; +import { type DispatchInput, withDispatchRequestPolicy } from "../socket/effect-rpc-client"; +import { makeBaseApiWithRpc } from "../socket/trustgraph-socket"; const input: DispatchInput = { scope: "global", @@ -13,11 +13,12 @@ const input: DispatchInput = { describe("Effect RPC request policy", () => { it("threads BaseApi timeout and retry options into dispatch calls", async () => { const dispatch = vi.fn(() => Promise.resolve({ ok: true })); - const api = Object.create(BaseApi.prototype) as BaseApi; - - (api as unknown as { rpc: { dispatch: typeof dispatch } }).rpc = { + const api = makeBaseApiWithRpc("alice", undefined, "ws://example.test/rpc", { dispatch, - }; + dispatchStream: vi.fn(() => Promise.resolve(undefined)), + close: vi.fn(() => Promise.resolve()), + subscribe: vi.fn(() => () => {}), + }); await api.makeRequest("config", { operation: "list" }, 25, 2); @@ -28,52 +29,33 @@ describe("Effect RPC request policy", () => { }); it("rejects stalled dispatch calls at the requested timeout", async () => { - const client = Object.create(EffectRpcClient.prototype) as EffectRpcClient; const startedAt = Date.now(); - setClientPromise(client, { - Dispatch: () => Effect.never, - DispatchStream: () => Stream.never, - }); - await expect( - client.dispatch(input, { timeoutMs: 20, retries: 1 }), + Effect.runPromise(withDispatchRequestPolicy(Effect.never, { timeoutMs: 20, retries: 1 })), ).rejects.toBeInstanceOf(DispatchError); expect(Date.now() - startedAt).toBeLessThan(1_000); }); it("retries dispatch failures up to the requested attempt count", async () => { - const client = Object.create(EffectRpcClient.prototype) as EffectRpcClient; let attempts = 0; - setClientPromise(client, { - Dispatch: () => - Effect.suspend(() => { - attempts += 1; - if (attempts < 3) { - return Effect.fail(new DispatchError({ message: String(attempts) })); - } - return Effect.succeed({ ok: true }); - }), - DispatchStream: () => Stream.never, - }); - - await expect(client.dispatch(input, { timeoutMs: 100, retries: 3 })).resolves.toEqual({ - ok: true, - }); + await expect( + Effect.runPromise( + withDispatchRequestPolicy( + Effect.suspend(() => { + attempts += 1; + if (attempts < 3) { + return Effect.fail(new DispatchError({ message: String(attempts) })); + } + return Effect.succeed({ ok: true }); + }), + { timeoutMs: 100, retries: 3 }, + ), + ), + ).resolves.toEqual({ ok: true }); expect(attempts).toBe(3); }); }); - -function setClientPromise( - client: EffectRpcClient, - fakeClient: { - Dispatch: (payload: unknown) => Effect.Effect; - DispatchStream: (payload: unknown) => Stream.Stream; - }, -): void { - (client as unknown as { clientPromise: Promise }).clientPromise = - Promise.resolve(fakeClient); -} diff --git a/ts/packages/client/src/socket/effect-rpc-client.ts b/ts/packages/client/src/socket/effect-rpc-client.ts index cb003f36..4a3ba88d 100644 --- a/ts/packages/client/src/socket/effect-rpc-client.ts +++ b/ts/packages/client/src/socket/effect-rpc-client.ts @@ -38,91 +38,55 @@ export interface DispatchOptions { const DEFAULT_REQUEST_TIMEOUT_MS = 10_000; const DEFAULT_REQUEST_ATTEMPTS = 3; -export class EffectRpcClient { - private readonly url: string; - private readonly onConnect: (() => void) | undefined; - private readonly onDisconnect: (() => void) | undefined; - private readonly scopePromise: Promise; - private readonly clientPromise: Promise; - private readonly listeners = new Set<(state: RpcConnectionState) => void>(); - private state: RpcConnectionState = { status: "connecting" }; - private closed = false; +type NewableFactory = { + new (...args: Args): A; + (...args: Args): A; + readonly prototype: A; +}; - constructor( - url: string, - onConnect?: () => void, - onDisconnect?: () => void, - ) { - this.url = url; - this.onConnect = onConnect; - this.onDisconnect = onDisconnect; - this.scopePromise = Effect.runPromise(Scope.make()); - this.clientPromise = this.scopePromise.then((scope) => - Effect.runPromise(this.makeClient().pipe(Scope.provide(scope))), - ); - this.clientPromise.catch((cause) => { - this.setState({ - status: "failed", - lastError: errorMessage(cause), - }); - }); +function newableFactory( + factory: (...args: Args) => A, +): NewableFactory { + function Constructor(...args: Args): A { + return factory(...args); } + return Constructor as unknown as NewableFactory; +} - subscribe(listener: (state: RpcConnectionState) => void): () => void { - this.listeners.add(listener); - listener(this.state); - return () => { - this.listeners.delete(listener); - }; - } - - async dispatch(input: DispatchInput, options: DispatchOptions = {}): Promise { - const client = await this.clientPromise; - return await Effect.runPromise( - this.withRequestPolicy(client.Dispatch(new DispatchPayload(input)), options), - ); - } - - async dispatchStream( +export interface EffectRpcClient { + readonly subscribe: (listener: (state: RpcConnectionState) => void) => () => void; + readonly dispatch: ( + input: DispatchInput, + options?: DispatchOptions, + ) => Promise; + readonly dispatchStream: ( input: DispatchInput, receiver: (chunk: DispatchStreamChunk) => boolean, - options: DispatchOptions = {}, - ): Promise { - const client = await this.clientPromise; - let last: DispatchStreamChunk | undefined; - await Effect.runPromise( - this.withRequestPolicy( - client.DispatchStream(new DispatchPayload(input)).pipe( - Stream.runForEach((chunk) => - Effect.suspend(() => { - last = chunk; - if (receiver(chunk)) return Effect.fail(new StopStreaming()); - return Effect.void; - }), - ), - Effect.catchIf( - (cause): cause is StopStreaming => cause instanceof StopStreaming, - () => Effect.void, - ), - ), - options, - ), - ); - return last; - } + options?: DispatchOptions, + ) => Promise; + readonly close: () => Promise; +} - async close(): Promise { - if (this.closed) return; - this.closed = true; - this.setState({ status: "closed" }); - const scope = await this.scopePromise; - await Effect.runPromise(Scope.close(scope, Exit.void)); - } +export function makeEffectRpcClient( + url: string, + onConnect?: () => void, + onDisconnect?: () => void, +): EffectRpcClient { + const listeners = new Set<(state: RpcConnectionState) => void>(); + let state: RpcConnectionState = { status: "connecting" }; + let closed = false; - private makeClient(): Effect.Effect { + const setState = (nextState: RpcConnectionState): void => { + state = nextState; + for (const listener of listeners) { + listener(nextState); + } + }; + + const makeClient = (): Effect.Effect => { const socketLayer = Layer.effect( Socket.Socket, - Socket.makeWebSocket(this.url, { + Socket.makeWebSocket(url, { closeCodeIsError: (code) => code !== 1000, openTimeout: "10 seconds", }), @@ -132,17 +96,17 @@ export class EffectRpcClient { RpcClient.ConnectionHooks, RpcClient.ConnectionHooks.of({ onConnect: Effect.sync(() => { - this.setState({ status: "connected" }); - this.onConnect?.(); + setState({ status: "connected" }); + onConnect?.(); }), onDisconnect: Effect.sync(() => { - if (!this.closed) { - this.setState({ + if (!closed) { + setState({ status: "connecting", lastError: "Disconnected from gateway", }); } - this.onDisconnect?.(); + onDisconnect?.(); }), }), ); @@ -164,35 +128,87 @@ export class EffectRpcClient { Layer.build(clientLayer), (context) => Context.get(context, TrustGraphRpcClientService), ); - } + }; - private setState(state: RpcConnectionState): void { - this.state = state; - for (const listener of this.listeners) { + const scopePromise = Effect.runPromise(Scope.make()); + const clientPromise = scopePromise.then((scope) => + Effect.runPromise(makeClient().pipe(Scope.provide(scope))), + ); + clientPromise.catch((cause) => { + setState({ + status: "failed", + lastError: errorMessage(cause), + }); + }); + + return { + subscribe: (listener) => { + listeners.add(listener); listener(state); - } - } - - private withRequestPolicy( - effect: Effect.Effect, - options: DispatchOptions, - ): Effect.Effect { - const timeoutMs = normalizeTimeoutMs(options.timeoutMs); - const retryTimes = normalizeAttempts(options.retries) - 1; - const timed = effect.pipe( - Effect.timeoutOrElse({ - duration: timeoutMs, - orElse: () => - Effect.fail( - new DispatchError({ - message: `Request timed out after ${timeoutMs}ms`, - }), + return () => { + listeners.delete(listener); + }; + }, + dispatch: async (input, options = {}) => { + const client = await clientPromise; + return await Effect.runPromise( + withDispatchRequestPolicy(client.Dispatch(new DispatchPayload(input)), options), + ); + }, + dispatchStream: async (input, receiver, options = {}) => { + const client = await clientPromise; + let last: DispatchStreamChunk | undefined; + await Effect.runPromise( + withDispatchRequestPolicy( + client.DispatchStream(new DispatchPayload(input)).pipe( + Stream.runForEach((chunk) => + Effect.suspend(() => { + last = chunk; + if (receiver(chunk)) return Effect.fail(new StopStreaming()); + return Effect.void; + }), + ), + Effect.catchIf( + (cause): cause is StopStreaming => cause instanceof StopStreaming, + () => Effect.void, + ), ), - }), - ); + options, + ), + ); + return last; + }, + close: async () => { + if (closed) return; + closed = true; + setState({ status: "closed" }); + const scope = await scopePromise; + await Effect.runPromise(Scope.close(scope, Exit.void)); + }, + }; +} - return retryTimes > 0 ? timed.pipe(Effect.retry({ times: retryTimes })) : timed; - } +export const EffectRpcClient = newableFactory(makeEffectRpcClient); + +export function withDispatchRequestPolicy( + effect: Effect.Effect, + options: DispatchOptions, +): Effect.Effect { + const timeoutMs = normalizeTimeoutMs(options.timeoutMs); + const retryTimes = normalizeAttempts(options.retries) - 1; + const timed = effect.pipe( + Effect.timeoutOrElse({ + duration: timeoutMs, + orElse: () => + Effect.fail( + new DispatchError({ + message: `Request timed out after ${timeoutMs}ms`, + }), + ), + }), + ); + + return retryTimes > 0 ? timed.pipe(Effect.retry({ times: retryTimes })) : timed; } class StopStreaming extends Data.TaggedError("StopStreaming")<{}> {} diff --git a/ts/packages/client/src/socket/trustgraph-socket.ts b/ts/packages/client/src/socket/trustgraph-socket.ts index 7bad5536..e7fb6ea1 100644 --- a/ts/packages/client/src/socket/trustgraph-socket.ts +++ b/ts/packages/client/src/socket/trustgraph-socket.ts @@ -5,6 +5,7 @@ import { type DispatchInput, type DispatchOptions, type RpcConnectionState, + makeEffectRpcClient, } from "./effect-rpc-client.js"; import { getDefaultSocketUrl, getRandomValues } from "./websocket-adapter.js"; @@ -167,7 +168,7 @@ function throwIfResponseError(error: ResponseError | undefined): void { } } -interface ConfigValueEntry { +export interface ConfigValueEntry { workspace?: string; type?: string; key: string; @@ -309,6 +310,21 @@ function makeid(length: number) { ); } +type NewableFactory = { + new (...args: Args): A; + (...args: Args): A; + readonly prototype: A; +}; + +function newableFactory( + factory: (...args: Args) => A, +): NewableFactory { + function Constructor(...args: Args): A { + return factory(...args); + } + return Constructor as unknown as NewableFactory; +} + /** * BaseApi - Core WebSocket client for TrustGraph API * Manages connection lifecycle, message routing, and provides base request @@ -330,180 +346,190 @@ export interface ConnectionState { lastError?: string; } -export class BaseApi { - tag: string; // Unique client identifier - id: number; // Counter for generating unique message IDs - token: string | undefined; // Optional authentication token - user: string; // User identifier for API requests - socketUrl: string; // WebSocket URL - private readonly rpc: EffectRpcClient; +export function makeBaseApi( + user: string, + token?: string, + socketUrl?: string, + rpcFactory: (url: string) => EffectRpcClient = makeEffectRpcClient, +) { + let rpc: EffectRpcClient; + const connectionStateListeners: ((state: ConnectionState) => void)[] = []; + let lastError: string | undefined = undefined; + let rpcState: RpcConnectionState = { status: "connecting" }; - // Connection state tracking for UI - private connectionStateListeners: ((state: ConnectionState) => void)[] = []; - private lastError: string | undefined = undefined; - private rpcState: RpcConnectionState = { status: "connecting" }; + const api = { + tag: makeid(16), // Unique client identifier + id: 1, // Counter for generating unique message IDs + token, // Optional authentication token + user, // User identifier for API requests + socketUrl: withDefault(socketUrl, SOCKET_URL), // Use provided URL or default - constructor(user: string, token?: string, socketUrl?: string) { - this.tag = makeid(16); // Generate unique client tag - this.id = 1; // Start message ID counter - this.token = token; // Store authentication token - this.user = user; // Store user identifier - this.socketUrl = withDefault(socketUrl, SOCKET_URL); // Use provided URL or default - this.rpc = new EffectRpcClient(this.socketUrlWithToken()); - this.rpc.subscribe((state) => { - this.rpcState = state; - this.lastError = state.lastError; - this.notifyStateChange(); - }); + /** + * Subscribe to connection state changes for UI updates + */ + onConnectionStateChange(listener: (state: ConnectionState) => void) { + connectionStateListeners.push(listener); + // Immediately send current state + listener(getConnectionState()); - console.log( - "SOCKET: opening socket...", - isNonEmptyString(token) ? "with auth" : "without auth", - "user:", - user, - ); - } + // Return unsubscribe function + return () => { + const index = connectionStateListeners.indexOf(listener); + if (index > -1) { + connectionStateListeners.splice(index, 1); + } + }; + }, - /** - * Subscribe to connection state changes for UI updates - */ - onConnectionStateChange(listener: (state: ConnectionState) => void) { - this.connectionStateListeners.push(listener); - // Immediately send current state - listener(this.getConnectionState()); + /** + * Closes the WebSocket connection and cleans up + */ + close() { + rpc.close().catch((err) => { + console.error("[socket close error]", err); + }); + }, - // Return unsubscribe function - return () => { - const index = this.connectionStateListeners.indexOf(listener); - if (index > -1) { - this.connectionStateListeners.splice(index, 1); - } - }; - } + /** + * Generates the next unique message ID for requests + * Format: {clientTag}-{incrementingNumber} + */ + getNextId() { + const mid = api.tag + "-" + api.id.toString(); + api.id++; + return mid; + }, + + /** + * Core method for making service requests over WebSocket + * @param service - Name of the service to call + * @param request - Request payload + * @param timeout - Request timeout in milliseconds (default: 10000) + * @param retries - Number of retry attempts (default: 3) + * @param flow - Optional flow identifier + * @returns Promise resolving to the service response + */ + makeRequest( + service: string, + request: RequestType, + timeout?: number, + retries?: number, + flow?: string, + ) { + return rpc + .dispatch(dispatchInput(service, request, flow), dispatchOptions(timeout, retries)) + .then((obj) => { + return obj as ResponseType; + }); + }, + + /** + * Makes a request that can receive multiple responses (streaming) + * Used for operations that return data in chunks + */ + makeRequestMulti( + service: string, + request: RequestType, + receiver: (resp: unknown) => boolean, // Callback to handle each response chunk + timeout?: number, + retries?: number, + flow?: string, + ) { + return rpc + .dispatchStream( + dispatchInput(service, request, flow), + (chunk) => { + return receiver({ response: chunk.response, complete: chunk.complete }); + }, + dispatchOptions(timeout, retries), + ) + .then((obj) => { + return obj as ResponseType; + }); + }, + + /** + * Convenience method for making flow-specific requests + * Defaults to "default" flow if none specified + */ + makeFlowRequest( + service: string, + request: RequestType, + timeout?: number, + retries?: number, + flow?: string, + ) { + const flowName = isNonEmptyString(flow) ? flow : "default"; + + return api.makeRequest( + service, + request, + timeout, + retries, + flowName, + ); + }, + + // Factory methods for creating specialized API instances + librarian() { + return new LibrarianApi(api); + }, + + flows() { + return new FlowsApi(api); + }, + + flow(id: string) { + return new FlowApi(api, id); + }, + + knowledge() { + return new KnowledgeApi(api); + }, + + config() { + return new ConfigApi(api); + }, + + collectionManagement() { + return new CollectionManagementApi(api); + }, + }; /** * Get current connection state */ - private getConnectionState(): ConnectionState { - const hasApiKey = isNonEmptyString(this.token); - const status = this.connectionStatusFromRpc(hasApiKey); + const getConnectionState = (): ConnectionState => { + const hasApiKey = isNonEmptyString(api.token); + const status = connectionStatusFromRpc(hasApiKey); const state: ConnectionState = { status, hasApiKey, }; - if (this.lastError !== undefined) { - state.lastError = this.lastError; + if (lastError !== undefined) { + state.lastError = lastError; } return state; - } + }; /** * Notify all listeners of connection state changes */ - private notifyStateChange() { - const state = this.getConnectionState(); - this.connectionStateListeners.forEach((listener) => { + const notifyStateChange = () => { + const state = getConnectionState(); + connectionStateListeners.forEach((listener) => { try { listener(state); } catch (error) { console.error("Error in connection state listener:", error); } }); - } + }; - /** - * Closes the WebSocket connection and cleans up - */ - close() { - this.rpc.close().catch((err) => { - console.error("[socket close error]", err); - }); - } - - /** - * Generates the next unique message ID for requests - * Format: {clientTag}-{incrementingNumber} - */ - getNextId() { - const mid = this.tag + "-" + this.id.toString(); - this.id++; - return mid; - } - - /** - * Core method for making service requests over WebSocket - * @param service - Name of the service to call - * @param request - Request payload - * @param timeout - Request timeout in milliseconds (default: 10000) - * @param retries - Number of retry attempts (default: 3) - * @param flow - Optional flow identifier - * @returns Promise resolving to the service response - */ - makeRequest( - service: string, - request: RequestType, - timeout?: number, - retries?: number, - flow?: string, - ) { - return this.rpc - .dispatch(this.dispatchInput(service, request, flow), dispatchOptions(timeout, retries)) - .then((obj) => { - return obj as ResponseType; - }); - } - - /** - * Makes a request that can receive multiple responses (streaming) - * Used for operations that return data in chunks - */ - makeRequestMulti( - service: string, - request: RequestType, - receiver: (resp: unknown) => boolean, // Callback to handle each response chunk - timeout?: number, - retries?: number, - flow?: string, - ) { - return this.rpc - .dispatchStream( - this.dispatchInput(service, request, flow), - (chunk) => { - return receiver({ response: chunk.response, complete: chunk.complete }); - }, - dispatchOptions(timeout, retries), - ) - .then((obj) => { - return obj as ResponseType; - }); - } - - /** - * Convenience method for making flow-specific requests - * Defaults to "default" flow if none specified - */ - makeFlowRequest( - service: string, - request: RequestType, - timeout?: number, - retries?: number, - flow?: string, - ) { - if (!isNonEmptyString(flow)) flow = "default"; - - return this.makeRequest( - service, - request, - timeout, - retries, - flow, - ); - } - - private connectionStatusFromRpc(hasApiKey: boolean): ConnectionState["status"] { - switch (this.rpcState.status) { + const connectionStatusFromRpc = (hasApiKey: boolean): ConnectionState["status"] => { + switch (rpcState.status) { case "connected": return hasApiKey ? "authenticated" : "unauthenticated"; case "failed": @@ -511,15 +537,15 @@ export class BaseApi { case "closed": return "failed"; case "connecting": - return this.lastError === undefined ? "connecting" : "reconnecting"; + return lastError === undefined ? "connecting" : "reconnecting"; } - } + }; - private dispatchInput( + const dispatchInput = ( service: string, request: RequestType, flow?: string, - ): DispatchInput { + ): DispatchInput => { if (isNonEmptyString(flow)) { return { scope: "flow", @@ -533,1811 +559,1950 @@ export class BaseApi { service, request: request as Record, }; - } + }; - private socketUrlWithToken(): string { - if (!isNonEmptyString(this.token)) return this.socketUrl; - const separator = this.socketUrl.includes("?") ? "&" : "?"; - return `${this.socketUrl}${separator}token=${encodeURIComponent(this.token)}`; - } + const socketUrlWithToken = (): string => { + if (!isNonEmptyString(api.token)) return api.socketUrl; + const separator = api.socketUrl.includes("?") ? "&" : "?"; + return `${api.socketUrl}${separator}token=${encodeURIComponent(api.token)}`; + }; - // Factory methods for creating specialized API instances - librarian() { - return new LibrarianApi(this); - } + rpc = rpcFactory(socketUrlWithToken()); + rpc.subscribe((state) => { + rpcState = state; + lastError = state.lastError; + notifyStateChange(); + }); - flows() { - return new FlowsApi(this); - } + console.log( + "SOCKET: opening socket...", + isNonEmptyString(token) ? "with auth" : "without auth", + "user:", + user, + ); - flow(id: string) { - return new FlowApi(this, id); - } + return api; +} - knowledge() { - return new KnowledgeApi(this); - } +export type BaseApi = ReturnType; +export const BaseApi = newableFactory(makeBaseApi); - config() { - return new ConfigApi(this); - } - - collectionManagement() { - return new CollectionManagementApi(this); - } +export function makeBaseApiWithRpc( + user: string, + token: string | undefined, + socketUrl: string | undefined, + rpc: EffectRpcClient, +): BaseApi { + return makeBaseApi(user, token, socketUrl, () => rpc); } /** * LibrarianApi - Manages document storage and retrieval * Handles document lifecycle including upload, processing, and removal */ -export class LibrarianApi { - api: BaseApi; +export function makeLibrarianApi(api: BaseApi) { + return { + api, - constructor(api: BaseApi) { - this.api = api; - } - /** - * Retrieves list of all documents in the system - */ - getDocuments() { - return this.api - .makeRequest( - "librarian", - { - operation: "list-documents", + + /** + * Retrieves list of all documents in the system + */ + getDocuments() { + return this.api + .makeRequest( + "librarian", + { + operation: "list-documents", + user: this.api.user, + }, + 60000, // 60 second timeout for potentially large lists + ) + .then((r) => r["document-metadatas"] ?? r.documents ?? []); + }, + + + + /** + * Retrieves list of documents currently being processed + */ + getProcessing() { + return this.api + .makeRequest( + "librarian", + { + operation: "list-processing", + user: this.api.user, + }, + 60000, + ) + .then((r) => r["processing-metadatas"] ?? r.processing ?? r["processing-metadata"] ?? []); + }, + + + + /** + * Retrieves metadata for a single document by ID + * @param documentId - Document URI/ID to fetch + * @returns Document metadata including title, comments, tags, and RDF metadata + */ + getDocumentMetadata(documentId: string): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "get-document-metadata", + "document-id": documentId, + documentId, + user: this.api.user, + }, + 30000, + ) + .then((r) => r["document-metadata"] ?? r.documentMetadata ?? null); + }, + + + + /** + * Uploads a document to the library with full metadata + * @param document - Base64-encoded document content + * @param id - Optional document identifier + * @param metadata - Optional metadata as triples + * @param mimeType - Document MIME type + * @param title - Document title + * @param comments - Additional comments + * @param tags - Document tags for categorization + */ + loadDocument( + document: string, // base64-encoded doc + mimeType: string, + title: string, + comments: string, + tags: string[], + id?: string, + metadata?: Triple[], + ) { + const documentMetadata: DocumentMetadata = { + time: Math.floor(Date.now() / 1000), // Unix timestamp + kind: mimeType, + title, + comments, user: this.api.user, - }, - 60000, // 60 second timeout for potentially large lists - ) - .then((r) => r["document-metadatas"] ?? r.documents ?? []); - } + tags, + "document-type": "source", + documentType: "source", + }; + if (id !== undefined) { + documentMetadata.id = id; + } + if (metadata !== undefined) { + documentMetadata.metadata = metadata; + } - /** - * Retrieves list of documents currently being processed - */ - getProcessing() { - return this.api - .makeRequest( - "librarian", - { - operation: "list-processing", - user: this.api.user, - }, - 60000, - ) - .then((r) => r["processing-metadatas"] ?? r.processing ?? r["processing-metadata"] ?? []); - } + return this.api.makeRequest( + "librarian", + { + operation: "add-document", + "document-metadata": documentMetadata, + documentMetadata, + content: document, + }, + 30000, // 30 second timeout for document upload + ); + }, - /** - * Retrieves metadata for a single document by ID - * @param documentId - Document URI/ID to fetch - * @returns Document metadata including title, comments, tags, and RDF metadata - */ - getDocumentMetadata(documentId: string): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "get-document-metadata", + + + /** + * Removes a document from the library + */ + removeDocument(id: string, collection?: string) { + return this.api.makeRequest( + "librarian", + { + operation: "remove-document", + "document-id": id, + documentId: id, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + ); + }, + + + + /** + * Adds a document to the processing queue + * @param id - Processing job identifier + * @param doc_id - Document to process + * @param flow - Processing flow to use + * @param collection - Collection to add processed data to + * @param tags - Tags for the processing job + */ + addProcessing( + id: string, + doc_id: string, + flow: string, + collection?: string, + tags?: string[], + ) { + return this.api.makeRequest( + "librarian", + { + operation: "add-processing", + "processing-metadata": { + id: id, + "document-id": doc_id, + documentId: doc_id, + time: Math.floor(Date.now() / 1000), + flow: flow, + user: this.api.user, + collection: withDefault(collection, "default"), + tags: tags ?? [], + }, + }, + 30000, + ); + }, + + + + // ========== Chunked Upload API ========== + + /** + * Initialize a chunked upload session for large documents (>2MB) + * @param metadata - Document metadata including id, title, kind (MIME type), etc. + * @param totalSize - Total size of the document in bytes + * @param chunkSize - Optional chunk size (default: 5MB) + * @returns Upload session info including upload-id and total-chunks + */ + beginUpload( + metadata: ChunkedUploadDocumentMetadata, + totalSize: number, + chunkSize?: number, + ): Promise { + const request: BeginUploadRequest = { + operation: "begin-upload", + "document-metadata": metadata, + documentMetadata: metadata, + "total-size": totalSize, + }; + if (chunkSize !== undefined) { + request["chunk-size"] = chunkSize; + } + + return this.api + .makeRequest( + "librarian", + request, + 30000, + ) + .then((r) => { + throwIfResponseError(r.error); + return r; + }); + }, + + + + /** + * Upload a single chunk of a document + * Chunks can be uploaded in any order and in parallel + * @param uploadId - Upload session ID from beginUpload + * @param chunkIndex - Zero-based chunk index + * @param content - Base64-encoded chunk content + * @returns Progress info including chunks-received and bytes-received + */ + uploadChunk( + uploadId: string, + chunkIndex: number, + content: string, + ): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "upload-chunk", + "upload-id": uploadId, + "chunk-index": chunkIndex, + content: content, + user: this.api.user, + }, + 60000, // Longer timeout for chunk uploads + ) + .then((r) => { + throwIfResponseError(r.error); + return r; + }); + }, + + + + /** + * Finalize a chunked upload after all chunks are received + * Triggers document processing + * @param uploadId - Upload session ID from beginUpload + * @returns Document ID and object ID + */ + completeUpload(uploadId: string): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "complete-upload", + "upload-id": uploadId, + user: this.api.user, + }, + 30000, + ) + .then((r) => { + throwIfResponseError(r.error); + return r; + }); + }, + + + + /** + * Check upload progress (useful for resuming interrupted uploads) + * @param uploadId - Upload session ID + * @returns Status including received/missing chunks + */ + getUploadStatus(uploadId: string): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "get-upload-status", + "upload-id": uploadId, + user: this.api.user, + }, + 30000, + ) + .then((r) => { + throwIfResponseError(r.error); + return r; + }); + }, + + + + /** + * Cancel an in-progress upload and clean up + * @param uploadId - Upload session ID to abort + */ + abortUpload(uploadId: string): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "abort-upload", + "upload-id": uploadId, + user: this.api.user, + }, + 30000, + ) + .then((r) => { + throwIfResponseError(r.error); + }); + }, + + + + /** + * List pending upload sessions for the current user + * @returns Array of upload sessions with metadata and progress + */ + listUploads(): Promise { + return this.api + .makeRequest( + "librarian", + { + operation: "list-uploads", + user: this.api.user, + }, + 30000, + ) + .then((r) => { + throwIfResponseError(r.error); + return r["upload-sessions"] ?? []; + }); + }, + + + + /** + * Stream a document in chunks for retrieval (streaming response) + * Sends one request, receives multiple chunk responses via callback + * @param documentId - Document ID to retrieve + * @param onChunk - Callback for each chunk: (content, chunkIndex, totalChunks, complete) => void + * @param onError - Callback for errors + * @param chunkSize - Optional chunk size (default: 1MB) + */ + streamDocument( + documentId: string, + onChunk: (content: string, chunkIndex: number, totalChunks: number, complete: boolean) => void, + onError: (error: string) => void, + chunkSize?: number, + ): void { + const receiver = (message: unknown): boolean => { + const msg = message as { response?: StreamDocumentResponse; complete?: boolean; error?: string }; + + // Check for top-level error + if (msg.error !== undefined) { + onError(msg.error); + return true; + } + + const resp = msg.response; + if (resp === undefined) { + return msg.complete === true; + } + + // Check for response-level error + if (resp.error !== undefined) { + onError(resp.error.message); + return true; + } + + const complete = msg.complete === true; + onChunk(resp.content, resp["chunk-index"], resp["total-chunks"], complete); + + return complete; + }; + + const request: StreamDocumentRequest = { + operation: "stream-document", "document-id": documentId, - documentId, user: this.api.user, - }, - 30000, - ) - .then((r) => r["document-metadata"] ?? r.documentMetadata ?? null); - } + }; + if (chunkSize !== undefined) { + request["chunk-size"] = chunkSize; + } - /** - * Uploads a document to the library with full metadata - * @param document - Base64-encoded document content - * @param id - Optional document identifier - * @param metadata - Optional metadata as triples - * @param mimeType - Document MIME type - * @param title - Document title - * @param comments - Additional comments - * @param tags - Document tags for categorization - */ - loadDocument( - document: string, // base64-encoded doc - mimeType: string, - title: string, - comments: string, - tags: string[], - id?: string, - metadata?: Triple[], - ) { - const documentMetadata: DocumentMetadata = { - time: Math.floor(Date.now() / 1000), // Unix timestamp - kind: mimeType, - title, - comments, - user: this.api.user, - tags, - "document-type": "source", - documentType: "source", - }; - if (id !== undefined) { - documentMetadata.id = id; - } - if (metadata !== undefined) { - documentMetadata.metadata = metadata; - } - - return this.api.makeRequest( - "librarian", - { - operation: "add-document", - "document-metadata": documentMetadata, - documentMetadata, - content: document, + this.api.makeRequestMulti( + "librarian", + request, + receiver, + 300000, // 5 minute timeout for full document stream + ); }, - 30000, // 30 second timeout for document upload - ); - } - - /** - * Removes a document from the library - */ - removeDocument(id: string, collection?: string) { - return this.api.makeRequest( - "librarian", - { - operation: "remove-document", - "document-id": id, - documentId: id, - user: this.api.user, - collection: withDefault(collection, "default"), - }, - 30000, - ); - } - - /** - * Adds a document to the processing queue - * @param id - Processing job identifier - * @param doc_id - Document to process - * @param flow - Processing flow to use - * @param collection - Collection to add processed data to - * @param tags - Tags for the processing job - */ - addProcessing( - id: string, - doc_id: string, - flow: string, - collection?: string, - tags?: string[], - ) { - return this.api.makeRequest( - "librarian", - { - operation: "add-processing", - "processing-metadata": { - id: id, - "document-id": doc_id, - documentId: doc_id, - time: Math.floor(Date.now() / 1000), - flow: flow, - user: this.api.user, - collection: withDefault(collection, "default"), - tags: tags ?? [], - }, - }, - 30000, - ); - } - - // ========== Chunked Upload API ========== - - /** - * Initialize a chunked upload session for large documents (>2MB) - * @param metadata - Document metadata including id, title, kind (MIME type), etc. - * @param totalSize - Total size of the document in bytes - * @param chunkSize - Optional chunk size (default: 5MB) - * @returns Upload session info including upload-id and total-chunks - */ - beginUpload( - metadata: ChunkedUploadDocumentMetadata, - totalSize: number, - chunkSize?: number, - ): Promise { - const request: BeginUploadRequest = { - operation: "begin-upload", - "document-metadata": metadata, - documentMetadata: metadata, - "total-size": totalSize, - }; - if (chunkSize !== undefined) { - request["chunk-size"] = chunkSize; - } - - return this.api - .makeRequest( - "librarian", - request, - 30000, - ) - .then((r) => { - throwIfResponseError(r.error); - return r; - }); - } - - /** - * Upload a single chunk of a document - * Chunks can be uploaded in any order and in parallel - * @param uploadId - Upload session ID from beginUpload - * @param chunkIndex - Zero-based chunk index - * @param content - Base64-encoded chunk content - * @returns Progress info including chunks-received and bytes-received - */ - uploadChunk( - uploadId: string, - chunkIndex: number, - content: string, - ): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "upload-chunk", - "upload-id": uploadId, - "chunk-index": chunkIndex, - content: content, - user: this.api.user, - }, - 60000, // Longer timeout for chunk uploads - ) - .then((r) => { - throwIfResponseError(r.error); - return r; - }); - } - - /** - * Finalize a chunked upload after all chunks are received - * Triggers document processing - * @param uploadId - Upload session ID from beginUpload - * @returns Document ID and object ID - */ - completeUpload(uploadId: string): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "complete-upload", - "upload-id": uploadId, - user: this.api.user, - }, - 30000, - ) - .then((r) => { - throwIfResponseError(r.error); - return r; - }); - } - - /** - * Check upload progress (useful for resuming interrupted uploads) - * @param uploadId - Upload session ID - * @returns Status including received/missing chunks - */ - getUploadStatus(uploadId: string): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "get-upload-status", - "upload-id": uploadId, - user: this.api.user, - }, - 30000, - ) - .then((r) => { - throwIfResponseError(r.error); - return r; - }); - } - - /** - * Cancel an in-progress upload and clean up - * @param uploadId - Upload session ID to abort - */ - abortUpload(uploadId: string): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "abort-upload", - "upload-id": uploadId, - user: this.api.user, - }, - 30000, - ) - .then((r) => { - throwIfResponseError(r.error); - }); - } - - /** - * List pending upload sessions for the current user - * @returns Array of upload sessions with metadata and progress - */ - listUploads(): Promise { - return this.api - .makeRequest( - "librarian", - { - operation: "list-uploads", - user: this.api.user, - }, - 30000, - ) - .then((r) => { - throwIfResponseError(r.error); - return r["upload-sessions"] ?? []; - }); - } - - /** - * Stream a document in chunks for retrieval (streaming response) - * Sends one request, receives multiple chunk responses via callback - * @param documentId - Document ID to retrieve - * @param onChunk - Callback for each chunk: (content, chunkIndex, totalChunks, complete) => void - * @param onError - Callback for errors - * @param chunkSize - Optional chunk size (default: 1MB) - */ - streamDocument( - documentId: string, - onChunk: (content: string, chunkIndex: number, totalChunks: number, complete: boolean) => void, - onError: (error: string) => void, - chunkSize?: number, - ): void { - const receiver = (message: unknown): boolean => { - const msg = message as { response?: StreamDocumentResponse; complete?: boolean; error?: string }; - - // Check for top-level error - if (msg.error !== undefined) { - onError(msg.error); - return true; - } - - const resp = msg.response; - if (resp === undefined) { - return msg.complete === true; - } - - // Check for response-level error - if (resp.error !== undefined) { - onError(resp.error.message); - return true; - } - - const complete = msg.complete === true; - onChunk(resp.content, resp["chunk-index"], resp["total-chunks"], complete); - - return complete; - }; - - const request: StreamDocumentRequest = { - operation: "stream-document", - "document-id": documentId, - user: this.api.user, - }; - if (chunkSize !== undefined) { - request["chunk-size"] = chunkSize; - } - - this.api.makeRequestMulti( - "librarian", - request, - receiver, - 300000, // 5 minute timeout for full document stream - ); - } + }; } +export type LibrarianApi = ReturnType; +export const LibrarianApi = newableFactory(makeLibrarianApi); + /** * FlowsApi - Manages processing flows and configuration * Flows define how documents and data are processed through the system */ -export class FlowsApi { - api: BaseApi; +export function makeFlowsApi(api: BaseApi) { + return { + api, - constructor(api: BaseApi) { - this.api = api; - } - /** - * Retrieves list of available flows - */ - getFlows() { - return this.api - .makeRequest( - "flow", - { - operation: "list-flows", - }, - 60000, - ) - .then((r) => r["flow-ids"] ?? []); - } - /** - * Retrieves definition of a specific flow - */ - getFlow(id: string) { - return this.api - .makeRequest( - "flow", - { - operation: "get-flow", + /** + * Retrieves list of available flows + */ + getFlows() { + return this.api + .makeRequest( + "flow", + { + operation: "list-flows", + }, + 60000, + ) + .then((r) => r["flow-ids"] ?? []); + }, + + + + /** + * Retrieves definition of a specific flow + */ + getFlow(id: string) { + return this.api + .makeRequest( + "flow", + { + operation: "get-flow", + "flow-id": id, + }, + 60000, + ) + .then((r) => JSON.parse(r.flow ?? "{}")); // Parse JSON flow definition + }, + + + + // Configuration management methods + + /** + * Retrieves all configuration settings + */ + getConfigAll() { + return this.api.makeRequest( + "config", + { + operation: "config", + }, + 60000, + ); + }, + + + + /** + * Retrieves specific configuration values by key + */ + getConfig(keys: { type: string; key: string }[]) { + return this.api.makeRequest( + "config", + { + operation: "get", + keys: keys, + }, + 60000, + ); + }, + + + + /** + * Updates configuration values using the Python-compatible values array. + */ + putConfig(items: { type: string; key: string; value: string }[]) { + return this.api.makeRequest( + "config", + { + operation: "put", + values: items, + }, + 60000, + ); + }, + + + + /** + * Deletes a configuration entry + */ + deleteConfig(target: { type: string; key: string }) { + return this.api.makeRequest( + "config", + { + operation: "delete", + keys: [target], + }, + 30000, + ); + }, + + + + // Prompt management - specialized config operations for AI prompts + + /** + * Retrieves list of available prompt templates from config.prompt. + * Each template is stored at `config.prompt.` as an object + * `{system, prompt}`. The reserved key `system` holds an optional + * global system prompt and is excluded from the template list. + */ + getPrompts() { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: Record } }; + const promptNs = config.config?.prompt ?? {}; + return Object.keys(promptNs) + .filter((k) => k !== "system") + .sort() + .map((id) => ({ id, name: id })); + }); + }, + + + + /** + * Retrieves a specific prompt template object: `{system, prompt}`. + */ + getPrompt(id: string) { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: Record } }; + return config.config?.prompt?.[id] ?? null; + }); + }, + + + + /** + * Retrieves the optional global system prompt at `config.prompt.system`. + * Returns "" if not configured. + */ + getSystemPrompt() { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: { system?: unknown } } }; + const raw = config.config?.prompt?.system; + if (raw == null) return ""; + return typeof raw === "string" ? raw : raw; + }); + }, + + + + // Flow blueprint management - templates for creating flows + + /** + * Retrieves list of available flow blueprints (templates) + */ + getFlowBlueprints() { + return this.api + .makeRequest( + "flow", + { + operation: "list-blueprints", + }, + 60000, + ) + .then((r) => r["blueprint-names"]); + }, + + + + /** + * Retrieves definition of a specific flow blueprint + */ + getFlowBlueprint(name: string) { + return this.api + .makeRequest( + "flow", + { + operation: "get-blueprint", + "blueprint-name": name, + }, + 60000, + ) + .then((r) => JSON.parse(r["blueprint-definition"] ?? "{}")); + }, + + + + /** + * Deletes a flow blueprint + */ + deleteFlowBlueprint(name: string) { + return this.api.makeRequest( + "flow", + { + operation: "delete-blueprint", + "blueprint-name": name, + }, + 30000, + ); + }, + + + + // Flow lifecycle management + + /** + * Starts a new flow instance + */ + startFlow( + id: string, + blueprint_name: string, + description: string, + parameters?: Record, + ) { + const request: FlowRequest = { + operation: "start-flow", "flow-id": id, - }, - 60000, - ) - .then((r) => JSON.parse(r.flow ?? "{}")); // Parse JSON flow definition - } + "blueprint-name": blueprint_name, + description: description, + }; - // Configuration management methods - - /** - * Retrieves all configuration settings - */ - getConfigAll() { - return this.api.makeRequest( - "config", - { - operation: "config", - }, - 60000, - ); - } - - /** - * Retrieves specific configuration values by key - */ - getConfig(keys: { type: string; key: string }[]) { - return this.api.makeRequest( - "config", - { - operation: "get", - keys: keys, - }, - 60000, - ); - } - - /** - * Updates configuration values using the Python-compatible values array. - */ - putConfig(items: { type: string; key: string; value: string }[]) { - return this.api.makeRequest( - "config", - { - operation: "put", - values: items, - }, - 60000, - ); - } - - /** - * Deletes a configuration entry - */ - deleteConfig(target: { type: string; key: string }) { - return this.api.makeRequest( - "config", - { - operation: "delete", - keys: [target], - }, - 30000, - ); - } - - // Prompt management - specialized config operations for AI prompts - - /** - * Retrieves list of available prompt templates from config.prompt. - * Each template is stored at `config.prompt.` as an object - * `{system, prompt}`. The reserved key `system` holds an optional - * global system prompt and is excluded from the template list. - */ - getPrompts() { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: Record } }; - const promptNs = config.config?.prompt ?? {}; - return Object.keys(promptNs) - .filter((k) => k !== "system") - .sort() - .map((id) => ({ id, name: id })); - }); - } - - /** - * Retrieves a specific prompt template object: `{system, prompt}`. - */ - getPrompt(id: string) { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: Record } }; - return config.config?.prompt?.[id] ?? null; - }); - } - - /** - * Retrieves the optional global system prompt at `config.prompt.system`. - * Returns "" if not configured. - */ - getSystemPrompt() { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: { system?: unknown } } }; - const raw = config.config?.prompt?.system; - if (raw == null) return ""; - return typeof raw === "string" ? raw : raw; - }); - } - - // Flow blueprint management - templates for creating flows - - /** - * Retrieves list of available flow blueprints (templates) - */ - getFlowBlueprints() { - return this.api - .makeRequest( - "flow", - { - operation: "list-blueprints", - }, - 60000, - ) - .then((r) => r["blueprint-names"]); - } - - /** - * Retrieves definition of a specific flow blueprint - */ - getFlowBlueprint(name: string) { - return this.api - .makeRequest( - "flow", - { - operation: "get-blueprint", - "blueprint-name": name, - }, - 60000, - ) - .then((r) => JSON.parse(r["blueprint-definition"] ?? "{}")); - } - - /** - * Deletes a flow blueprint - */ - deleteFlowBlueprint(name: string) { - return this.api.makeRequest( - "flow", - { - operation: "delete-blueprint", - "blueprint-name": name, - }, - 30000, - ); - } - - // Flow lifecycle management - - /** - * Starts a new flow instance - */ - startFlow( - id: string, - blueprint_name: string, - description: string, - parameters?: Record, - ) { - const request: FlowRequest = { - operation: "start-flow", - "flow-id": id, - "blueprint-name": blueprint_name, - description: description, - }; - - // Only include parameters if provided and not empty - if (parameters !== undefined && Object.keys(parameters).length > 0) { - request.parameters = parameters; - } - - return this.api - .makeRequest("flow", request, 30000) - .then((response) => { - if (response.error !== undefined) { - throw new Error(toErrorMessage(response.error, "Flow start failed")); + // Only include parameters if provided and not empty + if (parameters !== undefined && Object.keys(parameters).length > 0) { + request.parameters = parameters; } - return response; - }); - } - /** - * Stops a running flow instance - */ - stopFlow(id: string) { - return this.api.makeRequest( - "flow", - { - operation: "stop-flow", - "flow-id": id, + return this.api + .makeRequest("flow", request, 30000) + .then((response) => { + if (response.error !== undefined) { + throw new Error(toErrorMessage(response.error, "Flow start failed")); + } + return response; + }); }, - 30000, - ); - } + + + + /** + * Stops a running flow instance + */ + stopFlow(id: string) { + return this.api.makeRequest( + "flow", + { + operation: "stop-flow", + "flow-id": id, + }, + 30000, + ); + }, + }; } +export type FlowsApi = ReturnType; +export const FlowsApi = newableFactory(makeFlowsApi); + /** * FlowApi - Interface for interacting with a specific flow instance * Provides flow-specific versions of core AI/ML operations */ -export class FlowApi { - api: BaseApi; - flowId: string; +export function makeFlowApi(api: BaseApi, flowId: string) { + return { + api, - constructor(api: BaseApi, flowId: string) { - this.api = api; - this.flowId = flowId; // All requests will be routed through this flow - } + flowId, - /** - * Performs text completion using AI models within this flow - */ - textCompletion(system: string, text: string): Promise { - return this.api - .makeRequest( - "text-completion", - { - system: system, // System prompt/instructions - prompt: text, // User prompt - }, - 30000, - undefined, // Use default retries - this.flowId, // Route through this flow - ) - .then((r) => r.response); - } - /** - * Performs Graph RAG (Retrieval Augmented Generation) query - */ - graphRag(text: string, options?: GraphRagOptions, collection?: string) { - const request: GraphRagRequest = { - query: text, - user: this.api.user, - collection: withDefault(collection, "default"), - }; - if (options?.entityLimit !== undefined) { - request["entity-limit"] = options.entityLimit; - } - if (options?.tripleLimit !== undefined) { - request["triple-limit"] = options.tripleLimit; - } - if (options?.maxSubgraphSize !== undefined) { - request["max-subgraph-size"] = options.maxSubgraphSize; - } - if (options?.pathLength !== undefined) { - request["max-path-length"] = options.pathLength; - } - return this.api - .makeRequest( - "graph-rag", - request, - 60000, // Longer timeout for complex graph operations - undefined, - this.flowId, - ) - .then((r) => r.response); - } + /** + * Performs text completion using AI models within this flow + */ + textCompletion(system: string, text: string): Promise { + return this.api + .makeRequest( + "text-completion", + { + system: system, // System prompt/instructions + prompt: text, // User prompt + }, + 30000, + undefined, // Use default retries + this.flowId, // Route through this flow + ) + .then((r) => r.response); + }, - /** - * Performs Document RAG (Retrieval Augmented Generation) query - */ - documentRag(text: string, docLimit?: number, collection?: string) { - return this.api - .makeRequest( - "document-rag", - { + + + /** + * Performs Graph RAG (Retrieval Augmented Generation) query + */ + graphRag(text: string, options?: GraphRagOptions, collection?: string) { + const request: GraphRagRequest = { query: text, user: this.api.user, collection: withDefault(collection, "default"), - "doc-limit": docLimit ?? 20, - }, - 60000, // Longer timeout for document operations - undefined, - this.flowId, - ) - .then((r) => r.response); - } - - /** - * Interacts with an AI agent that provides streaming responses - * BREAKING CHANGE: Callbacks now receive (chunk, complete, metadata?) instead of full messages - */ - agent( - question: string, - think: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - observe: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - answer: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - error: (s: string) => void, - onExplain?: (event: ExplainEvent) => void, - collection?: string, - ) { - const receiver = (message: unknown) => { - const msg = message as { response?: AgentResponse; complete?: boolean; error?: string }; - - // Check for top-level error - if (msg.error !== undefined) { - error(msg.error); - return true; - } - - const resp = msg.response ?? {}; - - // Check for errors in response - if (resp.chunk_type === "error" || resp.error !== undefined) { - error(resp.error?.message ?? "Unknown agent error"); - return true; // End streaming on error - } - - // Handle explainability events (agent uses chunk_type="explain") - if ( - (resp.chunk_type === "explain" || resp.message_type === "explain") && - (resp.explain_id !== undefined || resp.explain_triples !== undefined) - ) { - const event: ExplainEvent = { - explainId: resp.explain_id ?? "", - explainGraph: resp.explain_graph ?? "", }; - if (resp.explain_triples !== undefined) { - event.explainTriples = resp.explain_triples as Triple[]; + if (options?.entityLimit !== undefined) { + request["entity-limit"] = options.entityLimit; + } + if (options?.tripleLimit !== undefined) { + request["triple-limit"] = options.tripleLimit; + } + if (options?.maxSubgraphSize !== undefined) { + request["max-subgraph-size"] = options.maxSubgraphSize; + } + if (options?.pathLength !== undefined) { + request["max-path-length"] = options.pathLength; } - onExplain?.(event); - return false; - } - // Handle streaming chunks by chunk_type - const content = resp.content ?? ""; - const messageComplete = resp.end_of_message === true; - const dialogComplete = msg.complete === true || resp.end_of_dialog === true; + return this.api + .makeRequest( + "graph-rag", + request, + 60000, // Longer timeout for complex graph operations + undefined, + this.flowId, + ) + .then((r) => r.response); + }, - // Extract metadata from final message - const metadata = dialogComplete ? streamingMetadataFrom(resp) : undefined; - switch (resp.chunk_type) { - case "thought": - think(content, messageComplete, metadata); - break; - case "observation": - observe(content, messageComplete, metadata); - break; - case "answer": - case "final-answer": - answer(content, messageComplete, metadata); - break; - case "action": - // Actions are typically not streamed incrementally, just logged - console.log("Agent action:", content); - break; - } - return dialogComplete; // End when backend signals complete - }; + /** + * Performs Document RAG (Retrieval Augmented Generation) query + */ + documentRag(text: string, docLimit?: number, collection?: string) { + return this.api + .makeRequest( + "document-rag", + { + query: text, + user: this.api.user, + collection: withDefault(collection, "default"), + "doc-limit": docLimit ?? 20, + }, + 60000, // Longer timeout for document operations + undefined, + this.flowId, + ) + .then((r) => r.response); + }, - return this.api - .makeRequestMulti( - "agent", - { - question: question, + + + /** + * Interacts with an AI agent that provides streaming responses + * BREAKING CHANGE: Callbacks now receive (chunk, complete, metadata?) instead of full messages + */ + agent( + question: string, + think: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + observe: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + answer: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + error: (s: string) => void, + onExplain?: (event: ExplainEvent) => void, + collection?: string, + ) { + const receiver = (message: unknown) => { + const msg = message as { response?: AgentResponse; complete?: boolean; error?: string }; + + // Check for top-level error + if (msg.error !== undefined) { + error(msg.error); + return true; + } + + const resp = msg.response ?? {}; + + // Check for errors in response + if (resp.chunk_type === "error" || resp.error !== undefined) { + error(resp.error?.message ?? "Unknown agent error"); + return true; // End streaming on error + } + + // Handle explainability events (agent uses chunk_type="explain") + if ( + (resp.chunk_type === "explain" || resp.message_type === "explain") && + (resp.explain_id !== undefined || resp.explain_triples !== undefined) + ) { + const event: ExplainEvent = { + explainId: resp.explain_id ?? "", + explainGraph: resp.explain_graph ?? "", + }; + if (resp.explain_triples !== undefined) { + event.explainTriples = resp.explain_triples as Triple[]; + } + onExplain?.(event); + return false; + } + + // Handle streaming chunks by chunk_type + const content = resp.content ?? ""; + const messageComplete = resp.end_of_message === true; + const dialogComplete = msg.complete === true || resp.end_of_dialog === true; + + // Extract metadata from final message + const metadata = dialogComplete ? streamingMetadataFrom(resp) : undefined; + + switch (resp.chunk_type) { + case "thought": + think(content, messageComplete, metadata); + break; + case "observation": + observe(content, messageComplete, metadata); + break; + case "answer": + case "final-answer": + answer(content, messageComplete, metadata); + break; + case "action": + // Actions are typically not streamed incrementally, just logged + console.log("Agent action:", content); + break; + } + + return dialogComplete; // End when backend signals complete + }; + + return this.api + .makeRequestMulti( + "agent", + { + question: question, + user: this.api.user, + collection: withDefault(collection, "default"), + streaming: true, // Always use streaming mode + }, + receiver, + 120000, + 2, + this.flowId, + ) + .catch((err) => { + const errorMessage = toErrorMessage(err, "Unknown error"); + error(`Agent request failed: ${errorMessage}`); + }); + }, + + + + /** + * Performs Graph RAG query with streaming response + * @param text - Query text + * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk + * @param onError - Called on error + * @param options - Graph RAG options (including explainable flag) + * @param collection - Collection name + * @param onExplain - Optional callback for explainability events + */ + graphRagStreaming( + text: string, + receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + onError: (error: string) => void, + options?: GraphRagOptions, + collection?: string, + onExplain?: (event: ExplainEvent) => void, + ): void { + const recv = (message: unknown): boolean => { + const msg = message as { response?: GraphRagResponse; complete?: boolean; error?: string }; + + // Check for top-level error + if (msg.error !== undefined) { + onError(msg.error); + return true; + } + + const resp = (msg.response ?? {}) as GraphRagResponse; + + // Check for response-level error + if (resp.error !== undefined) { + onError(resp.error.message); + return true; + } + + // Extract explain data if present (may be embedded in the answer message) + if ( + resp.message_type === "explain" && + (resp.explain_id !== undefined || resp.explain_triples !== undefined) + ) { + const event: ExplainEvent = { + explainId: resp.explain_id ?? "", + explainGraph: resp.explain_graph ?? "", + }; + if (resp.explain_triples !== undefined) { + event.explainTriples = resp.explain_triples as Triple[]; + } + onExplain?.(event); + // If this message also carries answer text, fall through to chunk handling. + // If it's a standalone explain event (no answer text), stop here. + if (resp.response === undefined && resp.endOfStream !== true && resp.end_of_session !== true) { + return false; + } + } + + // Handle chunk messages (default behavior) + const chunk = resp.response ?? resp.chunk ?? ""; + const complete = resp.end_of_session === true || resp.endOfStream === true || msg.complete === true; + + // Extract metadata from final message + const metadata = complete ? streamingMetadataFrom(resp) : undefined; + + receiver(chunk, complete, metadata); + + return complete; + }; + + const request: GraphRagRequest = { + query: text, user: this.api.user, - collection: withDefault(collection, "default"), - streaming: true, // Always use streaming mode - }, - receiver, - 120000, - 2, - this.flowId, - ) - .catch((err) => { - const errorMessage = toErrorMessage(err, "Unknown error"); - error(`Agent request failed: ${errorMessage}`); - }); - } - - /** - * Performs Graph RAG query with streaming response - * @param text - Query text - * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk - * @param onError - Called on error - * @param options - Graph RAG options (including explainable flag) - * @param collection - Collection name - * @param onExplain - Optional callback for explainability events - */ - graphRagStreaming( - text: string, - receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - onError: (error: string) => void, - options?: GraphRagOptions, - collection?: string, - onExplain?: (event: ExplainEvent) => void, - ): void { - const recv = (message: unknown): boolean => { - const msg = message as { response?: GraphRagResponse; complete?: boolean; error?: string }; - - // Check for top-level error - if (msg.error !== undefined) { - onError(msg.error); - return true; - } - - const resp = (msg.response ?? {}) as GraphRagResponse; - - // Check for response-level error - if (resp.error !== undefined) { - onError(resp.error.message); - return true; - } - - // Extract explain data if present (may be embedded in the answer message) - if ( - resp.message_type === "explain" && - (resp.explain_id !== undefined || resp.explain_triples !== undefined) - ) { - const event: ExplainEvent = { - explainId: resp.explain_id ?? "", - explainGraph: resp.explain_graph ?? "", + collection: withDefault(collection, "default"), + streaming: true, }; - if (resp.explain_triples !== undefined) { - event.explainTriples = resp.explain_triples as Triple[]; + if (options?.entityLimit !== undefined) { + request["entity-limit"] = options.entityLimit; } - onExplain?.(event); - // If this message also carries answer text, fall through to chunk handling. - // If it's a standalone explain event (no answer text), stop here. - if (resp.response === undefined && resp.endOfStream !== true && resp.end_of_session !== true) { - return false; + if (options?.tripleLimit !== undefined) { + request["triple-limit"] = options.tripleLimit; + } + if (options?.maxSubgraphSize !== undefined) { + request["max-subgraph-size"] = options.maxSubgraphSize; + } + if (options?.pathLength !== undefined) { + request["max-path-length"] = options.pathLength; } - } - // Handle chunk messages (default behavior) - const chunk = resp.response ?? resp.chunk ?? ""; - const complete = resp.end_of_session === true || resp.endOfStream === true || msg.complete === true; - - // Extract metadata from final message - const metadata = complete ? streamingMetadataFrom(resp) : undefined; - - receiver(chunk, complete, metadata); - - return complete; - }; - - const request: GraphRagRequest = { - query: text, - user: this.api.user, - collection: withDefault(collection, "default"), - streaming: true, - }; - if (options?.entityLimit !== undefined) { - request["entity-limit"] = options.entityLimit; - } - if (options?.tripleLimit !== undefined) { - request["triple-limit"] = options.tripleLimit; - } - if (options?.maxSubgraphSize !== undefined) { - request["max-subgraph-size"] = options.maxSubgraphSize; - } - if (options?.pathLength !== undefined) { - request["max-path-length"] = options.pathLength; - } - - this.api.makeRequestMulti( - "graph-rag", - request, - recv, - 60000, - undefined, - this.flowId, - ).catch((err) => { - const errorMessage = toErrorMessage(err, "Unknown error"); - onError(`Graph RAG request failed: ${errorMessage}`); - }); - } - - /** - * Performs Document RAG query with streaming response - * @param text - Query text - * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk - * @param onError - Called on error - * @param docLimit - Maximum documents to retrieve - * @param collection - Collection name - */ - documentRagStreaming( - text: string, - receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - onError: (error: string) => void, - docLimit?: number, - collection?: string, - onExplain?: (event: ExplainEvent) => void, - ): void { - const recv = (message: unknown): boolean => { - const msg = message as { response?: DocumentRagResponse; complete?: boolean; error?: string }; - - // Check for top-level error - if (msg.error !== undefined) { - onError(msg.error); - return true; - } - - const resp = (msg.response ?? {}) as DocumentRagResponse; - - // Check for response-level error - if (resp.error !== undefined) { - onError(resp.error.message); - return true; - } - - // Handle explainability events - if ( - resp.message_type === "explain" && - resp.explain_id !== undefined && - resp.explain_graph !== undefined - ) { - onExplain?.({ - explainId: resp.explain_id, - explainGraph: resp.explain_graph, + this.api.makeRequestMulti( + "graph-rag", + request, + recv, + 60000, + undefined, + this.flowId, + ).catch((err) => { + const errorMessage = toErrorMessage(err, "Unknown error"); + onError(`Graph RAG request failed: ${errorMessage}`); }); - return false; - } - - const chunk = resp.response ?? resp.chunk ?? ""; - const complete = resp.end_of_session === true || resp.endOfStream === true || msg.complete === true; - - // Extract metadata from final message - const metadata = complete ? streamingMetadataFrom(resp) : undefined; - - receiver(chunk, complete, metadata); - - return complete; - }; - - const request: DocumentRagRequest = { - query: text, - user: this.api.user, - collection: withDefault(collection, "default"), - streaming: true, - }; - if (docLimit !== undefined) { - request["doc-limit"] = docLimit; - } - - this.api.makeRequestMulti( - "document-rag", - request, - recv, - 60000, - undefined, - this.flowId, - ).catch((err) => { - const errorMessage = toErrorMessage(err, "Unknown error"); - onError(`Document RAG request failed: ${errorMessage}`); - }); - } - - /** - * Performs text completion with streaming response - * @param system - System prompt - * @param text - User prompt - * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk - * @param onError - Called on error - */ - textCompletionStreaming( - system: string, - text: string, - receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - onError: (error: string) => void, - ): void { - const recv = (message: unknown): boolean => { - const msg = message as { response?: TextCompletionResponse; complete?: boolean; error?: string }; - - // Check for top-level error - if (msg.error !== undefined) { - onError(msg.error); - return true; - } - - const resp = (msg.response ?? {}) as TextCompletionResponse; - - // Check for response-level error - if (resp.error !== undefined) { - onError(resp.error.message); - return true; - } - - // Text completion uses 'response' field for chunks - const chunk = resp.response ?? ""; - const complete = msg.complete === true; - - // Extract metadata from final message - const metadata = complete ? streamingMetadataFrom(resp) : undefined; - - receiver(chunk, complete, metadata); - - return complete; - }; - - this.api.makeRequestMulti( - "text-completion", - { - system: system, - prompt: text, - streaming: true, }, - recv, - 30000, - undefined, - this.flowId, - ); - } - /** - * Executes a prompt template with streaming response - * @param id - Prompt template ID - * @param terms - Template variables - * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk - * @param onError - Called on error - */ - promptStreaming( - id: string, - terms: Record, - receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, - onError: (error: string) => void, - ): void { - const recv = (message: unknown): boolean => { - const msg = message as { response?: PromptResponse; complete?: boolean; error?: string }; - // Check for top-level error - if (msg.error !== undefined) { - onError(msg.error); - return true; - } - const resp = (msg.response ?? {}) as PromptResponse; + /** + * Performs Document RAG query with streaming response + * @param text - Query text + * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk + * @param onError - Called on error + * @param docLimit - Maximum documents to retrieve + * @param collection - Collection name + */ + documentRagStreaming( + text: string, + receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + onError: (error: string) => void, + docLimit?: number, + collection?: string, + onExplain?: (event: ExplainEvent) => void, + ): void { + const recv = (message: unknown): boolean => { + const msg = message as { response?: DocumentRagResponse; complete?: boolean; error?: string }; - // Check for response-level error - if (resp.error !== undefined) { - onError(resp.error.message); - return true; - } + // Check for top-level error + if (msg.error !== undefined) { + onError(msg.error); + return true; + } - // Prompt service uses 'text' field for chunks - const chunk = resp.text ?? ""; - const complete = msg.complete === true; + const resp = (msg.response ?? {}) as DocumentRagResponse; - // Extract metadata from final message - const metadata = complete ? streamingMetadataFrom(resp) : undefined; + // Check for response-level error + if (resp.error !== undefined) { + onError(resp.error.message); + return true; + } - receiver(chunk, complete, metadata); + // Handle explainability events + if ( + resp.message_type === "explain" && + resp.explain_id !== undefined && + resp.explain_graph !== undefined + ) { + onExplain?.({ + explainId: resp.explain_id, + explainGraph: resp.explain_graph, + }); + return false; + } - return complete; - }; + const chunk = resp.response ?? resp.chunk ?? ""; + const complete = resp.end_of_session === true || resp.endOfStream === true || msg.complete === true; - this.api.makeRequestMulti( - "prompt", - { - id: id, - terms: terms, - streaming: true, - }, - recv, - 30000, - undefined, - this.flowId, - ); - } + // Extract metadata from final message + const metadata = complete ? streamingMetadataFrom(resp) : undefined; - /** - * Generates embeddings for multiple texts within this flow. - * Returns vectors[text_index][dimension_index] - one vector per input text. - */ - embeddings(texts: string[]) { - return this.api - .makeRequest( - "embeddings", - { - texts: texts, - }, - 30000, - undefined, - this.flowId, - ) - .then((r) => r.vectors); - } + receiver(chunk, complete, metadata); - /** - * Queries the knowledge graph using a single embedding vector - */ - graphEmbeddingsQuery( - vec: number[], - limit: number | undefined, - collection?: string, - ) { - return this.api - .makeRequest( - "graph-embeddings", - { - vector: vec, - limit: limit ?? 20, // Default to 20 results + return complete; + }; + + const request: DocumentRagRequest = { + query: text, user: this.api.user, collection: withDefault(collection, "default"), - }, - 30000, - undefined, - this.flowId, - ) - .then((r) => r.entities); - } - - /** - * Queries knowledge graph triples (subject-predicate-object relationships) - * All parameters are optional - omitted parameters act as wildcards - */ - triplesQuery( - s?: Term, - p?: Term, - o?: Term, - limit?: number, - collection?: string, - graph?: string, - ) { - const request: TriplesQueryRequest = { - limit: limit ?? 20, - user: this.api.user, - collection: withDefault(collection, "default"), - }; - if (s !== undefined) { - request.s = s; - } - if (p !== undefined) { - request.p = p; - } - if (o !== undefined) { - request.o = o; - } - if (graph !== undefined) { - request.g = graph; - } - - return this.api - .makeRequest( - "triples", - request, - 30000, - undefined, - this.flowId, - ) - .then((r) => r.triples ?? r.response ?? []); - } - - /** - * Loads a document into this flow for processing - */ - loadDocument( - document: string, // base64-encoded document - id?: string, - metadata?: Triple[], - ) { - const request: LoadDocumentRequest = { - data: document, - }; - if (id !== undefined) { - request.id = id; - } - if (metadata !== undefined) { - request.metadata = metadata; - } - - return this.api.makeRequest( - "document-load", - request, - 30000, - undefined, - this.flowId, - ); - } - - /** - * Loads plain text into this flow for processing - */ - loadText( - text: string, // Text content - id?: string, - metadata?: Triple[], - charset?: string, // Character encoding - ) { - const request: LoadTextRequest = { - text, - }; - if (id !== undefined) { - request.id = id; - } - if (metadata !== undefined) { - request.metadata = metadata; - } - if (charset !== undefined) { - request.charset = charset; - } - - return this.api.makeRequest( - "text-load", - request, - 30000, - undefined, - this.flowId, - ); - } - - /** - * Executes a GraphQL query against structured row data - */ - rowsQuery( - query: string, - collection?: string, - variables?: Record, - operationName?: string, - ) { - const request: RowsQueryRequest = { - query, - user: this.api.user, - collection: withDefault(collection, "default"), - }; - if (variables !== undefined) { - request.variables = variables; - } - if (operationName !== undefined) { - request.operation_name = operationName; - } - - return this.api - .makeRequest( - "rows", - request, - 30000, - undefined, - this.flowId, - ) - .then((r) => { - // Return the GraphQL response structure directly - const result: Record = {}; - if (r.data !== undefined) result.data = r.data; - if (r.errors !== undefined) result.errors = r.errors; - if (r.extensions !== undefined) result.extensions = r.extensions; - return result; - }); - } - - /** - * Converts a natural language question to a GraphQL query - */ - nlpQuery(question: string, maxResults?: number) { - return this.api - .makeRequest( - "nlp-query", - { - question: question, - max_results: maxResults ?? 100, - }, - 30000, - undefined, - this.flowId, - ) - .then((r) => r); - } - - /** - * Executes a natural language question against structured data - * Combines NLP query conversion and GraphQL execution - */ - structuredQuery(question: string, collection?: string) { - return this.api - .makeRequest( - "structured-query", - { - question: question, - user: this.api.user, - collection: withDefault(collection, "default"), - }, - 30000, - undefined, - this.flowId, - ) - .then((r) => { - // Return the response structure directly - const result: Record = {}; - if (r.data !== undefined) result.data = r.data; - if (r.errors !== undefined) result.errors = r.errors; - return result; - }); - } - - /** - * Performs semantic search on structured data indexes using embedding vectors - * @param vectors - Embedding vectors to search for - * @param schemaName - Name of the schema to search - * @param collection - Optional collection name - * @param indexName - Optional index name to filter results - * @param limit - Maximum number of results to return (default: 10) - */ - rowEmbeddingsQuery( - vector: number[], - schemaName: string, - collection?: string, - indexName?: string, - limit?: number, - ): Promise { - const request: RowEmbeddingsQueryRequest = { - vector: vector, - schema_name: schemaName, - user: this.api.user, - collection: withDefault(collection, "default"), - limit: limit ?? 10, - }; - - if (indexName !== undefined) { - request.index_name = indexName; - } - - return this.api - .makeRequest( - "row-embeddings", - request, - 30000, - undefined, - this.flowId, - ) - .then((r) => { - if (r.error !== undefined) { - throw new Error(r.error.message); + streaming: true, + }; + if (docLimit !== undefined) { + request["doc-limit"] = docLimit; } - return r.matches ?? []; - }); - } + + this.api.makeRequestMulti( + "document-rag", + request, + recv, + 60000, + undefined, + this.flowId, + ).catch((err) => { + const errorMessage = toErrorMessage(err, "Unknown error"); + onError(`Document RAG request failed: ${errorMessage}`); + }); + }, + + + + /** + * Performs text completion with streaming response + * @param system - System prompt + * @param text - User prompt + * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk + * @param onError - Called on error + */ + textCompletionStreaming( + system: string, + text: string, + receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + onError: (error: string) => void, + ): void { + const recv = (message: unknown): boolean => { + const msg = message as { response?: TextCompletionResponse; complete?: boolean; error?: string }; + + // Check for top-level error + if (msg.error !== undefined) { + onError(msg.error); + return true; + } + + const resp = (msg.response ?? {}) as TextCompletionResponse; + + // Check for response-level error + if (resp.error !== undefined) { + onError(resp.error.message); + return true; + } + + // Text completion uses 'response' field for chunks + const chunk = resp.response ?? ""; + const complete = msg.complete === true; + + // Extract metadata from final message + const metadata = complete ? streamingMetadataFrom(resp) : undefined; + + receiver(chunk, complete, metadata); + + return complete; + }; + + this.api.makeRequestMulti( + "text-completion", + { + system: system, + prompt: text, + streaming: true, + }, + recv, + 30000, + undefined, + this.flowId, + ); + }, + + + + /** + * Executes a prompt template with streaming response + * @param id - Prompt template ID + * @param terms - Template variables + * @param receiver - Called for each chunk with (chunk, complete) where complete=true on final chunk + * @param onError - Called on error + */ + promptStreaming( + id: string, + terms: Record, + receiver: (chunk: string, complete: boolean, metadata?: StreamingMetadata) => void, + onError: (error: string) => void, + ): void { + const recv = (message: unknown): boolean => { + const msg = message as { response?: PromptResponse; complete?: boolean; error?: string }; + + // Check for top-level error + if (msg.error !== undefined) { + onError(msg.error); + return true; + } + + const resp = (msg.response ?? {}) as PromptResponse; + + // Check for response-level error + if (resp.error !== undefined) { + onError(resp.error.message); + return true; + } + + // Prompt service uses 'text' field for chunks + const chunk = resp.text ?? ""; + const complete = msg.complete === true; + + // Extract metadata from final message + const metadata = complete ? streamingMetadataFrom(resp) : undefined; + + receiver(chunk, complete, metadata); + + return complete; + }; + + this.api.makeRequestMulti( + "prompt", + { + id: id, + terms: terms, + streaming: true, + }, + recv, + 30000, + undefined, + this.flowId, + ); + }, + + + + /** + * Generates embeddings for multiple texts within this flow. + * Returns vectors[text_index][dimension_index] - one vector per input text. + */ + embeddings(texts: string[]) { + return this.api + .makeRequest( + "embeddings", + { + texts: texts, + }, + 30000, + undefined, + this.flowId, + ) + .then((r) => r.vectors); + }, + + + + /** + * Queries the knowledge graph using a single embedding vector + */ + graphEmbeddingsQuery( + vec: number[], + limit: number | undefined, + collection?: string, + ) { + return this.api + .makeRequest( + "graph-embeddings", + { + vector: vec, + limit: limit ?? 20, // Default to 20 results + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + undefined, + this.flowId, + ) + .then((r) => r.entities); + }, + + + + /** + * Queries knowledge graph triples (subject-predicate-object relationships) + * All parameters are optional - omitted parameters act as wildcards + */ + triplesQuery( + s?: Term, + p?: Term, + o?: Term, + limit?: number, + collection?: string, + graph?: string, + ) { + const request: TriplesQueryRequest = { + limit: limit ?? 20, + user: this.api.user, + collection: withDefault(collection, "default"), + }; + if (s !== undefined) { + request.s = s; + } + if (p !== undefined) { + request.p = p; + } + if (o !== undefined) { + request.o = o; + } + if (graph !== undefined) { + request.g = graph; + } + + return this.api + .makeRequest( + "triples", + request, + 30000, + undefined, + this.flowId, + ) + .then((r) => r.triples ?? r.response ?? []); + }, + + + + /** + * Loads a document into this flow for processing + */ + loadDocument( + document: string, // base64-encoded document + id?: string, + metadata?: Triple[], + ) { + const request: LoadDocumentRequest = { + data: document, + }; + if (id !== undefined) { + request.id = id; + } + if (metadata !== undefined) { + request.metadata = metadata; + } + + return this.api.makeRequest( + "document-load", + request, + 30000, + undefined, + this.flowId, + ); + }, + + + + /** + * Loads plain text into this flow for processing + */ + loadText( + text: string, // Text content + id?: string, + metadata?: Triple[], + charset?: string, // Character encoding + ) { + const request: LoadTextRequest = { + text, + }; + if (id !== undefined) { + request.id = id; + } + if (metadata !== undefined) { + request.metadata = metadata; + } + if (charset !== undefined) { + request.charset = charset; + } + + return this.api.makeRequest( + "text-load", + request, + 30000, + undefined, + this.flowId, + ); + }, + + + + /** + * Executes a GraphQL query against structured row data + */ + rowsQuery( + query: string, + collection?: string, + variables?: Record, + operationName?: string, + ) { + const request: RowsQueryRequest = { + query, + user: this.api.user, + collection: withDefault(collection, "default"), + }; + if (variables !== undefined) { + request.variables = variables; + } + if (operationName !== undefined) { + request.operation_name = operationName; + } + + return this.api + .makeRequest( + "rows", + request, + 30000, + undefined, + this.flowId, + ) + .then((r) => { + // Return the GraphQL response structure directly + const result: Record = {}; + if (r.data !== undefined) result.data = r.data; + if (r.errors !== undefined) result.errors = r.errors; + if (r.extensions !== undefined) result.extensions = r.extensions; + return result; + }); + }, + + + + /** + * Converts a natural language question to a GraphQL query + */ + nlpQuery(question: string, maxResults?: number) { + return this.api + .makeRequest( + "nlp-query", + { + question: question, + max_results: maxResults ?? 100, + }, + 30000, + undefined, + this.flowId, + ) + .then((r) => r); + }, + + + + /** + * Executes a natural language question against structured data + * Combines NLP query conversion and GraphQL execution + */ + structuredQuery(question: string, collection?: string) { + return this.api + .makeRequest( + "structured-query", + { + question: question, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + undefined, + this.flowId, + ) + .then((r) => { + // Return the response structure directly + const result: Record = {}; + if (r.data !== undefined) result.data = r.data; + if (r.errors !== undefined) result.errors = r.errors; + return result; + }); + }, + + + + /** + * Performs semantic search on structured data indexes using embedding vectors + * @param vectors - Embedding vectors to search for + * @param schemaName - Name of the schema to search + * @param collection - Optional collection name + * @param indexName - Optional index name to filter results + * @param limit - Maximum number of results to return (default: 10) + */ + rowEmbeddingsQuery( + vector: number[], + schemaName: string, + collection?: string, + indexName?: string, + limit?: number, + ): Promise { + const request: RowEmbeddingsQueryRequest = { + vector: vector, + schema_name: schemaName, + user: this.api.user, + collection: withDefault(collection, "default"), + limit: limit ?? 10, + }; + + if (indexName !== undefined) { + request.index_name = indexName; + } + + return this.api + .makeRequest( + "row-embeddings", + request, + 30000, + undefined, + this.flowId, + ) + .then((r) => { + if (r.error !== undefined) { + throw new Error(r.error.message); + } + return r.matches ?? []; + }); + }, + }; } +export type FlowApi = ReturnType; +export const FlowApi = newableFactory(makeFlowApi); + /** * ConfigApi - Dedicated configuration management interface * Handles system configuration, prompts, and token cost tracking */ -export class ConfigApi { - api: BaseApi; +export function makeConfigApi(api: BaseApi) { + return { + api, - constructor(api: BaseApi) { - this.api = api; - } - /** - * Retrieves complete configuration - */ - getConfigAll() { - return this.api.makeRequest( - "config", - { - operation: "config", + + /** + * Retrieves complete configuration + */ + getConfigAll() { + return this.api.makeRequest( + "config", + { + operation: "config", + }, + 60000, + ); }, - 60000, - ); - } - /** - * Retrieves specific configuration entries - */ - getConfig(keys: { type: string; key: string }[]) { - return this.api.makeRequest( - "config", - { - operation: "get", - keys: keys, + + + /** + * Retrieves specific configuration entries + */ + getConfig(keys: { type: string; key: string }[]) { + return this.api.makeRequest( + "config", + { + operation: "get", + keys: keys, + }, + 60000, + ); }, - 60000, - ); - } - /** - * Updates configuration values using the Python-compatible values array. - */ - putConfig(items: { type: string; key: string; value: string }[]) { - return this.api.makeRequest( - "config", - { - operation: "put", - values: items, + + + /** + * Updates configuration values using the Python-compatible values array. + */ + putConfig(items: { type: string; key: string; value: string }[]) { + return this.api.makeRequest( + "config", + { + operation: "put", + values: items, + }, + 60000, + ); }, - 60000, - ); - } - /** - * Deletes a configuration entry - */ - deleteConfig(target: { type: string; key: string }) { - return this.api.makeRequest( - "config", - { - operation: "delete", - keys: [target], - }, - 30000, - ); - } - // Specialized prompt management methods - /** - * Retrieves list of available prompt templates from config.prompt. - * Each template is stored at `config.prompt.` as an object - * `{system, prompt}`. The reserved key `system` holds an optional - * global system prompt and is excluded from the template list. - */ - getPrompts() { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: Record } }; - const promptNs = config.config?.prompt ?? {}; - return Object.keys(promptNs) - .filter((k) => k !== "system") - .sort() - .map((id) => ({ id, name: id })); - }); - } + /** + * Deletes a configuration entry + */ + deleteConfig(target: { type: string; key: string }) { + return this.api.makeRequest( + "config", + { + operation: "delete", + keys: [target], + }, + 30000, + ); + }, - /** - * Retrieves a specific prompt template object: `{system, prompt}`. - */ - getPrompt(id: string) { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: Record } }; - return config.config?.prompt?.[id] ?? null; - }); - } - /** - * Retrieves the optional global system prompt at `config.prompt.system`. - * Returns "" if not configured. - */ - getSystemPrompt() { - return this.getConfigAll().then((r) => { - const config = r as { config?: { prompt?: { system?: unknown } } }; - const raw = config.config?.prompt?.system; - if (raw == null) return ""; - return typeof raw === "string" ? raw : raw; - }); - } - /** - * Lists available configuration types - */ - list(type: string) { - return this.api - .makeRequest( - "config", - { - operation: "list", - type: type, - }, - 60000, - ) - .then((r) => r); - } + // Specialized prompt management methods - /** - * Retrieves all key/values for a specific type - */ - getValues(type: string) { - return this.api - .makeRequest( - "config", - { - operation: "getvalues", - type: type, - }, - 60000, - ) - .then((r) => asConfigValues(r)); - } + /** + * Retrieves list of available prompt templates from config.prompt. + * Each template is stored at `config.prompt.` as an object + * `{system, prompt}`. The reserved key `system` holds an optional + * global system prompt and is excluded from the template list. + */ + getPrompts() { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: Record } }; + const promptNs = config.config?.prompt ?? {}; + return Object.keys(promptNs) + .filter((k) => k !== "system") + .sort() + .map((id) => ({ id, name: id })); + }); + }, - /** - * Retrieves token cost information for different AI models - * Useful for cost tracking and optimization - */ - getTokenCosts() { - return this.api - .makeRequest( - "config", - { - operation: "getvalues", - type: "token-cost", - }, - 60000, - ) - .then((r) => { - return asConfigValues(r).map((item) => ({ - key: item.key, - value: parseConfigJson(item.value), - })); - }) - .then((r) => - // Transform to more usable format - r.map((x: unknown) => { - const item = x as Record; - const value = item.value as Record; - return { - model: item.key, - input_price: value.input_price, // Cost per input token - output_price: value.output_price, // Cost per output token - }; - }), - ); - } + + + /** + * Retrieves a specific prompt template object: `{system, prompt}`. + */ + getPrompt(id: string) { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: Record } }; + return config.config?.prompt?.[id] ?? null; + }); + }, + + + + /** + * Retrieves the optional global system prompt at `config.prompt.system`. + * Returns "" if not configured. + */ + getSystemPrompt() { + return this.getConfigAll().then((r) => { + const config = r as { config?: { prompt?: { system?: unknown } } }; + const raw = config.config?.prompt?.system; + if (raw == null) return ""; + return typeof raw === "string" ? raw : raw; + }); + }, + + + + /** + * Lists available configuration types + */ + list(type: string) { + return this.api + .makeRequest( + "config", + { + operation: "list", + type: type, + }, + 60000, + ) + .then((r) => r); + }, + + + + /** + * Retrieves all key/values for a specific type + */ + getValues(type: string) { + return this.api + .makeRequest( + "config", + { + operation: "getvalues", + type: type, + }, + 60000, + ) + .then((r) => asConfigValues(r)); + }, + + + + /** + * Retrieves token cost information for different AI models + * Useful for cost tracking and optimization + */ + getTokenCosts() { + return this.api + .makeRequest( + "config", + { + operation: "getvalues", + type: "token-cost", + }, + 60000, + ) + .then((r) => { + return asConfigValues(r).map((item) => ({ + key: item.key, + value: parseConfigJson(item.value), + })); + }) + .then((r) => + // Transform to more usable format + r.map((x: unknown) => { + const item = x as Record; + const value = item.value as Record; + return { + model: item.key, + input_price: value.input_price, // Cost per input token + output_price: value.output_price, // Cost per output token + }; + }), + ); + }, + }; } +export type ConfigApi = ReturnType; +export const ConfigApi = newableFactory(makeConfigApi); + /** * KnowledgeApi - Manages knowledge graph cores and data * Knowledge cores appear to be collections of processed knowledge graph data */ -export class KnowledgeApi { - api: BaseApi; +export function makeKnowledgeApi(api: BaseApi) { + return { + api, - constructor(api: BaseApi) { - this.api = api; - } - /** - * Retrieves list of available knowledge graph cores - */ - getKnowledgeCores() { - return this.api - .makeRequest( - "knowledge", - { - operation: "list-kg-cores", - user: this.api.user, - }, - 60000, - ) - .then((r) => r.ids ?? []); - } - getDocumentEmbeddingCores() { - return this.api - .makeRequest( - "knowledge", - { - operation: "list-de-cores", - user: this.api.user, - }, - 60000, - ) - .then((r) => r.ids ?? []); - } - - /** - * Deletes a knowledge graph core - */ - deleteKgCore(id: string, collection?: string) { - return this.api.makeRequest( - "knowledge", - { - operation: "delete-kg-core", - id: id, - user: this.api.user, - collection: withDefault(collection, "default"), + /** + * Retrieves list of available knowledge graph cores + */ + getKnowledgeCores() { + return this.api + .makeRequest( + "knowledge", + { + operation: "list-kg-cores", + user: this.api.user, + }, + 60000, + ) + .then((r) => r.ids ?? []); }, - 30000, - ); - } - /** - * Deletes a knowledge graph core - */ - loadKgCore(id: string, flow: string, collection?: string) { - return this.api.makeRequest( - "knowledge", - { - operation: "load-kg-core", - id: id, - flow: flow, - user: this.api.user, - collection: withDefault(collection, "default"), + + + getDocumentEmbeddingCores() { + return this.api + .makeRequest( + "knowledge", + { + operation: "list-de-cores", + user: this.api.user, + }, + 60000, + ) + .then((r) => r.ids ?? []); }, - 30000, - ); - } - unloadKgCore(id: string, flow: string) { - return this.api.makeRequest( - "knowledge", - { - operation: "unload-kg-core", - id, - flow, - user: this.api.user, + + + /** + * Deletes a knowledge graph core + */ + deleteKgCore(id: string, collection?: string) { + return this.api.makeRequest( + "knowledge", + { + operation: "delete-kg-core", + id: id, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + ); }, - 30000, - ); - } - deleteDeCore(id: string) { - return this.api.makeRequest( - "knowledge", - { - operation: "delete-de-core", - id, - user: this.api.user, + + + /** + * Deletes a knowledge graph core + */ + loadKgCore(id: string, flow: string, collection?: string) { + return this.api.makeRequest( + "knowledge", + { + operation: "load-kg-core", + id: id, + flow: flow, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + ); }, - 30000, - ); - } - loadDeCore(id: string, flow: string, collection?: string) { - return this.api.makeRequest( - "knowledge", - { - operation: "load-de-core", - id, - flow, - user: this.api.user, - collection: withDefault(collection, "default"), + + + unloadKgCore(id: string, flow: string) { + return this.api.makeRequest( + "knowledge", + { + operation: "unload-kg-core", + id, + flow, + user: this.api.user, + }, + 30000, + ); }, - 30000, - ); - } - /** - * Retrieves a knowledge graph core with streaming data - * Uses multi-request pattern for large datasets - * @param receiver - Callback function to handle streaming data chunks - */ - getKgCore( - id: string, - collection: string | undefined, - receiver: (msg: unknown, eos: boolean) => void, - ) { - // Wrapper to handle end-of-stream detection - const recv = (msg: unknown) => { - const response = msg as Record; - if (response.eos === true) { - // End of stream - notify receiver and signal completion - receiver(msg, true); - return true; - } else { - // Regular message - continue streaming - receiver(msg, false); - return false; - } - }; - return this.api.makeRequestMulti( - "knowledge", - { - operation: "get-kg-core", - id: id, - user: this.api.user, - collection: withDefault(collection, "default"), + + deleteDeCore(id: string) { + return this.api.makeRequest( + "knowledge", + { + operation: "delete-de-core", + id, + user: this.api.user, + }, + 30000, + ); }, - recv, // Stream handler - 30000, - ); - } + + + + loadDeCore(id: string, flow: string, collection?: string) { + return this.api.makeRequest( + "knowledge", + { + operation: "load-de-core", + id, + flow, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + 30000, + ); + }, + + + + /** + * Retrieves a knowledge graph core with streaming data + * Uses multi-request pattern for large datasets + * @param receiver - Callback function to handle streaming data chunks + */ + getKgCore( + id: string, + collection: string | undefined, + receiver: (msg: unknown, eos: boolean) => void, + ) { + // Wrapper to handle end-of-stream detection + const recv = (msg: unknown) => { + const response = msg as Record; + if (response.eos === true) { + // End of stream - notify receiver and signal completion + receiver(msg, true); + return true; + } else { + // Regular message - continue streaming + receiver(msg, false); + return false; + } + }; + + return this.api.makeRequestMulti( + "knowledge", + { + operation: "get-kg-core", + id: id, + user: this.api.user, + collection: withDefault(collection, "default"), + }, + recv, // Stream handler + 30000, + ); + }, + }; } +export type KnowledgeApi = ReturnType; +export const KnowledgeApi = newableFactory(makeKnowledgeApi); + /** * CollectionManagementApi - Manages collections for organizing documents * Provides operations for listing, creating, updating, and deleting collections */ -export class CollectionManagementApi { - api: BaseApi; +export function makeCollectionManagementApi(api: BaseApi) { + return { + api, - constructor(api: BaseApi) { - this.api = api; - } - /** - * Lists all collections for the current user with optional tag filtering - * @param tagFilter - Optional array of tags to filter collections - * @returns Promise resolving to array of collection metadata - */ - listCollections(tagFilter?: string[]) { - const request: Record = { - operation: "list-collections", - user: this.api.user, - }; - if (tagFilter !== undefined && tagFilter.length > 0) { - request.tag_filter = tagFilter; - } + /** + * Lists all collections for the current user with optional tag filtering + * @param tagFilter - Optional array of tags to filter collections + * @returns Promise resolving to array of collection metadata + */ + listCollections(tagFilter?: string[]) { + const request: Record = { + operation: "list-collections", + user: this.api.user, + }; - return this.api - .makeRequest< - Record, - Record - >("collection-management", request, 30000) - .then((r) => r.collections ?? []); - } - - /** - * Creates or updates a collection for the current user - * @param collection - Collection ID (unique identifier) - * @param name - Display name for the collection - * @param description - Description of the collection - * @param tags - Array of tags for categorization - * @returns Promise resolving to updated collection metadata - */ - updateCollection( - collection: string, - name?: string, - description?: string, - tags?: string[], - ) { - const request: Record = { - operation: "update-collection", - user: this.api.user, - collection, - }; - - if (name !== undefined) { - request.name = name; - } - if (description !== undefined) { - request.description = description; - } - if (tags !== undefined) { - request.tags = tags; - } - - return this.api - .makeRequest< - Record, - Record - >("collection-management", request, 30000) - .then((r) => { - if ( - r.collections !== undefined && - Array.isArray(r.collections) && - r.collections.length > 0 - ) { - return r.collections[0]; + if (tagFilter !== undefined && tagFilter.length > 0) { + request.tag_filter = tagFilter; } - throw new Error("Failed to update collection"); - }); - } - /** - * Deletes a collection and all its data for the current user - * @param collection - Collection ID to delete - * @returns Promise resolving when deletion is complete - */ - deleteCollection(collection: string) { - return this.api.makeRequest< - Record, - Record - >( - "collection-management", - { - operation: "delete-collection", - user: this.api.user, - collection, + return this.api + .makeRequest< + Record, + Record + >("collection-management", request, 30000) + .then((r) => r.collections ?? []); }, - 30000, - ); - } + + + + /** + * Creates or updates a collection for the current user + * @param collection - Collection ID (unique identifier) + * @param name - Display name for the collection + * @param description - Description of the collection + * @param tags - Array of tags for categorization + * @returns Promise resolving to updated collection metadata + */ + updateCollection( + collection: string, + name?: string, + description?: string, + tags?: string[], + ) { + const request: Record = { + operation: "update-collection", + user: this.api.user, + collection, + }; + + if (name !== undefined) { + request.name = name; + } + if (description !== undefined) { + request.description = description; + } + if (tags !== undefined) { + request.tags = tags; + } + + return this.api + .makeRequest< + Record, + Record + >("collection-management", request, 30000) + .then((r) => { + if ( + r.collections !== undefined && + Array.isArray(r.collections) && + r.collections.length > 0 + ) { + return r.collections[0]; + } + throw new Error("Failed to update collection"); + }); + }, + + + + /** + * Deletes a collection and all its data for the current user + * @param collection - Collection ID to delete + * @returns Promise resolving when deletion is complete + */ + deleteCollection(collection: string) { + return this.api.makeRequest< + Record, + Record + >( + "collection-management", + { + operation: "delete-collection", + user: this.api.user, + collection, + }, + 30000, + ); + }, + }; } +export type CollectionManagementApi = ReturnType; +export const CollectionManagementApi = newableFactory(makeCollectionManagementApi); + /** * Factory function to create a new TrustGraph WebSocket connection * This is the main entry point for using the TrustGraph API diff --git a/ts/packages/flow/src/agent/mcp-tool/service.ts b/ts/packages/flow/src/agent/mcp-tool/service.ts index 1d0cc194..ae7f15b1 100644 --- a/ts/packages/flow/src/agent/mcp-tool/service.ts +++ b/ts/packages/flow/src/agent/mcp-tool/service.ts @@ -14,13 +14,14 @@ import { Client } from "@modelcontextprotocol/sdk/client/index.js"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, makeFlowProcessorProgram, errorMessage, type ProcessorConfig, type FlowContext, + type FlowProcessorRuntime, type ToolRequest, type ToolResponse, type EffectConfigHandler, @@ -281,41 +282,35 @@ const onMcpToolRequest = Effect.fn("McpToolService.onRequest")(function* ( }); export const makeMcpToolSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "mcp-tool-request", onMcpToolRequest, ), - new ProducerSpec("mcp-tool-response"), + makeProducerSpec("mcp-tool-response"), ]; export const makeMcpToolConfigHandlers = (): ReadonlyArray< EffectConfigHandler > => [onMcpConfig]; -export class McpToolService extends FlowProcessor { - private readonly runtime = Effect.runSync(makeMcpToolRuntime); +export type McpToolService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeMcpToolSpecs()) { - this.registerSpecification(spec); - } - - this.registerConfigHandler((config, version) => - Effect.runPromise(onMcpConfig(config, version).pipe( - Effect.provideService(McpToolRuntime, this.runtime), - )), - ); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(McpToolRuntime, this.runtime), - ); - } +export function makeMcpToolService(config: ProcessorConfig): McpToolService { + const runtime = Effect.runSync(makeMcpToolRuntime); + const service = makeFlowProcessor(config, { + specifications: makeMcpToolSpecs(), + provide: (effect) => effect.pipe(Effect.provideService(McpToolRuntime, runtime)), + }); + service.registerConfigHandler((pushedConfig, version) => + Effect.runPromise(onMcpConfig(pushedConfig, version).pipe( + Effect.provideService(McpToolRuntime, runtime), + )), + ); + return service; } +export const McpToolService = makeMcpToolService; + export const program = makeFlowProcessorProgram({ id: "mcp-tool", specs: () => makeMcpToolSpecs(), diff --git a/ts/packages/flow/src/agent/react/index.ts b/ts/packages/flow/src/agent/react/index.ts index 90eab368..f1c74aac 100644 --- a/ts/packages/flow/src/agent/react/index.ts +++ b/ts/packages/flow/src/agent/react/index.ts @@ -1,7 +1,7 @@ // ReAct agent -- barrel exports export { AgentService } from "./service.js"; -export { StreamingReActParser } from "./parser.js"; +export { makeStreamingReActParser, type StreamingReActParser } from "./parser.js"; export { buildReActPrompt } from "./prompt.js"; export { createKnowledgeQueryTool, diff --git a/ts/packages/flow/src/agent/react/parser.ts b/ts/packages/flow/src/agent/react/parser.ts index 809d7b41..01f4b9ab 100644 --- a/ts/packages/flow/src/agent/react/parser.ts +++ b/ts/packages/flow/src/agent/react/parser.ts @@ -22,57 +22,75 @@ const MARKERS = [ // Longest marker prefix for partial-match detection const MAX_MARKER_LEN = Math.max(...MARKERS.map((m) => m.prefix.length)); -export class StreamingReActParser { - private state: ReActState = "initial"; - private buffer = ""; - private onThought: (text: string) => void; - private onAction: (name: string) => void; - private onActionInput: (input: string) => void; - private onFinalAnswer: (text: string) => void; +export interface StreamingReActParser { + readonly feed: (text: string) => void; + readonly flush: () => void; +} - constructor( - onThought: (text: string) => void, - onAction: (name: string) => void, - onActionInput: (input: string) => void, - onFinalAnswer: (text: string) => void, - ) { - this.onThought = onThought; - this.onAction = onAction; - this.onActionInput = onActionInput; - this.onFinalAnswer = onFinalAnswer; - } +export function makeStreamingReActParser( + onThought: (text: string) => void, + onAction: (name: string) => void, + onActionInput: (input: string) => void, + onFinalAnswer: (text: string) => void, +): StreamingReActParser { + let state: ReActState = "initial"; + let buffer = ""; - /** - * Feed a chunk of LLM output text into the parser. - * Accumulates in a buffer and processes complete lines. - */ - feed(text: string): void { - this.buffer += text; - this.processBuffer(false); - } + const emitContent = (content: string): void => { + if (content.length === 0) return; - /** - * Flush any remaining buffered content at the end of output. - */ - flush(): void { - this.processBuffer(true); - // Emit any remaining buffer content in the current state - if (this.buffer.trim().length > 0) { - this.emitContent(this.buffer); - this.buffer = ""; + switch (state) { + case "thought": + onThought(content); + break; + case "action": + onAction(content); + break; + case "action_input": + onActionInput(content); + break; + case "final_answer": + onFinalAnswer(content); + break; + case "initial": + // Content before any marker -- treat as thought + state = "thought"; + onThought(content); + break; + case "complete": + break; } - } + }; - private processBuffer(isFinal: boolean): void { + const processLine = (line: string): void => { + const trimmed = line.trimStart(); + + // Check if this line starts a new section + for (const marker of MARKERS) { + if (trimmed.startsWith(marker.prefix)) { + const content = trimmed.slice(marker.prefix.length).trim(); + state = marker.state; + emitContent(content); + return; + } + } + + // Otherwise, this is continuation content for the current state + if (trimmed.length > 0) { + emitContent(trimmed); + } + }; + + const processBuffer = (isFinal: boolean): void => { // Process complete lines (terminated by newline) while (true) { - const newlineIdx = this.buffer.indexOf("\n"); + const newlineIdx = buffer.indexOf("\n"); if (newlineIdx === -1) { // No complete line yet. // If not final, check for partial marker match at the end and wait. if (!isFinal) { // If the remaining buffer could be the start of a marker, wait for more input. - const trimmed = this.buffer.trimStart(); + const trimmed = buffer.trimStart(); if (trimmed.length > 0 && trimmed.length < MAX_MARKER_LEN) { const couldBeMarker = MARKERS.some((m) => m.prefix.startsWith(trimmed), @@ -86,54 +104,29 @@ export class StreamingReActParser { break; } - const line = this.buffer.slice(0, newlineIdx); - this.buffer = this.buffer.slice(newlineIdx + 1); - this.processLine(line); + const line = buffer.slice(0, newlineIdx); + buffer = buffer.slice(newlineIdx + 1); + processLine(line); } - } + }; - private processLine(line: string): void { - const trimmed = line.trimStart(); + /** + * Feed a chunk of LLM output text into the parser. + * Accumulates in a buffer and processes complete lines. + */ + const feed = (text: string): void => { + buffer += text; + processBuffer(false); + }; - // Check if this line starts a new section - for (const marker of MARKERS) { - if (trimmed.startsWith(marker.prefix)) { - const content = trimmed.slice(marker.prefix.length).trim(); - this.state = marker.state; - this.emitContent(content); - return; - } + const flush = (): void => { + processBuffer(true); + // Emit any remaining buffer content in the current state + if (buffer.trim().length > 0) { + emitContent(buffer); + buffer = ""; } + }; - // Otherwise, this is continuation content for the current state - if (trimmed.length > 0) { - this.emitContent(trimmed); - } - } - - private emitContent(content: string): void { - if (content.length === 0) return; - - switch (this.state) { - case "thought": - this.onThought(content); - break; - case "action": - this.onAction(content); - break; - case "action_input": - this.onActionInput(content); - break; - case "final_answer": - this.onFinalAnswer(content); - break; - case "initial": - // Content before any marker -- treat as thought - this.state = "thought"; - this.onThought(content); - break; - case "complete": - break; - } - } + return { feed, flush }; } diff --git a/ts/packages/flow/src/agent/react/service.ts b/ts/packages/flow/src/agent/react/service.ts index d89dc6fe..8faf283a 100644 --- a/ts/packages/flow/src/agent/react/service.ts +++ b/ts/packages/flow/src/agent/react/service.ts @@ -17,14 +17,15 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, - RequestResponseSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, + makeRequestResponseSpec, makeFlowProcessorProgram, errorMessage, type ProcessorConfig, type FlowContext, + type FlowProcessorRuntime, type AgentRequest, type AgentResponse, type TextCompletionRequest, @@ -488,32 +489,32 @@ const onAgentRequest = Effect.fn("AgentService.onRequest")(function* ( }); export const makeAgentSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "agent-request", onAgentRequest, ), - new ProducerSpec("agent-response"), - new RequestResponseSpec( + makeProducerSpec("agent-response"), + makeRequestResponseSpec( "llm", "text-completion-request", "text-completion-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "graph-rag", "graph-rag-request", "graph-rag-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "doc-rag", "document-rag-request", "document-rag-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "triples", "triples-request", "triples-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "mcp-tool", "mcp-tool-request", "mcp-tool-response", @@ -524,32 +525,25 @@ export const makeAgentConfigHandlers = (): ReadonlyArray< EffectConfigHandler > => [onToolsConfig]; -export class AgentService extends FlowProcessor { - private readonly runtime = Effect.runSync(makeAgentRuntime); +export type AgentService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeAgentSpecs()) { - this.registerSpecification(spec); - } - - this.registerConfigHandler((config, version) => - Effect.runPromise(onToolsConfig(config, version).pipe( - Effect.provideService(AgentRuntime, this.runtime), - )), - ); - - console.log("[AgentService] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(AgentRuntime, this.runtime), - ); - } +export function makeAgentService(config: ProcessorConfig): AgentService { + const runtime = Effect.runSync(makeAgentRuntime); + const service = makeFlowProcessor(config, { + specifications: makeAgentSpecs(), + provide: (effect) => effect.pipe(Effect.provideService(AgentRuntime, runtime)), + }); + service.registerConfigHandler((pushedConfig, version) => + Effect.runPromise(onToolsConfig(pushedConfig, version).pipe( + Effect.provideService(AgentRuntime, runtime), + )), + ); + console.log("[AgentService] Service initialized"); + return service; } +export const AgentService = makeAgentService; + /** * Simple line-based parser for ReAct LLM output. * diff --git a/ts/packages/flow/src/chunking/service.ts b/ts/packages/flow/src/chunking/service.ts index 5d4b93fd..2226cf72 100644 --- a/ts/packages/flow/src/chunking/service.ts +++ b/ts/packages/flow/src/chunking/service.ts @@ -10,11 +10,12 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, - ParameterSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, + makeParameterSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type MessagingDeliveryError, @@ -74,28 +75,28 @@ const onChunkMessage = Effect.fn("ChunkingService.onMessage")(function* ( export const makeChunkingSpecs = (): ReadonlyArray< Spec > => [ - new ConsumerSpec( + makeConsumerSpec( "chunk-input", onChunkMessage, ), - new ProducerSpec("chunk-output"), - new ProducerSpec("chunk-triples"), - new ParameterSpec("chunk-size"), - new ParameterSpec("chunk-overlap"), + makeProducerSpec("chunk-output"), + makeProducerSpec("chunk-triples"), + makeParameterSpec("chunk-size"), + makeParameterSpec("chunk-overlap"), ]; -export class ChunkingService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); +export type ChunkingService = FlowProcessorRuntime; - for (const spec of makeChunkingSpecs()) { - this.registerSpecification(spec); - } - - console.log("[ChunkingService] Service initialized"); - } +export function makeChunkingService(config: ProcessorConfig): ChunkingService { + const service = makeFlowProcessor(config, { + specifications: makeChunkingSpecs(), + }); + console.log("[ChunkingService] Service initialized"); + return service; } +export const ChunkingService = makeChunkingService; + export const program = makeFlowProcessorProgram({ id: "chunking", specs: () => makeChunkingSpecs(), diff --git a/ts/packages/flow/src/config/service.ts b/ts/packages/flow/src/config/service.ts index 38784861..061a6515 100644 --- a/ts/packages/flow/src/config/service.ts +++ b/ts/packages/flow/src/config/service.ts @@ -14,8 +14,9 @@ import { Effect } from "effect"; import * as S from "effect/Schema"; import { - AsyncProcessor, + makeAsyncProcessor, type ProcessorConfig, + type AsyncProcessorRuntime, topics, ConfigRequest as ConfigRequestSchema, ConfigResponse as ConfigResponseSchema, @@ -27,7 +28,7 @@ import { makeProcessorProgram, optionalStringConfig, } from "@trustgraph/base"; -import type { BackendProducer, BackendConsumer, Message } from "@trustgraph/base"; +import type { Message } from "@trustgraph/base"; import { readTextFile, writeTextFile } from "../runtime/effect-files.js"; export interface ConfigServiceConfig extends ProcessorConfig { @@ -69,524 +70,596 @@ function optionalString(value: unknown): string | undefined { return typeof value === "string" && value.length > 0 ? value : undefined; } -export class ConfigService extends AsyncProcessor { - private store = new Map(); - private version = 0; - private readonly persistPath: string | null; - private consumer: BackendConsumer | null = null; - private responseProducer: BackendProducer | null = null; - private pushProducer: BackendProducer | null = null; +export type ConfigService = AsyncProcessorRuntime & Record; - constructor(config: ConfigServiceConfig) { - super(config); - this.persistPath = config.persistPath ?? null; - } +export function makeConfigService(config: ConfigServiceConfig): ConfigService { + const service = makeAsyncProcessor(config, { + run: async () => { + await service.run(); + }, + }) as ConfigService; + const baseStop = service.stop; + service.store = new Map(); + service.version = 0; + service.consumer = null; + service.responseProducer = null; + service.pushProducer = null; + service.persistPath = config.persistPath ?? null; + Object.assign(service, { - protected override async run(): Promise { - // Optionally load persisted state - if (this.persistPath !== null) { - await this.loadFromDisk(); - } - // Create producers - this.responseProducer = await this.pubsub.createProducer({ - topic: topics.configResponse, - schema: ConfigResponseSchema, - }); - this.pushProducer = await this.pubsub.createProducer({ - topic: topics.configPush, - schema: ConfigPushSchema, - }); - - // Create consumer for config requests - this.consumer = await this.pubsub.createConsumer({ - topic: topics.configRequest, - subscription: `${this.config.id}-config-request`, - schema: ConfigRequestSchema, - }); - - // Push initial config - await this.pushConfig(); - - console.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"); - - const msg = await consumer.receive(2000); - if (msg === null) continue; - - 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); - } - } - } - - private async handleMessage(msg: Message): Promise { - const request = await Effect.runPromise(S.decodeUnknownEffect(ConfigRequestSchema)(msg.value())); - const props = msg.properties(); - const requestId = props.id; - - if (requestId === undefined || requestId.length === 0) { - console.warn("[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 }, - ); - } - } - - private async handleOperation(request: ConfigRequest): Promise { - const op: ConfigOperation = request.operation; - - switch (op) { - case "get": - return this.handleGet(request); - - case "put": - return await this.handlePut(request); - - case "delete": - return await this.handleDelete(request); - - case "list": - return this.handleList(request); - - case "config": - return this.handleConfigDump(request); - - case "getvalues": - return this.handleGetValues(request); - - case "getvalues-all-ws": - return this.handleGetValuesAllWorkspaces(request); - - default: - throw new Error(`Unknown config operation: ${op as string}`); - } - } - - private requestRecord(request: ConfigRequest): Record { - return request as Record; - } - - private workspaceFor(request: ConfigRequest): string { - return optionalString(this.requestRecord(request).workspace) ?? DEFAULT_WORKSPACE; - } - - private workspaceStore(workspace: string, create: boolean): WorkspaceStore | undefined { - let store = this.store.get(workspace); - if (store === undefined && create) { - store = new Map(); - this.store.set(workspace, store); - } - return store; - } - - private namespaceStore( - workspace: string, - namespace: string, - create: boolean, - ): NamespaceStore | undefined { - const ws = this.workspaceStore(workspace, create); - if (ws === undefined) return undefined; - - let ns = ws.get(namespace); - if (ns === undefined && create) { - ns = new Map(); - ws.set(namespace, ns); - } - return ns; - } - - private rawKeys(request: ConfigRequest): unknown[] { - const keys = this.requestRecord(request).keys; - return Array.isArray(keys) ? keys : []; - } - - private stringKeys(request: ConfigRequest): string[] { - return this.rawKeys(request).filter((key): key is string => typeof key === "string"); - } - - private objectKeys(request: ConfigRequest): ConfigKeyLike[] { - return this.rawKeys(request).flatMap((key) => { - if (!isRecord(key)) return []; - const type = optionalString(key.type); - if (type === undefined) return []; - const item: ConfigKeyLike = { type }; - const keyValue = optionalString(key.key); - if (keyValue !== undefined) item.key = keyValue; - return [item]; - }); - } - - private requestType(request: ConfigRequest): string | undefined { - return optionalString(this.requestRecord(request).type) ?? this.stringKeys(request)[0]; - } - - private configValues(request: ConfigRequest): ConfigValueLike[] { - const req = this.requestRecord(request); - const rawValues = req.values; - const workspace = this.workspaceFor(request); - - if (Array.isArray(rawValues)) { - return rawValues.flatMap((value) => { - if (!isRecord(value)) return []; - const type = optionalString(value.type); - const key = optionalString(value.key); - if (type === undefined || key === undefined) return []; - return [{ - workspace: optionalString(value.workspace) ?? workspace, - type, - key, - value: value.value, - }]; - }); - } - - if (isRecord(rawValues)) { - const namespace = this.requestType(request); - if (namespace === undefined) return []; - return Object.entries(rawValues).map(([key, value]) => ({ - workspace, - type: namespace, - key, - value, - })); - } - - return []; - } - - private handleGet(request: ConfigRequest): ConfigResponse { - const workspace = this.workspaceFor(request); - const objectKeys = this.objectKeys(request); - - if (objectKeys.length > 0) { - const values = objectKeys.map((key) => ({ - type: key.type, - key: key.key ?? "", - value: key.key !== undefined - ? this.namespaceStore(workspace, key.type, false)?.get(key.key) - : undefined, - })); - return { version: this.version, values }; - } - - const keys = this.stringKeys(request); - if (keys.length === 0) { - return { version: this.version, values: {} }; - } - - const values: Record = {}; - const namespace = keys[0]; - const subMap = this.namespaceStore(workspace, namespace, false); - - if (subMap !== undefined) { - if (keys.length === 1) { - // Return entire namespace - for (const [k, v] of subMap) { - values[k] = v; + run: async function(this: ConfigService): Promise { + // Optionally load persisted state + if (this.persistPath !== null) { + await this.loadFromDisk(); } - } else { - // Return specific keys within namespace - for (let i = 1; i < keys.length; i++) { - const key = keys[i]; - if (subMap.has(key)) { - values[key] = subMap.get(key); + + // Create producers + this.responseProducer = await this.pubsub.createProducer({ + topic: topics.configResponse, + schema: ConfigResponseSchema, + }); + this.pushProducer = await this.pubsub.createProducer({ + topic: topics.configPush, + schema: ConfigPushSchema, + }); + + // Create consumer for config requests + this.consumer = await this.pubsub.createConsumer({ + topic: topics.configRequest, + subscription: `${this.config.id}-config-request`, + schema: ConfigRequestSchema, + }); + + // Push initial config + await this.pushConfig(); + + console.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"); + + const msg = await consumer.receive(2000); + if (msg === null) continue; + + 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); } } - } - } - return { version: this.version, values }; - } + }, - private async handlePut(request: ConfigRequest): Promise { - const values = this.configValues(request); - if (values.length === 0) throw new Error("Put requires config values"); - for (const item of values) { - this.namespaceStore(item.workspace ?? DEFAULT_WORKSPACE, item.type, true)?.set(item.key, item.value); - } - this.version++; - await this.persist(); - await this.pushConfig(); + handleMessage: async function(this: ConfigService, msg: Message): Promise { + const request = await Effect.runPromise(S.decodeUnknownEffect(ConfigRequestSchema)(msg.value())); + const props = msg.properties(); + const requestId = props.id; - return { version: this.version }; - } + if (requestId === undefined || requestId.length === 0) { + console.warn("[ConfigService] Received request without id, ignoring"); + return; + } - private async handleDelete(request: ConfigRequest): Promise { - const workspace = this.workspaceFor(request); - const objectKeys = this.objectKeys(request); - if (objectKeys.length > 0) { - for (const key of objectKeys) { + 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 }, + ); + } + + }, + + + + handleOperation: async function(this: ConfigService, request: ConfigRequest): Promise { + const op: ConfigOperation = request.operation; + + switch (op) { + case "get": + return this.handleGet(request); + + case "put": + return await this.handlePut(request); + + case "delete": + return await this.handleDelete(request); + + case "list": + return this.handleList(request); + + case "config": + return this.handleConfigDump(request); + + case "getvalues": + return this.handleGetValues(request); + + case "getvalues-all-ws": + return this.handleGetValuesAllWorkspaces(request); + + default: + throw new Error(`Unknown config operation: ${op as string}`); + } + + }, + + + + requestRecord: function(this: ConfigService, request: ConfigRequest): Record { + return request as Record; + + }, + + + + workspaceFor: function(this: ConfigService, request: ConfigRequest): string { + return optionalString(this.requestRecord(request).workspace) ?? DEFAULT_WORKSPACE; + + }, + + + + workspaceStore: function(this: ConfigService, workspace: string, create: boolean): WorkspaceStore | undefined { + let store = this.store.get(workspace); + if (store === undefined && create) { + store = new Map(); + this.store.set(workspace, store); + } + return store; + + }, + + + + namespaceStore: function(this: ConfigService, workspace: string, namespace: string, create: boolean): NamespaceStore | undefined { + const ws = this.workspaceStore(workspace, create); + if (ws === undefined) return undefined; + + let ns = ws.get(namespace); + if (ns === undefined && create) { + ns = new Map(); + ws.set(namespace, ns); + } + return ns; + + }, + + + + rawKeys: function(this: ConfigService, request: ConfigRequest): unknown[] { + const keys = this.requestRecord(request).keys; + return Array.isArray(keys) ? keys : []; + + }, + + + + stringKeys: function(this: ConfigService, request: ConfigRequest): string[] { + return (this.rawKeys(request) as unknown[]).filter((key: unknown): key is string => typeof key === "string"); + + }, + + + + objectKeys: function(this: ConfigService, request: ConfigRequest): ConfigKeyLike[] { + return (this.rawKeys(request) as unknown[]).flatMap((key: unknown) => { + if (!isRecord(key)) return []; + const type = optionalString(key.type); + if (type === undefined) return []; + const item: ConfigKeyLike = { type }; + const keyValue = optionalString(key.key); + if (keyValue !== undefined) item.key = keyValue; + return [item]; + }); + + }, + + + + requestType: function(this: ConfigService, request: ConfigRequest): string | undefined { + return optionalString(this.requestRecord(request).type) ?? this.stringKeys(request)[0]; + + }, + + + + configValues: function(this: ConfigService, request: ConfigRequest): ConfigValueLike[] { + const req = this.requestRecord(request); + const rawValues = req.values; + const workspace = this.workspaceFor(request); + + if (Array.isArray(rawValues)) { + return rawValues.flatMap((value) => { + if (!isRecord(value)) return []; + const type = optionalString(value.type); + const key = optionalString(value.key); + if (type === undefined || key === undefined) return []; + return [{ + workspace: optionalString(value.workspace) ?? workspace, + type, + key, + value: value.value, + }]; + }); + } + + if (isRecord(rawValues)) { + const namespace = this.requestType(request); + if (namespace === undefined) return []; + return Object.entries(rawValues).map(([key, value]) => ({ + workspace, + type: namespace, + key, + value, + })); + } + + return []; + + }, + + + + handleGet: function(this: ConfigService, request: ConfigRequest): ConfigResponse { + const workspace = this.workspaceFor(request); + const objectKeys = this.objectKeys(request) as ConfigKeyLike[]; + + if (objectKeys.length > 0) { + const values = objectKeys.map((key) => ({ + type: key.type, + key: key.key ?? "", + value: key.key !== undefined + ? this.namespaceStore(workspace, key.type, false)?.get(key.key) + : undefined, + })); + return { version: this.version, values }; + } + + const keys = this.stringKeys(request) as string[]; + if (keys.length === 0) { + return { version: this.version, values: {} }; + } + + const values: Record = {}; + const namespace = keys[0]; + const subMap = this.namespaceStore(workspace, namespace, false) as NamespaceStore | undefined; + + if (subMap !== undefined) { + if (keys.length === 1) { + // Return entire namespace + for (const [k, v] of subMap) { + values[k] = v; + } + } else { + // Return specific keys within namespace + for (let i = 1; i < keys.length; i++) { + const key = keys[i]; + if (key !== undefined && subMap.has(key)) { + values[key] = subMap.get(key); + } + } + } + } + + return { version: this.version, values }; + + }, + + + + handlePut: async function(this: ConfigService, request: ConfigRequest): Promise { + const values = this.configValues(request); + if (values.length === 0) throw new Error("Put requires config values"); + + for (const item of values) { + this.namespaceStore(item.workspace ?? DEFAULT_WORKSPACE, item.type, true)?.set(item.key, item.value); + } + + this.version++; + await this.persist(); + await this.pushConfig(); + + return { version: this.version }; + + }, + + + + handleDelete: async function(this: ConfigService, request: ConfigRequest): Promise { + const workspace = this.workspaceFor(request); + const objectKeys = this.objectKeys(request); + if (objectKeys.length > 0) { + for (const key of objectKeys) { + const ws = this.workspaceStore(workspace, false); + if (ws === undefined) continue; + if (key.key === undefined) { + ws.delete(key.type); + } else { + const ns = ws.get(key.type); + ns?.delete(key.key); + if (ns !== undefined && ns.size === 0) ws.delete(key.type); + } + } + + this.version++; + await this.persist(); + await this.pushConfig(); + return { version: this.version }; + } + + const keys = this.stringKeys(request); + if (keys.length === 0) { + throw new Error("Delete requires at least one key"); + } + + const namespace = keys[0]; const ws = this.workspaceStore(workspace, false); - if (ws === undefined) 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); - } - } + if (ws === undefined) return { version: this.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]); - } - if (subMap.size === 0) { + 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(); + this.version++; + await this.persist(); + await this.pushConfig(); - return { version: this.version }; - } + return { version: this.version }; - private handleList(request: ConfigRequest): ConfigResponse { - const workspace = this.workspaceFor(request); - const ws = this.workspaceStore(workspace, false); - const namespace = this.requestType(request); + }, - if (namespace === undefined) { - // List all namespaces - return { - version: this.version, - directory: ws !== undefined ? [...ws.keys()] : [], - }; - } - const subMap = ws?.get(namespace); - return { - version: this.version, - directory: subMap !== undefined ? [...subMap.keys()] : [], - }; - } + handleList: function(this: ConfigService, request: ConfigRequest): ConfigResponse { + const workspace = this.workspaceFor(request); + const ws = this.workspaceStore(workspace, false); + const namespace = this.requestType(request); - private handleGetValues(request: ConfigRequest): ConfigResponse { - const workspace = this.workspaceFor(request); - const type = this.requestType(request) ?? ""; - const ws = this.workspaceStore(workspace, false); - - const values: { type: string; key: string; value: unknown }[] = []; - - for (const [namespace, subMap] of ws ?? new Map()) { - if ( - type.length === 0 || - namespace === type - ) { - for (const [k, v] of subMap) { - values.push({ type: namespace, key: k, value: v }); + if (namespace === undefined) { + // List all namespaces + return { + version: this.version, + directory: ws !== undefined ? [...ws.keys()] : [], + }; } - } - } - return { version: this.version, values }; - } + const subMap = ws?.get(namespace); - private handleGetValuesAllWorkspaces(request: ConfigRequest): ConfigResponse { - const type = this.requestType(request) ?? ""; - const values: { workspace: string; type: string; key: string; value: unknown }[] = []; + return { + version: this.version, + directory: subMap !== undefined ? [...subMap.keys()] : [], + }; - for (const [workspace, ws] of this.store) { - for (const [namespace, subMap] of ws) { - if (type.length > 0 && namespace !== type) continue; - for (const [key, value] of subMap) { - values.push({ workspace, type: namespace, key, value }); + }, + + + + handleGetValues: function(this: ConfigService, request: ConfigRequest): ConfigResponse { + const workspace = this.workspaceFor(request); + const type = this.requestType(request) ?? ""; + const ws = this.workspaceStore(workspace, false); + + const values: { type: string; key: string; value: unknown }[] = []; + + for (const [namespace, subMap] of ws ?? new Map()) { + if ( + type.length === 0 || + namespace === type + ) { + for (const [k, v] of subMap) { + values.push({ type: namespace, key: k, value: v }); + } + } } - } - } - return { version: this.version, values }; - } + return { version: this.version, values }; - private handleConfigDump(request: ConfigRequest): ConfigResponse { - const workspace = this.workspaceFor(request); - const ws = this.workspaceStore(workspace, false); - const config: Record = {}; + }, - for (const [namespace, subMap] of ws ?? new Map()) { - const obj: Record = {}; - for (const [k, v] of subMap) { - obj[k] = v; - } - config[namespace] = obj; - } - return { - version: this.version, - config, - }; - } - private async pushConfig(): Promise { - const pushProducer = this.pushProducer; - if (pushProducer === null) return; + handleGetValuesAllWorkspaces: function(this: ConfigService, request: ConfigRequest): ConfigResponse { + const type = this.requestType(request) ?? ""; + const values: { workspace: string; type: string; key: string; value: unknown }[] = []; - const config: Record = {}; - const ws = this.workspaceStore(DEFAULT_WORKSPACE, false); - for (const [namespace, subMap] of ws ?? new Map()) { - const obj: Record = {}; - for (const [k, v] of subMap) { - obj[k] = v; - } - config[namespace] = obj; - } + for (const [workspace, ws] of this.store) { + for (const [namespace, subMap] of ws) { + if (type.length > 0 && namespace !== type) continue; + for (const [key, value] of subMap) { + values.push({ workspace, type: namespace, key, value }); + } + } + } - await pushProducer.send({ - version: this.version, - config, - }); + return { version: this.version, values }; - console.log(`[ConfigService] Pushed configuration version ${this.version}`); - } + }, - private async persist(): Promise { - const persistPath = this.persistPath; - if (persistPath === null) return; - try { - const workspaces: Record>> = {}; - for (const [workspace, ws] of this.store) { - const workspaceData: Record> = {}; - for (const [namespace, subMap] of ws) { + handleConfigDump: function(this: ConfigService, request: ConfigRequest): ConfigResponse { + const workspace = this.workspaceFor(request); + const ws = this.workspaceStore(workspace, false); + const config: Record = {}; + + for (const [namespace, subMap] of ws ?? new Map()) { const obj: Record = {}; 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, - ); + return { + version: this.version, + config, + }; - await writeTextFile(persistPath, json); - } catch (err) { - await Effect.runPromise(Effect.logError("[ConfigService] Failed to persist config", { error: errorMessage(err) })); - } - } + }, - private async loadFromDisk(): Promise { - const persistPath = this.persistPath; - if (persistPath === null) return; - try { - const raw = await readTextFile(persistPath); - const parsed = JSON.parse(raw) as { - version: number; - data?: Record>; - workspaces?: Record>>; - }; - this.version = parsed.version ?? 0; - this.store.clear(); + pushConfig: async function(this: ConfigService): Promise { + const pushProducer = this.pushProducer; + if (pushProducer === null) return; - if (parsed.workspaces !== undefined) { - for (const [workspace, namespaces] of Object.entries(parsed.workspaces)) { - const ws = new Map(); - for (const [namespace, obj] of Object.entries(namespaces)) { - const subMap = new Map(); - for (const [k, v] of Object.entries(obj)) { - subMap.set(k, v); + const config: Record = {}; + const ws = this.workspaceStore(DEFAULT_WORKSPACE, false); + for (const [namespace, subMap] of ws ?? new Map()) { + const obj: Record = {}; + 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 { + const persistPath = this.persistPath; + if (persistPath === null) return; + + try { + const workspaces: Record>> = {}; + + for (const [workspace, ws] of this.store) { + const workspaceData: Record> = {}; + for (const [namespace, subMap] of ws) { + const obj: Record = {}; + for (const [k, v] of subMap) { + obj[k] = v; + } + workspaceData[namespace] = obj; } - ws.set(namespace, subMap); + workspaces[workspace] = workspaceData; } - this.store.set(workspace, ws); + + const json = JSON.stringify( + { version: this.version, workspaces }, + null, + 2, + ); + + await writeTextFile(persistPath, json); + } catch (err) { + await Effect.runPromise(Effect.logError("[ConfigService] Failed to persist config", { error: errorMessage(err) })); } - } else { - const ws = new Map(); - for (const [namespace, obj] of Object.entries(parsed.data ?? {})) { - const subMap = new Map(); - for (const [k, v] of Object.entries(obj)) { - subMap.set(k, v); + + }, + + + + loadFromDisk: async function(this: ConfigService): Promise { + const persistPath = this.persistPath; + if (persistPath === null) return; + + try { + const raw = await readTextFile(persistPath); + const parsed = JSON.parse(raw) as { + version: number; + data?: Record>; + workspaces?: Record>>; + }; + + 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(); + for (const [namespace, obj] of Object.entries(namespaces)) { + const subMap = new Map(); + 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(); + for (const [namespace, obj] of Object.entries(parsed.data ?? {})) { + const subMap = new Map(); + for (const [k, v] of Object.entries(obj)) { + subMap.set(k, v); + } + ws.set(namespace, subMap); + } + this.store.set(DEFAULT_WORKSPACE, ws); } - ws.set(namespace, subMap); + + 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")); } - 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")); - } - } + }, - override async stop(): Promise { - 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 super.stop(); - } + + + stop: async function(this: ConfigService): Promise { + 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(); + + } + }); + return service; } +export const ConfigService = makeConfigService; + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } @@ -605,7 +678,7 @@ export const loadConfigServiceRuntimeConfig = Effect.fn("loadConfigServiceRuntim export const program = makeProcessorProgram({ id: "config-svc", loadConfig: loadConfigServiceRuntimeConfig(), - make: (config) => new ConfigService(config), + make: (config) => makeConfigService(config), }); export async function run(): Promise { diff --git a/ts/packages/flow/src/cores/service.ts b/ts/packages/flow/src/cores/service.ts index 48516225..321b0343 100644 --- a/ts/packages/flow/src/cores/service.ts +++ b/ts/packages/flow/src/cores/service.ts @@ -11,8 +11,9 @@ */ import { - AsyncProcessor, + makeAsyncProcessor, type ProcessorConfig, + type AsyncProcessorRuntime, topics, type KnowledgeRequest, type KnowledgeResponse, @@ -20,7 +21,7 @@ import { type Term, } from "@trustgraph/base"; import { makeProcessorProgram } from "@trustgraph/base"; -import type { BackendProducer, BackendConsumer, Message } from "@trustgraph/base"; +import type { Message } from "@trustgraph/base"; import { Effect } from "effect"; import { ensureDirectory, joinPath, readTextFile, writeTextFile } from "../runtime/effect-files.js"; @@ -39,394 +40,461 @@ interface DocumentEmbeddingsCore { [key: string]: unknown; } -export class KnowledgeCoreService extends AsyncProcessor { - /** Keyed by `${user}:${id}` */ - private cores = new Map(); - private deCores = new Map(); - private readonly dataDir: string; - private readonly persistPath: string; +export type KnowledgeCoreService = AsyncProcessorRuntime & Record; - private consumer: BackendConsumer | null = null; - private responseProducer: BackendProducer | null = null; +export function makeKnowledgeCoreService(config: KnowledgeCoreServiceConfig): KnowledgeCoreService { + const service = makeAsyncProcessor(config, { + run: async () => { + await service.run(); + }, + }) as KnowledgeCoreService; + const baseStop = service.stop; + service.cores = new Map(); + service.deCores = new Map(); + service.consumer = null; + service.responseProducer = null; + const dataDir = config.dataDir ?? process.env.KNOWLEDGE_DATA_DIR ?? "./data/knowledge"; + service.dataDir = dataDir; + service.persistPath = joinPath(dataDir, "knowledge-state.json"); + Object.assign(service, { - constructor(config: KnowledgeCoreServiceConfig) { - super(config); - const dataDir = config.dataDir ?? process.env.KNOWLEDGE_DATA_DIR ?? "./data/knowledge"; - this.dataDir = dataDir; - this.persistPath = joinPath(dataDir, "knowledge-state.json"); - } - private coreKey(user: string, id: string): string { - return `${user}:${id}`; - } + coreKey: function(this: KnowledgeCoreService, user: string, id: string): string { + return `${user}:${id}`; - protected override async run(): Promise { - await ensureDirectory(this.dataDir); - // Load persisted state - await this.loadFromDisk(); + }, - // Create producer - this.responseProducer = await this.pubsub.createProducer({ - topic: topics.knowledgeResponse, - }); - // Create consumer - this.consumer = await this.pubsub.createConsumer({ - topic: topics.knowledgeRequest, - subscription: `${this.config.id}-knowledge-request`, - }); - console.log(`[KnowledgeCoreService] Listening on ${topics.knowledgeRequest}`); + run: async function(this: KnowledgeCoreService): Promise { + await ensureDirectory(this.dataDir); + // Load persisted state + await this.loadFromDisk(); - // Main consume loop - while (this.running) { - try { - const msg = await this.consumer.receive(2000); - if (msg === null) continue; - - 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); - } - } - } - - private async handleMessage(msg: Message): Promise { - 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; - } - - 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 }, - ); - } - } - - private async handleOperation(request: KnowledgeRequest, requestId: string): Promise { - switch (request.operation) { - case "list-kg-cores": - return this.listKgCores(request, requestId); - case "get-kg-core": - return this.getKgCore(request, requestId); - case "delete-kg-core": - return this.deleteKgCore(request, requestId); - case "put-kg-core": - return this.putKgCore(request, requestId); - case "load-kg-core": - return this.loadKgCore(request, requestId); - case "unload-kg-core": - return this.unloadKgCore(request, requestId); - case "list-de-cores": - return this.listDeCores(request, requestId); - case "get-de-core": - return this.getDeCore(request, requestId); - case "delete-de-core": - return this.deleteDeCore(request, requestId); - case "put-de-core": - return this.putDeCore(request, requestId); - case "load-de-core": - return this.loadDeCore(request, requestId); - default: - throw new Error(`Unknown knowledge operation: ${request.operation as string}`); - } - } - - private requestRecord(request: KnowledgeRequest): Record { - return request as Record; - } - - private graphEmbeddings(request: KnowledgeRequest): { entity: Term; vectors: number[][] }[] { - const req = this.requestRecord(request); - const value = request.graphEmbeddings ?? req["graph-embeddings"]; - return Array.isArray(value) ? value as { entity: Term; vectors: number[][] }[] : []; - } - - private documentEmbeddings(request: KnowledgeRequest): DocumentEmbeddingsCore | undefined { - const req = this.requestRecord(request); - const value = request.documentEmbeddings ?? req["document-embeddings"]; - if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined; - return value as DocumentEmbeddingsCore; - } - - private async listKgCores(request: KnowledgeRequest, requestId: string): Promise { - const user = request.user ?? ""; - const prefix = user.length > 0 ? `${user}:` : ""; - - const ids: string[] = []; - for (const key of this.cores.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 }); - } - - private async getKgCore(request: KnowledgeRequest, requestId: string): Promise { - 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 }); - } - } - - private async deleteKgCore(request: KnowledgeRequest, requestId: string): Promise { - 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 }); - } - - private async putKgCore(request: KnowledgeRequest, requestId: string): Promise { - 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}`, - ); - await this.responseProducer!.send({}, { id: requestId }); - } - - private async loadKgCore(request: KnowledgeRequest, requestId: string): Promise { - 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}`); - } - - if (core.triples.length > 0) { - const producer = await this.pubsub.createProducer({ topic: "tg.flow.triples" }); - try { - await producer.send({ - metadata: { - id: coreId, - root: coreId, - user, - collection: request.collection ?? "default", - }, - triples: core.triples, + // Create producer + this.responseProducer = await this.pubsub.createProducer({ + topic: topics.knowledgeResponse, }); - } finally { - await producer.close(); - } - } - console.log( - `[KnowledgeCoreService] Loaded core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length})`, - ); - await this.responseProducer!.send({}, { id: requestId }); - } + // Create consumer + this.consumer = await this.pubsub.createConsumer({ + topic: topics.knowledgeRequest, + subscription: `${this.config.id}-knowledge-request`, + }); - private async unloadKgCore(_request: KnowledgeRequest, requestId: string): Promise { - await this.responseProducer!.send({}, { id: requestId }); - } + console.log(`[KnowledgeCoreService] Listening on ${topics.knowledgeRequest}`); - private async listDeCores(request: KnowledgeRequest, requestId: string): Promise { - 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 }); - } + // Main consume loop + while (this.running) { + try { + const msg = await this.consumer.receive(2000); + if (msg === null) continue; - private async getDeCore(request: KnowledgeRequest, requestId: string): Promise { - const user = request.user ?? ""; - const coreId = request.id ?? ""; - const key = this.coreKey(user, coreId); - const core = this.deCores.get(key); - if (core === undefined) throw new Error(`Document embeddings core not found: ${key}`); - - for (let i = 0; i < core.length; i++) { - const isLast = i === core.length - 1; - await this.responseProducer!.send( - { - documentEmbeddings: core[i], - "document-embeddings": core[i], - eos: isLast, - } as KnowledgeResponse, - { id: requestId }, - ); - } - if (core.length === 0) { - await this.responseProducer!.send({ eos: true }, { id: requestId }); - } - } - - private async deleteDeCore(request: KnowledgeRequest, requestId: string): Promise { - const user = request.user ?? ""; - const coreId = request.id ?? ""; - this.deCores.delete(this.coreKey(user, coreId)); - await this.persist(); - await this.responseProducer!.send({}, { id: requestId }); - } - - private async putDeCore(request: KnowledgeRequest, requestId: string): Promise { - const user = request.user ?? ""; - const coreId = request.id ?? ""; - const key = this.coreKey(user, coreId); - const item = this.documentEmbeddings(request); - if (item === undefined) throw new Error("put-de-core requires document-embeddings"); - const core = this.deCores.get(key) ?? []; - core.push(item); - this.deCores.set(key, core); - await this.persist(); - await this.responseProducer!.send({}, { id: requestId }); - } - - private async loadDeCore(request: KnowledgeRequest, requestId: string): Promise { - const user = request.user ?? ""; - const coreId = request.id ?? ""; - const key = this.coreKey(user, coreId); - if (!this.deCores.has(key)) throw new Error(`Document embeddings core not found: ${key}`); - await this.responseProducer!.send({}, { id: requestId }); - } - - // ---------- Persistence ---------- - - private async persist(): Promise { - try { - // Serialize Map to object - const data: { - kg: Record; - de: Record; - } = { 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); - } - } - - private async loadFromDisk(): Promise { - try { - const raw = await readTextFile(this.persistPath); - const parsed = JSON.parse(raw) as Record | { - kg?: Record; - de?: Record; - }; - - this.cores.clear(); - this.deCores.clear(); - const kg = "kg" in parsed && parsed.kg !== undefined ? parsed.kg : parsed as Record; - 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); + 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); + } } - } - console.log(`[KnowledgeCoreService] Loaded persisted state (kg=${this.cores.size}, de=${this.deCores.size})`); - } catch { - console.log("[KnowledgeCoreService] No persisted state found, starting fresh"); - } - } + }, - override async stop(): Promise { - if (this.consumer !== null) { - await this.consumer.close(); - this.consumer = null; - } - if (this.responseProducer !== null) { - await this.responseProducer.close(); - this.responseProducer = null; - } - await super.stop(); - } + + + handleMessage: async function(this: KnowledgeCoreService, msg: Message): Promise { + 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; + } + + 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 }, + ); + } + + }, + + + + handleOperation: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + switch (request.operation) { + case "list-kg-cores": + return this.listKgCores(request, requestId); + case "get-kg-core": + return this.getKgCore(request, requestId); + case "delete-kg-core": + return this.deleteKgCore(request, requestId); + case "put-kg-core": + return this.putKgCore(request, requestId); + case "load-kg-core": + return this.loadKgCore(request, requestId); + case "unload-kg-core": + return this.unloadKgCore(request, requestId); + case "list-de-cores": + return this.listDeCores(request, requestId); + case "get-de-core": + return this.getDeCore(request, requestId); + case "delete-de-core": + return this.deleteDeCore(request, requestId); + case "put-de-core": + return this.putDeCore(request, requestId); + case "load-de-core": + return this.loadDeCore(request, requestId); + default: + throw new Error(`Unknown knowledge operation: ${request.operation as string}`); + } + + }, + + + + requestRecord: function(this: KnowledgeCoreService, request: KnowledgeRequest): Record { + return request as Record; + + }, + + + + graphEmbeddings: function(this: KnowledgeCoreService, request: KnowledgeRequest): { entity: Term; vectors: number[][] }[] { + const req = this.requestRecord(request); + const value = request.graphEmbeddings ?? req["graph-embeddings"]; + return Array.isArray(value) ? value as { entity: Term; vectors: number[][] }[] : []; + + }, + + + + documentEmbeddings: function(this: KnowledgeCoreService, request: KnowledgeRequest): DocumentEmbeddingsCore | undefined { + const req = this.requestRecord(request); + const value = request.documentEmbeddings ?? req["document-embeddings"]; + if (typeof value !== "object" || value === null || Array.isArray(value)) return undefined; + return value as DocumentEmbeddingsCore; + + }, + + + + listKgCores: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + const user = request.user ?? ""; + const prefix = user.length > 0 ? `${user}:` : ""; + + const ids: string[] = []; + for (const key of (this.cores as Map).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 { + 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 { + 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 { + 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}`, + ); + await this.responseProducer!.send({}, { id: requestId }); + + }, + + + + loadKgCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + 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}`); + } + + if (core.triples.length > 0) { + const producer = await this.pubsub.createProducer({ topic: "tg.flow.triples" }); + try { + await producer.send({ + metadata: { + id: coreId, + root: coreId, + user, + collection: request.collection ?? "default", + }, + triples: core.triples, + }); + } finally { + await producer.close(); + } + } + + console.log( + `[KnowledgeCoreService] Loaded core ${key} (triples=${core.triples.length}, embeddings=${core.graphEmbeddings.length})`, + ); + await this.responseProducer!.send({}, { id: requestId }); + + }, + + + + unloadKgCore: async function(this: KnowledgeCoreService, _request: KnowledgeRequest, requestId: string): Promise { + await this.responseProducer!.send({}, { id: requestId }); + + }, + + + + listDeCores: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + 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 }); + + }, + + + + getDeCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + const user = request.user ?? ""; + const coreId = request.id ?? ""; + const key = this.coreKey(user, coreId); + const core = this.deCores.get(key); + if (core === undefined) throw new Error(`Document embeddings core not found: ${key}`); + + for (let i = 0; i < core.length; i++) { + const isLast = i === core.length - 1; + await this.responseProducer!.send( + { + documentEmbeddings: core[i], + "document-embeddings": core[i], + eos: isLast, + } as KnowledgeResponse, + { id: requestId }, + ); + } + if (core.length === 0) { + await this.responseProducer!.send({ eos: true }, { id: requestId }); + } + + }, + + + + deleteDeCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + const user = request.user ?? ""; + const coreId = request.id ?? ""; + this.deCores.delete(this.coreKey(user, coreId)); + await this.persist(); + await this.responseProducer!.send({}, { id: requestId }); + + }, + + + + putDeCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + 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 }); + + }, + + + + loadDeCore: async function(this: KnowledgeCoreService, request: KnowledgeRequest, requestId: string): Promise { + const user = request.user ?? ""; + const coreId = request.id ?? ""; + const key = this.coreKey(user, coreId); + if (!(this.deCores as Map).has(key)) throw new Error(`Document embeddings core not found: ${key}`); + await this.responseProducer!.send({}, { id: requestId }); + + }, + + + + // ---------- Persistence ---------- + + persist: async function(this: KnowledgeCoreService): Promise { + try { + // Serialize Map to object + const data: { + kg: Record; + de: Record; + } = { 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 { + try { + const raw = await readTextFile(this.persistPath); + const parsed = JSON.parse(raw) as Record | { + kg?: Record; + de?: Record; + }; + + this.cores.clear(); + this.deCores.clear(); + const kg = "kg" in parsed && parsed.kg !== undefined ? parsed.kg : parsed as Record; + for (const [key, core] of Object.entries(kg)) { + this.cores.set(key, core); + } + if ("de" in parsed && parsed.de !== undefined) { + for (const [key, core] of Object.entries(parsed.de)) { + this.deCores.set(key, core); + } + } + + console.log(`[KnowledgeCoreService] Loaded persisted state (kg=${this.cores.size}, de=${this.deCores.size})`); + } catch { + console.log("[KnowledgeCoreService] No persisted state found, starting fresh"); + } + + }, + + + + stop: async function(this: KnowledgeCoreService): Promise { + if (this.consumer !== null) { + await this.consumer.close(); + this.consumer = null; + } + if (this.responseProducer !== null) { + await this.responseProducer.close(); + this.responseProducer = null; + } + await baseStop(); + + } + }); + return service; } +export const KnowledgeCoreService = makeKnowledgeCoreService; + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } export const program = makeProcessorProgram({ id: "knowledge-svc", - make: (config) => new KnowledgeCoreService(config), + make: (config) => makeKnowledgeCoreService(config), }); export async function run(): Promise { diff --git a/ts/packages/flow/src/decoding/pdf-decoder.ts b/ts/packages/flow/src/decoding/pdf-decoder.ts index 01c3c319..96c579fb 100644 --- a/ts/packages/flow/src/decoding/pdf-decoder.ts +++ b/ts/packages/flow/src/decoding/pdf-decoder.ts @@ -16,11 +16,12 @@ import { getDocument } from "pdfjs-dist/legacy/build/pdf.mjs"; import type { TextItem } from "pdfjs-dist/types/src/display/api.js"; import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, - RequestResponseSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, + makeRequestResponseSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type Document, @@ -209,28 +210,28 @@ const onPdfDecodeMessage = Effect.fn("PdfDecoderService.onMessage")(function* ( }); export const makePdfDecoderSpecs = (): ReadonlyArray> => [ - new ConsumerSpec("decode-input", onPdfDecodeMessage), - new ProducerSpec("decode-output"), - new ProducerSpec("decode-triples"), - new RequestResponseSpec( + makeConsumerSpec("decode-input", onPdfDecodeMessage), + makeProducerSpec("decode-output"), + makeProducerSpec("decode-triples"), + makeRequestResponseSpec( "librarian-client", "librarian-request", "librarian-response", ), ]; -export class PdfDecoderService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); +export type PdfDecoderService = FlowProcessorRuntime; - for (const spec of makePdfDecoderSpecs()) { - this.registerSpecification(spec); - } - - console.log("[PdfDecoder] Service initialized"); - } +export function makePdfDecoderService(config: ProcessorConfig): PdfDecoderService { + const service = makeFlowProcessor(config, { + specifications: makePdfDecoderSpecs(), + }); + console.log("[PdfDecoder] Service initialized"); + return service; } +export const PdfDecoderService = makePdfDecoderService; + function iriTerm(iri: string): Term { return { type: "IRI", iri }; } diff --git a/ts/packages/flow/src/embeddings/ollama.ts b/ts/packages/flow/src/embeddings/ollama.ts index 6e884769..65925008 100644 --- a/ts/packages/flow/src/embeddings/ollama.ts +++ b/ts/packages/flow/src/embeddings/ollama.ts @@ -8,8 +8,8 @@ import { Effect, Layer } from "effect"; import * as S from "effect/Schema"; import { Embeddings, - EmbeddingsService, embeddingsError, + makeEmbeddingsService, makeEmbeddingsSpecs, type EmbeddingsServiceShape, type ProcessorConfig, @@ -84,25 +84,18 @@ export function OllamaEmbeddingsLive(config: OllamaEmbeddingsConfig): Layer.Laye ); } -export class OllamaEmbeddingsProcessor extends EmbeddingsService { - private readonly embeddings: EmbeddingsServiceShape; +export type OllamaEmbeddingsProcessor = ReturnType; - constructor(config: OllamaEmbeddingsConfig) { - super(config); - this.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"})`, - ); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(Embeddings, Embeddings.of(this.embeddings)), - ); - } +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"})`, + ); + return makeEmbeddingsService(config, embeddings); } +export const OllamaEmbeddingsProcessor = makeOllamaEmbeddingsProcessor; + export const program = makeFlowProcessorProgram({ id: "embeddings", specs: () => makeEmbeddingsSpecs(), diff --git a/ts/packages/flow/src/extract/knowledge-extract.ts b/ts/packages/flow/src/extract/knowledge-extract.ts index 238451b9..63a4d3c1 100644 --- a/ts/packages/flow/src/extract/knowledge-extract.ts +++ b/ts/packages/flow/src/extract/knowledge-extract.ts @@ -11,12 +11,13 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, - RequestResponseSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, + makeRequestResponseSpec, makeFlowProcessorProgram, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type Chunk, type Triples, @@ -264,36 +265,36 @@ const onKnowledgeExtractMessage = Effect.fn("KnowledgeExtractService.onMessage") }); export const makeKnowledgeExtractSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "extract-input", onKnowledgeExtractMessage, ), - new ProducerSpec("extract-triples"), - new ProducerSpec("extract-entity-contexts"), - new RequestResponseSpec( + makeProducerSpec("extract-triples"), + makeProducerSpec("extract-entity-contexts"), + makeRequestResponseSpec( "prompt-client", "prompt-request", "prompt-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "llm-client", "text-completion-request", "text-completion-response", ), ]; -export class KnowledgeExtractService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); +export type KnowledgeExtractService = FlowProcessorRuntime; - for (const spec of makeKnowledgeExtractSpecs()) { - this.registerSpecification(spec); - } - - console.log("[KnowledgeExtract] Service initialized"); - } +export function makeKnowledgeExtractService(config: ProcessorConfig): KnowledgeExtractService { + const service = makeFlowProcessor(config, { + specifications: makeKnowledgeExtractSpecs(), + }); + console.log("[KnowledgeExtract] Service initialized"); + return service; } +export const KnowledgeExtractService = makeKnowledgeExtractService; + // ---------- Helpers ---------- function toEntityIri(name: string): Term { diff --git a/ts/packages/flow/src/flow-manager/service.ts b/ts/packages/flow/src/flow-manager/service.ts index 3fc83fe4..077710c4 100644 --- a/ts/packages/flow/src/flow-manager/service.ts +++ b/ts/packages/flow/src/flow-manager/service.ts @@ -15,19 +15,16 @@ */ import { - AsyncProcessor, + makeAsyncProcessor, type ProcessorConfig, + type AsyncProcessorRuntime, topics, - RequestResponse, + makeRequestResponse, type ConfigRequest, type ConfigResponse, } from "@trustgraph/base"; import { makeProcessorProgram } from "@trustgraph/base"; -import type { - BackendProducer, - BackendConsumer, - Message, -} from "@trustgraph/base"; +import type { Message } from "@trustgraph/base"; import { Effect } from "effect"; // ---------- Internal state types ---------- @@ -136,451 +133,496 @@ const DEFAULT_BLUEPRINT: Blueprint = { // ---------- Service ---------- -export class FlowManagerService extends AsyncProcessor { - private flows = new Map(); - private blueprints = new Map(); +export type FlowManagerService = AsyncProcessorRuntime & Record; - private consumer: BackendConsumer> | null = null; - private responseProducer: BackendProducer> | null = null; - private configClient: RequestResponse | null = null; +export function makeFlowManagerService(config: ProcessorConfig): FlowManagerService { + const service = makeAsyncProcessor(config, { + run: async () => { + await service.run(); + }, + }) as FlowManagerService; + const baseStop = service.stop; + service.flows = new Map(); + service.blueprints = new Map(); + service.consumer = null; + service.responseProducer = null; + service.configClient = null; + service.blueprints.set("default", DEFAULT_BLUEPRINT); + Object.assign(service, { - constructor(config: ProcessorConfig) { - super(config); - this.blueprints.set("default", DEFAULT_BLUEPRINT); - } - protected override async run(): Promise { - // Create config client for pushing flow configs to the config service - this.configClient = new RequestResponse({ - pubsub: this.pubsub, - requestTopic: topics.configRequest, - responseTopic: topics.configResponse, - subscription: `${this.config.id}-config-client`, - }); - await this.configClient.start(); - await this.ensureDefaultBlueprint(); - await this.refreshBlueprintsFromConfig(); + run: async function(this: FlowManagerService): Promise { + // Create config client for pushing flow configs to the config service + this.configClient = makeRequestResponse({ + pubsub: this.pubsub, + requestTopic: topics.configRequest, + responseTopic: topics.configResponse, + subscription: `${this.config.id}-config-client`, + }); + await this.configClient.start(); + await this.ensureDefaultBlueprint(); + await this.refreshBlueprintsFromConfig(); - // Create producer for flow-response topic - this.responseProducer = await this.pubsub.createProducer>({ - topic: topics.flowResponse, - }); + // Create producer for flow-response topic + this.responseProducer = await this.pubsub.createProducer>({ + topic: topics.flowResponse, + }); - // Create consumer for flow-request topic - this.consumer = await this.pubsub.createConsumer>({ - topic: topics.flowRequest, - subscription: `${this.config.id}-flow-request`, - }); + // Create consumer for flow-request topic + this.consumer = await this.pubsub.createConsumer>({ + topic: topics.flowRequest, + subscription: `${this.config.id}-flow-request`, + }); - console.log(`[FlowManager] Listening on ${topics.flowRequest}`); + console.log(`[FlowManager] Listening on ${topics.flowRequest}`); - // Main consume loop (same pattern as ConfigService) - while (this.running) { - try { - const msg = await this.consumer.receive(2000); - if (msg === null) continue; + // Main consume loop (same pattern as ConfigService) + while (this.running) { + try { + const msg = await this.consumer.receive(2000); + if (msg === null) continue; - await this.handleMessage(msg); - await this.consumer.acknowledge(msg); - } catch (err) { - if (!this.running) break; - console.error("[FlowManager] Error in consume loop:", err); - await sleep(1000); - } - } - } + await this.handleMessage(msg); + await this.consumer.acknowledge(msg); + } catch (err) { + if (!this.running) break; + console.error("[FlowManager] Error in consume loop:", err); + await sleep(1000); + } + } - private async handleMessage( - msg: Message>, - ): Promise { - const request = msg.value(); - const props = msg.properties(); - const requestId = props.id; - - if (requestId === undefined || requestId.length === 0) { - console.warn("[FlowManager] Received request without id, ignoring"); - return; - } - - try { - const response = await this.handleOperation(request); - await this.responseProducer!.send(response, { id: requestId }); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - await this.responseProducer!.send( - { - error: { type: "flow-error", message }, }, - { id: requestId }, - ); - } - } - private async configRequest(request: ConfigRequest): Promise { - if (this.configClient === null) throw new Error("Config client not started"); - return this.configClient.request(request); - } - private async ensureDefaultBlueprint(): Promise { - const response = await this.configRequest({ - operation: "getvalues", - type: "flow-blueprint", - }); - if (configValues(response).some((value) => value.key === "default")) { - return; - } - await this.configRequest({ - operation: "put", - keys: ["flow-blueprint"], - values: { - default: JSON.stringify(DEFAULT_BLUEPRINT), - }, - }); - } + handleMessage: async function(this: FlowManagerService, msg: Message>): Promise { + const request = msg.value(); + const props = msg.properties(); + const requestId = props.id; - private async refreshBlueprintsFromConfig(): Promise { - const response = await this.configRequest({ - operation: "getvalues", - type: "flow-blueprint", - }); - const next = new Map(); + if (requestId === undefined || requestId.length === 0) { + console.warn("[FlowManager] Received request without id, ignoring"); + return; + } - for (const item of configValues(response)) { - const parsed = parseConfigRecord(item.value); - if (parsed === undefined) continue; - next.set(item.key, parsed as Blueprint); - } + try { + const response = await this.handleOperation(request); + await this.responseProducer!.send(response, { id: requestId }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + await this.responseProducer!.send( + { + error: { type: "flow-error", message }, + }, + { id: requestId }, + ); + } - if (!next.has("default")) { - next.set("default", DEFAULT_BLUEPRINT); - } - this.blueprints = next; - } + }, - private async refreshFlowsFromConfig(): Promise { - const response = await this.configRequest({ - operation: "getvalues", - type: "flow", - }); - const next = new Map(); - for (const item of configValues(response)) { - const parsed = parseConfigRecord(item.value); - if (parsed === undefined) continue; - const parameters = isRecord(parsed.parameters) ? parsed.parameters : {}; - next.set(item.key, { - id: item.key, - blueprintName: optionalString(parsed["blueprint-name"]) ?? optionalString(parsed.blueprintName) ?? "default", - description: optionalString(parsed.description) ?? "", - parameters, - status: "running", - }); - } - if (next.size === 0) { - const flowsResponse = await this.configRequest({ - operation: "getvalues", - type: "flows", - }); - for (const item of configValues(flowsResponse)) { - next.set(item.key, { - id: item.key, - blueprintName: "default", - description: "", - parameters: {}, + configRequest: async function(this: FlowManagerService, request: ConfigRequest): Promise { + if (this.configClient === null) throw new Error("Config client not started"); + return this.configClient.request(request); + + }, + + + + ensureDefaultBlueprint: async function(this: FlowManagerService): Promise { + const response = await this.configRequest({ + operation: "getvalues", + type: "flow-blueprint", + }); + if (configValues(response).some((value) => value.key === "default")) { + return; + } + + await this.configRequest({ + operation: "put", + keys: ["flow-blueprint"], + values: { + default: JSON.stringify(DEFAULT_BLUEPRINT), + }, + }); + + }, + + + + refreshBlueprintsFromConfig: async function(this: FlowManagerService): Promise { + const response = await this.configRequest({ + operation: "getvalues", + type: "flow-blueprint", + }); + const next = new Map(); + + for (const item of configValues(response)) { + const parsed = parseConfigRecord(item.value); + if (parsed === undefined) continue; + next.set(item.key, parsed as Blueprint); + } + + if (!next.has("default")) { + next.set("default", DEFAULT_BLUEPRINT); + } + this.blueprints = next; + + }, + + + + refreshFlowsFromConfig: async function(this: FlowManagerService): Promise { + const response = await this.configRequest({ + operation: "getvalues", + type: "flow", + }); + const next = new Map(); + + for (const item of configValues(response)) { + const parsed = parseConfigRecord(item.value); + if (parsed === undefined) continue; + const parameters = isRecord(parsed.parameters) ? parsed.parameters : {}; + next.set(item.key, { + id: item.key, + blueprintName: optionalString(parsed["blueprint-name"]) ?? optionalString(parsed.blueprintName) ?? "default", + description: optionalString(parsed.description) ?? "", + parameters, + status: "running", + }); + } + + if (next.size === 0) { + const flowsResponse = await this.configRequest({ + operation: "getvalues", + type: "flows", + }); + for (const item of configValues(flowsResponse)) { + next.set(item.key, { + id: item.key, + blueprintName: "default", + description: "", + parameters: {}, + status: "running", + }); + } + } + + this.flows = next; + + }, + + + + handleOperation: async function(this: FlowManagerService, request: Record): Promise> { + const op = request.operation as string; + await this.refreshBlueprintsFromConfig(); + await this.refreshFlowsFromConfig(); + + switch (op) { + case "list-blueprints": + return this.handleListBlueprints(); + + case "put-blueprint": + return await this.handlePutBlueprint(request); + + case "get-blueprint": + return this.handleGetBlueprint(request); + + case "delete-blueprint": + return this.handleDeleteBlueprint(request); + + case "list-flows": + return this.handleListFlows(); + + case "get-flow": + return this.handleGetFlow(request); + + case "start-flow": + return await this.handleStartFlow(request); + + case "stop-flow": + return await this.handleStopFlow(request); + + default: + throw new Error(`Unknown flow operation: ${op}`); + } + + }, + + + + // ---------- Blueprint operations ---------- + + handleListBlueprints: function(this: FlowManagerService): Record { + return { + "blueprint-names": [...this.blueprints.keys()], + }; + + }, + + + + handleGetBlueprint: function(this: FlowManagerService, request: Record): Record { + const name = request["blueprint-name"] as string | undefined; + if (name === undefined || name.length === 0) { + throw new Error("Missing blueprint-name"); + } + + const blueprint = this.blueprints.get(name); + if (blueprint === undefined) { + throw new Error(`Blueprint not found: ${name}`); + } + + return { + "blueprint-definition": JSON.stringify(blueprint), + }; + + }, + + + + handlePutBlueprint: async function(this: FlowManagerService, request: Record): Promise> { + const name = request["blueprint-name"] as string | undefined; + if (name === undefined || name.length === 0) { + throw new Error("Missing blueprint-name"); + } + const rawDefinition = request["blueprint-definition"]; + if (rawDefinition === undefined) { + throw new Error("Missing blueprint-definition"); + } + const definition = typeof rawDefinition === "string" + ? rawDefinition + : JSON.stringify(rawDefinition); + + await this.configRequest({ + operation: "put", + keys: ["flow-blueprint"], + values: { [name]: definition }, + }); + await this.refreshBlueprintsFromConfig(); + return {}; + + }, + + + + handleDeleteBlueprint: async function(this: FlowManagerService, request: Record): Promise> { + const name = request["blueprint-name"] as string | undefined; + if (name === undefined || name.length === 0) { + throw new Error("Missing blueprint-name"); + } + + if (name === "default") { + throw new Error("Cannot delete the default blueprint"); + } + + await this.configRequest({ + operation: "delete", + keys: ["flow-blueprint", name], + }); + this.blueprints.delete(name); + + return {}; + + }, + + + + // ---------- Flow operations ---------- + + handleListFlows: function(this: FlowManagerService): Record { + return { + "flow-ids": [...this.flows.keys()], + }; + + }, + + + + handleGetFlow: function(this: FlowManagerService, request: Record): Record { + const id = request["flow-id"] as string | undefined; + if (id === undefined || id.length === 0) { + throw new Error("Missing flow-id"); + } + + const inst = this.flows.get(id); + if (inst === undefined) { + throw new Error(`Flow not found: ${id}`); + } + + return { + flow: JSON.stringify({ + "blueprint-name": inst.blueprintName, + description: inst.description, + parameters: inst.parameters, + }), + }; + + }, + + + + handleStartFlow: async function(this: FlowManagerService, request: Record): Promise> { + const id = request["flow-id"] as string | undefined; + const blueprintName = (request["blueprint-name"] as string) ?? "default"; + const description = (request["description"] as string) ?? ""; + const parameters = (request["parameters"] as Record) ?? {}; + + if (id === undefined || id.length === 0) { + throw new Error("Missing flow-id"); + } + + if ((this.flows as Map).has(id)) { + throw new Error(`Flow already exists: ${id}`); + } + + const blueprint = this.blueprints.get(blueprintName); + if (blueprint === undefined) { + throw new Error(`Blueprint not found: ${blueprintName}`); + } + + // Create the flow instance + const inst: FlowInstance = { + id, + blueprintName, + description, + parameters, status: "running", + }; + this.flows.set(id, inst); + + console.log( + `[FlowManager] Started flow "${id}" with blueprint "${blueprintName}"`, + ); + + // Push updated flows config to the config service + await this.pushFlowsConfig(); + + return {}; + + }, + + + + handleStopFlow: async function(this: FlowManagerService, request: Record): Promise> { + const id = request["flow-id"] as string | undefined; + if (id === undefined || id.length === 0) { + throw new Error("Missing flow-id"); + } + + const inst = this.flows.get(id); + if (inst === undefined) { + throw new Error(`Flow not found: ${id}`); + } + + this.flows.delete(id); + + console.log(`[FlowManager] Stopped flow "${id}"`); + + await this.deleteFlowConfig(id); + + // Push updated flows config (without the removed flow) + await this.pushFlowsConfig(); + + return {}; + + }, + + + + // ---------- Config push ---------- + + /** + * Build the flows config object from all running flows and push it + * to the config service via a PUT operation. + */ + pushFlowsConfig: async function(this: FlowManagerService): Promise { + if (this.configClient === null) return; + + const flowsConfig: Record }> = {}; + const flowRecords: Record = {}; + for (const [id, inst] of this.flows) { + const blueprint = this.blueprints.get(inst.blueprintName); + if (blueprint !== undefined) { + flowsConfig[id] = { topics: blueprint.topics }; + flowRecords[id] = JSON.stringify({ + "blueprint-name": inst.blueprintName, + description: inst.description, + parameters: inst.parameters, + }); + } + } + + try { + await this.configClient.request({ + operation: "put", + keys: ["flows"], + values: flowsConfig, + }); + await this.configClient.request({ + operation: "put", + keys: ["flow"], + values: flowRecords, + }); + console.log( + `[FlowManager] Pushed flows config (${this.flows.size} active flows)`, + ); + } catch (err) { + console.error("[FlowManager] Failed to push flows config:", err); + } + + }, + + + + deleteFlowConfig: async function(this: FlowManagerService, id: string): Promise { + if (this.configClient === null) return; + await this.configClient.request({ + operation: "delete", + keys: ["flows", id], }); - } - } - - this.flows = next; - } - - private async handleOperation( - request: Record, - ): Promise> { - const op = request.operation as string; - await this.refreshBlueprintsFromConfig(); - await this.refreshFlowsFromConfig(); - - switch (op) { - case "list-blueprints": - return this.handleListBlueprints(); - - case "put-blueprint": - return await this.handlePutBlueprint(request); - - case "get-blueprint": - return this.handleGetBlueprint(request); - - case "delete-blueprint": - return this.handleDeleteBlueprint(request); - - case "list-flows": - return this.handleListFlows(); - - case "get-flow": - return this.handleGetFlow(request); - - case "start-flow": - return await this.handleStartFlow(request); - - case "stop-flow": - return await this.handleStopFlow(request); - - default: - throw new Error(`Unknown flow operation: ${op}`); - } - } - - // ---------- Blueprint operations ---------- - - private handleListBlueprints(): Record { - return { - "blueprint-names": [...this.blueprints.keys()], - }; - } - - private handleGetBlueprint( - request: Record, - ): Record { - const name = request["blueprint-name"] as string | undefined; - if (name === undefined || name.length === 0) { - throw new Error("Missing blueprint-name"); - } - - const blueprint = this.blueprints.get(name); - if (blueprint === undefined) { - throw new Error(`Blueprint not found: ${name}`); - } - - return { - "blueprint-definition": JSON.stringify(blueprint), - }; - } - - private async handlePutBlueprint( - request: Record, - ): Promise> { - const name = request["blueprint-name"] as string | undefined; - if (name === undefined || name.length === 0) { - throw new Error("Missing blueprint-name"); - } - const rawDefinition = request["blueprint-definition"]; - if (rawDefinition === undefined) { - throw new Error("Missing blueprint-definition"); - } - const definition = typeof rawDefinition === "string" - ? rawDefinition - : JSON.stringify(rawDefinition); - - await this.configRequest({ - operation: "put", - keys: ["flow-blueprint"], - values: { [name]: definition }, - }); - await this.refreshBlueprintsFromConfig(); - return {}; - } - - private async handleDeleteBlueprint( - request: Record, - ): Promise> { - const name = request["blueprint-name"] as string | undefined; - if (name === undefined || name.length === 0) { - throw new Error("Missing blueprint-name"); - } - - if (name === "default") { - throw new Error("Cannot delete the default blueprint"); - } - - await this.configRequest({ - operation: "delete", - keys: ["flow-blueprint", name], - }); - this.blueprints.delete(name); - - return {}; - } - - // ---------- Flow operations ---------- - - private handleListFlows(): Record { - return { - "flow-ids": [...this.flows.keys()], - }; - } - - private handleGetFlow( - request: Record, - ): Record { - const id = request["flow-id"] as string | undefined; - if (id === undefined || id.length === 0) { - throw new Error("Missing flow-id"); - } - - const inst = this.flows.get(id); - if (inst === undefined) { - throw new Error(`Flow not found: ${id}`); - } - - return { - flow: JSON.stringify({ - "blueprint-name": inst.blueprintName, - description: inst.description, - parameters: inst.parameters, - }), - }; - } - - private async handleStartFlow( - request: Record, - ): Promise> { - const id = request["flow-id"] as string | undefined; - const blueprintName = (request["blueprint-name"] as string) ?? "default"; - const description = (request["description"] as string) ?? ""; - const parameters = (request["parameters"] as Record) ?? {}; - - if (id === undefined || id.length === 0) { - throw new Error("Missing flow-id"); - } - - if (this.flows.has(id)) { - throw new Error(`Flow already exists: ${id}`); - } - - const blueprint = this.blueprints.get(blueprintName); - if (blueprint === undefined) { - throw new Error(`Blueprint not found: ${blueprintName}`); - } - - // Create the flow instance - const inst: FlowInstance = { - id, - blueprintName, - description, - parameters, - status: "running", - }; - this.flows.set(id, inst); - - console.log( - `[FlowManager] Started flow "${id}" with blueprint "${blueprintName}"`, - ); - - // Push updated flows config to the config service - await this.pushFlowsConfig(); - - return {}; - } - - private async handleStopFlow( - request: Record, - ): Promise> { - const id = request["flow-id"] as string | undefined; - if (id === undefined || id.length === 0) { - throw new Error("Missing flow-id"); - } - - const inst = this.flows.get(id); - if (inst === undefined) { - throw new Error(`Flow not found: ${id}`); - } - - this.flows.delete(id); - - console.log(`[FlowManager] Stopped flow "${id}"`); - - await this.deleteFlowConfig(id); - - // Push updated flows config (without the removed flow) - await this.pushFlowsConfig(); - - return {}; - } - - // ---------- Config push ---------- - - /** - * Build the flows config object from all running flows and push it - * to the config service via a PUT operation. - */ - private async pushFlowsConfig(): Promise { - if (this.configClient === null) return; - - const flowsConfig: Record }> = {}; - const flowRecords: Record = {}; - for (const [id, inst] of this.flows) { - const blueprint = this.blueprints.get(inst.blueprintName); - if (blueprint !== undefined) { - flowsConfig[id] = { topics: blueprint.topics }; - flowRecords[id] = JSON.stringify({ - "blueprint-name": inst.blueprintName, - description: inst.description, - parameters: inst.parameters, + await this.configClient.request({ + operation: "delete", + keys: ["flow", id], }); - } - } - try { - await this.configClient.request({ - operation: "put", - keys: ["flows"], - values: flowsConfig, - }); - await this.configClient.request({ - operation: "put", - keys: ["flow"], - values: flowRecords, - }); - console.log( - `[FlowManager] Pushed flows config (${this.flows.size} active flows)`, - ); - } catch (err) { - console.error("[FlowManager] Failed to push flows config:", err); - } - } + }, - private async deleteFlowConfig(id: string): Promise { - if (this.configClient === null) return; - await this.configClient.request({ - operation: "delete", - keys: ["flows", id], - }); - await this.configClient.request({ - operation: "delete", - keys: ["flow", id], - }); - } - // ---------- Lifecycle ---------- - override async stop(): Promise { - if (this.consumer !== null) { - await this.consumer.close(); - this.consumer = null; - } - if (this.responseProducer !== null) { - await this.responseProducer.close(); - this.responseProducer = null; - } - if (this.configClient !== null) { - await this.configClient.stop(); - this.configClient = null; - } - await super.stop(); - } + // ---------- Lifecycle ---------- + + stop: async function(this: FlowManagerService): Promise { + if (this.consumer !== null) { + await this.consumer.close(); + this.consumer = null; + } + if (this.responseProducer !== null) { + await this.responseProducer.close(); + this.responseProducer = null; + } + if (this.configClient !== null) { + await this.configClient.stop(); + this.configClient = null; + } + await baseStop(); + + } + }); + return service; } +export const FlowManagerService = makeFlowManagerService; + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } export const program = makeProcessorProgram({ id: "flow-manager", - make: (config) => new FlowManagerService(config), + make: (config) => makeFlowManagerService(config), }); export async function run(): Promise { diff --git a/ts/packages/flow/src/gateway/dispatch/manager.ts b/ts/packages/flow/src/gateway/dispatch/manager.ts index bc44a256..fb612b5e 100644 --- a/ts/packages/flow/src/gateway/dispatch/manager.ts +++ b/ts/packages/flow/src/gateway/dispatch/manager.ts @@ -8,7 +8,7 @@ * Python reference: trustgraph-flow/trustgraph/gateway/dispatch/manager.py */ -import { NatsBackend, RequestResponse, type PubSubBackend } from "@trustgraph/base"; +import { makeNatsBackend, makeRequestResponse, type PubSubBackend, type RequestResponse } from "@trustgraph/base"; import type { GatewayConfig } from "../server.js"; import { translateRequest, translateResponse } from "./serialize.js"; @@ -66,38 +66,75 @@ function topicName(name: string): string { // ---------- Manager ---------- -export class DispatcherManager { - private readonly pubsub: PubSubBackend; - private requestors = new Map>>(); +export interface DispatcherManager { + readonly start: () => Promise; + readonly stop: () => Promise; + readonly dispatchGlobalService: ( + kind: string, + request: Record, + ) => Promise; + readonly dispatchGlobalServiceStreaming: ( + kind: string, + request: Record, + responder: Responder, + ) => Promise; + readonly dispatchFlowService: ( + flow: string, + kind: string, + request: Record, + ) => Promise; + readonly dispatchFlowServiceStreaming: ( + flow: string, + kind: string, + request: Record, + responder: Responder, + ) => Promise; + readonly publishToTopic: ( + topic: string, + message: unknown, + id?: string, + ) => Promise; +} - constructor(config: GatewayConfig) { - this.pubsub = new NatsBackend(config.natsUrl ?? "nats://localhost:4222"); - } +export const dispatcherManagerFlowServiceNames = (): readonly string[] => [ + ...FLOW_SERVICES.keys(), +]; - async start(): Promise { +export const dispatcherManagerGlobalServiceNames = (): readonly string[] => [ + ...GLOBAL_SERVICES.keys(), +]; + +export const dispatcherManagerIsStreamingService = (kind: string): boolean => + STREAMING_SERVICES.has(kind); + +export function makeDispatcherManager(config: GatewayConfig): DispatcherManager { + const pubsub: PubSubBackend = makeNatsBackend(config.natsUrl ?? "nats://localhost:4222"); + const requestors = new Map>>(); + + const start = async (): Promise => { // Requestors are created on demand when first accessed - } + }; - async stop(): Promise { - for (const pending of this.requestors.values()) { + const stop = async (): Promise => { + for (const pending of requestors.values()) { const rr = await pending; await rr.stop(); } - await this.pubsub.close(); - } + await pubsub.close(); + }; // ---------- Internal helpers ---------- - private getRequestor( + const getRequestor = ( requestTopic: string, responseTopic: string, key: string, - ): Promise> { - let pending = this.requestors.get(key); + ): Promise> => { + let pending = requestors.get(key); if (pending === undefined) { pending = (async () => { - const rr = new RequestResponse({ - pubsub: this.pubsub, + const rr = makeRequestResponse({ + pubsub, requestTopic, responseTopic, subscription: `gateway-${key}`, @@ -105,14 +142,14 @@ export class DispatcherManager { await rr.start(); return rr; })(); - this.requestors.set(key, pending); + requestors.set(key, pending); } return pending; - } + }; - private resolveGlobalTopics( + const resolveGlobalTopics = ( kind: string, - ): { requestTopic: string; responseTopic: string } { + ): { requestTopic: string; responseTopic: string } => { const entry = GLOBAL_SERVICES.get(kind); if (entry !== undefined) { return { @@ -125,11 +162,11 @@ export class DispatcherManager { requestTopic: topicName(`${kind}-request`), responseTopic: topicName(`${kind}-response`), }; - } + }; - private resolveFlowTopics( + const resolveFlowTopics = ( kind: string, - ): { requestTopic: string; responseTopic: string } { + ): { requestTopic: string; responseTopic: string } => { const entry = FLOW_SERVICES.get(kind); if (entry !== undefined) { return { @@ -142,13 +179,13 @@ export class DispatcherManager { requestTopic: topicName(`${kind}-request`), responseTopic: topicName(`${kind}-response`), }; - } + }; /** * Determine whether a response is the final one in a streaming sequence. * Checks for various end-of-stream markers used by different services. */ - private isComplete(response: unknown): boolean { + const isComplete = (response: unknown): boolean => { if (typeof response !== "object" || response === null) return true; const res = response as Record; return ( @@ -162,50 +199,50 @@ export class DispatcherManager { // error responses are always final (res.error !== undefined && res.error !== null) ); - } + }; // ---------- Global service dispatch ---------- - async dispatchGlobalService( + const dispatchGlobalService = async ( kind: string, request: Record, - ): Promise { - const { requestTopic, responseTopic } = this.resolveGlobalTopics(kind); - const rr = await this.getRequestor(requestTopic, responseTopic, `global:${kind}`); + ): Promise => { + const { requestTopic, responseTopic } = resolveGlobalTopics(kind); + const rr = await getRequestor(requestTopic, responseTopic, `global:${kind}`); const translated = translateRequest(kind, request); const response = await rr.request(translated); return translateResponse(kind, response); - } + }; - async dispatchGlobalServiceStreaming( + const dispatchGlobalServiceStreaming = async ( kind: string, request: Record, responder: Responder, - ): Promise { - const { requestTopic, responseTopic } = this.resolveGlobalTopics(kind); - const rr = await this.getRequestor(requestTopic, responseTopic, `global:${kind}`); + ): Promise => { + const { requestTopic, responseTopic } = resolveGlobalTopics(kind); + const rr = await getRequestor(requestTopic, responseTopic, `global:${kind}`); const translated = translateRequest(kind, request); await rr.request(translated, { recipient: async (response) => { const translatedRes = translateResponse(kind, response); - const complete = this.isComplete(translatedRes); + const complete = isComplete(translatedRes); await responder(translatedRes, complete); return complete; }, }); - } + }; // ---------- Flow-scoped service dispatch ---------- - async dispatchFlowService( + const dispatchFlowService = async ( flow: string, kind: string, request: Record, - ): Promise { - const { requestTopic, responseTopic } = this.resolveFlowTopics(kind); - const rr = await this.getRequestor( + ): Promise => { + const { requestTopic, responseTopic } = resolveFlowTopics(kind); + const rr = await getRequestor( requestTopic, responseTopic, `flow:${flow}:${kind}`, @@ -214,16 +251,16 @@ export class DispatcherManager { const translated = translateRequest(kind, request); const response = await rr.request(translated); return translateResponse(kind, response); - } + }; - async dispatchFlowServiceStreaming( + const dispatchFlowServiceStreaming = async ( flow: string, kind: string, request: Record, responder: Responder, - ): Promise { - const { requestTopic, responseTopic } = this.resolveFlowTopics(kind); - const rr = await this.getRequestor( + ): Promise => { + const { requestTopic, responseTopic } = resolveFlowTopics(kind); + const rr = await getRequestor( requestTopic, responseTopic, `flow:${flow}:${kind}`, @@ -233,12 +270,12 @@ export class DispatcherManager { await rr.request(translated, { recipient: async (response) => { const translatedRes = translateResponse(kind, response); - const complete = this.isComplete(translatedRes); + const complete = isComplete(translatedRes); await responder(translatedRes, complete); return complete; }, }); - } + }; // ---------- Fire-and-forget publish ---------- @@ -246,24 +283,20 @@ export class DispatcherManager { * Publish a single message to an arbitrary topic (no request/response). * Used for injecting documents into the processing pipeline. */ - async publishToTopic(topic: string, message: unknown, id?: string): Promise { - const producer = await this.pubsub.createProducer({ topic }); + const publishToTopic = async (topic: string, message: unknown, id?: string): Promise => { + const producer = await pubsub.createProducer({ topic }); const messageId = id ?? `pub-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; await producer.send(message, { id: messageId }); await producer.close(); - } + }; - // ---------- Static introspection ---------- - - static get flowServiceNames(): readonly string[] { - return [...FLOW_SERVICES.keys()]; - } - - static get globalServiceNames(): readonly string[] { - return [...GLOBAL_SERVICES.keys()]; - } - - static isStreamingService(kind: string): boolean { - return STREAMING_SERVICES.has(kind); - } + return { + start, + stop, + dispatchGlobalService, + dispatchGlobalServiceStreaming, + dispatchFlowService, + dispatchFlowServiceStreaming, + publishToTopic, + }; } diff --git a/ts/packages/flow/src/gateway/index.ts b/ts/packages/flow/src/gateway/index.ts index 9ebbfff4..b6fc4b44 100644 --- a/ts/packages/flow/src/gateway/index.ts +++ b/ts/packages/flow/src/gateway/index.ts @@ -1,5 +1,11 @@ export { createGateway, run, type GatewayConfig } from "./server.js"; -export { DispatcherManager } from "./dispatch/manager.js"; +export { + dispatcherManagerFlowServiceNames, + dispatcherManagerGlobalServiceNames, + dispatcherManagerIsStreamingService, + makeDispatcherManager, + type DispatcherManager, +} from "./dispatch/manager.js"; export { clientTermToInternal, clientTripleToInternal, diff --git a/ts/packages/flow/src/gateway/server.ts b/ts/packages/flow/src/gateway/server.ts index 2a0d3dcf..ac6c3ca9 100644 --- a/ts/packages/flow/src/gateway/server.ts +++ b/ts/packages/flow/src/gateway/server.ts @@ -14,7 +14,7 @@ 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 } from "@trustgraph/base"; -import { DispatcherManager } from "./dispatch/manager.js"; +import { makeDispatcherManager } from "./dispatch/manager.js"; import { makeGatewayRpcServer } from "./rpc-server.js"; export interface GatewayConfig { @@ -28,7 +28,7 @@ export async function createGateway(config: GatewayConfig) { const app = Fastify({ logger: true }); await app.register(websocketPlugin); - const dispatcher = new DispatcherManager(config); + const dispatcher = makeDispatcherManager(config); await dispatcher.start(); const rpcScope = await Effect.runPromise(Scope.make()); const rpcServer = await Effect.runPromise( diff --git a/ts/packages/flow/src/index.ts b/ts/packages/flow/src/index.ts index d8654619..134610bc 100644 --- a/ts/packages/flow/src/index.ts +++ b/ts/packages/flow/src/index.ts @@ -4,39 +4,43 @@ export { createGateway, type GatewayConfig } from "./gateway/index.js"; export { OpenAIProcessor } from "./model/text-completion/openai.js"; export { ClaudeProcessor } from "./model/text-completion/claude.js"; export { - GraphRag, GraphRagEngine, GraphRagLive, + makeGraphRag, makeGraphRagEngine, normalizeGraphRagConfig, stringToTerm, termToString, + type GraphRag, type GraphRagConfig, type GraphRagClients, type GraphRagEngineShape, type GraphRagQueryOptions, } from "./retrieval/graph-rag.js"; export { - DocumentRag, DocumentRagEngine, DocumentRagLive, + makeDocumentRag, makeDocumentRagEngine, + type DocumentRag, type DocumentRagClients, type DocumentRagEngineShape, type DocumentRagQueryOptions, } from "./retrieval/document-rag.js"; -export { FalkorDBTriplesStore, type FalkorDBConfig } from "./storage/triples/falkordb.js"; -export { FalkorDBTriplesQuery, type FalkorDBQueryConfig } from "./query/triples/falkordb.js"; +export { makeFalkorDBTriplesStore, type FalkorDBTriplesStore, type FalkorDBConfig } from "./storage/triples/falkordb.js"; +export { makeFalkorDBTriplesQuery, type FalkorDBTriplesQuery, type FalkorDBQueryConfig } from "./query/triples/falkordb.js"; // Qdrant embeddings storage export { - QdrantDocEmbeddingsStore, + makeQdrantDocEmbeddingsStore, + type QdrantDocEmbeddingsStore, type QdrantDocEmbeddingsConfig, type DocEmbeddingsMessage, type DocEmbeddingChunk, } from "./storage/embeddings/qdrant-doc.js"; export { - QdrantGraphEmbeddingsStore, + makeQdrantGraphEmbeddingsStore, + type QdrantGraphEmbeddingsStore, type QdrantGraphEmbeddingsConfig, type GraphEmbeddingsMessage, type GraphEmbeddingEntity, @@ -44,13 +48,15 @@ export { // Qdrant embeddings query export { - QdrantDocEmbeddingsQuery, + makeQdrantDocEmbeddingsQuery, + type QdrantDocEmbeddingsQuery, type QdrantDocQueryConfig, type ChunkMatch, type DocEmbeddingsQueryRequest, } from "./query/embeddings/qdrant-doc.js"; export { - QdrantGraphEmbeddingsQuery, + makeQdrantGraphEmbeddingsQuery, + type QdrantGraphEmbeddingsQuery, type QdrantGraphQueryConfig, type EntityMatch, type GraphEmbeddingsQueryRequest, @@ -81,7 +87,7 @@ export { filterToolsByGroupAndState, getNextState } from "./agent/tool-filter.js // Librarian service export { LibrarianService, type LibrarianServiceConfig } from "./librarian/service.js"; -export { CollectionManager, type CollectionEntry } from "./librarian/collection-manager.js"; +export { makeCollectionManager, type CollectionEntry, type CollectionManager } from "./librarian/collection-manager.js"; // Chunking service export { recursiveSplit } from "./chunking/recursive-splitter.js"; diff --git a/ts/packages/flow/src/librarian/collection-manager.ts b/ts/packages/flow/src/librarian/collection-manager.ts index c92953fa..46dfb537 100644 --- a/ts/packages/flow/src/librarian/collection-manager.ts +++ b/ts/packages/flow/src/librarian/collection-manager.ts @@ -14,60 +14,66 @@ export interface CollectionEntry { tags: string[]; } -export class CollectionManager { - /** keyed by `${user}:${collection}` */ - private collections = new Map(); - - private key(user: string, collection: string): string { - return `${user}:${collection}`; - } - - listCollections(user: string): CollectionEntry[] { - const result: CollectionEntry[] = []; - for (const entry of this.collections.values()) { - if (entry.user === user) { - result.push(entry); - } - } - return result; - } - - getCollection(user: string, collection: string): CollectionEntry | undefined { - return this.collections.get(this.key(user, collection)); - } - - updateCollection( +export interface CollectionManager { + readonly listCollections: (user: string) => CollectionEntry[]; + readonly getCollection: (user: string, collection: string) => CollectionEntry | undefined; + readonly updateCollection: ( user: string, collection: string, name: string, description: string, tags: string[], - ): CollectionEntry { - const entry: CollectionEntry = { user, collection, name, description, tags }; - this.collections.set(this.key(user, collection), entry); - return entry; - } - - deleteCollection(user: string, collection: string): boolean { - return this.collections.delete(this.key(user, collection)); - } - - ensureCollectionExists(user: string, collection: string): CollectionEntry { - const existing = this.getCollection(user, collection); - if (existing !== undefined) return existing; - return this.updateCollection(user, collection, collection, "", []); - } - - /** Serialize to a plain array for JSON persistence. */ - toJSON(): CollectionEntry[] { - return [...this.collections.values()]; - } - - /** Restore from a serialized array. */ - loadFromJSON(entries: CollectionEntry[]): void { - this.collections.clear(); - for (const entry of entries) { - this.collections.set(this.key(entry.user, entry.collection), entry); - } - } + ) => CollectionEntry; + readonly deleteCollection: (user: string, collection: string) => boolean; + readonly ensureCollectionExists: (user: string, collection: string) => CollectionEntry; + readonly toJSON: () => CollectionEntry[]; + readonly loadFromJSON: (entries: CollectionEntry[]) => void; +} + +export function makeCollectionManager(): CollectionManager { + /** keyed by `${user}:${collection}` */ + const collections = new Map(); + + const key = (user: string, collection: string): string => `${user}:${collection}`; + + const updateCollection = ( + user: string, + collection: string, + name: string, + description: string, + tags: string[], + ): CollectionEntry => { + const entry: CollectionEntry = { user, collection, name, description, tags }; + collections.set(key(user, collection), entry); + return entry; + }; + + return { + listCollections: (user) => { + const result: CollectionEntry[] = []; + for (const entry of collections.values()) { + if (entry.user === user) { + result.push(entry); + } + } + return result; + }, + getCollection: (user, collection) => collections.get(key(user, collection)), + updateCollection, + deleteCollection: (user, collection) => collections.delete(key(user, collection)), + ensureCollectionExists: (user, collection) => { + const existing = collections.get(key(user, collection)); + if (existing !== undefined) return existing; + return updateCollection(user, collection, collection, "", []); + }, + /** Serialize to a plain array for JSON persistence. */ + toJSON: () => [...collections.values()], + /** Restore from a serialized array. */ + loadFromJSON: (entries) => { + collections.clear(); + for (const entry of entries) { + collections.set(key(entry.user, entry.collection), entry); + } + }, + }; } diff --git a/ts/packages/flow/src/librarian/service.ts b/ts/packages/flow/src/librarian/service.ts index bedb5f16..1a4e4523 100644 --- a/ts/packages/flow/src/librarian/service.ts +++ b/ts/packages/flow/src/librarian/service.ts @@ -11,8 +11,9 @@ */ import { - AsyncProcessor, + makeAsyncProcessor, type ProcessorConfig, + type AsyncProcessorRuntime, topics, type LibrarianRequest, type LibrarianResponse, @@ -22,9 +23,9 @@ import { type ProcessingMetadata, } from "@trustgraph/base"; import { makeProcessorProgram } from "@trustgraph/base"; -import type { BackendProducer, BackendConsumer, Message } from "@trustgraph/base"; +import type { Message } from "@trustgraph/base"; import { Effect } from "effect"; -import { CollectionManager } from "./collection-manager.js"; +import { makeCollectionManager } from "./collection-manager.js"; import { ensureDirectory, joinPath, @@ -58,786 +59,899 @@ function optionalString(value: unknown): string | undefined { return typeof value === "string" && value.length > 0 ? value : undefined; } -export class LibrarianService extends AsyncProcessor { - private documents = new Map(); - private processing = new Map(); - private uploads = new Map(); - private collectionManager = new CollectionManager(); - private readonly dataDir: string; - private readonly persistPath: string; +export type LibrarianService = AsyncProcessorRuntime & Record; - // Librarian topic consumers/producers - private libConsumer: BackendConsumer | null = null; - private libProducer: BackendProducer | null = null; +export function makeLibrarianService(config: LibrarianServiceConfig): LibrarianService { + const service = makeAsyncProcessor(config, { + run: async () => { + await service.run(); + }, + }) as LibrarianService; + const baseStop = service.stop; + service.documents = new Map(); + service.processing = new Map(); + service.uploads = new Map(); + service.collectionManager = makeCollectionManager(); + service.libConsumer = null; + service.libProducer = null; + service.colConsumer = null; + service.colProducer = null; + service.dataDir = config.dataDir ?? process.env.LIBRARIAN_DATA_DIR ?? "./data/librarian"; + service.persistPath = joinPath(service.dataDir, "librarian-state.json"); + Object.assign(service, { - // Collection management topic consumers/producers - private colConsumer: BackendConsumer | null = null; - private colProducer: BackendProducer | null = null; - constructor(config: LibrarianServiceConfig) { - super(config); - this.dataDir = config.dataDir ?? process.env.LIBRARIAN_DATA_DIR ?? "./data/librarian"; - this.persistPath = joinPath(this.dataDir, "librarian-state.json"); - } + run: async function(this: LibrarianService): Promise { + // Ensure directories exist + await ensureDirectory(joinPath(this.dataDir, "docs")); - protected override async run(): Promise { - // Ensure directories exist - await ensureDirectory(joinPath(this.dataDir, "docs")); + // Load persisted state + await this.loadFromDisk(); - // Load persisted state - await this.loadFromDisk(); + // Create producers + this.libProducer = await this.pubsub.createProducer({ + topic: topics.librarianResponse, + }); + this.colProducer = await this.pubsub.createProducer({ + topic: topics.collectionManagementResponse, + }); - // Create producers - this.libProducer = await this.pubsub.createProducer({ - topic: topics.librarianResponse, - }); - this.colProducer = await this.pubsub.createProducer({ - topic: topics.collectionManagementResponse, - }); + // Create consumers + this.libConsumer = await this.pubsub.createConsumer({ + topic: topics.librarianRequest, + subscription: `${this.config.id}-librarian-request`, + }); + this.colConsumer = await this.pubsub.createConsumer({ + topic: topics.collectionManagementRequest, + subscription: `${this.config.id}-collection-management-request`, + }); - // Create consumers - this.libConsumer = await this.pubsub.createConsumer({ - topic: topics.librarianRequest, - subscription: `${this.config.id}-librarian-request`, - }); - this.colConsumer = await this.pubsub.createConsumer({ - topic: topics.collectionManagementRequest, - subscription: `${this.config.id}-collection-management-request`, - }); + console.log(`[LibrarianService] Listening on ${topics.librarianRequest} and ${topics.collectionManagementRequest}`); - console.log(`[LibrarianService] Listening on ${topics.librarianRequest} and ${topics.collectionManagementRequest}`); + // Main consume loop — poll both consumers + while (this.running) { + try { + // Poll librarian requests + const libMsg = await this.libConsumer.receive(2000); + if (libMsg !== null) { + await this.handleLibrarianMessage(libMsg); + await this.libConsumer.acknowledge(libMsg); + } - // Main consume loop — poll both consumers - while (this.running) { - try { - // Poll librarian requests - const libMsg = await this.libConsumer.receive(2000); - if (libMsg !== null) { - await this.handleLibrarianMessage(libMsg); - await this.libConsumer.acknowledge(libMsg); + // Poll collection management requests + const colMsg = await this.colConsumer.receive(2000); + if (colMsg !== null) { + await this.handleCollectionMessage(colMsg); + await this.colConsumer.acknowledge(colMsg); + } + } catch (err) { + if (!this.running) break; + console.error("[LibrarianService] Error in consume loop:", err); + await sleep(1000); + } } - // Poll collection management requests - const colMsg = await this.colConsumer.receive(2000); - if (colMsg !== null) { - await this.handleCollectionMessage(colMsg); - await this.colConsumer.acknowledge(colMsg); + }, + + + + // ---------- Librarian message handling ---------- + + requestRecord: function(this: LibrarianService, request: LibrarianRequest): Record { + return request as Record; + + }, + + + + documentId: function(this: LibrarianService, request: LibrarianRequest): string | undefined { + const req = this.requestRecord(request); + return optionalString(req.documentId) ?? optionalString(req["document-id"]); + + }, + + + + processingId: function(this: LibrarianService, request: LibrarianRequest): string | undefined { + const req = this.requestRecord(request); + return optionalString(req.processingId) ?? optionalString(req["processing-id"]); + + }, + + + + documentMetadata: function(this: LibrarianService, request: LibrarianRequest): DocumentMetadata | undefined { + const req = this.requestRecord(request); + const value = req.documentMetadata ?? req["document-metadata"]; + return isRecord(value) ? this.normaliseDocumentMetadata(value) : undefined; + + }, + + + + processingMetadata: function(this: LibrarianService, request: LibrarianRequest): ProcessingMetadata | undefined { + const req = this.requestRecord(request); + const value = req.processingMetadata ?? req["processing-metadata"]; + if (!isRecord(value)) return undefined; + const documentId = optionalString(value.documentId) ?? optionalString(value["document-id"]) ?? ""; + return { + id: optionalString(value.id) ?? crypto.randomUUID(), + documentId, + "document-id": documentId, + time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000), + flow: optionalString(value.flow) ?? "default", + user: optionalString(value.user) ?? optionalString(this.requestRecord(request).user) ?? "default", + collection: optionalString(value.collection) ?? optionalString(this.requestRecord(request).collection) ?? "default", + tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [], + }; + + }, + + + + normaliseDocumentMetadata: function(this: LibrarianService, value: Record): DocumentMetadata { + const id = optionalString(value.id) ?? crypto.randomUUID(); + const parentId = optionalString(value.parentId) ?? optionalString(value["parent-id"]); + const documentType = optionalString(value.documentType) ?? optionalString(value["document-type"]) ?? "source"; + return { + id, + time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000), + kind: optionalString(value.kind) ?? "application/octet-stream", + title: optionalString(value.title) ?? "", + comments: optionalString(value.comments) ?? "", + user: optionalString(value.user) ?? "default", + tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [], + ...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}), + documentType, + "document-type": documentType, + ...(Array.isArray(value.metadata) ? { metadata: value.metadata as NonNullable } : {}), + }; + + }, + + + + publicDocument: function(this: LibrarianService, doc: DocumentMetadata): DocumentMetadata { + const parentId = doc.parentId ?? doc["parent-id"]; + const documentType = doc.documentType ?? doc["document-type"] ?? "source"; + return { + ...doc, + ...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}), + documentType, + "document-type": documentType, + }; + + }, + + + + publicProcessing: function(this: LibrarianService, proc: ProcessingMetadata): ProcessingMetadata { + const documentId = proc.documentId ?? proc["document-id"] ?? ""; + return { + ...proc, + documentId, + "document-id": documentId, + }; + + }, + + + + documentResponse: function(this: LibrarianService, doc: DocumentMetadata): LibrarianResponse { + const publicDoc = this.publicDocument(doc); + return { + documentMetadata: publicDoc, + "document-metadata": publicDoc, + }; + + }, + + + + documentsResponse: function(this: LibrarianService, docs: DocumentMetadata[]): LibrarianResponse { + const publicDocs = docs.map((doc) => this.publicDocument(doc)); + return { + documents: publicDocs, + "document-metadatas": publicDocs, + }; + + }, + + + + processingResponse: function(this: LibrarianService, records: ProcessingMetadata[]): LibrarianResponse { + const publicRecords = records.map((proc) => this.publicProcessing(proc)); + return { + processing: publicRecords, + "processing-metadata": publicRecords, + "processing-metadatas": publicRecords, + }; + + }, + + + + handleLibrarianMessage: async function(this: LibrarianService, msg: Message): Promise { + const request = msg.value(); + const props = msg.properties(); + const requestId = props.id; + + if (requestId === undefined || requestId.length === 0) { + console.warn("[LibrarianService] Received request without id, ignoring"); + return; } - } catch (err) { - if (!this.running) break; - console.error("[LibrarianService] Error in consume loop:", err); - await sleep(1000); - } - } - } - // ---------- Librarian message handling ---------- - - private requestRecord(request: LibrarianRequest): Record { - return request as Record; - } - - private documentId(request: LibrarianRequest): string | undefined { - const req = this.requestRecord(request); - return optionalString(req.documentId) ?? optionalString(req["document-id"]); - } - - private processingId(request: LibrarianRequest): string | undefined { - const req = this.requestRecord(request); - return optionalString(req.processingId) ?? optionalString(req["processing-id"]); - } - - private documentMetadata(request: LibrarianRequest): DocumentMetadata | undefined { - const req = this.requestRecord(request); - const value = req.documentMetadata ?? req["document-metadata"]; - return isRecord(value) ? this.normaliseDocumentMetadata(value) : undefined; - } - - private processingMetadata(request: LibrarianRequest): ProcessingMetadata | undefined { - const req = this.requestRecord(request); - const value = req.processingMetadata ?? req["processing-metadata"]; - if (!isRecord(value)) return undefined; - const documentId = optionalString(value.documentId) ?? optionalString(value["document-id"]) ?? ""; - return { - id: optionalString(value.id) ?? crypto.randomUUID(), - documentId, - "document-id": documentId, - time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000), - flow: optionalString(value.flow) ?? "default", - user: optionalString(value.user) ?? optionalString(this.requestRecord(request).user) ?? "default", - collection: optionalString(value.collection) ?? optionalString(this.requestRecord(request).collection) ?? "default", - tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [], - }; - } - - private normaliseDocumentMetadata(value: Record): DocumentMetadata { - const id = optionalString(value.id) ?? crypto.randomUUID(); - const parentId = optionalString(value.parentId) ?? optionalString(value["parent-id"]); - const documentType = optionalString(value.documentType) ?? optionalString(value["document-type"]) ?? "source"; - return { - id, - time: typeof value.time === "number" ? value.time : Math.floor(Date.now() / 1000), - kind: optionalString(value.kind) ?? "application/octet-stream", - title: optionalString(value.title) ?? "", - comments: optionalString(value.comments) ?? "", - user: optionalString(value.user) ?? "default", - tags: Array.isArray(value.tags) ? value.tags.filter((tag): tag is string => typeof tag === "string") : [], - ...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}), - documentType, - "document-type": documentType, - ...(Array.isArray(value.metadata) ? { metadata: value.metadata as NonNullable } : {}), - }; - } - - private publicDocument(doc: DocumentMetadata): DocumentMetadata { - const parentId = doc.parentId ?? doc["parent-id"]; - const documentType = doc.documentType ?? doc["document-type"] ?? "source"; - return { - ...doc, - ...(parentId !== undefined ? { parentId, "parent-id": parentId } : {}), - documentType, - "document-type": documentType, - }; - } - - private publicProcessing(proc: ProcessingMetadata): ProcessingMetadata { - const documentId = proc.documentId ?? proc["document-id"] ?? ""; - return { - ...proc, - documentId, - "document-id": documentId, - }; - } - - private documentResponse(doc: DocumentMetadata): LibrarianResponse { - const publicDoc = this.publicDocument(doc); - return { - documentMetadata: publicDoc, - "document-metadata": publicDoc, - }; - } - - private documentsResponse(docs: DocumentMetadata[]): LibrarianResponse { - const publicDocs = docs.map((doc) => this.publicDocument(doc)); - return { - documents: publicDocs, - "document-metadatas": publicDocs, - }; - } - - private processingResponse(records: ProcessingMetadata[]): LibrarianResponse { - const publicRecords = records.map((proc) => this.publicProcessing(proc)); - return { - processing: publicRecords, - "processing-metadata": publicRecords, - "processing-metadatas": publicRecords, - }; - } - - private async handleLibrarianMessage(msg: Message): Promise { - const request = msg.value(); - const props = msg.properties(); - const requestId = props.id; - - if (requestId === undefined || requestId.length === 0) { - console.warn("[LibrarianService] Received request without id, ignoring"); - return; - } - - try { - if (request.operation === "stream-document") { - for (const response of await this.streamDocument(request)) { + try { + if (request.operation === "stream-document") { + for (const response of await this.streamDocument(request)) { + await this.libProducer!.send(response, { id: requestId }); + } + return; + } + const response = await this.handleLibrarianOperation(request); await this.libProducer!.send(response, { id: requestId }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + await this.libProducer!.send( + { error: { type: "librarian-error", message } }, + { id: requestId }, + ); } - return; - } - const response = await this.handleLibrarianOperation(request); - await this.libProducer!.send(response, { id: requestId }); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - await this.libProducer!.send( - { error: { type: "librarian-error", message } }, - { id: requestId }, - ); - } - } - private async handleLibrarianOperation(request: LibrarianRequest): Promise { - switch (request.operation) { - case "add-document": - return this.addDocument(request); - case "remove-document": - return this.removeDocument(request); - case "update-document": - return this.updateDocument(request); - case "list-documents": - return this.listDocuments(request); - case "get-document-metadata": - return this.getDocumentMetadata(request); - case "get-document-content": - return this.getDocumentContent(request); - case "add-child-document": - return this.addChildDocument(request); - case "list-children": - return this.listChildren(request); - case "add-processing": - return this.addProcessing(request); - case "remove-processing": - return this.removeProcessing(request); - case "list-processing": - return this.listProcessing(request); - case "begin-upload": - return this.beginUpload(request); - case "upload-chunk": - return this.uploadChunk(request); - case "complete-upload": - return this.completeUpload(request); - case "get-upload-status": - return this.getUploadStatus(request); - case "abort-upload": - return this.abortUpload(request); - case "list-uploads": - return this.listUploads(request); - case "stream-document": - throw new Error("stream-document must be handled as a streaming operation"); - default: - throw new Error(`Unknown librarian operation: ${request.operation as string}`); - } - } - - private async addDocument(request: LibrarianRequest): Promise { - const meta = this.documentMetadata(request); - if (meta === undefined) throw new Error("add-document requires documentMetadata"); - - const id = meta.id; - const now = Math.floor(Date.now() / 1000); - - const doc: DocumentMetadata = { - ...meta, - id, - time: now, - }; - - this.documents.set(id, doc); - - // Store file content if provided - if (request.content !== undefined && request.content.length > 0) { - const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); - const buf = Buffer.from(request.content, "base64"); - await writeBinaryFile(filePath, buf); - } - - await this.persist(); - console.log(`[LibrarianService] Added document ${id}: ${doc.title}`); - - return this.documentResponse(doc); - } - - private async removeDocument(request: LibrarianRequest): Promise { - const id = this.documentId(request); - if (id === undefined || id.length === 0) { - throw new Error("remove-document requires documentId"); - } - - // Remove the document itself - this.documents.delete(id); - - // Remove the file - try { - await removePath(joinPath(this.dataDir, "docs", `${id}.bin`)); - } catch { - // File may not exist — that's fine - } - - // Cascade: remove children - const childIds = [...this.documents.entries()] - .filter(([, doc]) => doc.parentId === id) - .map(([childId]) => childId); - - for (const childId of childIds) { - this.documents.delete(childId); - try { - await removePath(joinPath(this.dataDir, "docs", `${childId}.bin`)); - } catch { - // ignore - } - } - - // Remove associated processing records - const procIds = [...this.processing.entries()] - .filter(([, proc]) => proc.documentId === id) - .map(([procId]) => procId); - - for (const procId of procIds) { - this.processing.delete(procId); - } - - await this.persist(); - console.log(`[LibrarianService] Removed document ${id} (cascade: ${childIds.length} children, ${procIds.length} processing)`); - - return {}; - } - - private async updateDocument(request: LibrarianRequest): Promise { - const id = this.documentId(request) ?? this.documentMetadata(request)?.id; - if (id === undefined || id.length === 0) { - throw new Error("update-document requires documentId"); - } - const existing = this.documents.get(id); - if (existing === undefined) throw new Error(`Document not found: ${id}`); - const meta = this.documentMetadata(request); - if (meta === undefined) throw new Error("update-document requires documentMetadata"); - - const doc: DocumentMetadata = this.publicDocument({ - ...existing, - ...meta, - id, - time: meta.time ?? existing.time, - }); - this.documents.set(id, doc); - await this.persist(); - return this.documentResponse(doc); - } - - private listDocuments(request: LibrarianRequest): LibrarianResponse { - const user = request.user ?? ""; - const includeChildren = this.requestRecord(request)["include-children"] === true; - const docs: DocumentMetadata[] = []; - - for (const doc of this.documents.values()) { - // Filter by user - if (user.length > 0 && doc.user !== user) continue; - // Exclude children (only top-level documents) unless explicitly requested - if (!includeChildren && doc.parentId !== undefined && doc.parentId.length > 0) continue; - docs.push(doc); - } - - return this.documentsResponse(docs); - } - - private getDocumentMetadata(request: LibrarianRequest): LibrarianResponse { - const id = this.documentId(request); - if (id === undefined || id.length === 0) { - throw new Error("get-document-metadata requires documentId"); - } - - const doc = this.documents.get(id); - if (doc === undefined) throw new Error(`Document not found: ${id}`); - - return this.documentResponse(doc); - } - - private async getDocumentContent(request: LibrarianRequest): Promise { - const id = this.documentId(request); - if (id === undefined || id.length === 0) { - throw new Error("get-document-content requires documentId"); - } - - const doc = this.documents.get(id); - if (doc === undefined) throw new Error(`Document not found: ${id}`); - - try { - const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); - const buf = await readBinaryFile(filePath); - const content = Buffer.from(buf).toString("base64"); - return { ...this.documentResponse(doc), content }; - } catch { - throw new Error(`Document content not found on disk: ${id}`); - } - } - - private async addChildDocument(request: LibrarianRequest): Promise { - const meta = this.documentMetadata(request); - if (meta === undefined) { - throw new Error("add-child-document requires documentMetadata"); - } - if (meta.parentId === undefined || meta.parentId.length === 0) { - throw new Error("add-child-document requires parentId in metadata"); - } - - // Verify parent exists - if (!this.documents.has(meta.parentId)) { - throw new Error(`Parent document not found: ${meta.parentId}`); - } - - const id = meta.id; - const now = Math.floor(Date.now() / 1000); - - const doc: DocumentMetadata = { - ...meta, - id, - time: now, - }; - - this.documents.set(id, doc); - - // Store file content if provided - if (request.content !== undefined && request.content.length > 0) { - const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); - const buf = Buffer.from(request.content, "base64"); - await writeBinaryFile(filePath, buf); - } - - await this.persist(); - console.log(`[LibrarianService] Added child document ${id} (parent: ${meta.parentId})`); - - return this.documentResponse(doc); - } - - private listChildren(request: LibrarianRequest): LibrarianResponse { - const parentId = this.documentId(request); - if (parentId === undefined || parentId.length === 0) { - throw new Error("list-children requires documentId"); - } - - const children: DocumentMetadata[] = []; - for (const doc of this.documents.values()) { - if (doc.parentId === parentId) { - children.push(doc); - } - } - - return this.documentsResponse(children); - } - - private async addProcessing(request: LibrarianRequest): Promise { - const proc = this.processingMetadata(request); - if (proc === undefined) throw new Error("add-processing requires processingMetadata"); - - const id = proc.id; - const now = Math.floor(Date.now() / 1000); - - const record: ProcessingMetadata = { - ...proc, - id, - time: now, - }; - - this.processing.set(id, record); - await this.persist(); - - console.log(`[LibrarianService] Added processing ${id} for document ${proc.documentId}`); - return this.processingResponse([record]); - } - - private async removeProcessing(request: LibrarianRequest): Promise { - const id = this.processingId(request); - if (id === undefined || id.length === 0) { - throw new Error("remove-processing requires processingId"); - } - - this.processing.delete(id); - await this.persist(); - - return {}; - } - - private listProcessing(request: LibrarianRequest): LibrarianResponse { - const documentId = this.documentId(request); - const records: ProcessingMetadata[] = []; - - for (const proc of this.processing.values()) { - const procDocumentId = proc.documentId ?? proc["document-id"]; - if (documentId !== undefined && documentId.length > 0 && procDocumentId !== documentId) { - continue; - } - records.push(proc); - } - - return this.processingResponse(records); - } - - private beginUpload(request: LibrarianRequest): LibrarianResponse { - const meta = this.documentMetadata(request); - if (meta === undefined) throw new Error("begin-upload requires documentMetadata"); - const req = this.requestRecord(request); - const totalSize = typeof req["total-size"] === "number" ? req["total-size"] : 0; - if (totalSize <= 0) throw new Error("begin-upload requires total-size"); - const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0 - ? req["chunk-size"] - : 3 * 1024 * 1024; - const totalChunks = Math.max(1, Math.ceil(totalSize / chunkSize)); - const uploadId = crypto.randomUUID(); - - this.uploads.set(uploadId, { - id: uploadId, - documentMetadata: meta, - totalSize, - chunkSize, - totalChunks, - createdAt: new Date().toISOString(), - chunks: new Map(), - user: meta.user ?? optionalString(req.user) ?? "default", - }); - - return { - "upload-id": uploadId, - "chunk-size": chunkSize, - "total-chunks": totalChunks, - } as LibrarianResponse; - } - - private uploadChunk(request: LibrarianRequest): LibrarianResponse { - const req = this.requestRecord(request); - const uploadId = optionalString(req["upload-id"]); - if (uploadId === undefined) throw new Error("upload-chunk requires upload-id"); - const session = this.uploads.get(uploadId); - if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); - const chunkIndex = typeof req["chunk-index"] === "number" ? req["chunk-index"] : -1; - if (!Number.isInteger(chunkIndex) || chunkIndex < 0 || chunkIndex >= session.totalChunks) { - throw new Error("upload-chunk requires a valid chunk-index"); - } - const content = optionalString(req.content); - if (content === undefined) throw new Error("upload-chunk requires content"); - session.chunks.set(chunkIndex, content); - - const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0); - return { - "upload-id": uploadId, - "chunk-index": chunkIndex, - "chunks-received": session.chunks.size, - "total-chunks": session.totalChunks, - "bytes-received": bytesReceived, - "total-bytes": session.totalSize, - } as LibrarianResponse; - } - - private async completeUpload(request: LibrarianRequest): Promise { - const uploadId = optionalString(this.requestRecord(request)["upload-id"]); - if (uploadId === undefined) throw new Error("complete-upload requires upload-id"); - const session = this.uploads.get(uploadId); - if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); - if (session.chunks.size !== session.totalChunks) { - throw new Error(`Upload incomplete: ${session.chunks.size}/${session.totalChunks} chunks received`); - } - - const content = Array.from({ length: session.totalChunks }, (_, i) => session.chunks.get(i) ?? "").join(""); - const response = await this.addDocument({ - operation: "add-document", - documentMetadata: session.documentMetadata, - "document-metadata": session.documentMetadata, - content, - user: session.user, - } as LibrarianRequest); - this.uploads.delete(uploadId); - const documentId = response.documentMetadata?.id ?? response["document-metadata"]?.id ?? session.documentMetadata.id; - return { - ...response, - "document-id": documentId, - "object-id": documentId, - } as LibrarianResponse; - } - - private getUploadStatus(request: LibrarianRequest): LibrarianResponse { - const uploadId = optionalString(this.requestRecord(request)["upload-id"]); - if (uploadId === undefined) throw new Error("get-upload-status requires upload-id"); - const session = this.uploads.get(uploadId); - if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); - const receivedChunks = [...session.chunks.keys()].sort((a, b) => a - b); - const receivedSet = new Set(receivedChunks); - const missingChunks = Array.from({ length: session.totalChunks }, (_, i) => i).filter((i) => !receivedSet.has(i)); - const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0); - return { - "upload-id": uploadId, - "upload-state": "in-progress", - "chunks-received": session.chunks.size, - "total-chunks": session.totalChunks, - "received-chunks": receivedChunks, - "missing-chunks": missingChunks, - "bytes-received": bytesReceived, - "total-bytes": session.totalSize, - } as LibrarianResponse; - } - - private abortUpload(request: LibrarianRequest): LibrarianResponse { - const uploadId = optionalString(this.requestRecord(request)["upload-id"]); - if (uploadId === undefined) throw new Error("abort-upload requires upload-id"); - this.uploads.delete(uploadId); - return {}; - } - - private listUploads(request: LibrarianRequest): LibrarianResponse { - const user = optionalString(this.requestRecord(request).user); - const sessions = [...this.uploads.values()] - .filter((session) => user === undefined || session.user === user) - .map((session) => ({ - "upload-id": session.id, - "document-id": session.documentMetadata.id, - "document-metadata-json": JSON.stringify(this.publicDocument(session.documentMetadata)), - "total-size": session.totalSize, - "chunk-size": session.chunkSize, - "total-chunks": session.totalChunks, - "chunks-received": session.chunks.size, - "created-at": session.createdAt, - })); - return { "upload-sessions": sessions } as LibrarianResponse; - } - - private async streamDocument(request: LibrarianRequest): Promise { - const id = this.documentId(request); - if (id === undefined) throw new Error("stream-document requires documentId"); - const req = this.requestRecord(request); - const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0 - ? req["chunk-size"] - : 1024 * 1024; - const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); - const buf = await readBinaryFile(filePath); - const base64 = Buffer.from(buf).toString("base64"); - const totalChunks = Math.max(1, Math.ceil(base64.length / chunkSize)); - return Array.from({ length: totalChunks }, (_, index) => { - const start = index * chunkSize; - const content = base64.slice(start, start + chunkSize); - return { - content, - "chunk-index": index, - "total-chunks": totalChunks, - eos: index === totalChunks - 1, - } as LibrarianResponse; - }); - } - - // ---------- Collection management ---------- - - private async handleCollectionMessage(msg: Message): Promise { - const request = msg.value(); - const props = msg.properties(); - const requestId = props.id; - - if (requestId === undefined || requestId.length === 0) { - console.warn("[LibrarianService] Received collection request without id, ignoring"); - return; - } - - try { - const response = this.handleCollectionOperation(request); - await this.colProducer!.send(response, { id: requestId }); - } catch (err) { - const message = err instanceof Error ? err.message : String(err); - await this.colProducer!.send( - { error: { type: "collection-error", message } }, - { id: requestId }, - ); - } - } - - private handleCollectionOperation(request: CollectionManagementRequest): CollectionManagementResponse { - switch (request.operation) { - case "list-collections": { - const user = request.user ?? ""; - const collections = this.collectionManager.listCollections(user); - return { collections }; - } - - case "update-collection": { - const user = request.user ?? ""; - const collection = request.collection ?? ""; - const name = request.name ?? collection; - const description = request.description ?? ""; - const tags = request.tags ?? []; - - this.collectionManager.updateCollection(user, collection, name, description, tags); - // Persist after mutation - this.persist().catch((err) => console.error("[LibrarianService] Persist failed:", err)); - - const collections = this.collectionManager.listCollections(user); - return { collections }; - } - - case "delete-collection": { - const user = request.user ?? ""; - const collection = request.collection ?? ""; - - this.collectionManager.deleteCollection(user, collection); - this.persist().catch((err) => console.error("[LibrarianService] Persist failed:", err)); + }, + + + + handleLibrarianOperation: async function(this: LibrarianService, request: LibrarianRequest): Promise { + switch (request.operation) { + case "add-document": + return this.addDocument(request); + case "remove-document": + return this.removeDocument(request); + case "update-document": + return this.updateDocument(request); + case "list-documents": + return this.listDocuments(request); + case "get-document-metadata": + return this.getDocumentMetadata(request); + case "get-document-content": + return this.getDocumentContent(request); + case "add-child-document": + return this.addChildDocument(request); + case "list-children": + return this.listChildren(request); + case "add-processing": + return this.addProcessing(request); + case "remove-processing": + return this.removeProcessing(request); + case "list-processing": + return this.listProcessing(request); + case "begin-upload": + return this.beginUpload(request); + case "upload-chunk": + return this.uploadChunk(request); + case "complete-upload": + return this.completeUpload(request); + case "get-upload-status": + return this.getUploadStatus(request); + case "abort-upload": + return this.abortUpload(request); + case "list-uploads": + return this.listUploads(request); + case "stream-document": + throw new Error("stream-document must be handled as a streaming operation"); + default: + throw new Error(`Unknown librarian operation: ${request.operation as string}`); + } + + }, + + + + addDocument: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const meta = this.documentMetadata(request); + if (meta === undefined) throw new Error("add-document requires documentMetadata"); + + const id = meta.id; + const now = Math.floor(Date.now() / 1000); + + const doc: DocumentMetadata = { + ...meta, + id, + time: now, + }; + + this.documents.set(id, doc); + + // Store file content if provided + if (request.content !== undefined && request.content.length > 0) { + const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); + const buf = Buffer.from(request.content, "base64"); + await writeBinaryFile(filePath, buf); + } + + await this.persist(); + console.log(`[LibrarianService] Added document ${id}: ${doc.title}`); + + return this.documentResponse(doc); + + }, + + + + removeDocument: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const id = this.documentId(request); + if (id === undefined || id.length === 0) { + throw new Error("remove-document requires documentId"); + } + + // Remove the document itself + this.documents.delete(id); + + // Remove the file + try { + await removePath(joinPath(this.dataDir, "docs", `${id}.bin`)); + } catch { + // File may not exist — that's fine + } + + // Cascade: remove children + const childIds = [...this.documents.entries()] + .filter(([, doc]) => doc.parentId === id) + .map(([childId]) => childId); + + for (const childId of childIds) { + this.documents.delete(childId); + try { + await removePath(joinPath(this.dataDir, "docs", `${childId}.bin`)); + } catch { + // ignore + } + } + + // Remove associated processing records + const procIds = [...this.processing.entries()] + .filter(([, proc]) => proc.documentId === id) + .map(([procId]) => procId); + + for (const procId of procIds) { + this.processing.delete(procId); + } + + await this.persist(); + console.log(`[LibrarianService] Removed document ${id} (cascade: ${childIds.length} children, ${procIds.length} processing)`); return {}; - } - default: - throw new Error(`Unknown collection operation: ${request.operation as string}`); - } - } + }, - // ---------- Persistence ---------- - private async persist(): Promise { - try { - const data = { - documents: Object.fromEntries(this.documents), - processing: Object.fromEntries(this.processing), - collections: this.collectionManager.toJSON(), - }; - const json = JSON.stringify(data, null, 2); - await writeTextFile(this.persistPath, json); - } catch (err) { - console.error("[LibrarianService] Failed to persist state:", err); - } - } - - private async loadFromDisk(): Promise { - try { - const raw = await readTextFile(this.persistPath); - const parsed = JSON.parse(raw) as { - documents?: Record; - processing?: Record; - collections?: Array<{ user: string; collection: string; name: string; description: string; tags: string[] }>; - }; - - this.documents.clear(); - if (parsed.documents !== undefined) { - for (const [id, doc] of Object.entries(parsed.documents)) { - this.documents.set(id, this.publicDocument(doc)); + updateDocument: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const id = this.documentId(request) ?? this.documentMetadata(request)?.id; + if (id === undefined || id.length === 0) { + throw new Error("update-document requires documentId"); } - } + const existing = this.documents.get(id); + if (existing === undefined) throw new Error(`Document not found: ${id}`); + const meta = this.documentMetadata(request); + if (meta === undefined) throw new Error("update-document requires documentMetadata"); - this.processing.clear(); - if (parsed.processing !== undefined) { - for (const [id, proc] of Object.entries(parsed.processing)) { - this.processing.set(id, this.publicProcessing(proc)); + const doc: DocumentMetadata = this.publicDocument({ + ...existing, + ...meta, + id, + time: meta.time ?? existing.time, + }); + this.documents.set(id, doc); + await this.persist(); + return this.documentResponse(doc); + + }, + + + + listDocuments: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const user = request.user ?? ""; + const includeChildren = this.requestRecord(request)["include-children"] === true; + const docs: DocumentMetadata[] = []; + + for (const doc of this.documents.values()) { + // Filter by user + if (user.length > 0 && doc.user !== user) continue; + // Exclude children (only top-level documents) unless explicitly requested + if (!includeChildren && doc.parentId !== undefined && doc.parentId.length > 0) continue; + docs.push(doc); } - } - if (parsed.collections !== undefined) { - this.collectionManager.loadFromJSON(parsed.collections); - } + return this.documentsResponse(docs); - console.log( - `[LibrarianService] Loaded persisted state (documents=${this.documents.size}, processing=${this.processing.size})`, - ); - } catch { - console.log("[LibrarianService] No persisted state found, starting fresh"); - } - } + }, - override async stop(): Promise { - if (this.libConsumer !== null) { - await this.libConsumer.close(); - this.libConsumer = null; - } - if (this.libProducer !== null) { - await this.libProducer.close(); - this.libProducer = null; - } - if (this.colConsumer !== null) { - await this.colConsumer.close(); - this.colConsumer = null; - } - if (this.colProducer !== null) { - await this.colProducer.close(); - this.colProducer = null; - } - await super.stop(); - } + + + getDocumentMetadata: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const id = this.documentId(request); + if (id === undefined || id.length === 0) { + throw new Error("get-document-metadata requires documentId"); + } + + const doc = this.documents.get(id); + if (doc === undefined) throw new Error(`Document not found: ${id}`); + + return this.documentResponse(doc); + + }, + + + + getDocumentContent: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const id = this.documentId(request); + if (id === undefined || id.length === 0) { + throw new Error("get-document-content requires documentId"); + } + + const doc = this.documents.get(id); + if (doc === undefined) throw new Error(`Document not found: ${id}`); + + try { + const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); + const buf = await readBinaryFile(filePath); + const content = Buffer.from(buf).toString("base64"); + return { ...this.documentResponse(doc), content }; + } catch { + throw new Error(`Document content not found on disk: ${id}`); + } + + }, + + + + addChildDocument: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const meta = this.documentMetadata(request); + if (meta === undefined) { + throw new Error("add-child-document requires documentMetadata"); + } + if (meta.parentId === undefined || meta.parentId.length === 0) { + throw new Error("add-child-document requires parentId in metadata"); + } + + // Verify parent exists + if (!(this.documents as Map).has(meta.parentId)) { + throw new Error(`Parent document not found: ${meta.parentId}`); + } + + const id = meta.id; + const now = Math.floor(Date.now() / 1000); + + const doc: DocumentMetadata = { + ...meta, + id, + time: now, + }; + + this.documents.set(id, doc); + + // Store file content if provided + if (request.content !== undefined && request.content.length > 0) { + const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); + const buf = Buffer.from(request.content, "base64"); + await writeBinaryFile(filePath, buf); + } + + await this.persist(); + console.log(`[LibrarianService] Added child document ${id} (parent: ${meta.parentId})`); + + return this.documentResponse(doc); + + }, + + + + listChildren: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const parentId = this.documentId(request); + if (parentId === undefined || parentId.length === 0) { + throw new Error("list-children requires documentId"); + } + + const children: DocumentMetadata[] = []; + for (const doc of this.documents.values()) { + if (doc.parentId === parentId) { + children.push(doc); + } + } + + return this.documentsResponse(children); + + }, + + + + addProcessing: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const proc = this.processingMetadata(request); + if (proc === undefined) throw new Error("add-processing requires processingMetadata"); + + const id = proc.id; + const now = Math.floor(Date.now() / 1000); + + const record: ProcessingMetadata = { + ...proc, + id, + time: now, + }; + + this.processing.set(id, record); + await this.persist(); + + console.log(`[LibrarianService] Added processing ${id} for document ${proc.documentId}`); + return this.processingResponse([record]); + + }, + + + + removeProcessing: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const id = this.processingId(request); + if (id === undefined || id.length === 0) { + throw new Error("remove-processing requires processingId"); + } + + this.processing.delete(id); + await this.persist(); + + return {}; + + }, + + + + listProcessing: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const documentId = this.documentId(request); + const records: ProcessingMetadata[] = []; + + for (const proc of this.processing.values()) { + const procDocumentId = proc.documentId ?? proc["document-id"]; + if (documentId !== undefined && documentId.length > 0 && procDocumentId !== documentId) { + continue; + } + records.push(proc); + } + + return this.processingResponse(records); + + }, + + + + beginUpload: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const meta = this.documentMetadata(request); + if (meta === undefined) throw new Error("begin-upload requires documentMetadata"); + const req = this.requestRecord(request); + const totalSize = typeof req["total-size"] === "number" ? req["total-size"] : 0; + if (totalSize <= 0) throw new Error("begin-upload requires total-size"); + const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0 + ? req["chunk-size"] + : 3 * 1024 * 1024; + const totalChunks = Math.max(1, Math.ceil(totalSize / chunkSize)); + const uploadId = crypto.randomUUID(); + + this.uploads.set(uploadId, { + id: uploadId, + documentMetadata: meta, + totalSize, + chunkSize, + totalChunks, + createdAt: new Date().toISOString(), + chunks: new Map(), + user: meta.user ?? optionalString(req.user) ?? "default", + }); + + return { + "upload-id": uploadId, + "chunk-size": chunkSize, + "total-chunks": totalChunks, + } as LibrarianResponse; + + }, + + + + uploadChunk: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const req = this.requestRecord(request); + const uploadId = optionalString(req["upload-id"]); + if (uploadId === undefined) throw new Error("upload-chunk requires upload-id"); + const session = this.uploads.get(uploadId); + if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); + const chunkIndex = typeof req["chunk-index"] === "number" ? req["chunk-index"] : -1; + if (!Number.isInteger(chunkIndex) || chunkIndex < 0 || chunkIndex >= session.totalChunks) { + throw new Error("upload-chunk requires a valid chunk-index"); + } + const content = optionalString(req.content); + if (content === undefined) throw new Error("upload-chunk requires content"); + session.chunks.set(chunkIndex, content); + + const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0); + return { + "upload-id": uploadId, + "chunk-index": chunkIndex, + "chunks-received": session.chunks.size, + "total-chunks": session.totalChunks, + "bytes-received": bytesReceived, + "total-bytes": session.totalSize, + } as LibrarianResponse; + + }, + + + + completeUpload: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const uploadId = optionalString(this.requestRecord(request)["upload-id"]); + if (uploadId === undefined) throw new Error("complete-upload requires upload-id"); + const session = this.uploads.get(uploadId); + if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); + if (session.chunks.size !== session.totalChunks) { + throw new Error(`Upload incomplete: ${session.chunks.size}/${session.totalChunks} chunks received`); + } + + const content = Array.from({ length: session.totalChunks }, (_, i) => session.chunks.get(i) ?? "").join(""); + const response = await this.addDocument({ + operation: "add-document", + documentMetadata: session.documentMetadata, + "document-metadata": session.documentMetadata, + content, + user: session.user, + } as LibrarianRequest); + this.uploads.delete(uploadId); + const documentId = response.documentMetadata?.id ?? response["document-metadata"]?.id ?? session.documentMetadata.id; + return { + ...response, + "document-id": documentId, + "object-id": documentId, + } as LibrarianResponse; + + }, + + + + getUploadStatus: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const uploadId = optionalString(this.requestRecord(request)["upload-id"]); + if (uploadId === undefined) throw new Error("get-upload-status requires upload-id"); + const session = this.uploads.get(uploadId); + if (session === undefined) throw new Error(`Upload not found: ${uploadId}`); + const receivedChunks = [...session.chunks.keys()].sort((a, b) => a - b); + const receivedSet = new Set(receivedChunks); + const missingChunks = Array.from({ length: session.totalChunks }, (_, i) => i).filter((i) => !receivedSet.has(i)); + const bytesReceived = [...session.chunks.values()].reduce((sum, chunk) => sum + chunk.length, 0); + return { + "upload-id": uploadId, + "upload-state": "in-progress", + "chunks-received": session.chunks.size, + "total-chunks": session.totalChunks, + "received-chunks": receivedChunks, + "missing-chunks": missingChunks, + "bytes-received": bytesReceived, + "total-bytes": session.totalSize, + } as LibrarianResponse; + + }, + + + + abortUpload: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const uploadId = optionalString(this.requestRecord(request)["upload-id"]); + if (uploadId === undefined) throw new Error("abort-upload requires upload-id"); + this.uploads.delete(uploadId); + return {}; + + }, + + + + listUploads: function(this: LibrarianService, request: LibrarianRequest): LibrarianResponse { + const user = optionalString(this.requestRecord(request).user); + const sessions = [...this.uploads.values()] + .filter((session) => user === undefined || session.user === user) + .map((session) => ({ + "upload-id": session.id, + "document-id": session.documentMetadata.id, + "document-metadata-json": JSON.stringify(this.publicDocument(session.documentMetadata)), + "total-size": session.totalSize, + "chunk-size": session.chunkSize, + "total-chunks": session.totalChunks, + "chunks-received": session.chunks.size, + "created-at": session.createdAt, + })); + return { "upload-sessions": sessions } as LibrarianResponse; + + }, + + + + streamDocument: async function(this: LibrarianService, request: LibrarianRequest): Promise { + const id = this.documentId(request); + if (id === undefined) throw new Error("stream-document requires documentId"); + const req = this.requestRecord(request); + const chunkSize = typeof req["chunk-size"] === "number" && req["chunk-size"] > 0 + ? req["chunk-size"] + : 1024 * 1024; + const filePath = joinPath(this.dataDir, "docs", `${id}.bin`); + const buf = await readBinaryFile(filePath); + const base64 = Buffer.from(buf).toString("base64"); + const totalChunks = Math.max(1, Math.ceil(base64.length / chunkSize)); + return Array.from({ length: totalChunks }, (_, index) => { + const start = index * chunkSize; + const content = base64.slice(start, start + chunkSize); + return { + content, + "chunk-index": index, + "total-chunks": totalChunks, + eos: index === totalChunks - 1, + } as LibrarianResponse; + }); + + }, + + + + // ---------- Collection management ---------- + + handleCollectionMessage: async function(this: LibrarianService, msg: Message): Promise { + const request = msg.value(); + const props = msg.properties(); + const requestId = props.id; + + if (requestId === undefined || requestId.length === 0) { + console.warn("[LibrarianService] Received collection request without id, ignoring"); + return; + } + + try { + const response = this.handleCollectionOperation(request); + await this.colProducer!.send(response, { id: requestId }); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + await this.colProducer!.send( + { error: { type: "collection-error", message } }, + { id: requestId }, + ); + } + + }, + + + + handleCollectionOperation: function(this: LibrarianService, request: CollectionManagementRequest): CollectionManagementResponse { + switch (request.operation) { + case "list-collections": { + const user = request.user ?? ""; + const collections = this.collectionManager.listCollections(user); + return { collections }; + } + + case "update-collection": { + const user = request.user ?? ""; + const collection = request.collection ?? ""; + const name = request.name ?? collection; + const description = request.description ?? ""; + const tags = request.tags ?? []; + + this.collectionManager.updateCollection(user, collection, name, description, tags); + // Persist after mutation + this.persist().catch((err: unknown) => console.error("[LibrarianService] Persist failed:", err)); + + const collections = this.collectionManager.listCollections(user); + return { collections }; + } + + case "delete-collection": { + const user = request.user ?? ""; + const collection = request.collection ?? ""; + + this.collectionManager.deleteCollection(user, collection); + this.persist().catch((err: unknown) => console.error("[LibrarianService] Persist failed:", err)); + + return {}; + } + + default: + throw new Error(`Unknown collection operation: ${request.operation as string}`); + } + + }, + + + + // ---------- Persistence ---------- + + persist: async function(this: LibrarianService): Promise { + try { + const data = { + documents: Object.fromEntries(this.documents), + processing: Object.fromEntries(this.processing), + collections: this.collectionManager.toJSON(), + }; + + const json = JSON.stringify(data, null, 2); + await writeTextFile(this.persistPath, json); + } catch (err) { + console.error("[LibrarianService] Failed to persist state:", err); + } + + }, + + + + loadFromDisk: async function(this: LibrarianService): Promise { + try { + const raw = await readTextFile(this.persistPath); + const parsed = JSON.parse(raw) as { + documents?: Record; + processing?: Record; + collections?: Array<{ user: string; collection: string; name: string; description: string; tags: string[] }>; + }; + + this.documents.clear(); + if (parsed.documents !== undefined) { + for (const [id, doc] of Object.entries(parsed.documents)) { + this.documents.set(id, this.publicDocument(doc)); + } + } + + this.processing.clear(); + if (parsed.processing !== undefined) { + for (const [id, proc] of Object.entries(parsed.processing)) { + this.processing.set(id, this.publicProcessing(proc)); + } + } + + if (parsed.collections !== undefined) { + this.collectionManager.loadFromJSON(parsed.collections); + } + + console.log( + `[LibrarianService] Loaded persisted state (documents=${this.documents.size}, processing=${this.processing.size})`, + ); + } catch { + console.log("[LibrarianService] No persisted state found, starting fresh"); + } + + }, + + + + stop: async function(this: LibrarianService): Promise { + if (this.libConsumer !== null) { + await this.libConsumer.close(); + this.libConsumer = null; + } + if (this.libProducer !== null) { + await this.libProducer.close(); + this.libProducer = null; + } + if (this.colConsumer !== null) { + await this.colConsumer.close(); + this.colConsumer = null; + } + if (this.colProducer !== null) { + await this.colProducer.close(); + this.colProducer = null; + } + await baseStop(); + + } + }); + return service; } +export const LibrarianService = makeLibrarianService; + function sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } export const program = makeProcessorProgram({ id: "librarian-svc", - make: (config) => new LibrarianService(config), + make: (config) => makeLibrarianService(config), }); export async function run(): Promise { diff --git a/ts/packages/flow/src/model/text-completion/azure-openai.ts b/ts/packages/flow/src/model/text-completion/azure-openai.ts index db843016..cf4f598a 100644 --- a/ts/packages/flow/src/model/text-completion/azure-openai.ts +++ b/ts/packages/flow/src/model/text-completion/azure-openai.ts @@ -11,10 +11,11 @@ import { AzureOpenAI } from "openai"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, @@ -22,27 +23,19 @@ import { } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class AzureOpenAIProcessor extends LlmService { - private client: AzureOpenAI; - private readonly defaultModel: string; - private readonly defaultTemperature: number; - private readonly maxOutput: number; +export type AzureOpenAIProcessorConfig = ProcessorConfig & { + model?: string; + apiKey?: string; + endpoint?: string; + apiVersion?: string; + temperature?: number; + maxOutput?: number; +}; - constructor( - config: ProcessorConfig & { - model?: string; - apiKey?: string; - endpoint?: string; - apiVersion?: string; - temperature?: number; - maxOutput?: number; - }, - ) { - super(config); - - this.defaultModel = config.model ?? process.env.AZURE_MODEL ?? "gpt-4o"; - this.defaultTemperature = config.temperature ?? 0.0; - this.maxOutput = config.maxOutput ?? 4096; +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) { @@ -59,115 +52,122 @@ export class AzureOpenAIProcessor extends LlmService { process.env.AZURE_API_VERSION ?? "2024-12-01-preview"; - this.client = new AzureOpenAI({ apiKey, apiVersion, endpoint }); + const client = new AzureOpenAI({ apiKey, apiVersion, endpoint }); console.log("[AzureOpenAI] LLM service initialized"); - } - async generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - try { - const resp = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_completion_tokens: this.maxOutput, - }); + 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(); + 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; } - throw err; - } - } + }, + supportsStreaming: () => true, + generateContentStream: async function* ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - override supportsStreaming(): boolean { - return true; - } + 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 }, + }); - async *generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + let totalInputTokens = 0; + let totalOutputTokens = 0; - try { - const stream = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_completion_tokens: this.maxOutput, - stream: true, - stream_options: { include_usage: true }, - }); + 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, + }; + } - 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, - }; + if (chunk.usage !== null && chunk.usage !== undefined) { + totalInputTokens = chunk.usage.prompt_tokens; + totalOutputTokens = chunk.usage.completion_tokens; + } } - if (chunk.usage !== null && chunk.usage !== undefined) { - totalInputTokens = chunk.usage.prompt_tokens; - totalOutputTokens = chunk.usage.completion_tokens; + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + } catch (err) { + if ((err as any)?.status === 429) { + throw tooManyRequestsError(); } + throw err; } - - yield { - text: "", - inToken: totalInputTokens, - outToken: totalOutputTokens, - model: modelName, - isFinal: true, - }; - } catch (err) { - if ((err as any)?.status === 429) { - throw tooManyRequestsError(); - } - throw err; - } - } + }, + }; } +export type AzureOpenAIProcessor = ReturnType; + +export function makeAzureOpenAIProcessor( + config: AzureOpenAIProcessorConfig, +): ReturnType { + return makeLlmService(config, makeAzureOpenAIProvider(config)); +} + +export const AzureOpenAIProcessor = makeAzureOpenAIProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new AzureOpenAIProcessor(config))), + Llm.of(makeLlmServiceShape(makeAzureOpenAIProvider(config))), ), }); diff --git a/ts/packages/flow/src/model/text-completion/claude.ts b/ts/packages/flow/src/model/text-completion/claude.ts index 39ced278..4d43dd4c 100644 --- a/ts/packages/flow/src/model/text-completion/claude.ts +++ b/ts/packages/flow/src/model/text-completion/claude.ts @@ -7,10 +7,11 @@ import Anthropic from "@anthropic-ai/sdk"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, @@ -18,132 +19,130 @@ import { } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class ClaudeProcessor extends LlmService { - private client: Anthropic; - private readonly defaultModel: string; - private readonly defaultTemperature: number; - private readonly maxOutput: number; - - constructor(config: ProcessorConfig & { - model?: string; - apiKey?: string; - temperature?: number; - maxOutput?: number; - }) { - super(config); - - this.defaultModel = config.model ?? "claude-sonnet-4-20250514"; - this.defaultTemperature = config.temperature ?? 0.0; - this.maxOutput = config.maxOutput ?? 8192; +export type ClaudeProcessorConfig = ProcessorConfig & { + model?: string; + apiKey?: string; + temperature?: number; + maxOutput?: number; +}; +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"); } - this.client = new Anthropic({ apiKey }); + const client = new Anthropic({ apiKey }); console.log("[Claude] LLM service initialized"); - } - async generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - try { - const response = await this.client.messages.create({ - model: modelName, - max_tokens: this.maxOutput, - temperature: temp, - system, - messages: [ - { role: "user", content: prompt }, - ], - }); + try { + const response = await client.messages.create({ + model: modelName, + max_tokens: maxOutput, + temperature: temp, + system, + messages: [ + { role: "user", content: prompt }, + ], + }); - const text = response.content[0].type === "text" - ? response.content[0].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; - } - } - - override supportsStreaming(): boolean { - return true; - } - - async *generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; - - try { - const stream = this.client.messages.stream({ - model: modelName, - max_tokens: this.maxOutput, - temperature: temp, - system, - messages: [ - { role: "user", content: prompt }, - ], - }); - - for await (const event of stream) { - if (event.type === "content_block_delta" && event.delta.type === "text_delta") { - yield { - text: event.delta.text, - inToken: null, - outToken: null, - model: modelName, - isFinal: false, - }; + 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; } + }, + supportsStreaming: () => true, + generateContentStream: async function* ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - 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(); + 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, + model: modelName, + isFinal: false, + }; + } + } + + const finalMessage = await stream.finalMessage(); + yield { + text: "", + inToken: finalMessage.usage.input_tokens, + outToken: finalMessage.usage.output_tokens, + model: modelName, + isFinal: true, + }; + } catch (err) { + if (err instanceof Anthropic.RateLimitError) { + throw tooManyRequestsError(); + } + throw err; } - throw err; - } - } + }, + }; } +export type ClaudeProcessor = ReturnType; + +export function makeClaudeProcessor(config: ClaudeProcessorConfig): ReturnType { + return makeLlmService(config, makeClaudeProvider(config)); +} + +export const ClaudeProcessor = makeClaudeProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new ClaudeProcessor(config))), + Llm.of(makeLlmServiceShape(makeClaudeProvider(config))), ), }); diff --git a/ts/packages/flow/src/model/text-completion/mistral.ts b/ts/packages/flow/src/model/text-completion/mistral.ts index 7708a498..8968919a 100644 --- a/ts/packages/flow/src/model/text-completion/mistral.ts +++ b/ts/packages/flow/src/model/text-completion/mistral.ts @@ -9,10 +9,11 @@ import { Mistral } from "@mistralai/mistralai"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, @@ -20,140 +21,136 @@ import { } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class MistralProcessor extends LlmService { - private client: Mistral; - private readonly defaultModel: string; - private readonly defaultTemperature: number; - private readonly maxOutput: number; - - constructor( - config: ProcessorConfig & { - model?: string; - apiKey?: string; - temperature?: number; - maxOutput?: number; - }, - ) { - super(config); - - this.defaultModel = - config.model ?? process.env.MISTRAL_MODEL ?? "ministral-8b-latest"; - this.defaultTemperature = config.temperature ?? 0.0; - this.maxOutput = config.maxOutput ?? 4096; +export type MistralProcessorConfig = ProcessorConfig & { + model?: string; + apiKey?: string; + temperature?: number; + maxOutput?: number; +}; +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"); } - this.client = new Mistral({ apiKey }); + const client = new Mistral({ apiKey }); console.log("[Mistral] LLM service initialized"); - } - async generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - try { - const resp = await this.client.chat.complete({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - maxTokens: this.maxOutput, - }); + 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(); + 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; } - throw err; - } - } + }, + supportsStreaming: () => true, + generateContentStream: async function* ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - override supportsStreaming(): boolean { - return true; - } + try { + const stream = await client.chat.stream({ + model: modelName, + messages: [ + { role: "system", content: system }, + { role: "user", content: prompt }, + ], + temperature: temp, + maxTokens: maxOutput, + }); - async *generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + let totalInputTokens = 0; + let totalOutputTokens = 0; - try { - const stream = await this.client.chat.stream({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - maxTokens: this.maxOutput, - }); + 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, + }; + } - 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, - }; + if (chunk.data?.usage !== undefined) { + totalInputTokens = chunk.data.usage.promptTokens ?? 0; + totalOutputTokens = chunk.data.usage.completionTokens ?? 0; + } } - if (chunk.data?.usage !== undefined) { - totalInputTokens = chunk.data.usage.promptTokens ?? 0; - totalOutputTokens = chunk.data.usage.completionTokens ?? 0; + 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; } - - 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; - } - } + }, + }; } +export type MistralProcessor = ReturnType; + +export function makeMistralProcessor(config: MistralProcessorConfig): ReturnType { + return makeLlmService(config, makeMistralProvider(config)); +} + +export const MistralProcessor = makeMistralProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new MistralProcessor(config))), + Llm.of(makeLlmServiceShape(makeMistralProvider(config))), ), }); diff --git a/ts/packages/flow/src/model/text-completion/ollama.ts b/ts/packages/flow/src/model/text-completion/ollama.ts index 428a3817..5ba792b7 100644 --- a/ts/packages/flow/src/model/text-completion/ollama.ts +++ b/ts/packages/flow/src/model/text-completion/ollama.ts @@ -9,27 +9,24 @@ import { Ollama } from "ollama"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class OllamaProcessor extends LlmService { - private client: Ollama; - private readonly defaultModel: string; +export type OllamaProcessorConfig = ProcessorConfig & { + model?: string; + ollamaUrl?: string; +}; - constructor(config: ProcessorConfig & { - model?: string; - ollamaUrl?: string; - }) { - super(config); - - this.defaultModel = +export function makeOllamaProvider(config: OllamaProcessorConfig): LlmProvider { + const defaultModel = config.model ?? process.env.OLLAMA_MODEL ?? "qwen2.5:0.5b"; @@ -39,96 +36,101 @@ export class OllamaProcessor extends LlmService { process.env.OLLAMA_URL ?? "http://localhost:11434"; - this.client = new Ollama({ host }); + const client = new Ollama({ host }); console.log( - `[Ollama] LLM service initialized (host=${host}, model=${this.defaultModel})`, + `[Ollama] LLM service initialized (host=${host}, model=${defaultModel})`, ); - } - async generateContent( - system: string, - prompt: string, - model?: string, - _temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const fullPrompt = system + "\n\n" + prompt; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + _temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const fullPrompt = system + "\n\n" + prompt; - const resp = await this.client.generate({ - model: modelName, - prompt: fullPrompt, - stream: false, - }); + 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 { + text: resp.response, + inToken: resp.prompt_eval_count ?? 0, + outToken: resp.eval_count ?? 0, + model: modelName, + }; + }, + supportsStreaming: () => true, + generateContentStream: async function* ( + system: string, + prompt: string, + model?: string, + _temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const fullPrompt = system + "\n\n" + prompt; - override supportsStreaming(): boolean { - return true; - } + const stream = await client.generate({ + model: modelName, + prompt: fullPrompt, + stream: true, + }); - async *generateContentStream( - system: string, - prompt: string, - model?: string, - _temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const fullPrompt = system + "\n\n" + prompt; + let totalInputTokens = 0; + let totalOutputTokens = 0; - const stream = await this.client.generate({ - model: modelName, - prompt: fullPrompt, - stream: true, - }); - - let totalInputTokens = 0; - let totalOutputTokens = 0; - - for await (const chunk of stream) { + 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; - } + 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) { - yield { - text: chunk.response, - inToken: null, - outToken: null, - model: modelName, - isFinal: false, - }; + if (chunk.response.length > 0) { + yield { + text: chunk.response, + inToken: null, + outToken: null, + model: modelName, + isFinal: false, + }; + } } - } // Final chunk with accumulated token counts - yield { - text: "", - inToken: totalInputTokens, - outToken: totalOutputTokens, - model: modelName, - isFinal: true, - }; - } + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + }, + }; } +export type OllamaProcessor = ReturnType; + +export function makeOllamaProcessor(config: OllamaProcessorConfig): ReturnType { + return makeLlmService(config, makeOllamaProvider(config)); +} + +export const OllamaProcessor = makeOllamaProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new OllamaProcessor(config))), + Llm.of(makeLlmServiceShape(makeOllamaProvider(config))), ), }); diff --git a/ts/packages/flow/src/model/text-completion/openai-compatible.ts b/ts/packages/flow/src/model/text-completion/openai-compatible.ts index bd3f8fc9..a5e4755b 100644 --- a/ts/packages/flow/src/model/text-completion/openai-compatible.ts +++ b/ts/packages/flow/src/model/text-completion/openai-compatible.ts @@ -12,37 +12,32 @@ import OpenAI from "openai"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class OpenAICompatibleProcessor extends LlmService { - private client: OpenAI; - private readonly defaultModel: string; - private readonly defaultTemperature: number; - private readonly maxOutput: number; +export type OpenAICompatibleProcessorConfig = ProcessorConfig & { + model?: string; + apiKey?: string; + baseUrl?: string; + temperature?: number; + maxOutput?: number; +}; - constructor( - config: ProcessorConfig & { - model?: string; - apiKey?: string; - baseUrl?: string; - temperature?: number; - maxOutput?: number; - }, - ) { - super(config); - - this.defaultModel = - config.model ?? process.env.OPENAI_COMPAT_MODEL ?? "default"; - this.defaultTemperature = config.temperature ?? 0.0; - this.maxOutput = config.maxOutput ?? 4096; +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) { @@ -54,100 +49,107 @@ export class OpenAICompatibleProcessor extends LlmService { const apiKey = config.apiKey ?? process.env.OPENAI_COMPAT_KEY ?? "sk-no-key-required"; - this.client = new OpenAI({ baseURL, apiKey }); + const client = new OpenAI({ baseURL, apiKey }); console.log("[OpenAI-Compatible] LLM service initialized"); - } - async generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - const resp = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_tokens: this.maxOutput, - }); + 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 { + 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* ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - override supportsStreaming(): boolean { - return true; - } + 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, + }); - async *generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + let totalInputTokens = 0; + let totalOutputTokens = 0; - const stream = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_tokens: this.maxOutput, - stream: true, - }); + 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, + }; + } - 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, - }; + if (chunk.usage !== null && chunk.usage !== undefined) { + totalInputTokens = chunk.usage.prompt_tokens; + totalOutputTokens = chunk.usage.completion_tokens; + } } - if (chunk.usage !== null && chunk.usage !== undefined) { - totalInputTokens = chunk.usage.prompt_tokens; - totalOutputTokens = chunk.usage.completion_tokens; - } - } - - yield { - text: "", - inToken: totalInputTokens, - outToken: totalOutputTokens, - model: modelName, - isFinal: true, - }; - } + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + }, + }; } +export type OpenAICompatibleProcessor = ReturnType; + +export function makeOpenAICompatibleProcessor( + config: OpenAICompatibleProcessorConfig, +): ReturnType { + return makeLlmService(config, makeOpenAICompatibleProvider(config)); +} + +export const OpenAICompatibleProcessor = makeOpenAICompatibleProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new OpenAICompatibleProcessor(config))), + Llm.of(makeLlmServiceShape(makeOpenAICompatibleProvider(config))), ), }); diff --git a/ts/packages/flow/src/model/text-completion/openai.ts b/ts/packages/flow/src/model/text-completion/openai.ts index 8c8dd2ad..a6547915 100644 --- a/ts/packages/flow/src/model/text-completion/openai.ts +++ b/ts/packages/flow/src/model/text-completion/openai.ts @@ -7,10 +7,11 @@ import OpenAI from "openai"; import { Llm, - LlmService, + makeLlmService, makeFlowProcessorProgram, makeLlmServiceShape, makeLlmSpecs, + type LlmProvider, type ProcessorConfig, type LlmResult, type LlmChunk, @@ -18,142 +19,140 @@ import { } from "@trustgraph/base"; import { Effect, Layer } from "effect"; -export class OpenAIProcessor extends LlmService { - private client: OpenAI; - private readonly defaultModel: string; - private readonly defaultTemperature: number; - private readonly maxOutput: number; - - constructor(config: ProcessorConfig & { - model?: string; - apiKey?: string; - baseUrl?: string; - temperature?: number; - maxOutput?: number; - }) { - super(config); - - this.defaultModel = config.model ?? "gpt-4o"; - this.defaultTemperature = config.temperature ?? 0.0; - this.maxOutput = config.maxOutput ?? 4096; +export type OpenAIProcessorConfig = ProcessorConfig & { + model?: string; + apiKey?: string; + baseUrl?: string; + temperature?: number; + maxOutput?: number; +}; +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"); } - this.client = new OpenAI({ + const client = new OpenAI({ apiKey, baseURL: config.baseUrl ?? process.env.OPENAI_BASE_URL, }); console.log("[OpenAI] LLM service initialized"); - } - async generateContent( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): Promise { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + return { + generateContent: async ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): Promise => { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - try { - const resp = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_completion_tokens: this.maxOutput, - }); + 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(); + 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; } - throw err; - } - } + }, + supportsStreaming: () => true, + generateContentStream: async function* ( + system: string, + prompt: string, + model?: string, + temperature?: number, + ): AsyncGenerator { + const modelName = model ?? defaultModel; + const temp = temperature ?? defaultTemperature; - override supportsStreaming(): boolean { - return true; - } + 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 }, + }); - async *generateContentStream( - system: string, - prompt: string, - model?: string, - temperature?: number, - ): AsyncGenerator { - const modelName = model ?? this.defaultModel; - const temp = temperature ?? this.defaultTemperature; + let totalInputTokens = 0; + let totalOutputTokens = 0; - try { - const stream = await this.client.chat.completions.create({ - model: modelName, - messages: [ - { role: "system", content: system }, - { role: "user", content: prompt }, - ], - temperature: temp, - max_completion_tokens: this.maxOutput, - stream: true, - stream_options: { include_usage: true }, - }); + 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, + }; + } - 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, - }; + if (chunk.usage !== null && chunk.usage !== undefined) { + totalInputTokens = chunk.usage.prompt_tokens; + totalOutputTokens = chunk.usage.completion_tokens; + } } - if (chunk.usage !== null && chunk.usage !== undefined) { - totalInputTokens = chunk.usage.prompt_tokens; - totalOutputTokens = chunk.usage.completion_tokens; + yield { + text: "", + inToken: totalInputTokens, + outToken: totalOutputTokens, + model: modelName, + isFinal: true, + }; + } catch (err) { + if (err instanceof OpenAI.RateLimitError) { + throw tooManyRequestsError(); } + throw err; } - - yield { - text: "", - inToken: totalInputTokens, - outToken: totalOutputTokens, - model: modelName, - isFinal: true, - }; - } catch (err) { - if (err instanceof OpenAI.RateLimitError) { - throw tooManyRequestsError(); - } - throw err; - } - } + }, + }; } +export type OpenAIProcessor = ReturnType; + +export function makeOpenAIProcessor(config: OpenAIProcessorConfig): ReturnType { + return makeLlmService(config, makeOpenAIProvider(config)); +} + +export const OpenAIProcessor = makeOpenAIProcessor; + export const program = makeFlowProcessorProgram({ id: "text-completion", specs: () => makeLlmSpecs(), layer: (config) => Layer.succeed( Llm, - Llm.of(makeLlmServiceShape(new OpenAIProcessor(config))), + Llm.of(makeLlmServiceShape(makeOpenAIProvider(config))), ), }); diff --git a/ts/packages/flow/src/prompt/template.ts b/ts/packages/flow/src/prompt/template.ts index 47e408d0..4c5b1873 100644 --- a/ts/packages/flow/src/prompt/template.ts +++ b/ts/packages/flow/src/prompt/template.ts @@ -25,12 +25,13 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, type ProcessorConfig, type EffectConfigHandler, type FlowContext, + type FlowProcessorRuntime, type FlowResourceNotFoundError, type MessagingDeliveryError, type PromptRequest, @@ -136,11 +137,11 @@ const makePromptTemplateRuntime = (config: PromptTemplateConfig): PromptTemplate return { specs: [ - new ConsumerSpec( + makeConsumerSpec( "prompt-request", onRequest, ), - new ProducerSpec("prompt-response"), + makeProducerSpec("prompt-response"), ], configHandlers: [onPromptConfig], }; @@ -154,27 +155,24 @@ const promptTemplateRuntime = (config: PromptTemplateConfig): PromptTemplateRunt return runtime; }; -export class PromptTemplateService extends FlowProcessor { - private readonly runtime: PromptTemplateRuntime; +export type PromptTemplateService = FlowProcessorRuntime; - constructor(config: PromptTemplateConfig) { - super(config); - - this.runtime = makePromptTemplateRuntime(config); - - for (const spec of this.runtime.specs) { - this.registerSpecification(spec); - } - for (const handler of this.runtime.configHandlers) { - this.registerConfigHandler((pushedConfig, version) => - Effect.runPromise(handler(pushedConfig, version)), - ); - } - - console.log("[PromptTemplate] Service initialized"); +export function makePromptTemplateService(config: PromptTemplateConfig): PromptTemplateService { + const runtime = makePromptTemplateRuntime(config); + const service = makeFlowProcessor(config, { + specifications: runtime.specs, + }); + for (const handler of runtime.configHandlers) { + service.registerConfigHandler((pushedConfig, version) => + Effect.runPromise(handler(pushedConfig, version)), + ); } + console.log("[PromptTemplate] Service initialized"); + return service; } +export const PromptTemplateService = makePromptTemplateService; + /** * Simple template rendering: replaces {variable} placeholders with values. * Unmatched placeholders are left as-is. diff --git a/ts/packages/flow/src/query/embeddings/qdrant-doc-service.ts b/ts/packages/flow/src/query/embeddings/qdrant-doc-service.ts index 209f1c78..6666d17c 100644 --- a/ts/packages/flow/src/query/embeddings/qdrant-doc-service.ts +++ b/ts/packages/flow/src/query/embeddings/qdrant-doc-service.ts @@ -8,10 +8,11 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type MessagingDeliveryError, @@ -78,37 +79,34 @@ const onDocEmbeddingsQueryMessage = Effect.fn("DocEmbeddingsQueryService.onMessa }); export const makeDocEmbeddingsQuerySpecs = (): ReadonlyArray> => [ - new ConsumerSpec< + makeConsumerSpec< DocumentEmbeddingsRequest, FlowResourceNotFoundError | MessagingDeliveryError, QdrantDocEmbeddingsQueryService >("document-embeddings-request", onDocEmbeddingsQueryMessage), - new ProducerSpec("document-embeddings-response"), + makeProducerSpec("document-embeddings-response"), ]; -export class DocEmbeddingsQueryService extends FlowProcessor { - private readonly query = makeQdrantDocEmbeddingsQueryService(); +export type DocEmbeddingsQueryService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeDocEmbeddingsQuerySpecs()) { - this.registerSpecification(spec); - } - - console.log("[DocEmbeddingsQuery] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService( - QdrantDocEmbeddingsQueryService, - QdrantDocEmbeddingsQueryService.of(this.query), +export function makeDocEmbeddingsQueryService(config: ProcessorConfig): DocEmbeddingsQueryService { + const query = makeQdrantDocEmbeddingsQueryService(); + const service = makeFlowProcessor(config, { + specifications: makeDocEmbeddingsQuerySpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService( + QdrantDocEmbeddingsQueryService, + QdrantDocEmbeddingsQueryService.of(query), + ), ), - ); - } + }); + console.log("[DocEmbeddingsQuery] Service initialized"); + return service; } +export const DocEmbeddingsQueryService = makeDocEmbeddingsQueryService; + export const program = makeFlowProcessorProgram({ id: "doc-embeddings-query", specs: () => makeDocEmbeddingsQuerySpecs(), diff --git a/ts/packages/flow/src/query/embeddings/qdrant-doc.ts b/ts/packages/flow/src/query/embeddings/qdrant-doc.ts index 28d17b65..8aa91387 100644 --- a/ts/packages/flow/src/query/embeddings/qdrant-doc.ts +++ b/ts/packages/flow/src/query/embeddings/qdrant-doc.ts @@ -30,22 +30,24 @@ export interface DocEmbeddingsQueryRequest { limit: number; } -export class QdrantDocEmbeddingsQuery { - private client: QdrantClient; +export interface QdrantDocEmbeddingsQuery { + readonly query: (request: DocEmbeddingsQueryRequest) => Promise; +} - constructor(config: QdrantDocQueryConfig = {}) { - const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333"; - const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY; +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; - this.client = new QdrantClient({ - url, - ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), - }); + const client = new QdrantClient({ + url, + ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), + }); - console.log("[QdrantDocQuery] Query service initialized"); - } + console.log("[QdrantDocQuery] Query service initialized"); - async query(request: DocEmbeddingsQueryRequest): Promise { + const query = async (request: DocEmbeddingsQueryRequest): Promise => { const { vector, user, collection, limit } = request; if (vector.length === 0) { @@ -56,7 +58,7 @@ export class QdrantDocEmbeddingsQuery { const collectionName = `d_${user}_${collection}_${dim}`; // Check if collection exists -- return empty if not - const exists = await this.client.collectionExists(collectionName); + const exists = await client.collectionExists(collectionName); if (!exists.exists) { console.log( `[QdrantDocQuery] Collection ${collectionName} does not exist, returning empty results`, @@ -64,7 +66,7 @@ export class QdrantDocEmbeddingsQuery { return []; } - const searchResult = await this.client.search(collectionName, { + const searchResult = await client.search(collectionName, { vector, limit, with_payload: true, @@ -84,7 +86,9 @@ export class QdrantDocEmbeddingsQuery { } return chunks; - } + }; + + return { query }; } export class QdrantDocEmbeddingsQueryError extends S.TaggedErrorClass()( @@ -119,7 +123,7 @@ const qdrantDocEmbeddingsQueryError = (operation: string, cause: unknown) => export const makeQdrantDocEmbeddingsQueryService = ( config: QdrantDocQueryConfig = {}, ): QdrantDocEmbeddingsQueryServiceShape => { - const query = new QdrantDocEmbeddingsQuery(config); + const query = makeQdrantDocEmbeddingsQuery(config); return { query: Effect.fn("QdrantDocEmbeddingsQuery.query")(function* (request) { return yield* Effect.tryPromise({ diff --git a/ts/packages/flow/src/query/embeddings/qdrant-graph-service.ts b/ts/packages/flow/src/query/embeddings/qdrant-graph-service.ts index 5425e08a..7140f30e 100644 --- a/ts/packages/flow/src/query/embeddings/qdrant-graph-service.ts +++ b/ts/packages/flow/src/query/embeddings/qdrant-graph-service.ts @@ -8,10 +8,11 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type MessagingDeliveryError, @@ -79,37 +80,34 @@ const onGraphEmbeddingsQueryMessage = Effect.fn("GraphEmbeddingsQueryService.onM }); export const makeGraphEmbeddingsQuerySpecs = (): ReadonlyArray> => [ - new ConsumerSpec< + makeConsumerSpec< GraphEmbeddingsRequest, FlowResourceNotFoundError | MessagingDeliveryError, QdrantGraphEmbeddingsQueryService >("graph-embeddings-request", onGraphEmbeddingsQueryMessage), - new ProducerSpec("graph-embeddings-response"), + makeProducerSpec("graph-embeddings-response"), ]; -export class GraphEmbeddingsQueryService extends FlowProcessor { - private readonly query = makeQdrantGraphEmbeddingsQueryService(); +export type GraphEmbeddingsQueryService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeGraphEmbeddingsQuerySpecs()) { - this.registerSpecification(spec); - } - - console.log("[GraphEmbeddingsQuery] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService( - QdrantGraphEmbeddingsQueryService, - QdrantGraphEmbeddingsQueryService.of(this.query), +export function makeGraphEmbeddingsQueryService(config: ProcessorConfig): GraphEmbeddingsQueryService { + const query = makeQdrantGraphEmbeddingsQueryService(); + const service = makeFlowProcessor(config, { + specifications: makeGraphEmbeddingsQuerySpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService( + QdrantGraphEmbeddingsQueryService, + QdrantGraphEmbeddingsQueryService.of(query), + ), ), - ); - } + }); + console.log("[GraphEmbeddingsQuery] Service initialized"); + return service; } +export const GraphEmbeddingsQueryService = makeGraphEmbeddingsQueryService; + export const program = makeFlowProcessorProgram({ id: "graph-embeddings-query", specs: () => makeGraphEmbeddingsQuerySpecs(), diff --git a/ts/packages/flow/src/query/embeddings/qdrant-graph.ts b/ts/packages/flow/src/query/embeddings/qdrant-graph.ts index 0f2cf411..5a71d1c7 100644 --- a/ts/packages/flow/src/query/embeddings/qdrant-graph.ts +++ b/ts/packages/flow/src/query/embeddings/qdrant-graph.ts @@ -39,22 +39,24 @@ function createTerm(value: string): Term { return { type: "LITERAL", value }; } -export class QdrantGraphEmbeddingsQuery { - private client: QdrantClient; +export interface QdrantGraphEmbeddingsQuery { + readonly query: (request: GraphEmbeddingsQueryRequest) => Promise; +} - constructor(config: QdrantGraphQueryConfig = {}) { - const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333"; - const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY; +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; - this.client = new QdrantClient({ - url, - ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), - }); + const client = new QdrantClient({ + url, + ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), + }); - console.log("[QdrantGraphQuery] Query service initialized"); - } + console.log("[QdrantGraphQuery] Query service initialized"); - async query(request: GraphEmbeddingsQueryRequest): Promise { + const query = async (request: GraphEmbeddingsQueryRequest): Promise => { const { vector, user, collection, limit } = request; if (vector.length === 0) { @@ -65,7 +67,7 @@ export class QdrantGraphEmbeddingsQuery { const collectionName = `t_${user}_${collection}_${dim}`; // Check if collection exists -- return empty if not - const exists = await this.client.collectionExists(collectionName); + const exists = await client.collectionExists(collectionName); if (!exists.exists) { console.log( `[QdrantGraphQuery] Collection ${collectionName} does not exist, returning empty results`, @@ -75,7 +77,7 @@ export class QdrantGraphEmbeddingsQuery { // 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 this.client.search(collectionName, { + const searchResult = await client.search(collectionName, { vector, limit: limit * 2, with_payload: true, @@ -104,7 +106,9 @@ export class QdrantGraphEmbeddingsQuery { } return entities; - } + }; + + return { query }; } export class QdrantGraphEmbeddingsQueryError extends S.TaggedErrorClass()( @@ -139,7 +143,7 @@ const qdrantGraphEmbeddingsQueryError = (operation: string, cause: unknown) => export const makeQdrantGraphEmbeddingsQueryService = ( config: QdrantGraphQueryConfig = {}, ): QdrantGraphEmbeddingsQueryServiceShape => { - const query = new QdrantGraphEmbeddingsQuery(config); + const query = makeQdrantGraphEmbeddingsQuery(config); return { query: Effect.fn("QdrantGraphEmbeddingsQuery.query")(function* (request) { return yield* Effect.tryPromise({ diff --git a/ts/packages/flow/src/query/triples/falkordb-service.ts b/ts/packages/flow/src/query/triples/falkordb-service.ts index d32b7237..1be6a406 100644 --- a/ts/packages/flow/src/query/triples/falkordb-service.ts +++ b/ts/packages/flow/src/query/triples/falkordb-service.ts @@ -8,10 +8,11 @@ */ import { - FlowProcessor, - ConsumerSpec, - ProducerSpec, + makeFlowProcessor, + makeConsumerSpec, + makeProducerSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type MessagingDeliveryError, @@ -65,37 +66,34 @@ const onTriplesQueryMessage = Effect.fn("TriplesQueryService.onMessage")(functio }); export const makeTriplesQuerySpecs = (): ReadonlyArray> => [ - new ConsumerSpec< + makeConsumerSpec< TriplesQueryRequest, FlowResourceNotFoundError | MessagingDeliveryError, FalkorDBTriplesQueryService >("triples-request", onTriplesQueryMessage), - new ProducerSpec("triples-response"), + makeProducerSpec("triples-response"), ]; -export class TriplesQueryService extends FlowProcessor { - private readonly query = makeFalkorDBTriplesQueryService(); +export type TriplesQueryService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeTriplesQuerySpecs()) { - this.registerSpecification(spec); - } - - console.log("[TriplesQuery] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService( - FalkorDBTriplesQueryService, - FalkorDBTriplesQueryService.of(this.query), +export function makeTriplesQueryService(config: ProcessorConfig): TriplesQueryService { + const query = makeFalkorDBTriplesQueryService(); + const service = makeFlowProcessor(config, { + specifications: makeTriplesQuerySpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService( + FalkorDBTriplesQueryService, + FalkorDBTriplesQueryService.of(query), + ), ), - ); - } + }); + console.log("[TriplesQuery] Service initialized"); + return service; } +export const TriplesQueryService = makeTriplesQueryService; + export const program = makeFlowProcessorProgram({ id: "triples-query", specs: () => makeTriplesQuerySpecs(), diff --git a/ts/packages/flow/src/query/triples/falkordb.ts b/ts/packages/flow/src/query/triples/falkordb.ts index 56c4134d..ec4322ed 100644 --- a/ts/packages/flow/src/query/triples/falkordb.ts +++ b/ts/packages/flow/src/query/triples/falkordb.ts @@ -41,35 +41,194 @@ function field(row: unknown, key: string): string { return (row as Record)?.[key] as string ?? ""; } -export class FalkorDBTriplesQuery { - private graph: Graph; - private connectPromise: Promise; +export interface FalkorDBTriplesQuery { + readonly queryTriples: ( + s?: Term, + p?: Term, + o?: Term, + limit?: number, + ) => Promise; +} - constructor(config: FalkorDBQueryConfig = {}) { - const url = config.url ?? process.env.FALKORDB_URL ?? "redis://localhost:6379"; - const database = config.database ?? "falkordb"; +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 }); - this.graph = new Graph(client, database); - this.connectPromise = client.connect().then(() => { - console.log(`[FalkorDBTriplesQuery] Connected to ${url}, graph: ${database}`); - }).catch((err) => { - console.error(`[FalkorDBTriplesQuery] Connection failed:`, err); - throw err; - }); - } + 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; + }); - private async ensureConnected(): Promise { - await this.connectPromise; - } + const ensureConnected = async (): Promise => { + await connectPromise; + }; - async queryTriples( + const matchPattern = async ( + out: [string, string, string][], + sv: string, pv: string, ov: string, limit: number, + ): Promise => { + 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 => { + // 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 => { + 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 { - await this.ensureConnected(); + ): Promise => { + await ensureConnected(); const sv = termToValue(s); const pv = termToValue(p); const ov = termToValue(o); @@ -79,28 +238,28 @@ export class FalkorDBTriplesQuery { // Query both Node and Literal targets for each pattern if (sv !== null && pv !== null && ov !== null) { // SPO — exact match - await this.matchPattern(rawTriples, sv, pv, ov, limit); + await matchPattern(rawTriples, sv, pv, ov, limit); } else if (sv !== null && pv !== null) { // SP — known subject + predicate - await this.matchSP(rawTriples, sv, pv, limit); + await matchSP(rawTriples, sv, pv, limit); } else if (sv !== null && ov !== null) { // SO — known subject + object - await this.matchSO(rawTriples, sv, ov, limit); + await matchSO(rawTriples, sv, ov, limit); } else if (pv !== null && ov !== null) { // PO — known predicate + object - await this.matchPO(rawTriples, pv, ov, limit); + await matchPO(rawTriples, pv, ov, limit); } else if (sv !== null) { // S only - await this.matchS(rawTriples, sv, limit); + await matchS(rawTriples, sv, limit); } else if (pv !== null) { // P only - await this.matchP(rawTriples, pv, limit); + await matchP(rawTriples, pv, limit); } else if (ov !== null) { // O only - await this.matchO(rawTriples, ov, limit); + await matchO(rawTriples, ov, limit); } else { // Wildcard — all triples - await this.matchAll(rawTriples, limit); + await matchAll(rawTriples, limit); } return rawTriples @@ -111,160 +270,9 @@ export class FalkorDBTriplesQuery { p: createTerm(p), o: createTerm(o), })); - } + }; - private async matchPattern( - out: [string, string, string][], - sv: string, pv: string, ov: string, limit: number, - ): Promise { - for (const destType of ["Literal", "Node"] as const) { - const destKey = destType === "Literal" ? "value" : "uri"; - const result = await this.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]); - } - } - } - - private async matchSP( - out: [string, string, string][], - sv: string, pv: string, limit: number, - ): Promise { - // Literals - const litResult = await this.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 this.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")]); - } - } - - private async matchSO( - out: [string, string, string][], - sv: string, ov: string, limit: number, - ): Promise { - for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) { - const result = await this.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]); - } - } - } - - private async matchPO( - out: [string, string, string][], - pv: string, ov: string, limit: number, - ): Promise { - for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) { - const result = await this.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]); - } - } - } - - private async matchS( - out: [string, string, string][], - sv: string, limit: number, - ): Promise { - const litResult = await this.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 this.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")]); - } - } - - private async matchP( - out: [string, string, string][], - pv: string, limit: number, - ): Promise { - const litResult = await this.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 this.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")]); - } - } - - private async matchO( - out: [string, string, string][], - ov: string, limit: number, - ): Promise { - for (const [destType, destKey] of [["Literal", "value"], ["Node", "uri"]] as const) { - const result = await this.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]); - } - } - } - - private async matchAll( - out: [string, string, string][], - limit: number, - ): Promise { - const litResult = await this.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 this.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")]); - } - } + return { queryTriples }; } export class FalkorDBTriplesQueryError extends S.TaggedErrorClass()( @@ -302,7 +310,7 @@ const falkorDBTriplesQueryError = (operation: string, cause: unknown) => export const makeFalkorDBTriplesQueryService = ( config: FalkorDBQueryConfig = {}, ): FalkorDBTriplesQueryServiceShape => { - const query = new FalkorDBTriplesQuery(config); + const query = makeFalkorDBTriplesQuery(config); return { queryTriples: Effect.fn("FalkorDBTriplesQuery.queryTriples")(( s: Term | undefined, diff --git a/ts/packages/flow/src/retrieval/document-rag-service.ts b/ts/packages/flow/src/retrieval/document-rag-service.ts index ac02a055..fd7d372d 100644 --- a/ts/packages/flow/src/retrieval/document-rag-service.ts +++ b/ts/packages/flow/src/retrieval/document-rag-service.ts @@ -8,10 +8,10 @@ */ import { - ConsumerSpec, - FlowProcessor, - ProducerSpec, - RequestResponseSpec, + makeConsumerSpec, + makeFlowProcessor, + makeProducerSpec, + makeRequestResponseSpec, makeFlowProcessorProgram, type DocumentEmbeddingsRequest, type DocumentEmbeddingsResponse, @@ -22,6 +22,7 @@ import { type EmbeddingsRequest, type EmbeddingsResponse, type FlowContext, + type FlowProcessorRuntime, type FlowRequestOptions, type FlowRequestor, type FlowResourceNotFoundError, @@ -113,48 +114,47 @@ const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function* }); export const makeDocumentRagSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "document-rag-request", onDocumentRagRequest, ), - new ProducerSpec("document-rag-response"), - new RequestResponseSpec( + makeProducerSpec("document-rag-response"), + makeRequestResponseSpec( "llm", "text-completion-request", "text-completion-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "embeddings", "embeddings-request", "embeddings-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "doc-embeddings", "document-embeddings-request", "document-embeddings-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "prompt", "prompt-request", "prompt-response", ), ]; -export class DocumentRagService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); - for (const spec of makeDocumentRagSpecs()) { - this.registerSpecification(spec); - } - } +export type DocumentRagService = FlowProcessorRuntime; - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(DocumentRagEngine, DocumentRagEngine.of(makeDocumentRagEngine())), - ); - } +export function makeDocumentRagService(config: ProcessorConfig): DocumentRagService { + return makeFlowProcessor(config, { + specifications: makeDocumentRagSpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService(DocumentRagEngine, DocumentRagEngine.of(makeDocumentRagEngine())), + ), + }); } +export const DocumentRagService = makeDocumentRagService; + export const program = makeFlowProcessorProgram({ id: "document-rag", specs: makeDocumentRagSpecs, diff --git a/ts/packages/flow/src/retrieval/document-rag.ts b/ts/packages/flow/src/retrieval/document-rag.ts index 400a1e40..d09137c1 100644 --- a/ts/packages/flow/src/retrieval/document-rag.ts +++ b/ts/packages/flow/src/retrieval/document-rag.ts @@ -82,20 +82,19 @@ export const DocumentRagLive: Layer.Layer = Layer.succeed( DocumentRagEngine.of(makeDocumentRagEngine()), ); -export class DocumentRag { - private readonly engine = makeDocumentRagEngine(); - private readonly clients: DocumentRagClients; - - constructor(clients: DocumentRagClients) { - this.clients = clients; - } - - query( +export interface DocumentRag { + readonly query: ( queryText: string, options?: DocumentRagQueryOptions, - ): Promise { - return Effect.runPromise(this.engine.query(this.clients, queryText, options)); - } + ) => Promise; +} + +export function makeDocumentRag(clients: DocumentRagClients): DocumentRag { + const engine = makeDocumentRagEngine(); + return { + query: (queryText, options) => + Effect.runPromise(engine.query(clients, queryText, options)), + }; } async function queryDocumentRag( diff --git a/ts/packages/flow/src/retrieval/graph-rag-service.ts b/ts/packages/flow/src/retrieval/graph-rag-service.ts index 949546c0..48b4051e 100644 --- a/ts/packages/flow/src/retrieval/graph-rag-service.ts +++ b/ts/packages/flow/src/retrieval/graph-rag-service.ts @@ -8,14 +8,15 @@ */ import { - ConsumerSpec, - FlowProcessor, - ProducerSpec, - RequestResponseSpec, + makeConsumerSpec, + makeFlowProcessor, + makeProducerSpec, + makeRequestResponseSpec, makeFlowProcessorProgram, type EffectRequestOptions, type EffectRequestResponse, type FlowContext, + type FlowProcessorRuntime, type FlowRequestOptions, type FlowRequestor, type FlowResourceNotFoundError, @@ -139,53 +140,52 @@ const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* ( }); export const makeGraphRagSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "graph-rag-request", onGraphRagRequest, ), - new ProducerSpec("graph-rag-response"), - new RequestResponseSpec( + makeProducerSpec("graph-rag-response"), + makeRequestResponseSpec( "llm", "text-completion-request", "text-completion-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "embeddings", "embeddings-request", "embeddings-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "graph-embeddings", "graph-embeddings-request", "graph-embeddings-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "triples", "triples-request", "triples-response", ), - new RequestResponseSpec( + makeRequestResponseSpec( "prompt", "prompt-request", "prompt-response", ), ]; -export class GraphRagService extends FlowProcessor { - constructor(config: ProcessorConfig) { - super(config); - for (const spec of makeGraphRagSpecs()) { - this.registerSpecification(spec); - } - } +export type GraphRagService = FlowProcessorRuntime; - override startEffect() { - return super.startEffect().pipe( - Effect.provideService(GraphRagEngine, GraphRagEngine.of(makeGraphRagEngine())), - ); - } +export function makeGraphRagService(config: ProcessorConfig): GraphRagService { + return makeFlowProcessor(config, { + specifications: makeGraphRagSpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService(GraphRagEngine, GraphRagEngine.of(makeGraphRagEngine())), + ), + }); } +export const GraphRagService = makeGraphRagService; + export const program = makeFlowProcessorProgram({ id: "graph-rag", specs: makeGraphRagSpecs, diff --git a/ts/packages/flow/src/retrieval/graph-rag.ts b/ts/packages/flow/src/retrieval/graph-rag.ts index 2c9c4e32..11e24a19 100644 --- a/ts/packages/flow/src/retrieval/graph-rag.ts +++ b/ts/packages/flow/src/retrieval/graph-rag.ts @@ -124,27 +124,22 @@ export const GraphRagLive: Layer.Layer = Layer.succeed( GraphRagEngine.of(makeGraphRagEngine()), ); -export class GraphRag { - private readonly engine = makeGraphRagEngine(); - private readonly clients: GraphRagClients; - private readonly config: GraphRagConfig; - - constructor( - clients: GraphRagClients, - config: GraphRagConfig = {}, - ) { - this.clients = clients; - this.config = config; - } - - query( +export interface GraphRag { + readonly query: ( queryText: string, options?: GraphRagQueryOptions, - ): Promise { - return Effect.runPromise( - this.engine.query(this.clients, queryText, options, this.config), - ); - } + ) => Promise; +} + +export function makeGraphRag( + clients: GraphRagClients, + config: GraphRagConfig = {}, +): GraphRag { + const engine = makeGraphRagEngine(); + return { + query: (queryText, options) => + Effect.runPromise(engine.query(clients, queryText, options, config)), + }; } async function queryGraphRag( diff --git a/ts/packages/flow/src/storage/embeddings/graph-embeddings-service.ts b/ts/packages/flow/src/storage/embeddings/graph-embeddings-service.ts index adb53426..55eb0b8a 100644 --- a/ts/packages/flow/src/storage/embeddings/graph-embeddings-service.ts +++ b/ts/packages/flow/src/storage/embeddings/graph-embeddings-service.ts @@ -10,10 +10,11 @@ */ import { - FlowProcessor, - ConsumerSpec, - RequestResponseSpec, + makeFlowProcessor, + makeConsumerSpec, + makeRequestResponseSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type FlowResourceNotFoundError, type MessagingDeliveryError, @@ -77,40 +78,37 @@ const onGraphEmbeddingsStoreMessage = Effect.fn("GraphEmbeddingsStoreService.onM }); export const makeGraphEmbeddingsStoreSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "store-graph-embeddings-input", onGraphEmbeddingsStoreMessage, ), - new RequestResponseSpec( + makeRequestResponseSpec( "embeddings-client", "embeddings-request", "embeddings-response", ), ]; -export class GraphEmbeddingsStoreService extends FlowProcessor { - private readonly store = makeQdrantGraphEmbeddingsStoreService(); +export type GraphEmbeddingsStoreService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeGraphEmbeddingsStoreSpecs()) { - this.registerSpecification(spec); - } - - console.log("[GraphEmbeddingsStore] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService( - QdrantGraphEmbeddingsStoreService, - QdrantGraphEmbeddingsStoreService.of(this.store), +export function makeGraphEmbeddingsStoreService(config: ProcessorConfig): GraphEmbeddingsStoreService { + const store = makeQdrantGraphEmbeddingsStoreService(); + const service = makeFlowProcessor(config, { + specifications: makeGraphEmbeddingsStoreSpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService( + QdrantGraphEmbeddingsStoreService, + QdrantGraphEmbeddingsStoreService.of(store), + ), ), - ); - } + }); + console.log("[GraphEmbeddingsStore] Service initialized"); + return service; } +export const GraphEmbeddingsStoreService = makeGraphEmbeddingsStoreService; + export const program = makeFlowProcessorProgram< ProcessorConfig & QdrantGraphEmbeddingsConfig, never, diff --git a/ts/packages/flow/src/storage/embeddings/qdrant-doc.ts b/ts/packages/flow/src/storage/embeddings/qdrant-doc.ts index 58648262..662c0120 100644 --- a/ts/packages/flow/src/storage/embeddings/qdrant-doc.ts +++ b/ts/packages/flow/src/storage/embeddings/qdrant-doc.ts @@ -27,51 +27,53 @@ export interface DocEmbeddingsMessage { chunks: DocEmbeddingChunk[]; } -export class QdrantDocEmbeddingsStore { - private client: QdrantClient; - private knownCollections = new Set(); +export interface QdrantDocEmbeddingsStore { + readonly store: (message: DocEmbeddingsMessage) => Promise; + readonly deleteCollection: (user: string, collection: string) => Promise; +} - constructor(config: QdrantDocEmbeddingsConfig = {}) { - const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333"; - const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY; +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; - this.client = new QdrantClient({ - url, - ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), - }); + const client = new QdrantClient({ + url, + ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), + }); + const knownCollections = new Set(); - console.log("[QdrantDocEmbeddings] Store initialized"); - } + console.log("[QdrantDocEmbeddings] Store initialized"); - private collectionName(user: string, collection: string, dim: number): string { - return `d_${user}_${collection}_${dim}`; - } + const collectionName = (user: string, collection: string, dim: number): string => + `d_${user}_${collection}_${dim}`; - private async ensureCollection(name: string, dim: number): Promise { - if (this.knownCollections.has(name)) return; + const ensureCollection = async (name: string, dim: number): Promise => { + if (knownCollections.has(name)) return; - const exists = await this.client.collectionExists(name); + const exists = await client.collectionExists(name); if (!exists.exists) { console.log(`[QdrantDocEmbeddings] Creating collection ${name} (dim=${dim})`); - await this.client.createCollection(name, { + await client.createCollection(name, { vectors: { size: dim, distance: "Cosine" }, }); } - this.knownCollections.add(name); - } + knownCollections.add(name); + }; - async store(message: DocEmbeddingsMessage): Promise { + const store = async (message: DocEmbeddingsMessage): Promise => { for (const chunk of message.chunks) { if (chunk.chunkId.length === 0) continue; if (chunk.vector.length === 0) continue; const dim = chunk.vector.length; - const name = this.collectionName(message.user, message.collection, dim); + const name = collectionName(message.user, message.collection, dim); - await this.ensureCollection(name, dim); + await ensureCollection(name, dim); - await this.client.upsert(name, { + await client.upsert(name, { points: [ { id: crypto.randomUUID(), @@ -86,12 +88,12 @@ export class QdrantDocEmbeddingsStore { ], }); } - } + }; - async deleteCollection(user: string, collection: string): Promise { + const deleteCollection = async (user: string, collection: string): Promise => { const prefix = `d_${user}_${collection}_`; - const allCollections = await this.client.getCollections(); + const allCollections = await client.getCollections(); const matching = allCollections.collections.filter((c) => c.name.startsWith(prefix), ); @@ -102,13 +104,15 @@ export class QdrantDocEmbeddingsStore { } for (const coll of matching) { - await this.client.deleteCollection(coll.name); - this.knownCollections.delete(coll.name); + await client.deleteCollection(coll.name); + knownCollections.delete(coll.name); console.log(`[QdrantDocEmbeddings] Deleted collection: ${coll.name}`); } console.log( `[QdrantDocEmbeddings] Deleted ${matching.length} collection(s) for ${user}/${collection}`, ); - } + }; + + return { store, deleteCollection }; } diff --git a/ts/packages/flow/src/storage/embeddings/qdrant-graph.ts b/ts/packages/flow/src/storage/embeddings/qdrant-graph.ts index 1006f49b..c1e3bbe2 100644 --- a/ts/packages/flow/src/storage/embeddings/qdrant-graph.ts +++ b/ts/packages/flow/src/storage/embeddings/qdrant-graph.ts @@ -43,57 +43,59 @@ function getTermValue(term: Term): string | null { } } -export class QdrantGraphEmbeddingsStore { - private client: QdrantClient; - private knownCollections = new Set(); +export interface QdrantGraphEmbeddingsStore { + readonly store: (message: GraphEmbeddingsMessage) => Promise; + readonly deleteCollection: (user: string, collection: string) => Promise; +} - constructor(config: QdrantGraphEmbeddingsConfig = {}) { - const url = config.url ?? process.env.QDRANT_URL ?? "http://localhost:6333"; - const apiKey = config.apiKey ?? process.env.QDRANT_API_KEY; +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; - this.client = new QdrantClient({ - url, - ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), - }); + const client = new QdrantClient({ + url, + ...(apiKey !== undefined && apiKey.length > 0 ? { apiKey } : {}), + }); + const knownCollections = new Set(); - console.log("[QdrantGraphEmbeddings] Store initialized"); - } + console.log("[QdrantGraphEmbeddings] Store initialized"); - private collectionName(user: string, collection: string, dim: number): string { - return `t_${user}_${collection}_${dim}`; - } + const collectionName = (user: string, collection: string, dim: number): string => + `t_${user}_${collection}_${dim}`; - private async ensureCollection(name: string, dim: number): Promise { - if (this.knownCollections.has(name)) return; + const ensureCollection = async (name: string, dim: number): Promise => { + if (knownCollections.has(name)) return; - const exists = await this.client.collectionExists(name); + const exists = await client.collectionExists(name); if (!exists.exists) { console.log(`[QdrantGraphEmbeddings] Creating collection ${name} (dim=${dim})`); - await this.client.createCollection(name, { + await client.createCollection(name, { vectors: { size: dim, distance: "Cosine" }, }); } - this.knownCollections.add(name); - } + knownCollections.add(name); + }; - async store(message: GraphEmbeddingsMessage): Promise { + const store = async (message: GraphEmbeddingsMessage): Promise => { for (const entry of message.entities) { const entityValue = getTermValue(entry.entity); if (entityValue === null || entityValue.length === 0) continue; if (entry.vector.length === 0) continue; const dim = entry.vector.length; - const name = this.collectionName(message.user, message.collection, dim); + const name = collectionName(message.user, message.collection, dim); - await this.ensureCollection(name, dim); + await ensureCollection(name, dim); const payload: Record = { entity: entityValue }; if (entry.chunkId !== undefined && entry.chunkId.length > 0) { payload.chunk_id = entry.chunkId; } - await this.client.upsert(name, { + await client.upsert(name, { points: [ { id: crypto.randomUUID(), @@ -103,12 +105,12 @@ export class QdrantGraphEmbeddingsStore { ], }); } - } + }; - async deleteCollection(user: string, collection: string): Promise { + const deleteCollection = async (user: string, collection: string): Promise => { const prefix = `t_${user}_${collection}_`; - const allCollections = await this.client.getCollections(); + const allCollections = await client.getCollections(); const matching = allCollections.collections.filter((c) => c.name.startsWith(prefix), ); @@ -119,15 +121,17 @@ export class QdrantGraphEmbeddingsStore { } for (const coll of matching) { - await this.client.deleteCollection(coll.name); - this.knownCollections.delete(coll.name); + await client.deleteCollection(coll.name); + knownCollections.delete(coll.name); console.log(`[QdrantGraphEmbeddings] Deleted collection: ${coll.name}`); } console.log( `[QdrantGraphEmbeddings] Deleted ${matching.length} collection(s) for ${user}/${collection}`, ); - } + }; + + return { store, deleteCollection }; } export class QdrantGraphEmbeddingsStoreError extends S.TaggedErrorClass()( @@ -166,7 +170,7 @@ const qdrantGraphEmbeddingsStoreError = (operation: string, cause: unknown) => export const makeQdrantGraphEmbeddingsStoreService = ( config: QdrantGraphEmbeddingsConfig = {}, ): QdrantGraphEmbeddingsStoreServiceShape => { - const store = new QdrantGraphEmbeddingsStore(config); + const store = makeQdrantGraphEmbeddingsStore(config); return { store: Effect.fn("QdrantGraphEmbeddingsStore.store")(function* (message) { return yield* Effect.tryPromise({ diff --git a/ts/packages/flow/src/storage/triples/falkordb-service.ts b/ts/packages/flow/src/storage/triples/falkordb-service.ts index 5457d3c5..713deaa9 100644 --- a/ts/packages/flow/src/storage/triples/falkordb-service.ts +++ b/ts/packages/flow/src/storage/triples/falkordb-service.ts @@ -9,9 +9,10 @@ */ import { - FlowProcessor, - ConsumerSpec, + makeFlowProcessor, + makeConsumerSpec, type ProcessorConfig, + type FlowProcessorRuntime, type FlowContext, type Triples, type Spec, @@ -45,35 +46,32 @@ const onStoreTriplesMessage = Effect.fn("TriplesStoreService.onMessage")(functio }); export const makeTriplesStoreSpecs = (): ReadonlyArray> => [ - new ConsumerSpec( + makeConsumerSpec( "store-triples-input", onStoreTriplesMessage, ), ]; -export class TriplesStoreService extends FlowProcessor { - private readonly store = makeFalkorDBTriplesStoreService(); +export type TriplesStoreService = FlowProcessorRuntime; - constructor(config: ProcessorConfig) { - super(config); - - for (const spec of makeTriplesStoreSpecs()) { - this.registerSpecification(spec); - } - - console.log("[TriplesStore] Service initialized"); - } - - override startEffect() { - return super.startEffect().pipe( - Effect.provideService( - FalkorDBTriplesStoreService, - FalkorDBTriplesStoreService.of(this.store), +export function makeTriplesStoreService(config: ProcessorConfig): TriplesStoreService { + const store = makeFalkorDBTriplesStoreService(); + const service = makeFlowProcessor(config, { + specifications: makeTriplesStoreSpecs(), + provide: (effect) => + effect.pipe( + Effect.provideService( + FalkorDBTriplesStoreService, + FalkorDBTriplesStoreService.of(store), + ), ), - ); - } + }); + console.log("[TriplesStore] Service initialized"); + return service; } +export const TriplesStoreService = makeTriplesStoreService; + export const program = makeFlowProcessorProgram({ id: "triples-store", specs: () => makeTriplesStoreSpecs(), diff --git a/ts/packages/flow/src/storage/triples/falkordb.ts b/ts/packages/flow/src/storage/triples/falkordb.ts index 13b75cb3..38a162c1 100644 --- a/ts/packages/flow/src/storage/triples/falkordb.ts +++ b/ts/packages/flow/src/storage/triples/falkordb.ts @@ -30,107 +30,136 @@ function getTermValue(term: Term): string { } } -export class FalkorDBTriplesStore { - private graph: Graph; - private connectPromise: Promise; +export interface FalkorDBTriplesStore { + readonly createNode: (uri: string, user: string, collection: string) => Promise; + readonly createLiteral: (value: string, user: string, collection: string) => Promise; + readonly relateNode: ( + src: string, + uri: string, + dest: string, + user: string, + collection: string, + ) => Promise; + readonly relateLiteral: ( + src: string, + uri: string, + dest: string, + user: string, + collection: string, + ) => Promise; + readonly storeTriples: ( + triples: Triple[], + user?: string, + collection?: string, + ) => Promise; + readonly deleteCollection: (user: string, collection: string) => Promise; +} - constructor(config: FalkorDBConfig = {}) { - const url = config.url ?? process.env.FALKORDB_URL ?? "redis://localhost:6379"; - const database = config.database ?? "falkordb"; +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 }); - this.graph = new Graph(client, database); - this.connectPromise = client.connect().then(() => { - console.log(`[FalkorDBTriplesStore] Connected to ${url}, graph: ${database}`); - }).catch((err) => { - console.error(`[FalkorDBTriplesStore] Connection failed:`, err); - throw err; - }); - } + 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; + }); - private async ensureConnected(): Promise { - await this.connectPromise; - } + const ensureConnected = async (): Promise => { + await connectPromise; + }; - async createNode(uri: string, user: string, collection: string): Promise { - await this.ensureConnected(); - await this.graph.query( + const createNode = async (uri: string, user: string, collection: string): Promise => { + await ensureConnected(); + await graph.query( "MERGE (n:Node {uri: $uri, user: $user, collection: $collection})", { params: { uri, user, collection } }, ); - } + }; - async createLiteral(value: string, user: string, collection: string): Promise { - await this.ensureConnected(); - await this.graph.query( + const createLiteral = async (value: string, user: string, collection: string): Promise => { + await ensureConnected(); + await graph.query( "MERGE (n:Literal {value: $value, user: $user, collection: $collection})", { params: { value, user, collection } }, ); - } + }; - async relateNode( + const relateNode = async ( src: string, uri: string, dest: string, user: string, collection: string, - ): Promise { - await this.ensureConnected(); - await this.graph.query( + ): Promise => { + 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 } }, ); - } + }; - async relateLiteral( + const relateLiteral = async ( src: string, uri: string, dest: string, user: string, collection: string, - ): Promise { - await this.ensureConnected(); - await this.graph.query( + ): Promise => { + 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 } }, ); - } + }; - async storeTriples( + const storeTriples = async ( triples: Triple[], user = "default", collection = "default", - ): Promise { + ): Promise => { for (const t of triples) { const s = getTermValue(t.s); const p = getTermValue(t.p); const o = getTermValue(t.o); - await this.createNode(s, user, collection); + await createNode(s, user, collection); if (t.o.type === "IRI") { - await this.createNode(o, user, collection); - await this.relateNode(s, p, o, user, collection); + await createNode(o, user, collection); + await relateNode(s, p, o, user, collection); } else { - await this.createLiteral(o, user, collection); - await this.relateLiteral(s, p, o, user, collection); + await createLiteral(o, user, collection); + await relateLiteral(s, p, o, user, collection); } } - } + }; - async deleteCollection(user: string, collection: string): Promise { - await this.ensureConnected(); - await this.graph.query( + const deleteCollection = async (user: string, collection: string): Promise => { + await ensureConnected(); + await graph.query( "MATCH (n:Node {user: $user, collection: $collection}) DETACH DELETE n", { params: { user, collection } }, ); - await this.graph.query( + await graph.query( "MATCH (n:Literal {user: $user, collection: $collection}) DETACH DELETE n", { params: { user, collection } }, ); - await this.graph.query( + 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()( @@ -171,7 +200,7 @@ const falkorDBTriplesStoreError = (operation: string, cause: unknown) => export const makeFalkorDBTriplesStoreService = ( config: FalkorDBConfig = {}, ): FalkorDBTriplesStoreServiceShape => { - const store = new FalkorDBTriplesStore(config); + const store = makeFalkorDBTriplesStore(config); return { storeTriples: Effect.fn("FalkorDBTriplesStore.storeTriples")(( triples: ReadonlyArray, diff --git a/ts/packages/workbench/package.json b/ts/packages/workbench/package.json index ef51f843..8ea785ef 100644 --- a/ts/packages/workbench/package.json +++ b/ts/packages/workbench/package.json @@ -10,19 +10,6 @@ "qa:browser": "playwright test" }, "dependencies": { - "@tanstack/react-query": "^5.75.0", - "@trustgraph/client": "workspace:*", - "clsx": "^2.1.0", - "lucide-react": "^0.513.0", - "react": "^19.1.0", - "react-dom": "^19.1.0", - "react-force-graph-2d": "^1.29.1", - "react-markdown": "^10.1.0", - "react-router": "^7.6.0", - "tailwind-merge": "^3.3.0", - "zustand": "^5.0.0", - "@effect/platform-node": "4.0.0-beta.74", - "@effect/platform-node-shared": "4.0.0-beta.74", "@effect/ai-anthropic": "4.0.0-beta.74", "@effect/ai-openai": "4.0.0-beta.74", "@effect/ai-openrouter": "4.0.0-beta.74", @@ -31,8 +18,22 @@ "@effect/opentelemetry": "4.0.0-beta.74", "@effect/platform-browser": "4.0.0-beta.74", "@effect/platform-bun": "4.0.0-beta.74", + "@effect/platform-node": "4.0.0-beta.74", + "@effect/platform-node-shared": "4.0.0-beta.74", "@effect/tsgo": "0.13.0", - "@effect/vitest": "4.0.0-beta.74" + "@effect/vitest": "4.0.0-beta.74", + "@tanstack/react-query": "^5.75.0", + "@trustgraph/client": "workspace:*", + "clsx": "^2.1.0", + "lucide-react": "^0.513.0", + "react": "^19.1.0", + "react-dom": "^19.1.0", + "react-error-boundary": "^6.1.2", + "react-force-graph-2d": "^1.29.1", + "react-markdown": "^10.1.0", + "react-router": "^7.6.0", + "tailwind-merge": "^3.3.0", + "zustand": "^5.0.0" }, "devDependencies": { "@effect/vitest": "4.0.0-beta.74", diff --git a/ts/packages/workbench/src/components/error-boundary.tsx b/ts/packages/workbench/src/components/error-boundary.tsx index 6d2586bd..c41e9c76 100644 --- a/ts/packages/workbench/src/components/error-boundary.tsx +++ b/ts/packages/workbench/src/components/error-boundary.tsx @@ -1,4 +1,8 @@ -import { Component, type ErrorInfo, type ReactNode } from "react"; +import type { ReactNode } from "react"; +import { + ErrorBoundary as ReactErrorBoundary, + type FallbackProps, +} from "react-error-boundary"; import { AlertTriangle, RefreshCw } from "lucide-react"; interface Props { @@ -7,55 +11,41 @@ interface Props { fallback?: ReactNode; } -interface State { - hasError: boolean; - error: Error | null; +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : "An unexpected error occurred."; + +function DefaultFallback({ error, resetErrorBoundary }: FallbackProps) { + return ( +
+
+ +

+ Something went wrong +

+

+ {errorMessage(error)} +

+ +
+
+ ); } -export class ErrorBoundary extends Component { - constructor(props: Props) { - super(props); - this.state = { hasError: false, error: null }; - } - - static getDerivedStateFromError(error: Error): State { - return { hasError: true, error }; - } - - override componentDidCatch(error: Error, info: ErrorInfo) { - console.error("[ErrorBoundary]", error, info.componentStack); - } - - handleReset = () => { - this.setState({ hasError: false, error: null }); - }; - - override render() { - if (this.state.hasError) { - if (this.props.fallback !== undefined) return this.props.fallback; - - return ( -
-
- -

- Something went wrong -

-

- {this.state.error?.message ?? "An unexpected error occurred."} -

- -
-
- ); - } - - return this.props.children; - } +export function ErrorBoundary({ children, fallback }: Props) { + return ( + fallback ?? } + onError={(error, info) => { + console.error("[ErrorBoundary]", error, info.componentStack); + }} + > + {children} + + ); } diff --git a/ts/packages/workbench/src/qa/mock-api.ts b/ts/packages/workbench/src/qa/mock-api.ts index f344400f..3e5099f4 100644 --- a/ts/packages/workbench/src/qa/mock-api.ts +++ b/ts/packages/workbench/src/qa/mock-api.ts @@ -1,4 +1,4 @@ -import { BaseApi, type ConnectionState, type DocumentMetadata, type ProcessingMetadata, type StreamingMetadata, type Triple } from "@trustgraph/client"; +import { makeBaseApiWithRpc, type BaseApi, type DocumentMetadata, type ProcessingMetadata, type StreamingMetadata, type Triple } from "@trustgraph/client"; import { Option, Schema as S } from "effect"; type ConfigValues = Record>; @@ -80,24 +80,6 @@ interface MockState { }; } -interface MockBaseApi extends BaseApi { - makeRequest( - service: string, - request: RequestType, - timeout?: number, - retries?: number, - flow?: string, - ): Promise; - makeRequestMulti( - service: string, - request: RequestType, - receiver: (resp: unknown) => boolean, - timeout?: number, - retries?: number, - flow?: string, - ): Promise; -} - const encodeJsonUnknown = S.encodeUnknownOption(S.fromJsonString(S.Unknown)); const decodeJsonUnknown = S.decodeUnknownOption(S.UnknownFromJsonString); @@ -533,40 +515,33 @@ function dispatchStream( export function makeMockBaseApi(fixture: MockWorkbenchFixture = {}): BaseApi { const state = createState(fixture); - const api = Object.create(BaseApi.prototype) as MockBaseApi; - api.tag = "mock-workbench"; - api.id = 1; - api.token = state.settings.apiKey.length > 0 ? state.settings.apiKey : undefined; - api.user = state.settings.user; - api.socketUrl = state.settings.gatewayUrl; - api.makeRequest = function makeRequest( - service: string, - request: RequestType, - _timeout?: number, - _retries?: number, - flow?: string, - ) { - return Promise.resolve(dispatchRequest(state, service, request as Record, flow) as ResponseType); - }; - api.makeRequestMulti = function makeRequestMulti( - service: string, - _request: RequestType, - receiver: (resp: unknown) => boolean, - _timeout?: number, - _retries?: number, - _flow?: string, - ) { - return dispatchStream(state, service, receiver); - }; - api.onConnectionStateChange = function onConnectionStateChange(listener: (state: ConnectionState) => void) { - listener({ - status: api.token === undefined ? "unauthenticated" : "authenticated", - hasApiKey: api.token !== undefined, - }); - return () => {}; - }; - api.close = function close() {}; - return api; + const token = state.settings.apiKey.length > 0 ? state.settings.apiKey : undefined; + return makeBaseApiWithRpc(state.settings.user, token, state.settings.gatewayUrl, { + dispatch: (input) => + Promise.resolve( + dispatchRequest( + state, + input.service, + input.request, + input.flow, + ), + ), + dispatchStream: async (input, receiver) => { + await dispatchStream(state, input.service, (message) => { + const chunk = message as { response?: unknown; complete?: boolean }; + return receiver({ + response: chunk.response, + complete: chunk.complete === true, + }); + }); + return undefined; + }, + subscribe: (listener) => { + listener({ status: token === undefined ? "connected" : "connected" }); + return () => {}; + }, + close: () => Promise.resolve(), + }); } export function qaSettingsFromFixture(fixture: MockWorkbenchFixture = {}) { diff --git a/ts/scripts/inventory-native-classes.ts b/ts/scripts/inventory-native-classes.ts new file mode 100644 index 00000000..7b2860b8 --- /dev/null +++ b/ts/scripts/inventory-native-classes.ts @@ -0,0 +1,170 @@ +import { existsSync, readdirSync, readFileSync, statSync } from "node:fs"; +import { join, relative, sep } from "node:path"; +import ts from "typescript"; + +type Scope = "production" | "non-production"; +type Classification = "blocking" | "candidate-effect-exemption" | "non-blocking"; + +interface ClassFinding { + readonly file: string; + readonly line: number; + readonly column: number; + readonly name: string; + readonly scope: Scope; + readonly classification: Classification; + readonly reason: string; + readonly extendsText?: string; +} + +const root = process.cwd(); +const packagesSrc = join(root, "packages"); +const scriptsDir = join(root, "scripts"); + +const sourceExtensions = new Set([".ts", ".tsx", ".mts", ".cts"]); +const effectClassPatterns = [ + /\bS\.(Class|TaggedClass|TaggedErrorClass|ErrorClass)\b/, + /\bSchema\.(Class|TaggedClass|TaggedErrorClass|ErrorClass)\b/, + /\bData\.TaggedError\b/, + /\bContext\.Service\b/, + /\bRpc\.make\b/, + /\bHttpApi\.make\b/, + /\bEffect\.Service\b/, +]; + +function extensionOf(path: string): string { + const match = path.match(/\.[cm]?tsx?$/); + return match?.[0] ?? ""; +} + +function isSourceFile(path: string): boolean { + return sourceExtensions.has(extensionOf(path)); +} + +function walk(dir: string): string[] { + if (!existsSync(dir)) { + return []; + } + + const files: string[] = []; + for (const entry of readdirSync(dir)) { + if (entry === "dist" || entry === "node_modules" || entry === ".turbo") { + continue; + } + + const path = join(dir, entry); + const stat = statSync(path); + if (stat.isDirectory()) { + files.push(...walk(path)); + } else if (stat.isFile() && isSourceFile(path)) { + files.push(path); + } + } + return files; +} + +function isProductionPackageSource(path: string): boolean { + const rel = relative(root, path).split(sep).join("/"); + return ( + rel.startsWith("packages/") && + rel.includes("/src/") && + !rel.includes("/__tests__/") && + !rel.endsWith(".test.ts") && + !rel.endsWith(".test.tsx") && + !rel.endsWith(".spec.ts") && + !rel.endsWith(".spec.tsx") + ); +} + +function getClassName(node: ts.ClassLikeDeclarationBase): string { + return node.name?.getText() ?? ""; +} + +function getExtendsText(node: ts.ClassLikeDeclarationBase, source: ts.SourceFile): string | undefined { + const heritage = node.heritageClauses?.find((clause) => clause.token === ts.SyntaxKind.ExtendsKeyword); + return heritage?.types.map((type) => type.expression.getText(source)).join(", "); +} + +function classify(scope: Scope, extendsText?: string): Pick { + if (scope === "non-production") { + return { + classification: "non-blocking", + reason: "outside production runtime source", + }; + } + + if (extendsText !== undefined && effectClassPatterns.some((pattern) => pattern.test(extendsText))) { + return { + classification: "candidate-effect-exemption", + reason: "Effect class-shaped API requires proof or functional replacement", + }; + } + + return { + classification: "blocking", + reason: "native class syntax in production runtime source", + }; +} + +function inspectFile(path: string): ClassFinding[] { + const sourceText = readFileSync(path, "utf8"); + const source = ts.createSourceFile( + path, + sourceText, + ts.ScriptTarget.Latest, + true, + path.endsWith(".tsx") ? ts.ScriptKind.TSX : ts.ScriptKind.TS, + ); + const scope: Scope = isProductionPackageSource(path) ? "production" : "non-production"; + const findings: ClassFinding[] = []; + + function visit(node: ts.Node): void { + if (ts.isClassDeclaration(node) || ts.isClassExpression(node)) { + const position = source.getLineAndCharacterOfPosition(node.getStart(source)); + const extendsText = getExtendsText(node, source); + const { classification, reason } = classify(scope, extendsText); + findings.push({ + file: relative(root, path).split(sep).join("/"), + line: position.line + 1, + column: position.character + 1, + name: getClassName(node), + scope, + classification, + reason, + extendsText, + }); + } + + ts.forEachChild(node, visit); + } + + visit(source); + return findings; +} + +const files = [...walk(packagesSrc), ...walk(scriptsDir)].sort(); +const findings = files.flatMap(inspectFile); +const productionFindings = findings.filter((finding) => finding.scope === "production"); +const blocking = productionFindings.filter((finding) => finding.classification === "blocking"); +const candidates = productionFindings.filter((finding) => finding.classification === "candidate-effect-exemption"); +const nonProduction = findings.filter((finding) => finding.scope === "non-production"); + +function printGroup(title: string, group: ClassFinding[]): void { + console.log(`${title}: ${group.length}`); + for (const finding of group) { + const extendsPart = finding.extendsText === undefined ? "" : ` extends ${finding.extendsText}`; + console.log( + ` ${finding.file}:${finding.line}:${finding.column} ${finding.name}${extendsPart} - ${finding.reason}`, + ); + } +} + +printGroup("Blocking production native classes", blocking); +printGroup("Candidate Effect class-shaped exemptions", candidates); +printGroup("Non-production class declarations", nonProduction); + +if (blocking.length > 0) { + console.error(`\nFound ${blocking.length} blocking production native class declarations.`); + process.exit(1); +} + +console.log("\nNo blocking production native class declarations found.");