Advance TS port Effect workbench

ts/packages/flow/src/retrieval/document-rag.ts

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