From 0b0bdb98df8dd19e869a14bcfdb88afd7b3e53a8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CSebastian?= <64795732+slegarraga@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:49:52 -0400 Subject: [PATCH] fix: emit retriever resources as sources --- src/dify-chat-language-model.test.ts | 80 ++++++++++++++++++++++++++-- src/dify-chat-language-model.ts | 61 ++++++++++++++++++++- 2 files changed, 136 insertions(+), 5 deletions(-) diff --git a/src/dify-chat-language-model.test.ts b/src/dify-chat-language-model.test.ts index 05b9173..ab95ce4 100644 --- a/src/dify-chat-language-model.test.ts +++ b/src/dify-chat-language-model.test.ts @@ -314,10 +314,14 @@ describe("DifyChatLanguageModel", () => { headers: { "Custom-Header": "custom-value" }, } as any); - // @ai-sdk/provider-utils v3 normalizes headers through the Headers API - // which lowercases all header names (per HTTP/2 spec) - expect(capturedOptions.headers["authorization"]).toBe("Bearer test"); - expect(capturedOptions.headers["custom-header"]).toBe("custom-value"); + const getHeader = (name: string) => + typeof capturedOptions.headers.get === "function" + ? capturedOptions.headers.get(name) + : Object.entries(capturedOptions.headers).find( + ([key]) => key.toLowerCase() === name + )?.[1]; + expect(getHeader("authorization")).toBe("Bearer test"); + expect(getHeader("custom-header")).toBe("custom-value"); }); it("should use default user-id when no user-id is provided", async () => { @@ -449,6 +453,74 @@ describe("DifyChatLanguageModel", () => { expect(finishPart?.usage?.outputTokens).toBe(25); // message_end doesn't have data.data.total_tokens }); + it("should emit retriever resources as source parts", async () => { + const retrieverResource = { + segment_id: "segment-1", + document_id: "document-1", + document_name: "Knowledge Base.md", + dataset_name: "Support docs", + content: "Relevant source excerpt", + score: 0.98, + }; + const mockResponseText = + `data: {"event":"agent_message","answer":"Agent response","id":"agent1","conversation_id":"conv1","message_id":"msg1"}\n\n` + + `data: ${JSON.stringify({ + event: "message_end", + id: "msg1", + metadata: { + usage: { prompt_tokens: 10, completion_tokens: 25, total_tokens: 35 }, + retriever_resources: [retrieverResource], + }, + conversation_id: "conv1", + message_id: "msg1", + task_id: "task1", + })}\n\n`; + + const mockFetch = createMockFetch({ + ok: true, + headers: new Map([["Content-Type", "text/event-stream"]]), + body: new ReadableStream({ + start(controller) { + const encoder = new TextEncoder(); + controller.enqueue(encoder.encode(mockResponseText)); + controller.close(); + }, + }), + status: 200, + }); + + const model = makeModel({ fetch: mockFetch }); + const { stream: resultStream } = await model.doStream({ + prompt: [{ role: "user", content: [{ type: "text", text: "Hi" }] }], + } as any); + + const parts: any[] = []; + const reader = resultStream.getReader(); + while (true) { + const { value, done } = await reader.read(); + if (done) break; + parts.push(value); + } + + const sourcePart = parts.find((p) => p.type === "source"); + expect(sourcePart).toMatchObject({ + type: "source", + sourceType: "document", + id: "segment-1", + mediaType: "text/plain", + title: "Knowledge Base.md", + filename: "Knowledge Base.md", + providerMetadata: { + difyWorkflowData: { + conversationId: "conv1", + messageId: "msg1", + taskId: "task1", + retrieverResource, + }, + }, + }); + }); + it("should handle message_end with usage tokens from data field", async () => { // Test message_end event with usage tokens in data.total_tokens (like workflow_finished) const mockResponseText = `data: {"event":"message_end","id":"msg1","data":{"total_tokens":50},"metadata":{"usage":{"prompt_tokens":10,"completion_tokens":25,"total_tokens":35}},"conversation_id":"conv1","message_id":"msg1"}\n\n`; diff --git a/src/dify-chat-language-model.ts b/src/dify-chat-language-model.ts index da3e89c..a943298 100644 --- a/src/dify-chat-language-model.ts +++ b/src/dify-chat-language-model.ts @@ -81,6 +81,51 @@ interface Message { [key: string]: unknown; } +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function getStringField(record: Record, keys: string[]): string | undefined { + for (const key of keys) { + const value = record[key]; + if (typeof value === "string" && value) return value; + } + return undefined; +} + +function buildRetrieverResourceSources( + resources: unknown, + metadata: Record +): LanguageModelV2StreamPart[] { + if (!Array.isArray(resources)) return []; + + return resources.map((resource, index) => { + const record = isRecord(resource) ? resource : {}; + const id = + getStringField(record, ["segment_id", "document_id", "id"]) ?? + `retriever-resource-${index + 1}`; + const title = + getStringField(record, ["document_name", "title", "dataset_name", "segment_id", "id"]) ?? + `Retriever resource ${index + 1}`; + const filename = getStringField(record, ["document_name", "filename"]); + + return { + type: "source", + sourceType: "document", + id, + mediaType: "text/plain", + title, + ...(filename ? { filename } : {}), + providerMetadata: { + difyWorkflowData: { + ...metadata, + retrieverResource: resource as JSONValue, + }, + }, + }; + }); +} + export class DifyChatLanguageModel implements LanguageModelV2 { readonly specificationVersion = "v2" as const; readonly modelId: string; @@ -233,11 +278,16 @@ export class DifyChatLanguageModel implements LanguageModelV2 { const logger = this.logger; const logMessages = this.settings.logMessages; - function buildDifyMetadata() { + function buildDifyWorkflowData() { const meta: Record = {}; if (conversationId) meta.conversationId = conversationId as JSONValue; if (messageId) meta.messageId = messageId as JSONValue; if (taskId) meta.taskId = taskId as JSONValue; + return meta; + } + + function buildDifyMetadata() { + const meta = buildDifyWorkflowData(); return Object.keys(meta).length ? { providerMetadata: { difyWorkflowData: meta } } : {}; } @@ -341,6 +391,15 @@ export class DifyChatLanguageModel implements LanguageModelV2 { controller.enqueue({ type: "text-end", id: "0", ...buildDifyMetadata() }); } + if (data.event === "message_end") { + for (const sourcePart of buildRetrieverResourceSources( + (data as any).metadata?.retriever_resources, + buildDifyWorkflowData() + )) { + controller.enqueue(sourcePart); + } + } + // Don't emit cleanText here since it was already emitted via stream for (const tc of toolCalls) {