Emit retriever resources as source parts

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`;

src/dify-chat-language-model.ts

@@ -81,6 +81,51 @@ interface Message {
   [key: string]: unknown;
 }
 
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function getStringField(record: Record<string, unknown>, 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<string, JSONValue>
+): 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<string, JSONValue> = {};
       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) {