mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-06 03:42:11 +02:00
Advance TS port Effect workbench
This commit is contained in:
elpresidank
parent
92dae8c374
commit
3515106670
116 changed files with 12286 additions and 9584 deletions
|
|
@ -1,118 +1,166 @@
|
|||
/**
|
||||
* Document RAG service — FlowProcessor wrapper around the DocumentRag class.
|
||||
* Document RAG service.
|
||||
*
|
||||
* Consumes DocumentRagRequest messages, runs the document retrieval pipeline
|
||||
* (embed query → find similar chunks → synthesize answer), emits DocumentRagResponse.
|
||||
*
|
||||
* Each request gets its own DocumentRag instance for security isolation.
|
||||
* Consumes DocumentRagRequest messages, runs the document retrieval pipeline,
|
||||
* and emits DocumentRagResponse.
|
||||
*
|
||||
* Python reference: trustgraph-flow/trustgraph/retrieval/document_rag/
|
||||
*/
|
||||
|
||||
import {
|
||||
FlowProcessor,
|
||||
ConsumerSpec,
|
||||
FlowProcessor,
|
||||
ProducerSpec,
|
||||
RequestResponseSpec,
|
||||
type ProcessorConfig,
|
||||
type FlowContext,
|
||||
type DocumentRagRequest,
|
||||
type DocumentRagResponse,
|
||||
type TextCompletionRequest,
|
||||
type TextCompletionResponse,
|
||||
type EmbeddingsRequest,
|
||||
type EmbeddingsResponse,
|
||||
makeFlowProcessorProgram,
|
||||
type DocumentEmbeddingsRequest,
|
||||
type DocumentEmbeddingsResponse,
|
||||
type DocumentRagRequest,
|
||||
type DocumentRagResponse,
|
||||
type EffectRequestOptions,
|
||||
type EffectRequestResponse,
|
||||
type EmbeddingsRequest,
|
||||
type EmbeddingsResponse,
|
||||
type FlowContext,
|
||||
type FlowRequestOptions,
|
||||
type FlowRequestor,
|
||||
type FlowResourceNotFoundError,
|
||||
type MessagingDeliveryError,
|
||||
type ProcessorConfig,
|
||||
type PromptRequest,
|
||||
type PromptResponse,
|
||||
type Spec,
|
||||
type TextCompletionRequest,
|
||||
type TextCompletionResponse,
|
||||
} from "@trustgraph/base";
|
||||
import { makeProcessorProgram } from "@trustgraph/base";
|
||||
import { DocumentRag } from "./document-rag.js";
|
||||
import { Effect } from "effect";
|
||||
import {
|
||||
DocumentRagEngine,
|
||||
DocumentRagEngineError,
|
||||
DocumentRagLive,
|
||||
makeDocumentRagEngine,
|
||||
type DocumentRagClients,
|
||||
} from "./document-rag.js";
|
||||
|
||||
export class DocumentRagService extends FlowProcessor {
|
||||
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>,
|
||||
flowCtx: FlowContext<DocumentRagEngine>,
|
||||
) {
|
||||
const requestId = properties.id;
|
||||
if (requestId === undefined || requestId.length === 0) return;
|
||||
|
||||
const producer = yield* flowCtx.flow.producerEffect<DocumentRagResponse>("document-rag-response");
|
||||
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")),
|
||||
};
|
||||
|
||||
const response = yield* engine.query(
|
||||
clients,
|
||||
msg.query,
|
||||
{
|
||||
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
|
||||
},
|
||||
).pipe(
|
||||
Effect.catch((error: DocumentRagEngineError) =>
|
||||
Effect.logError("[DocumentRag] Query failed", {
|
||||
error: error.message,
|
||||
operation: error.operation,
|
||||
}).pipe(
|
||||
Effect.flatMap(() =>
|
||||
producer.send(requestId, {
|
||||
response: "",
|
||||
error: { type: "rag-error", message: error.message },
|
||||
}),
|
||||
),
|
||||
Effect.as(undefined),
|
||||
),
|
||||
),
|
||||
);
|
||||
|
||||
if (response === undefined) return;
|
||||
yield* producer.send(requestId, { response, endOfStream: true });
|
||||
});
|
||||
|
||||
export const makeDocumentRagSpecs = (): ReadonlyArray<Spec<DocumentRagEngine>> => [
|
||||
new ConsumerSpec<DocumentRagRequest, FlowResourceNotFoundError | MessagingDeliveryError, DocumentRagEngine>(
|
||||
"document-rag-request",
|
||||
onDocumentRagRequest,
|
||||
),
|
||||
new ProducerSpec<DocumentRagResponse>("document-rag-response"),
|
||||
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
|
||||
"llm",
|
||||
"text-completion-request",
|
||||
"text-completion-response",
|
||||
),
|
||||
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
|
||||
"embeddings",
|
||||
"embeddings-request",
|
||||
"embeddings-response",
|
||||
),
|
||||
new RequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
|
||||
"doc-embeddings",
|
||||
"document-embeddings-request",
|
||||
"document-embeddings-response",
|
||||
),
|
||||
new RequestResponseSpec<PromptRequest, PromptResponse>(
|
||||
"prompt",
|
||||
"prompt-request",
|
||||
"prompt-response",
|
||||
),
|
||||
];
|
||||
|
||||
export class DocumentRagService extends FlowProcessor<DocumentRagEngine> {
|
||||
constructor(config: ProcessorConfig) {
|
||||
super(config);
|
||||
|
||||
// Consumer: document RAG requests
|
||||
this.registerSpecification(
|
||||
ConsumerSpec.fromPromise<DocumentRagRequest>("document-rag-request", this.onRequest.bind(this)),
|
||||
);
|
||||
|
||||
// Producer: document RAG responses
|
||||
this.registerSpecification(new ProducerSpec<DocumentRagResponse>("document-rag-response"));
|
||||
|
||||
// Request-response clients
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
|
||||
"llm",
|
||||
"text-completion-request",
|
||||
"text-completion-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
|
||||
"embeddings",
|
||||
"embeddings-request",
|
||||
"embeddings-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
|
||||
"doc-embeddings",
|
||||
"document-embeddings-request",
|
||||
"document-embeddings-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<PromptRequest, PromptResponse>(
|
||||
"prompt",
|
||||
"prompt-request",
|
||||
"prompt-response",
|
||||
),
|
||||
);
|
||||
|
||||
console.log("[DocumentRag] Service initialized");
|
||||
for (const spec of makeDocumentRagSpecs()) {
|
||||
this.registerSpecification(spec);
|
||||
}
|
||||
}
|
||||
|
||||
private async onRequest(
|
||||
msg: DocumentRagRequest,
|
||||
properties: Record<string, string>,
|
||||
flowCtx: FlowContext,
|
||||
): Promise<void> {
|
||||
const requestId = properties.id;
|
||||
if (requestId === undefined || requestId.length === 0) return;
|
||||
|
||||
const producer = flowCtx.flow.producer<DocumentRagResponse>("document-rag-response");
|
||||
|
||||
try {
|
||||
const documentRag = new DocumentRag({
|
||||
llm: flowCtx.flow.requestor<TextCompletionRequest, TextCompletionResponse>("llm"),
|
||||
embeddings: flowCtx.flow.requestor<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
|
||||
docEmbeddings: flowCtx.flow.requestor<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>("doc-embeddings"),
|
||||
prompt: flowCtx.flow.requestor<PromptRequest, PromptResponse>("prompt"),
|
||||
});
|
||||
|
||||
const response = await documentRag.query(msg.query, {
|
||||
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
|
||||
});
|
||||
|
||||
await producer.send(requestId, { response, endOfStream: true });
|
||||
} catch (err) {
|
||||
console.error("[DocumentRag] Query failed:", err);
|
||||
await producer.send(requestId, {
|
||||
response: "",
|
||||
error: { type: "rag-error", message: String(err) },
|
||||
});
|
||||
}
|
||||
override startEffect() {
|
||||
return super.startEffect().pipe(
|
||||
Effect.provideService(DocumentRagEngine, DocumentRagEngine.of(makeDocumentRagEngine())),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export const program = makeProcessorProgram({
|
||||
export const program = makeFlowProcessorProgram({
|
||||
id: "document-rag",
|
||||
make: (config) => new DocumentRagService(config),
|
||||
specs: makeDocumentRagSpecs,
|
||||
layer: () => DocumentRagLive,
|
||||
});
|
||||
|
||||
export async function run(): Promise<void> {
|
||||
await DocumentRagService.launch("document-rag");
|
||||
await Effect.runPromise(program);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,158 +1,197 @@
|
|||
/**
|
||||
* Graph RAG service — FlowProcessor wrapper around the GraphRag class.
|
||||
* Graph RAG service.
|
||||
*
|
||||
* Consumes GraphRagRequest messages from the agent/gateway, runs the full
|
||||
* Graph RAG pipeline (concept extraction → entity lookup → graph traversal →
|
||||
* edge scoring → answer synthesis), and emits GraphRagResponse.
|
||||
*
|
||||
* Each request gets its own GraphRag instance to prevent data leakage
|
||||
* across requests (security requirement from the Python implementation).
|
||||
* Graph RAG pipeline, and emits GraphRagResponse.
|
||||
*
|
||||
* Python reference: trustgraph-flow/trustgraph/retrieval/graph_rag/rag.py
|
||||
*/
|
||||
|
||||
import {
|
||||
FlowProcessor,
|
||||
ConsumerSpec,
|
||||
FlowProcessor,
|
||||
ProducerSpec,
|
||||
RequestResponseSpec,
|
||||
type ProcessorConfig,
|
||||
makeFlowProcessorProgram,
|
||||
type EffectRequestOptions,
|
||||
type EffectRequestResponse,
|
||||
type FlowContext,
|
||||
type GraphRagRequest,
|
||||
type GraphRagResponse,
|
||||
type TextCompletionRequest,
|
||||
type TextCompletionResponse,
|
||||
type EmbeddingsRequest,
|
||||
type EmbeddingsResponse,
|
||||
type FlowRequestOptions,
|
||||
type FlowRequestor,
|
||||
type FlowResourceNotFoundError,
|
||||
type GraphEmbeddingsRequest,
|
||||
type GraphEmbeddingsResponse,
|
||||
type TriplesQueryRequest,
|
||||
type TriplesQueryResponse,
|
||||
type GraphRagRequest,
|
||||
type GraphRagResponse,
|
||||
type EmbeddingsRequest,
|
||||
type EmbeddingsResponse,
|
||||
type MessagingDeliveryError,
|
||||
type ProcessorConfig,
|
||||
type PromptRequest,
|
||||
type PromptResponse,
|
||||
type Spec,
|
||||
type TextCompletionRequest,
|
||||
type TextCompletionResponse,
|
||||
type TriplesQueryRequest,
|
||||
type TriplesQueryResponse,
|
||||
} from "@trustgraph/base";
|
||||
import { makeProcessorProgram } from "@trustgraph/base";
|
||||
import { GraphRag } from "./graph-rag.js";
|
||||
import { Effect } from "effect";
|
||||
import {
|
||||
GraphRagEngine,
|
||||
GraphRagEngineError,
|
||||
GraphRagLive,
|
||||
makeGraphRagEngine,
|
||||
type GraphRagClients,
|
||||
type GraphRagConfig,
|
||||
} from "./graph-rag.js";
|
||||
|
||||
export class GraphRagService extends FlowProcessor {
|
||||
constructor(config: ProcessorConfig) {
|
||||
super(config);
|
||||
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)),
|
||||
}),
|
||||
};
|
||||
};
|
||||
|
||||
// Consumer: graph RAG requests
|
||||
this.registerSpecification(
|
||||
ConsumerSpec.fromPromise<GraphRagRequest>("graph-rag-request", this.onRequest.bind(this)),
|
||||
);
|
||||
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),
|
||||
});
|
||||
|
||||
// Producer: graph RAG responses
|
||||
this.registerSpecification(new ProducerSpec<GraphRagResponse>("graph-rag-response"));
|
||||
const graphRagConfigFromRequest = (msg: GraphRagRequest): GraphRagConfig => ({
|
||||
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
|
||||
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
|
||||
...(msg.maxSubgraphSize !== undefined ? { maxSubgraphSize: msg.maxSubgraphSize } : {}),
|
||||
...(msg.maxPathLength !== undefined ? { maxPathLength: msg.maxPathLength } : {}),
|
||||
});
|
||||
|
||||
// Request-response clients for the pipeline
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
|
||||
"llm",
|
||||
"text-completion-request",
|
||||
"text-completion-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
|
||||
"embeddings",
|
||||
"embeddings-request",
|
||||
"embeddings-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
|
||||
"graph-embeddings",
|
||||
"graph-embeddings-request",
|
||||
"graph-embeddings-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
|
||||
"triples",
|
||||
"triples-request",
|
||||
"triples-response",
|
||||
),
|
||||
);
|
||||
this.registerSpecification(
|
||||
new RequestResponseSpec<PromptRequest, PromptResponse>(
|
||||
"prompt",
|
||||
"prompt-request",
|
||||
"prompt-response",
|
||||
),
|
||||
);
|
||||
const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* (
|
||||
msg: GraphRagRequest,
|
||||
properties: Record<string, string>,
|
||||
flowCtx: FlowContext<GraphRagEngine>,
|
||||
) {
|
||||
const requestId = properties.id;
|
||||
if (requestId === undefined || requestId.length === 0) return;
|
||||
|
||||
console.log("[GraphRag] Service initialized");
|
||||
const producer = yield* flowCtx.flow.producerEffect<GraphRagResponse>("graph-rag-response");
|
||||
const engine = yield* GraphRagEngine;
|
||||
|
||||
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")),
|
||||
};
|
||||
|
||||
const result = yield* engine.query(
|
||||
clients,
|
||||
msg.query,
|
||||
{
|
||||
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
|
||||
},
|
||||
graphRagConfigFromRequest(msg),
|
||||
).pipe(
|
||||
Effect.catch((error: GraphRagEngineError) =>
|
||||
Effect.logError("[GraphRag] Query failed", {
|
||||
error: error.message,
|
||||
operation: error.operation,
|
||||
}).pipe(
|
||||
Effect.flatMap(() =>
|
||||
producer.send(requestId, {
|
||||
response: "",
|
||||
error: { type: "rag-error", message: error.message },
|
||||
}),
|
||||
),
|
||||
Effect.as(undefined),
|
||||
),
|
||||
),
|
||||
);
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
private async onRequest(
|
||||
msg: GraphRagRequest,
|
||||
properties: Record<string, string>,
|
||||
flowCtx: FlowContext,
|
||||
): Promise<void> {
|
||||
const requestId = properties.id;
|
||||
if (requestId === undefined || requestId.length === 0) return;
|
||||
yield* producer.send(requestId, response);
|
||||
});
|
||||
|
||||
const producer = flowCtx.flow.producer<GraphRagResponse>("graph-rag-response");
|
||||
console.log(`[GraphRagService] Received request ${requestId}: "${msg.query?.slice(0, 60)}..." collection=${msg.collection}`);
|
||||
export const makeGraphRagSpecs = (): ReadonlyArray<Spec<GraphRagEngine>> => [
|
||||
new ConsumerSpec<GraphRagRequest, FlowResourceNotFoundError | MessagingDeliveryError, GraphRagEngine>(
|
||||
"graph-rag-request",
|
||||
onGraphRagRequest,
|
||||
),
|
||||
new ProducerSpec<GraphRagResponse>("graph-rag-response"),
|
||||
new RequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
|
||||
"llm",
|
||||
"text-completion-request",
|
||||
"text-completion-response",
|
||||
),
|
||||
new RequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
|
||||
"embeddings",
|
||||
"embeddings-request",
|
||||
"embeddings-response",
|
||||
),
|
||||
new RequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
|
||||
"graph-embeddings",
|
||||
"graph-embeddings-request",
|
||||
"graph-embeddings-response",
|
||||
),
|
||||
new RequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
|
||||
"triples",
|
||||
"triples-request",
|
||||
"triples-response",
|
||||
),
|
||||
new RequestResponseSpec<PromptRequest, PromptResponse>(
|
||||
"prompt",
|
||||
"prompt-request",
|
||||
"prompt-response",
|
||||
),
|
||||
];
|
||||
|
||||
try {
|
||||
// Create a per-request GraphRag instance with flow clients
|
||||
const graphRag = new GraphRag(
|
||||
{
|
||||
llm: flowCtx.flow.requestor<TextCompletionRequest, TextCompletionResponse>("llm"),
|
||||
embeddings: flowCtx.flow.requestor<EmbeddingsRequest, EmbeddingsResponse>("embeddings"),
|
||||
graphEmbeddings: flowCtx.flow.requestor<GraphEmbeddingsRequest, GraphEmbeddingsResponse>("graph-embeddings"),
|
||||
triples: flowCtx.flow.requestor<TriplesQueryRequest, TriplesQueryResponse>("triples"),
|
||||
prompt: flowCtx.flow.requestor<PromptRequest, PromptResponse>("prompt"),
|
||||
},
|
||||
{
|
||||
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
|
||||
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
|
||||
...(msg.maxSubgraphSize !== undefined
|
||||
? { maxSubgraphSize: msg.maxSubgraphSize }
|
||||
: {}),
|
||||
...(msg.maxPathLength !== undefined ? { maxPathLength: msg.maxPathLength } : {}),
|
||||
},
|
||||
);
|
||||
|
||||
const result = await graphRag.query(msg.query, {
|
||||
...(msg.collection !== undefined ? { collection: msg.collection } : {}),
|
||||
});
|
||||
|
||||
// Send answer with explain data embedded in a SINGLE message.
|
||||
// Non-streaming callers (agent's RequestResponse) return the first
|
||||
// response — so the answer must be in that first (and only) message.
|
||||
// Streaming callers (gateway) extract explain data + answer from
|
||||
// the same message.
|
||||
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;
|
||||
}
|
||||
|
||||
await producer.send(requestId, response);
|
||||
} catch (err) {
|
||||
console.error("[GraphRag] Query failed:", err);
|
||||
await producer.send(requestId, {
|
||||
response: "",
|
||||
error: { type: "rag-error", message: String(err) },
|
||||
});
|
||||
export class GraphRagService extends FlowProcessor<GraphRagEngine> {
|
||||
constructor(config: ProcessorConfig) {
|
||||
super(config);
|
||||
for (const spec of makeGraphRagSpecs()) {
|
||||
this.registerSpecification(spec);
|
||||
}
|
||||
}
|
||||
|
||||
override startEffect() {
|
||||
return super.startEffect().pipe(
|
||||
Effect.provideService(GraphRagEngine, GraphRagEngine.of(makeGraphRagEngine())),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export const program = makeProcessorProgram({
|
||||
export const program = makeFlowProcessorProgram({
|
||||
id: "graph-rag",
|
||||
make: (config) => new GraphRagService(config),
|
||||
specs: makeGraphRagSpecs,
|
||||
layer: () => GraphRagLive,
|
||||
});
|
||||
|
||||
export async function run(): Promise<void> {
|
||||
await GraphRagService.launch("graph-rag");
|
||||
await Effect.runPromise(program);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,22 +1,15 @@
|
|||
/**
|
||||
* Graph RAG retrieval pipeline.
|
||||
*
|
||||
* This is the core RAG pipeline that:
|
||||
* 1. Extracts concepts from the query
|
||||
* 2. Embeds concepts to find matching entities
|
||||
* 3. Traverses the knowledge graph from those entities
|
||||
* 4. Scores and filters edges
|
||||
* 5. Synthesizes an answer with the selected context
|
||||
*
|
||||
* Python reference: trustgraph-flow/trustgraph/retrieval/graph_rag/graph_rag.py
|
||||
*/
|
||||
|
||||
import type {
|
||||
EmbeddingsRequest,
|
||||
EmbeddingsResponse,
|
||||
FlowRequestor,
|
||||
GraphEmbeddingsRequest,
|
||||
GraphEmbeddingsResponse,
|
||||
FlowRequestor,
|
||||
PromptRequest,
|
||||
PromptResponse,
|
||||
Term,
|
||||
|
|
@ -26,6 +19,10 @@ import type {
|
|||
TriplesQueryRequest,
|
||||
TriplesQueryResponse,
|
||||
} from "@trustgraph/base";
|
||||
import { errorMessage } from "@trustgraph/base";
|
||||
import { Context, Effect, Layer } from "effect";
|
||||
import * as O from "effect/Option";
|
||||
import * as S from "effect/Schema";
|
||||
|
||||
export interface GraphRagConfig {
|
||||
entityLimit?: number;
|
||||
|
|
@ -46,321 +43,373 @@ export interface GraphRagClients {
|
|||
|
||||
export type ChunkCallback = (text: string, endOfStream: boolean) => Promise<void>;
|
||||
|
||||
export interface GraphRagQueryOptions {
|
||||
readonly collection?: string;
|
||||
readonly streaming?: boolean;
|
||||
readonly chunkCallback?: ChunkCallback;
|
||||
}
|
||||
|
||||
export interface GraphRagResult {
|
||||
answer: string;
|
||||
subgraph: Triple[];
|
||||
}
|
||||
|
||||
interface NormalizedGraphRagConfig {
|
||||
entityLimit: number;
|
||||
tripleLimit: number;
|
||||
maxSubgraphSize: number;
|
||||
maxPathLength: number;
|
||||
edgeScoreLimit: number;
|
||||
edgeLimit: number;
|
||||
}
|
||||
|
||||
export class GraphRagEngineError extends S.TaggedErrorClass<GraphRagEngineError>()(
|
||||
"GraphRagEngineError",
|
||||
{
|
||||
message: S.String,
|
||||
operation: S.String,
|
||||
cause: S.DefectWithStack,
|
||||
},
|
||||
) {}
|
||||
|
||||
export interface GraphRagEngineShape {
|
||||
readonly query: (
|
||||
clients: GraphRagClients,
|
||||
queryText: string,
|
||||
options?: GraphRagQueryOptions,
|
||||
config?: GraphRagConfig,
|
||||
) => Effect.Effect<GraphRagResult, GraphRagEngineError>;
|
||||
}
|
||||
|
||||
export class GraphRagEngine extends Context.Service<GraphRagEngine, GraphRagEngineShape>()(
|
||||
"@trustgraph/flow/retrieval/graph-rag/GraphRagEngine",
|
||||
) {}
|
||||
|
||||
const graphRagError = (operation: string, cause: unknown) =>
|
||||
new GraphRagEngineError({
|
||||
operation,
|
||||
cause,
|
||||
message: errorMessage(cause),
|
||||
});
|
||||
|
||||
export function normalizeGraphRagConfig(config: GraphRagConfig = {}): NormalizedGraphRagConfig {
|
||||
return {
|
||||
entityLimit: config.entityLimit ?? 50,
|
||||
tripleLimit: config.tripleLimit ?? 30,
|
||||
maxSubgraphSize: config.maxSubgraphSize ?? 1000,
|
||||
maxPathLength: config.maxPathLength ?? 2,
|
||||
edgeScoreLimit: config.edgeScoreLimit ?? 50,
|
||||
edgeLimit: config.edgeLimit ?? 25,
|
||||
};
|
||||
}
|
||||
|
||||
export function makeGraphRagEngine(): GraphRagEngineShape {
|
||||
return {
|
||||
query: Effect.fn("GraphRagEngine.query")((
|
||||
clients: GraphRagClients,
|
||||
queryText: string,
|
||||
options?: GraphRagQueryOptions,
|
||||
config?: GraphRagConfig,
|
||||
) =>
|
||||
Effect.tryPromise({
|
||||
try: () => queryGraphRag(clients, queryText, options, config),
|
||||
catch: (cause) => graphRagError("query", cause),
|
||||
}),
|
||||
),
|
||||
};
|
||||
}
|
||||
|
||||
export const GraphRagLive: Layer.Layer<GraphRagEngine> = Layer.succeed(
|
||||
GraphRagEngine,
|
||||
GraphRagEngine.of(makeGraphRagEngine()),
|
||||
);
|
||||
|
||||
export class GraphRag {
|
||||
private readonly engine = makeGraphRagEngine();
|
||||
private readonly clients: GraphRagClients;
|
||||
private config: Required<GraphRagConfig>;
|
||||
private readonly config: GraphRagConfig;
|
||||
|
||||
constructor(
|
||||
clients: GraphRagClients,
|
||||
config: GraphRagConfig = {},
|
||||
) {
|
||||
this.clients = clients;
|
||||
this.config = {
|
||||
entityLimit: config.entityLimit ?? 50,
|
||||
tripleLimit: config.tripleLimit ?? 30,
|
||||
maxSubgraphSize: config.maxSubgraphSize ?? 1000,
|
||||
maxPathLength: config.maxPathLength ?? 2,
|
||||
edgeScoreLimit: config.edgeScoreLimit ?? 50,
|
||||
edgeLimit: config.edgeLimit ?? 25,
|
||||
};
|
||||
this.config = config;
|
||||
}
|
||||
|
||||
async query(
|
||||
query(
|
||||
queryText: string,
|
||||
options?: {
|
||||
collection?: string;
|
||||
streaming?: boolean;
|
||||
chunkCallback?: ChunkCallback;
|
||||
},
|
||||
options?: GraphRagQueryOptions,
|
||||
): Promise<GraphRagResult> {
|
||||
console.log(`[GraphRag] Query: "${queryText.slice(0, 80)}..."`);
|
||||
|
||||
// Step 1: Extract concepts from the query via prompt + LLM
|
||||
const concepts = await this.extractConcepts(queryText);
|
||||
console.log(`[GraphRag] Step 1: extracted ${concepts.length} concepts: ${concepts.slice(0, 5).join(", ")}`);
|
||||
|
||||
// Step 2: Embed concepts concurrently
|
||||
const vectors = await this.getVectors(concepts);
|
||||
console.log(`[GraphRag] Step 2: got ${vectors.length} vectors (dim=${vectors[0]?.length ?? 0})`);
|
||||
|
||||
// Step 3: Find matching entities via graph embeddings
|
||||
const entities = await this.getEntities(vectors, options?.collection);
|
||||
console.log(`[GraphRag] Step 3: found ${entities.length} matching entities`);
|
||||
|
||||
// Step 4: Traverse the knowledge graph from entities
|
||||
const subgraph = await this.followEdges(entities, options?.collection);
|
||||
console.log(`[GraphRag] Step 4: traversed graph, ${subgraph.length} triples in subgraph`);
|
||||
|
||||
// Step 5: Score and filter edges via LLM
|
||||
const scoredEdges = await this.scoreEdges(queryText, subgraph);
|
||||
console.log(`[GraphRag] Step 5: scored down to ${scoredEdges.length} edges`);
|
||||
|
||||
// Step 6: Synthesize answer
|
||||
console.log(`[GraphRag] Step 6: synthesizing answer from ${scoredEdges.length} edges...`);
|
||||
const answer = await this.synthesize(
|
||||
queryText,
|
||||
scoredEdges,
|
||||
options?.chunkCallback,
|
||||
return Effect.runPromise(
|
||||
this.engine.query(this.clients, queryText, options, this.config),
|
||||
);
|
||||
console.log(`[GraphRag] Step 6: done (${answer.length} chars)`);
|
||||
|
||||
return { answer, subgraph: scoredEdges };
|
||||
}
|
||||
|
||||
private async extractConcepts(query: string): Promise<string[]> {
|
||||
const promptResp = await this.clients.prompt.request({
|
||||
name: "extract-concepts",
|
||||
variables: { query },
|
||||
});
|
||||
|
||||
const llmResp = await this.clients.llm.request({
|
||||
system: (promptResp as PromptResponse).system,
|
||||
prompt: (promptResp as PromptResponse).prompt,
|
||||
});
|
||||
|
||||
// Parse concepts from LLM response (newline-separated)
|
||||
return (llmResp as TextCompletionResponse).response
|
||||
.split("\n")
|
||||
.map((c) => c.trim())
|
||||
.filter((c) => c.length > 0);
|
||||
}
|
||||
|
||||
private async getVectors(concepts: string[]): Promise<number[][]> {
|
||||
const resp = await this.clients.embeddings.request({ text: concepts });
|
||||
return (resp as EmbeddingsResponse).vectors;
|
||||
}
|
||||
|
||||
private async getEntities(vectors: number[][], collection?: string): Promise<Term[]> {
|
||||
const resp = await this.clients.graphEmbeddings.request({
|
||||
vectors,
|
||||
user: "default",
|
||||
collection: collection ?? "default",
|
||||
limit: this.config.entityLimit,
|
||||
});
|
||||
return (resp as GraphEmbeddingsResponse).entities;
|
||||
}
|
||||
|
||||
private async followEdges(entities: Term[], collection?: string): Promise<Triple[]> {
|
||||
// BFS multi-hop traversal up to maxPathLength
|
||||
const visited = new Set<string>();
|
||||
const subgraph: Triple[] = [];
|
||||
|
||||
// Current frontier: the set of entities to expand at this depth level
|
||||
let currentLevel = new Set<string>(
|
||||
entities.map((e) => termToString(e)),
|
||||
);
|
||||
|
||||
for (let depth = 0; depth < this.config.maxPathLength; depth++) {
|
||||
if (currentLevel.size === 0 || subgraph.length >= this.config.maxSubgraphSize) {
|
||||
break;
|
||||
}
|
||||
|
||||
// Filter out already-visited entities
|
||||
const unvisited = [...currentLevel].filter((e) => !visited.has(e));
|
||||
if (unvisited.length === 0) break;
|
||||
|
||||
// Batch triple queries for all unvisited entities at this depth
|
||||
// Query each entity as subject to get outgoing edges
|
||||
const queries = unvisited.map((entityStr) => {
|
||||
const term = stringToTerm(entityStr);
|
||||
const request: TriplesQueryRequest = {
|
||||
s: term,
|
||||
limit: this.config.tripleLimit,
|
||||
...(collection !== undefined ? { collection } : {}),
|
||||
};
|
||||
return this.clients.triples.request(request);
|
||||
});
|
||||
|
||||
const results = await Promise.all(queries);
|
||||
|
||||
const nextLevel = new Set<string>();
|
||||
|
||||
for (const result of results) {
|
||||
const triples = (result as TriplesQueryResponse).triples;
|
||||
for (const triple of triples) {
|
||||
subgraph.push(triple);
|
||||
|
||||
// Collect objects as next-level entities for further expansion
|
||||
// (only if we have more depth levels remaining)
|
||||
if (depth < this.config.maxPathLength - 1) {
|
||||
const objStr = termToString(triple.o);
|
||||
if (!visited.has(objStr)) {
|
||||
nextLevel.add(objStr);
|
||||
}
|
||||
}
|
||||
|
||||
if (subgraph.length >= this.config.maxSubgraphSize) {
|
||||
return subgraph;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Mark current level as visited and move to next
|
||||
for (const e of currentLevel) {
|
||||
visited.add(e);
|
||||
}
|
||||
currentLevel = nextLevel;
|
||||
}
|
||||
|
||||
return subgraph.slice(0, this.config.maxSubgraphSize);
|
||||
}
|
||||
|
||||
private async scoreEdges(query: string, triples: Triple[]): Promise<Triple[]> {
|
||||
if (triples.length === 0) return [];
|
||||
|
||||
// If the subgraph is small enough, skip LLM scoring entirely
|
||||
// 500 triples is well within LLM context limits and avoids lossy scoring
|
||||
if (triples.length <= 500) {
|
||||
console.log(`[GraphRag] Skipping edge scoring — ${triples.length} triples fits in context directly`);
|
||||
return triples;
|
||||
}
|
||||
|
||||
// Build a numbered list of edges for the LLM to score
|
||||
const edgeDescriptions = triples.map((t, i) => ({
|
||||
id: String(i),
|
||||
s: termToString(t.s),
|
||||
p: termToString(t.p),
|
||||
o: termToString(t.o),
|
||||
}));
|
||||
|
||||
// Limit how many edges we send for scoring to avoid overflowing context
|
||||
const toScore = edgeDescriptions.slice(0, this.config.edgeScoreLimit);
|
||||
|
||||
const knowledgeJson = JSON.stringify(toScore, null, 2);
|
||||
|
||||
// Ask the LLM to score each edge for relevance to the query
|
||||
const promptResp = await this.clients.prompt.request({
|
||||
name: "kg-edge-scoring",
|
||||
variables: {
|
||||
query,
|
||||
knowledge: knowledgeJson,
|
||||
},
|
||||
});
|
||||
|
||||
const llmResp = await this.clients.llm.request({
|
||||
system: (promptResp as PromptResponse).system,
|
||||
prompt: (promptResp as PromptResponse).prompt,
|
||||
});
|
||||
|
||||
const responseText = (llmResp as TextCompletionResponse).response;
|
||||
console.log(`[GraphRag] Edge scoring LLM response (first 500 chars): ${responseText.slice(0, 500)}`);
|
||||
|
||||
// Parse scores from LLM response
|
||||
// Expected format: JSON array of { id: string, score: number }
|
||||
// or newline-separated JSON objects
|
||||
const scored: Array<{ id: string; score: number }> = [];
|
||||
|
||||
try {
|
||||
// Try parsing as a JSON array first
|
||||
const parsed = JSON.parse(responseText) as Array<{ id: string; score: number }>;
|
||||
if (Array.isArray(parsed)) {
|
||||
for (const item of parsed) {
|
||||
if (
|
||||
typeof item === "object" &&
|
||||
item !== null &&
|
||||
typeof item.id === "string" &&
|
||||
typeof item.score === "number"
|
||||
) {
|
||||
scored.push({ id: item.id, score: item.score });
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
// Fall back to parsing line-by-line JSON objects
|
||||
for (const line of responseText.split("\n")) {
|
||||
const trimmed = line.trim();
|
||||
if (trimmed.length === 0) continue;
|
||||
try {
|
||||
const obj = JSON.parse(trimmed) as { id?: string; score?: number };
|
||||
if (
|
||||
typeof obj === "object" &&
|
||||
obj !== null &&
|
||||
typeof obj.id === "string" &&
|
||||
typeof obj.score === "number"
|
||||
) {
|
||||
scored.push({ id: obj.id, score: obj.score });
|
||||
}
|
||||
} catch {
|
||||
// Skip unparseable lines
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sort by score descending and keep top N
|
||||
scored.sort((a, b) => b.score - a.score);
|
||||
const topN = scored.slice(0, this.config.edgeLimit);
|
||||
// Map back to triples
|
||||
const result: Triple[] = [];
|
||||
for (const entry of topN) {
|
||||
const idx = parseInt(entry.id, 10);
|
||||
if (!isNaN(idx) && idx >= 0 && idx < triples.length) {
|
||||
result.push(triples[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`[GraphRag] Edge scoring: LLM returned ${scored.length} scores, keeping top ${topN.length}, mapped ${result.length} triples`);
|
||||
|
||||
// If scoring failed entirely, fall back to returning the first edgeLimit triples
|
||||
if (result.length === 0) {
|
||||
return triples.slice(0, this.config.edgeLimit);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
private async synthesize(
|
||||
query: string,
|
||||
edges: Triple[],
|
||||
chunkCallback?: ChunkCallback,
|
||||
): Promise<string> {
|
||||
// Format edges as context
|
||||
const context = edges
|
||||
.map((t) => `${termToString(t.s)} -> ${termToString(t.p)} -> ${termToString(t.o)}`)
|
||||
.join("\n");
|
||||
|
||||
const promptResp = await this.clients.prompt.request({
|
||||
name: "graph-rag-synthesize",
|
||||
variables: { query, context },
|
||||
});
|
||||
|
||||
if (chunkCallback !== undefined) {
|
||||
// Streaming response
|
||||
let fullText = "";
|
||||
await this.clients.llm.request(
|
||||
{
|
||||
system: (promptResp as PromptResponse).system,
|
||||
prompt: (promptResp as PromptResponse).prompt,
|
||||
streaming: true,
|
||||
},
|
||||
{
|
||||
recipient: async (resp) => {
|
||||
const r = resp as TextCompletionResponse;
|
||||
if (r.response.length > 0) {
|
||||
fullText += r.response;
|
||||
await chunkCallback(r.response, r.endOfStream === true);
|
||||
}
|
||||
return r.endOfStream === true;
|
||||
},
|
||||
},
|
||||
);
|
||||
return fullText;
|
||||
}
|
||||
|
||||
const resp = await this.clients.llm.request({
|
||||
system: (promptResp as PromptResponse).system,
|
||||
prompt: (promptResp as PromptResponse).prompt,
|
||||
});
|
||||
|
||||
return (resp as TextCompletionResponse).response;
|
||||
}
|
||||
}
|
||||
|
||||
function termToString(term: Term): string {
|
||||
async function queryGraphRag(
|
||||
clients: GraphRagClients,
|
||||
queryText: string,
|
||||
options?: GraphRagQueryOptions,
|
||||
rawConfig?: GraphRagConfig,
|
||||
): Promise<GraphRagResult> {
|
||||
const config = normalizeGraphRagConfig(rawConfig);
|
||||
console.log(`[GraphRag] Query: "${queryText.slice(0, 80)}..."`);
|
||||
|
||||
const concepts = await extractConcepts(clients, queryText);
|
||||
console.log(`[GraphRag] Step 1: extracted ${concepts.length} concepts: ${concepts.slice(0, 5).join(", ")}`);
|
||||
|
||||
const vectors = await getVectors(clients, concepts);
|
||||
console.log(`[GraphRag] Step 2: got ${vectors.length} vectors (dim=${vectors[0]?.length ?? 0})`);
|
||||
|
||||
const entities = await getEntities(clients, config, vectors, options?.collection);
|
||||
console.log(`[GraphRag] Step 3: found ${entities.length} matching entities`);
|
||||
|
||||
const subgraph = await followEdges(clients, config, entities, options?.collection);
|
||||
console.log(`[GraphRag] Step 4: traversed graph, ${subgraph.length} triples in subgraph`);
|
||||
|
||||
const scoredEdges = await scoreEdges(clients, config, queryText, subgraph);
|
||||
console.log(`[GraphRag] Step 5: scored down to ${scoredEdges.length} edges`);
|
||||
|
||||
console.log(`[GraphRag] Step 6: synthesizing answer from ${scoredEdges.length} edges...`);
|
||||
const answer = await synthesize(
|
||||
clients,
|
||||
queryText,
|
||||
scoredEdges,
|
||||
options?.chunkCallback,
|
||||
);
|
||||
console.log(`[GraphRag] Step 6: done (${answer.length} chars)`);
|
||||
|
||||
return { answer, subgraph: scoredEdges };
|
||||
}
|
||||
|
||||
async function extractConcepts(clients: GraphRagClients, query: string): Promise<string[]> {
|
||||
const promptResp = await clients.prompt.request({
|
||||
name: "extract-concepts",
|
||||
variables: { query },
|
||||
});
|
||||
|
||||
const llmResp = await clients.llm.request({
|
||||
system: promptResp.system,
|
||||
prompt: promptResp.prompt,
|
||||
});
|
||||
|
||||
return llmResp.response
|
||||
.split("\n")
|
||||
.map((concept) => concept.trim())
|
||||
.filter((concept) => concept.length > 0);
|
||||
}
|
||||
|
||||
async function getVectors(clients: GraphRagClients, concepts: string[]): Promise<number[][]> {
|
||||
const resp = await clients.embeddings.request({ text: concepts });
|
||||
return resp.vectors;
|
||||
}
|
||||
|
||||
async function getEntities(
|
||||
clients: GraphRagClients,
|
||||
config: NormalizedGraphRagConfig,
|
||||
vectors: number[][],
|
||||
collection?: string,
|
||||
): Promise<Term[]> {
|
||||
const resp = await clients.graphEmbeddings.request({
|
||||
vectors,
|
||||
user: "default",
|
||||
collection: collection ?? "default",
|
||||
limit: config.entityLimit,
|
||||
});
|
||||
return resp.entities;
|
||||
}
|
||||
|
||||
async function followEdges(
|
||||
clients: GraphRagClients,
|
||||
config: NormalizedGraphRagConfig,
|
||||
entities: Term[],
|
||||
collection?: string,
|
||||
): Promise<Triple[]> {
|
||||
const visited = new Set<string>();
|
||||
const subgraph: Triple[] = [];
|
||||
let currentLevel = new Set<string>(
|
||||
entities.map((entity) => termToString(entity)),
|
||||
);
|
||||
|
||||
for (let depth = 0; depth < config.maxPathLength; depth++) {
|
||||
if (currentLevel.size === 0 || subgraph.length >= config.maxSubgraphSize) {
|
||||
break;
|
||||
}
|
||||
|
||||
const unvisited = [...currentLevel].filter((entity) => !visited.has(entity));
|
||||
if (unvisited.length === 0) break;
|
||||
|
||||
const queries = unvisited.map((entityStr) => {
|
||||
const term = stringToTerm(entityStr);
|
||||
const request: TriplesQueryRequest = {
|
||||
s: term,
|
||||
limit: config.tripleLimit,
|
||||
...(collection !== undefined ? { collection } : {}),
|
||||
};
|
||||
return clients.triples.request(request);
|
||||
});
|
||||
|
||||
const results = await Promise.all(queries);
|
||||
const nextLevel = new Set<string>();
|
||||
|
||||
for (const result of results) {
|
||||
for (const triple of result.triples) {
|
||||
subgraph.push(triple);
|
||||
|
||||
if (depth < config.maxPathLength - 1) {
|
||||
const objStr = termToString(triple.o);
|
||||
if (!visited.has(objStr)) {
|
||||
nextLevel.add(objStr);
|
||||
}
|
||||
}
|
||||
|
||||
if (subgraph.length >= config.maxSubgraphSize) {
|
||||
return subgraph;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for (const entity of currentLevel) {
|
||||
visited.add(entity);
|
||||
}
|
||||
currentLevel = nextLevel;
|
||||
}
|
||||
|
||||
return subgraph.slice(0, config.maxSubgraphSize);
|
||||
}
|
||||
|
||||
async function scoreEdges(
|
||||
clients: GraphRagClients,
|
||||
config: NormalizedGraphRagConfig,
|
||||
query: string,
|
||||
triples: Triple[],
|
||||
): Promise<Triple[]> {
|
||||
if (triples.length === 0) return [];
|
||||
|
||||
if (triples.length <= 500) {
|
||||
console.log(`[GraphRag] Skipping edge scoring - ${triples.length} triples fits in context directly`);
|
||||
return triples;
|
||||
}
|
||||
|
||||
const edgeDescriptions = triples.map((triple, index) => ({
|
||||
id: String(index),
|
||||
s: termToString(triple.s),
|
||||
p: termToString(triple.p),
|
||||
o: termToString(triple.o),
|
||||
}));
|
||||
|
||||
const toScore = edgeDescriptions.slice(0, config.edgeScoreLimit);
|
||||
const knowledgeJson = JSON.stringify(toScore, null, 2);
|
||||
|
||||
const promptResp = await clients.prompt.request({
|
||||
name: "kg-edge-scoring",
|
||||
variables: {
|
||||
query,
|
||||
knowledge: knowledgeJson,
|
||||
},
|
||||
});
|
||||
|
||||
const llmResp = await clients.llm.request({
|
||||
system: promptResp.system,
|
||||
prompt: promptResp.prompt,
|
||||
});
|
||||
|
||||
console.log(`[GraphRag] Edge scoring LLM response (first 500 chars): ${llmResp.response.slice(0, 500)}`);
|
||||
|
||||
const scored = parseScoredEdges(llmResp.response);
|
||||
scored.sort((a, b) => b.score - a.score);
|
||||
const topN = scored.slice(0, config.edgeLimit);
|
||||
|
||||
const result: Triple[] = [];
|
||||
for (const entry of topN) {
|
||||
const idx = Number.parseInt(entry.id, 10);
|
||||
if (!Number.isNaN(idx) && idx >= 0 && idx < triples.length) {
|
||||
result.push(triples[idx]);
|
||||
}
|
||||
}
|
||||
|
||||
console.log(`[GraphRag] Edge scoring: LLM returned ${scored.length} scores, keeping top ${topN.length}, mapped ${result.length} triples`);
|
||||
|
||||
if (result.length === 0) {
|
||||
return triples.slice(0, config.edgeLimit);
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
async function synthesize(
|
||||
clients: GraphRagClients,
|
||||
query: string,
|
||||
edges: Triple[],
|
||||
chunkCallback?: ChunkCallback,
|
||||
): Promise<string> {
|
||||
const context = edges
|
||||
.map((triple) => `${termToString(triple.s)} -> ${termToString(triple.p)} -> ${termToString(triple.o)}`)
|
||||
.join("\n");
|
||||
|
||||
const promptResp = await clients.prompt.request({
|
||||
name: "graph-rag-synthesize",
|
||||
variables: { query, context },
|
||||
});
|
||||
|
||||
if (chunkCallback !== undefined) {
|
||||
let fullText = "";
|
||||
await clients.llm.request(
|
||||
{
|
||||
system: promptResp.system,
|
||||
prompt: promptResp.prompt,
|
||||
streaming: true,
|
||||
},
|
||||
{
|
||||
recipient: async (resp) => {
|
||||
if (resp.response.length > 0) {
|
||||
fullText += resp.response;
|
||||
await chunkCallback(resp.response, resp.endOfStream === true);
|
||||
}
|
||||
return resp.endOfStream === true;
|
||||
},
|
||||
},
|
||||
);
|
||||
return fullText;
|
||||
}
|
||||
|
||||
const resp = await clients.llm.request({
|
||||
system: promptResp.system,
|
||||
prompt: promptResp.prompt,
|
||||
});
|
||||
|
||||
return resp.response;
|
||||
}
|
||||
|
||||
const ScoredEdge = S.Struct({
|
||||
id: S.String,
|
||||
score: S.Number,
|
||||
});
|
||||
const ScoredEdgesFromJson = S.Array(ScoredEdge).pipe(S.fromJsonString);
|
||||
const ScoredEdgeFromJson = ScoredEdge.pipe(S.fromJsonString);
|
||||
const decodeScoredEdges = S.decodeUnknownOption(ScoredEdgesFromJson);
|
||||
const decodeScoredEdge = S.decodeUnknownOption(ScoredEdgeFromJson);
|
||||
|
||||
function parseScoredEdges(responseText: string): Array<typeof ScoredEdge.Type> {
|
||||
const parsedArray = decodeScoredEdges(responseText);
|
||||
if (O.isSome(parsedArray)) {
|
||||
return Array.from(parsedArray.value);
|
||||
}
|
||||
|
||||
const scored: Array<typeof ScoredEdge.Type> = [];
|
||||
for (const line of responseText.split("\n")) {
|
||||
const trimmed = line.trim();
|
||||
if (trimmed.length === 0) continue;
|
||||
const parsedLine = decodeScoredEdge(trimmed);
|
||||
if (O.isSome(parsedLine)) {
|
||||
scored.push(parsedLine.value);
|
||||
}
|
||||
}
|
||||
return scored;
|
||||
}
|
||||
|
||||
export function termToString(term: Term): string {
|
||||
switch (term.type) {
|
||||
case "IRI":
|
||||
return term.iri;
|
||||
|
|
@ -373,7 +422,7 @@ function termToString(term: Term): string {
|
|||
}
|
||||
}
|
||||
|
||||
function stringToTerm(value: string): Term {
|
||||
export function stringToTerm(value: string): Term {
|
||||
if (value.startsWith("http://") || value.startsWith("https://")) {
|
||||
return { type: "IRI", iri: value };
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue