Advance TS port Effect workbench

ts/packages/flow/src/__tests__/retrieval-rag.test.ts

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