apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-01 01:19:38 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Remove RAG requestor Promise bridges

Browse source
This commit is contained in:
elpresidank 2026-06-02 00:54:47 -05:00
parent 88db18fbda
commit 5979d38b99
11 changed files with 249 additions and 293 deletions
Download patch file Download diff file Expand all files Collapse all files

48
ts/packages/flow/src/retrieval/document-rag-service.ts
View file

@ -7,6 +7,7 @@
* Python reference: trustgraph-flow/trustgraph/retrieval/document_rag/
*/
import {NodeRuntime} from "@effect/platform-node";
import {
makeConsumerSpec,
makeFlowProcessor,
@ -17,14 +18,10 @@ import {
type DocumentEmbeddingsResponse,
type DocumentRagRequest,
type DocumentRagResponse,
type EffectRequestOptions,
type EffectRequestResponse,
type EmbeddingsRequest,
type EmbeddingsResponse,
type FlowContext,
type FlowProcessorRuntime,
type FlowRequestOptions,
type FlowRequestor,
type FlowResourceNotFoundError,
type MessagingDeliveryError,
type ProcessorConfig,
@ -34,7 +31,7 @@ import {
type TextCompletionRequest,
type TextCompletionResponse,
} from "@trustgraph/base";
import { Effect } from "effect";
import {Effect, Layer, ManagedRuntime} from "effect";
import {
DocumentRagEngine,
DocumentRagEngineError,
@ -43,29 +40,6 @@ import {
type DocumentRagClients,
} from "./document-rag.js";
const toEffectRequestOptions = <TRes>(
options: FlowRequestOptions<TRes> | undefined,
): EffectRequestOptions<TRes> | undefined => {
if (options === undefined) return undefined;
return {
...(options.timeoutMs === undefined ? {} : { timeoutMs: options.timeoutMs }),
...(options.recipient === undefined
? {}
: {
recipient: (response: TRes) =>
Effect.promise(() => options.recipient?.(response) ?? Promise.resolve(true)),
}),
};
};
const toPromiseRequestor = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
): FlowRequestor<TReq, TRes> => ({
request: (request, options) =>
Effect.runPromise(requestor.request(request, toEffectRequestOptions(options))),
stop: () => Effect.runPromise(requestor.stop),
});
const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function* (
msg: DocumentRagRequest,
properties: Record<string, string>,
@ -78,12 +52,10 @@ const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function*
const engine = yield* DocumentRagEngine;
const clients: DocumentRagClients = {
llm: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm")),
embeddings: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings")),
docEmbeddings: toPromiseRequestor(
yield* flowCtx.flow.requestorEffect<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>("doc-embeddings"),
),
prompt: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt")),
llm: yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm"),
embeddings: yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
docEmbeddings: yield* flowCtx.flow.requestorEffect<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>("doc-embeddings"),
prompt: yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt"),
};
const response = yield* engine.query(
@ -161,6 +133,12 @@ export const program = makeFlowProcessorProgram({
layer: () => DocumentRagLive,
});
const documentRagRuntime = ManagedRuntime.make(Layer.empty);
export function run(): Promise<void> {
return Effect.runPromise(program);
return documentRagRuntime.runPromise(program);
}
export function runMain(): void {
NodeRuntime.runMain(program);
}

44
ts/packages/flow/src/retrieval/document-rag.ts
View file

@ -9,7 +9,7 @@ import type {
DocumentEmbeddingsResponse,
EmbeddingsRequest,
EmbeddingsResponse,
FlowRequestor,
EffectRequestResponse,
PromptRequest,
PromptResponse,
TextCompletionRequest,
@ -20,10 +20,10 @@ import { Context, Effect, Layer } from "effect";
import * as S from "effect/Schema";
export interface DocumentRagClients {
llm: FlowRequestor<TextCompletionRequest, TextCompletionResponse>;
embeddings: FlowRequestor<EmbeddingsRequest, EmbeddingsResponse>;
docEmbeddings: FlowRequestor<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>;
prompt: FlowRequestor<PromptRequest, PromptResponse>;
llm: EffectRequestResponse<TextCompletionRequest, TextCompletionResponse>;
embeddings: EffectRequestResponse<EmbeddingsRequest, EmbeddingsResponse>;
docEmbeddings: EffectRequestResponse<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>;
prompt: EffectRequestResponse<PromptRequest, PromptResponse>;
}
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
@ -101,21 +101,19 @@ function queryDocumentRag(
return Effect.gen(function* () {
const collection = options?.collection ?? "default";
const embResp = yield* Effect.tryPromise({
try: () => clients.embeddings.request({ text: [queryText] }),
catch: (cause) => documentRagError("embeddings", cause),
});
const embResp = yield* clients.embeddings.request({ text: [queryText] }).pipe(
Effect.mapError((cause) => documentRagError("embeddings", cause)),
);
const vectors = embResp.vectors;
const docResp = yield* Effect.tryPromise({
try: () => clients.docEmbeddings.request({
const docResp = yield* clients.docEmbeddings.request({
vectors,
limit: 10,
collection,
user: "default",
}),
catch: (cause) => documentRagError("document-embeddings", cause),
});
}).pipe(
Effect.mapError((cause) => documentRagError("document-embeddings", cause)),
);
const chunks = docResp.chunks ?? [];
yield* Effect.log(`[DocumentRag] Found ${chunks.length} matching chunks`);
@ -125,21 +123,19 @@ function queryDocumentRag(
)
.join("\n\n---\n\n");
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
const promptResp = yield* clients.prompt.request({
name: "document-rag-synthesize",
variables: { query: queryText, context },
}),
catch: (cause) => documentRagError("prompt", cause),
});
}).pipe(
Effect.mapError((cause) => documentRagError("prompt", cause)),
);
const resp = yield* Effect.tryPromise({
try: () => clients.llm.request({
const resp = yield* clients.llm.request({
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => documentRagError("llm", cause),
});
}).pipe(
Effect.mapError((cause) => documentRagError("llm", cause)),
);
return resp.response;
});

72
ts/packages/flow/src/retrieval/graph-rag-service.ts
View file

@ -7,18 +7,15 @@
* Python reference: trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
*/
import {NodeRuntime} from "@effect/platform-node";
import {
makeConsumerSpec,
makeFlowProcessor,
makeProducerSpec,
makeRequestResponseSpec,
makeFlowProcessorProgram,
type EffectRequestOptions,
type EffectRequestResponse,
type FlowContext,
type FlowProcessorRuntime,
type FlowRequestOptions,
type FlowRequestor,
type FlowResourceNotFoundError,
type GraphEmbeddingsRequest,
type GraphEmbeddingsResponse,
@ -36,7 +33,7 @@ import {
type TriplesQueryRequest,
type TriplesQueryResponse,
} from "@trustgraph/base";
import { Effect } from "effect";
import {Effect, Layer, ManagedRuntime} from "effect";
import {
GraphRagEngine,
GraphRagEngineError,
@ -46,29 +43,6 @@ import {
type GraphRagConfig,
} from "./graph-rag.js";
const toEffectRequestOptions = <TRes>(
options: FlowRequestOptions<TRes> | undefined,
): EffectRequestOptions<TRes> | undefined => {
if (options === undefined) return undefined;
return {
...(options.timeoutMs === undefined ? {} : { timeoutMs: options.timeoutMs }),
...(options.recipient === undefined
? {}
: {
recipient: (response: TRes) =>
Effect.promise(() => options.recipient?.(response) ?? Promise.resolve(true)),
}),
};
};
const toPromiseRequestor = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
): FlowRequestor<TReq, TRes> => ({
request: (request, options) =>
Effect.runPromise(requestor.request(request, toEffectRequestOptions(options))),
stop: () => Effect.runPromise(requestor.stop),
});
const graphRagConfigFromRequest = (msg: GraphRagRequest): GraphRagConfig => ({
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
@ -90,13 +64,11 @@ const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* (
yield* Effect.log(`[GraphRagService] Received request ${requestId}: "${msg.query?.slice(0, 60)}..." collection=${msg.collection}`);
const clients: GraphRagClients = {
llm: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm")),
embeddings: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings")),
graphEmbeddings: toPromiseRequestor(
yield* flowCtx.flow.requestorEffect<GraphEmbeddingsRequest, GraphEmbeddingsResponse>("graph-embeddings"),
),
triples: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<TriplesQueryRequest, TriplesQueryResponse>("triples")),
prompt: toPromiseRequestor(yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt")),
llm: yield* flowCtx.flow.requestorEffect<TextCompletionRequest, TextCompletionResponse>("llm"),
embeddings: yield* flowCtx.flow.requestorEffect<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
graphEmbeddings: yield* flowCtx.flow.requestorEffect<GraphEmbeddingsRequest, GraphEmbeddingsResponse>("graph-embeddings"),
triples: yield* flowCtx.flow.requestorEffect<TriplesQueryRequest, TriplesQueryResponse>("triples"),
prompt: yield* flowCtx.flow.requestorEffect<PromptRequest, PromptResponse>("prompt"),
};
const result = yield* engine.query(
@ -125,16 +97,18 @@ const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* (
if (result === undefined) return;
const response: GraphRagResponse = {
response: result.answer,
endOfStream: true,
};
if (result.subgraph.length > 0) {
(response as Record<string, unknown>).message_type = "explain";
(response as Record<string, unknown>).explain_id = `explain-${requestId}`;
(response as Record<string, unknown>).explain_triples = result.subgraph;
}
const response: GraphRagResponse = result.subgraph.length === 0
? {
response: result.answer,
endOfStream: true,
}
: {
response: result.answer,
endOfStream: true,
message_type: "explain",
explain_id: `explain-${requestId}`,
explain_triples: result.subgraph,
};
yield* producer.send(requestId, response);
});
@ -192,6 +166,12 @@ export const program = makeFlowProcessorProgram({
layer: () => GraphRagLive,
});
const graphRagRuntime = ManagedRuntime.make(Layer.empty);
export function run(): Promise<void> {
return Effect.runPromise(program);
return graphRagRuntime.runPromise(program);
}
export function runMain(): void {
NodeRuntime.runMain(program);
}

146
ts/packages/flow/src/retrieval/graph-rag.ts
View file

@ -7,7 +7,8 @@
import type {
EmbeddingsRequest,
EmbeddingsResponse,
FlowRequestor,
EffectRequestOptions,
EffectRequestResponse,
GraphEmbeddingsRequest,
GraphEmbeddingsResponse,
PromptRequest,
@ -34,11 +35,11 @@ export interface GraphRagConfig {
}
export interface GraphRagClients {
llm: FlowRequestor<TextCompletionRequest, TextCompletionResponse>;
embeddings: FlowRequestor<EmbeddingsRequest, EmbeddingsResponse>;
graphEmbeddings: FlowRequestor<GraphEmbeddingsRequest, GraphEmbeddingsResponse>;
triples: FlowRequestor<TriplesQueryRequest, TriplesQueryResponse>;
prompt: FlowRequestor<PromptRequest, PromptResponse>;
llm: EffectRequestResponse<TextCompletionRequest, TextCompletionResponse>;
embeddings: EffectRequestResponse<EmbeddingsRequest, EmbeddingsResponse>;
graphEmbeddings: EffectRequestResponse<GraphEmbeddingsRequest, GraphEmbeddingsResponse>;
triples: EffectRequestResponse<TriplesQueryRequest, TriplesQueryResponse>;
prompt: EffectRequestResponse<PromptRequest, PromptResponse>;
}
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
@ -92,6 +93,16 @@ const graphRagError = (operation: string, cause: unknown) =>
message: errorMessage(cause),
});
const requestClient = <TReq, TRes>(
requestor: EffectRequestResponse<TReq, TRes>,
operation: string,
request: TReq,
options?: EffectRequestOptions<TRes, GraphRagEngineError>,
): Effect.Effect<TRes, GraphRagEngineError> =>
requestor.request(request, options).pipe(
Effect.mapError((cause) => graphRagError(operation, cause)),
);
export function normalizeGraphRagConfig(config: GraphRagConfig = {}): NormalizedGraphRagConfig {
return {
entityLimit: config.entityLimit ?? 50,
@ -178,21 +189,23 @@ function queryGraphRag(
function extractConcepts(clients: GraphRagClients, query: string): Effect.Effect<string[], GraphRagEngineError> {
return Effect.gen(function* () {
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
const promptResp = yield* requestClient(
clients.prompt,
"extract-concepts-prompt",
{
name: "extract-concepts",
variables: { query },
}),
catch: (cause) => graphRagError("extract-concepts-prompt", cause),
});
},
);
const llmResp = yield* Effect.tryPromise({
try: () => clients.llm.request({
const llmResp = yield* requestClient(
clients.llm,
"extract-concepts-llm",
{
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => graphRagError("extract-concepts-llm", cause),
});
},
);
return llmResp.response
.split("\n")
@ -203,10 +216,7 @@ function extractConcepts(clients: GraphRagClients, query: string): Effect.Effect
function getVectors(clients: GraphRagClients, concepts: string[]): Effect.Effect<number[][], GraphRagEngineError> {
return Effect.gen(function* () {
const resp = yield* Effect.tryPromise({
try: () => clients.embeddings.request({ text: concepts }),
catch: (cause) => graphRagError("get-vectors", cause),
});
const resp = yield* requestClient(clients.embeddings, "get-vectors", { text: concepts });
return resp.vectors;
});
}
@ -218,15 +228,16 @@ function getEntities(
collection?: string,
): Effect.Effect<Term[], GraphRagEngineError> {
return Effect.gen(function* () {
const resp = yield* Effect.tryPromise({
try: () => clients.graphEmbeddings.request({
const resp = yield* requestClient(
clients.graphEmbeddings,
"get-entities",
{
vectors,
user: "default",
collection: collection ?? "default",
limit: config.entityLimit,
}),
catch: (cause) => graphRagError("get-entities", cause),
});
},
);
return resp.entities;
});
}
@ -259,10 +270,7 @@ function followEdges(
limit: config.tripleLimit,
...(collection !== undefined ? { collection } : {}),
};
return Effect.tryPromise({
try: () => clients.triples.request(request),
catch: (cause) => graphRagError("follow-edges-query", cause),
});
return requestClient(clients.triples, "follow-edges-query", request);
});
const results = yield* Effect.all(queries);
@ -321,24 +329,26 @@ function scoreEdges(
Effect.mapError((cause) => graphRagError("edge-score-encode", cause)),
);
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
const promptResp = yield* requestClient(
clients.prompt,
"edge-score-prompt",
{
name: "kg-edge-scoring",
variables: {
query,
knowledge: knowledgeJson,
},
}),
catch: (cause) => graphRagError("edge-score-prompt", cause),
});
},
);
const llmResp = yield* Effect.tryPromise({
try: () => clients.llm.request({
const llmResp = yield* requestClient(
clients.llm,
"edge-score-llm",
{
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => graphRagError("edge-score-llm", cause),
});
},
);
yield* Effect.log(`[GraphRag] Edge scoring LLM response (first 500 chars): ${llmResp.response.slice(0, 500)}`);
@ -375,43 +385,49 @@ function synthesize(
.map((triple) => `${termToString(triple.s)} -> ${termToString(triple.p)} -> ${termToString(triple.o)}`)
.join("\n");
const promptResp = yield* Effect.tryPromise({
try: () => clients.prompt.request({
const promptResp = yield* requestClient(
clients.prompt,
"synthesize-prompt",
{
name: "graph-rag-synthesize",
variables: { query, context },
}),
catch: (cause) => graphRagError("synthesize-prompt", cause),
});
},
);
if (chunkCallback !== undefined) {
let fullText = "";
yield* Effect.tryPromise({
try: () => clients.llm.request(
{
system: promptResp.system,
prompt: promptResp.prompt,
streaming: true,
yield* requestClient(
clients.llm,
"synthesize-stream",
{
system: promptResp.system,
prompt: promptResp.prompt,
streaming: true,
},
{
recipient: (resp) => {
if (resp.response.length === 0) {
return Effect.succeed(resp.endOfStream === true);
}
fullText += resp.response;
return Effect.tryPromise({
try: () => chunkCallback(resp.response, resp.endOfStream === true).then(() => resp.endOfStream === true),
catch: (cause) => graphRagError("synthesize-stream-callback", cause),
});
},
{
recipient: (resp) => {
if (resp.response.length === 0) return Promise.resolve(resp.endOfStream === true);
fullText += resp.response;
return chunkCallback(resp.response, resp.endOfStream === true).then(() => resp.endOfStream === true);
},
},
),
catch: (cause) => graphRagError("synthesize-stream", cause),
});
},
);
return fullText;
}
const resp = yield* Effect.tryPromise({
try: () => clients.llm.request({
const resp = yield* requestClient(
clients.llm,
"synthesize-llm",
{
system: promptResp.system,
prompt: promptResp.prompt,
}),
catch: (cause) => graphRagError("synthesize-llm", cause),
});
},
);
return resp.response;
});