Add typed flow spec accessors

This commit is contained in:
elpresidank 2026-06-02 03:23:23 -05:00
parent abb6f3aed0
commit 44110c5bb4
19 changed files with 457 additions and 223 deletions

View file

@ -40,6 +40,28 @@ import {
type DocumentRagClients,
} from "./document-rag.js";
const DocumentRagResponseProducer = makeProducerSpec<DocumentRagResponse>("document-rag-response");
const DocumentRagLlmClient = makeRequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
);
const DocumentRagEmbeddingsClient = makeRequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
);
const DocumentRagDocEmbeddingsClient = makeRequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
"doc-embeddings",
"document-embeddings-request",
"document-embeddings-response",
);
const DocumentRagPromptClient = makeRequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
);
const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function* (
msg: DocumentRagRequest,
properties: Record<string, string>,
@ -48,14 +70,14 @@ const onDocumentRagRequest = Effect.fn("DocumentRagService.onRequest")(function*
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<DocumentRagResponse>("document-rag-response");
const producer = yield* flowCtx.flow.producerEffect(DocumentRagResponseProducer);
const engine = yield* DocumentRagEngine;
const clients: DocumentRagClients = {
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"),
llm: yield* flowCtx.flow.requestorEffect(DocumentRagLlmClient),
embeddings: yield* flowCtx.flow.requestorEffect(DocumentRagEmbeddingsClient),
docEmbeddings: yield* flowCtx.flow.requestorEffect(DocumentRagDocEmbeddingsClient),
prompt: yield* flowCtx.flow.requestorEffect(DocumentRagPromptClient),
};
const response = yield* engine.query(
@ -90,27 +112,11 @@ export const makeDocumentRagSpecs = (): ReadonlyArray<Spec<DocumentRagEngine>> =
"document-rag-request",
onDocumentRagRequest,
),
makeProducerSpec<DocumentRagResponse>("document-rag-response"),
makeRequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
makeRequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
makeRequestResponseSpec<DocumentEmbeddingsRequest, DocumentEmbeddingsResponse>(
"doc-embeddings",
"document-embeddings-request",
"document-embeddings-response",
),
makeRequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
DocumentRagResponseProducer,
DocumentRagLlmClient,
DocumentRagEmbeddingsClient,
DocumentRagDocEmbeddingsClient,
DocumentRagPromptClient,
];
export type DocumentRagService = FlowProcessorRuntime<DocumentRagEngine>;

View file

@ -43,6 +43,33 @@ import {
type GraphRagConfig,
} from "./graph-rag.js";
const GraphRagResponseProducer = makeProducerSpec<GraphRagResponse>("graph-rag-response");
const GraphRagLlmClient = makeRequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
);
const GraphRagEmbeddingsClient = makeRequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
);
const GraphRagGraphEmbeddingsClient = makeRequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
"graph-embeddings",
"graph-embeddings-request",
"graph-embeddings-response",
);
const GraphRagTriplesClient = makeRequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
);
const GraphRagPromptClient = makeRequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
);
const graphRagConfigFromRequest = (msg: GraphRagRequest): GraphRagConfig => ({
...(msg.entityLimit !== undefined ? { entityLimit: msg.entityLimit } : {}),
...(msg.tripleLimit !== undefined ? { tripleLimit: msg.tripleLimit } : {}),
@ -58,17 +85,17 @@ const onGraphRagRequest = Effect.fn("GraphRagService.onRequest")(function* (
const requestId = properties.id;
if (requestId === undefined || requestId.length === 0) return;
const producer = yield* flowCtx.flow.producerEffect<GraphRagResponse>("graph-rag-response");
const producer = yield* flowCtx.flow.producerEffect(GraphRagResponseProducer);
const engine = yield* GraphRagEngine;
yield* Effect.log(`[GraphRagService] Received request ${requestId}: "${msg.query?.slice(0, 60)}..." collection=${msg.collection}`);
const clients: GraphRagClients = {
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"),
llm: yield* flowCtx.flow.requestorEffect(GraphRagLlmClient),
embeddings: yield* flowCtx.flow.requestorEffect(GraphRagEmbeddingsClient),
graphEmbeddings: yield* flowCtx.flow.requestorEffect(GraphRagGraphEmbeddingsClient),
triples: yield* flowCtx.flow.requestorEffect(GraphRagTriplesClient),
prompt: yield* flowCtx.flow.requestorEffect(GraphRagPromptClient),
};
const result = yield* engine.query(
@ -118,32 +145,12 @@ export const makeGraphRagSpecs = (): ReadonlyArray<Spec<GraphRagEngine>> => [
"graph-rag-request",
onGraphRagRequest,
),
makeProducerSpec<GraphRagResponse>("graph-rag-response"),
makeRequestResponseSpec<TextCompletionRequest, TextCompletionResponse>(
"llm",
"text-completion-request",
"text-completion-response",
),
makeRequestResponseSpec<EmbeddingsRequest, EmbeddingsResponse>(
"embeddings",
"embeddings-request",
"embeddings-response",
),
makeRequestResponseSpec<GraphEmbeddingsRequest, GraphEmbeddingsResponse>(
"graph-embeddings",
"graph-embeddings-request",
"graph-embeddings-response",
),
makeRequestResponseSpec<TriplesQueryRequest, TriplesQueryResponse>(
"triples",
"triples-request",
"triples-response",
),
makeRequestResponseSpec<PromptRequest, PromptResponse>(
"prompt",
"prompt-request",
"prompt-response",
),
GraphRagResponseProducer,
GraphRagLlmClient,
GraphRagEmbeddingsClient,
GraphRagGraphEmbeddingsClient,
GraphRagTriplesClient,
GraphRagPromptClient,
];
export type GraphRagService = FlowProcessorRuntime<GraphRagEngine>;