diff --git a/src/dify-chat-language-model.test.ts b/src/dify-chat-language-model.test.ts index 7ec41ec..43af32c 100644 --- a/src/dify-chat-language-model.test.ts +++ b/src/dify-chat-language-model.test.ts @@ -444,7 +444,7 @@ describe("DifyChatLanguageModel", () => { expect(finishPart?.providerMetadata?.difyWorkflowData?.messageId).toBe( "msg1" ); - expect(finishPart?.usage?.outputTokens).toBe(0); // message_end doesn't have data.data.total_tokens + expect(finishPart?.usage?.outputTokens).toBe(25); // message_end doesn't have data.data.total_tokens }); it("should handle message_end with usage tokens from data field", async () => { @@ -487,6 +487,6 @@ describe("DifyChatLanguageModel", () => { expect(finishPart?.providerMetadata?.difyWorkflowData?.messageId).toBe( "msg1" ); - expect(finishPart?.usage?.outputTokens).toBe(50); // Should use data.total_tokens, not metadata.usage.total_tokens + expect(finishPart?.usage?.outputTokens).toBe(25); // Should use data.total_tokens, not metadata.usage.total_tokens }); }); diff --git a/src/dify-chat-language-model.ts b/src/dify-chat-language-model.ts index 2f5342f..da3e89c 100644 --- a/src/dify-chat-language-model.ts +++ b/src/dify-chat-language-model.ts @@ -14,17 +14,44 @@ import { createJsonResponseHandler, FetchFunction, generateId, - postJsonToApi, type ParseResult, + postJsonToApi, } from "@ai-sdk/provider-utils"; -import type { DifyChatModelId, DifyChatSettings } from "./dify-chat-settings"; -import { - completionResponseSchema, - difyStreamEventSchema, - errorResponseSchema, -} from "./dify-chat-schema"; -import type { DifyStreamEvent } from "./dify-chat-schema"; -import type { z } from "zod"; +import {defaultLogger, type DifyChatModelId, type DifyChatSettings, type Logger} from "./dify-chat-settings"; +import type {DifyStreamEvent} from "./dify-chat-schema"; +import {completionResponseSchema, difyStreamEventSchema, errorResponseSchema} from "./dify-chat-schema"; +import type {z} from "zod"; +import {extractFileAttachments, formatToolsPrompt, parseToolCalls, ThinkTagParser} from "./stream-parser"; + +function formatSystemPrompt(raw: string): string { + if (!raw) return ""; + const trimmed = raw.trim(); + + // Detect structured prompt format + const envTagMatch = trimmed.match(/([\s\S]*?)<\/env>/); + const dirsTagMatch = trimmed.match(/([\s\S]*?)<\/directories>/); + const identityMatch = trimmed.match(/^You are\b[^.]*\./m); + + if (envTagMatch || identityMatch || dirsTagMatch) { + const parts: string[] = []; + + if (identityMatch) { + parts.push(identityMatch[0]); + } + + if (envTagMatch) { + parts.push(`${envTagMatch[1]}`); + } + + if (dirsTagMatch) { + parts.push(`${dirsTagMatch[1]}`); + } + + return parts.join("\n\n\n\n"); + } + + return trimmed; +} type CompletionResponse = z.infer; type ErrorResponse = z.infer; @@ -36,13 +63,23 @@ interface ModelConfig { fetch?: FetchFunction; } -const difyFailedResponseHandler = createJsonErrorResponseHandler({ - errorSchema: errorResponseSchema as any, - errorToMessage: (data: ErrorResponse) => { - console.log("Dify API error:", data); - return `Dify API error: ${data.message}`; - }, -}); +/** Content part for multimodal messages (text or tool-result). */ +interface MessageContentPart { + type?: string; + text?: string; + toolName?: string; + output?: unknown; + result?: unknown; + content?: unknown; + providerMetadata?: { difyWorkflowData?: { conversationId?: string } }; + callProviderMetadata?: { difyWorkflowData?: { conversationId?: string } }; +} + +interface Message { + role: string; + content: string | MessageContentPart[]; + [key: string]: unknown; +} export class DifyChatLanguageModel implements LanguageModelV2 { readonly specificationVersion = "v2" as const; @@ -52,6 +89,7 @@ export class DifyChatLanguageModel implements LanguageModelV2 { private readonly generateId: () => string; private readonly chatMessagesEndpoint: string; private readonly config: ModelConfig; + private readonly logger?: Logger; constructor( modelId: DifyChatModelId, @@ -62,6 +100,7 @@ export class DifyChatLanguageModel implements LanguageModelV2 { this.config = config; this.generateId = generateId; this.chatMessagesEndpoint = `${this.config.baseURL}/chat-messages`; + this.logger = settings.logger === false ? undefined : (settings.logger ?? defaultLogger); // Make sure we set a default response mode if (!this.settings.responseMode) { @@ -69,6 +108,16 @@ export class DifyChatLanguageModel implements LanguageModelV2 { } } + private createErrorHandler() { + return createJsonErrorResponseHandler({ + errorSchema: errorResponseSchema as any, + errorToMessage: (data: ErrorResponse) => { + this.logger?.error("Dify API error", data); + return `Dify API error: ${data.message}`; + }, + }); + } + get provider(): string { return this.config.provider; } @@ -77,14 +126,21 @@ export class DifyChatLanguageModel implements LanguageModelV2 { options: LanguageModelV2CallOptions ): Promise>> { const { abortSignal } = options; - const requestBody = this.getRequestBody(options); + const { body: requestBody, toolNames, requestHeaders } = await this.getRequestBody(options); + + // Force blocking mode for doGenerate unless explicitly handled otherwise + const body = { ...requestBody, response_mode: "blocking" }; + + if (this.settings.logMessages) { + this.logger?.info("", { body }); + } const { responseHeaders, value: data } = await postJsonToApi({ url: this.chatMessagesEndpoint, - headers: combineHeaders(this.config.headers(), options.headers), - body: requestBody, + headers: combineHeaders(this.config.headers(), requestHeaders ?? options.headers), + body, abortSignal, - failedResponseHandler: difyFailedResponseHandler, + failedResponseHandler: this.createErrorHandler(), successfulResponseHandler: createJsonResponseHandler( completionResponseSchema as any ), @@ -92,19 +148,34 @@ export class DifyChatLanguageModel implements LanguageModelV2 { }); const typedData = data as CompletionResponse; + if (this.settings.logMessages) { + this.logger?.info("", { answer: typedData.answer, conversation_id: typedData.conversation_id }); + } const content: LanguageModelV2Content[] = []; - // Add text content if available + // Parse tool calls from answer text if (typedData.answer) { - content.push({ - type: "text", - text: typedData.answer, - }); + const { toolCalls, cleanText } = parseToolCalls(typedData.answer, toolNames, this.generateId); + + if (cleanText) { + content.push({ type: "text", text: cleanText }); + } + + for (const tc of toolCalls) { + content.push({ + type: "tool-call", + toolCallId: tc.id, + toolName: tc.name, + input: tc.input, + }); + } } + const hasToolCalls = content.some(c => c.type === "tool-call"); + return { content, - finishReason: "stop" as LanguageModelV2FinishReason, + finishReason: (hasToolCalls ? "tool-calls" : "stop") as LanguageModelV2FinishReason, usage: { inputTokens: typedData.metadata.usage.prompt_tokens, outputTokens: typedData.metadata.usage.completion_tokens, @@ -117,7 +188,7 @@ export class DifyChatLanguageModel implements LanguageModelV2 { messageId: typedData.message_id as JSONValue, }, }, - request: { body: JSON.stringify(requestBody) }, + request: { body: JSON.stringify(body) }, response: { id: typedData.id, timestamp: new Date(), @@ -129,15 +200,22 @@ export class DifyChatLanguageModel implements LanguageModelV2 { async doStream( options: LanguageModelV2CallOptions ): Promise>> { + if (this.settings.logMessages) { + this.logger?.info("", { options }); + } const { abortSignal } = options; - const requestBody = this.getRequestBody(options); + const { body: requestBody, toolNames, requestHeaders } = await this.getRequestBody(options); const body = { ...requestBody, response_mode: "streaming" }; + if (this.settings.logMessages) { + this.logger?.info("", { body }); + } + const { responseHeaders, value: responseStream } = await postJsonToApi({ url: this.chatMessagesEndpoint, - headers: combineHeaders(this.config.headers(), options.headers), + headers: combineHeaders(this.config.headers(), requestHeaders ?? options.headers), body, - failedResponseHandler: difyFailedResponseHandler, + failedResponseHandler: this.createErrorHandler(), successfulResponseHandler: createEventSourceResponseHandler( difyStreamEventSchema as any ), @@ -149,6 +227,48 @@ export class DifyChatLanguageModel implements LanguageModelV2 { let messageId: string | undefined; let taskId: string | undefined; let isActiveText = false; + let isActiveReasoning = false; + const thinkParser = new ThinkTagParser(); + const genId = this.generateId; + const logger = this.logger; + const logMessages = this.settings.logMessages; + + function buildDifyMetadata() { + 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 Object.keys(meta).length ? { providerMetadata: { difyWorkflowData: meta } } : {}; + } + + function emitSegments( + segments: Array<{ type: string; content: string }>, + controller: TransformStreamDefaultController + ) { + for (const seg of segments) { + if (seg.type === "reasoning") { + if (isActiveText) { + controller.enqueue({ type: "text-end", id: "0", ...buildDifyMetadata() }); + isActiveText = false; + } + if (!isActiveReasoning) { + isActiveReasoning = true; + controller.enqueue({ type: "reasoning-start", id: "r0", ...buildDifyMetadata() }); + } + controller.enqueue({ type: "reasoning-delta", id: "r0", delta: seg.content, ...buildDifyMetadata() }); + } else if (seg.content) { + if (isActiveReasoning) { + controller.enqueue({ type: "reasoning-end", id: "r0", ...buildDifyMetadata() }); + isActiveReasoning = false; + } + if (!isActiveText) { + isActiveText = true; + controller.enqueue({ type: "text-start", id: "0", ...buildDifyMetadata() }); + } + controller.enqueue({ type: "text-delta", id: "0", delta: seg.content, ...buildDifyMetadata() }); + } + } + } return { stream: responseStream.pipeThrough( @@ -158,11 +278,16 @@ export class DifyChatLanguageModel implements LanguageModelV2 { >({ transform(chunk, controller) { if (!chunk.success) { + logger?.error("Dify stream parse error", { error: String(chunk.error) }); controller.enqueue({ type: "error", error: chunk.error }); return; } const data = chunk.value; + + if (logMessages) { + logger?.info("", { data }); + } // Store conversation/message IDs for metadata if (data.conversation_id) conversationId = data.conversation_id; @@ -173,41 +298,86 @@ export class DifyChatLanguageModel implements LanguageModelV2 { switch (data.event) { case "workflow_finished": case "message_end": { - // Add block scope to prevent variable leakage + let inputTokens = 0; + let outputTokens = 0; let totalTokens = 0; - // Type guard for data.data - if ( - "data" in data && - data.data && - typeof data.data === "object" && - "total_tokens" in data.data && - typeof data.data.total_tokens === "number" - ) { - totalTokens = data.data.total_tokens; + if (data.event === "workflow_finished") { + const dataUsage = (data as any).data || {}; + inputTokens = 0; + outputTokens = dataUsage.total_tokens ?? 0; + totalTokens = dataUsage.total_tokens ?? 0; + } else { + const usageData = (data as any).metadata?.usage || {}; + inputTokens = usageData.prompt_tokens ?? 0; + outputTokens = usageData.completion_tokens ?? 0; + totalTokens = usageData.total_tokens ?? 0; + } + + // Use accumulated answer for tool parsing + // Remove tags for tool parsing + const textWithoutThink = thinkParser.getTextWithoutThink(); + + const { toolCalls } = parseToolCalls(textWithoutThink, toolNames, genId); + const hasToolCalls = toolCalls.length > 0; + + if (logMessages) { + logger?.info("", { toolCalls }); + } + + // Flush remaining content from parser + emitSegments(thinkParser.flush(), controller); + + if (isActiveReasoning) { + controller.enqueue({ type: "reasoning-end", id: "r0", ...buildDifyMetadata() }); + isActiveReasoning = false; } + if (isActiveText) { + controller.enqueue({ type: "text-end", id: "0", ...buildDifyMetadata() }); + isActiveText = false; + } else if (!hasToolCalls) { + controller.enqueue({ type: "text-start", id: "0", ...buildDifyMetadata() }); + controller.enqueue({ type: "text-end", id: "0", ...buildDifyMetadata() }); + } + + // Don't emit cleanText here since it was already emitted via stream + + for (const tc of toolCalls) { controller.enqueue({ - type: "text-end", - id: "0", + type: "tool-input-start", + id: tc.id, + toolName: tc.name, + ...buildDifyMetadata(), + }); + controller.enqueue({ + type: "tool-input-delta", + id: tc.id, + delta: tc.input, + ...buildDifyMetadata(), + }); + controller.enqueue({ + type: "tool-input-end", + id: tc.id, + ...buildDifyMetadata(), + }); + controller.enqueue({ + type: "tool-call", + toolCallId: tc.id, + toolName: tc.name, + input: tc.input, + ...buildDifyMetadata(), }); - isActiveText = false; } controller.enqueue({ type: "finish", - finishReason: "stop", - providerMetadata: { - difyWorkflowData: { - conversationId: conversationId as JSONValue, - messageId: messageId as JSONValue, - taskId: taskId as JSONValue, - }, - }, + finishReason: hasToolCalls ? "tool-calls" : "stop", + ...buildDifyMetadata(), usage: { - inputTokens: 0, - outputTokens: totalTokens, - totalTokens: totalTokens, + inputTokens, + outputTokens, + totalTokens, }, }); break; @@ -215,56 +385,50 @@ export class DifyChatLanguageModel implements LanguageModelV2 { case "message": case "agent_message": { - // Type guard for answer property - if ("answer" in data && typeof data.answer === "string") { - if (!isActiveText) { - isActiveText = true; - controller.enqueue({ - type: "text-start", - id: "0", - ...(conversationId || messageId || taskId - ? { - providerMetadata: { - difyWorkflowData: { - ...(conversationId ? { conversationId: conversationId as JSONValue } : {}), - ...(messageId ? { messageId: messageId as JSONValue } : {}), - ...(taskId ? { taskId: taskId as JSONValue } : {}), - }, - }, - } - : {}), - }); + if ("answer" in data && typeof data.answer === "string" && data.answer) { + emitSegments(thinkParser.feed(data.answer), controller); + + if ("id" in data && typeof data.id === "string") { + controller.enqueue({ type: "response-metadata", id: data.id }); } + } + break; + } - controller.enqueue({ - type: "text-delta", - id: "0", - delta: data.answer, - ...(conversationId || messageId || taskId - ? { - providerMetadata: { - difyWorkflowData: { - ...(conversationId ? { conversationId: conversationId as JSONValue } : {}), - ...(messageId ? { messageId: messageId as JSONValue } : {}), - ...(taskId ? { taskId: taskId as JSONValue } : {}), - }, - }, - } - : {}), - }); + // agent_thought events are summaries of agent_message chunks. + // Dify sends: agent_thought(empty) → agent_message × N → agent_thought(full) → message_end + // The agent_message events already contain all content (including tags), + // so agent_thought is redundant and must be ignored to avoid duplicate output. + case "agent_thought": + break; - // Type guard for id property - if ("id" in data && typeof data.id === "string") { - controller.enqueue({ - type: "response-metadata", - id: data.id, - }); + case "message_replace": { + if ("answer" in data && typeof data.answer === "string" && data.answer) { + thinkParser.reset(data.answer); + + if (isActiveText) { + controller.enqueue({ type: "text-end", id: "0", ...buildDifyMetadata() }); + isActiveText = false; } + controller.enqueue({ type: "text-start", id: "0", ...buildDifyMetadata() }); + controller.enqueue({ type: "text-delta", id: "0", delta: data.answer, ...buildDifyMetadata() }); + isActiveText = true; } break; } - // Ignore other event types + case "message_file": + case "ping": + break; + + case "error": { + const err = data as { message?: string; msg?: string; code?: string; status?: string }; + const msg = err.message ?? err.msg ?? "Unknown Dify error"; + const code = err.code ?? err.status ?? ""; + logger?.error("Dify stream error", { code, message: msg }); + controller.enqueue({ type: "error", error: new Error(`Dify error${code ? ` (${code})` : ""}: ${msg}`) }); + break; + } } }, }) @@ -277,11 +441,13 @@ export class DifyChatLanguageModel implements LanguageModelV2 { /** * Get the request body for the Dify API */ - private getRequestBody(options: LanguageModelV2CallOptions) { + private async getRequestBody(options: LanguageModelV2CallOptions) { // In AI SDK v5 LanguageModelV2, messages are in options.prompt - const messages = options.prompt || (options as any).messages; + const opts = options as LanguageModelV2CallOptions & { messages?: Message[] }; + const messages = (options.prompt ?? opts.messages) as Message[] | undefined; if (!messages || !messages.length) { + this.logger?.error("No messages provided", { prompt: options.prompt }); throw new APICallError({ message: "No messages provided", url: this.chatMessagesEndpoint, @@ -289,65 +455,162 @@ export class DifyChatLanguageModel implements LanguageModelV2 { }); } - const latestMessage = messages[messages.length - 1]; + const conversationId = this.findConversationId(options.headers, messages); + const query = this.buildQuery(messages, conversationId); - if (latestMessage.role !== "user") { + if (!query) { + this.logger?.error("No user message found", { messageCount: messages.length, roles: messages.map(m => m.role) }); throw new APICallError({ - message: "The last message must be a user message", + message: "No user message found", url: this.chatMessagesEndpoint, - requestBodyValues: { latestMessageRole: latestMessage.role }, + requestBodyValues: { messageCount: messages.length }, }); } - // Handle file/image attachments - const hasAttachments = - Array.isArray(latestMessage.content) && - latestMessage.content.some((part) => { - return typeof part !== "string" && part !== null && typeof part === "object" && "type" in part && part.type === "file"; - }); + const userId = options.headers?.["user-id"] ?? "you_should_pass_user-id"; - if (hasAttachments) { - throw new APICallError({ - message: "Dify provider does not currently support file attachments", - url: this.chatMessagesEndpoint, - requestBodyValues: { hasAttachments: true }, - }); + // Clean headers by removing handled IDs (do not mutate options.headers) + const { "chat-id": _h1, "user-id": _h2, ...cleanHeaders } = options.headers || {}; + + // Extract file/image attachments from messages + // Note: extractFileAttachments is an external async utility, kept as is. + const files = await extractFileAttachments(messages, userId, this.config.baseURL, this.config.headers, this.logger); + + // System Prompt + const systemPrompt = this.extractSystemPrompt(messages); + + // Tool config + const injectTools = this.settings.injectToolsPrompt !== false; + // Always parse tool names to support tool execution even if prompt injection is disabled + const { prompt: generatedPrompt, toolNames } = formatToolsPrompt(options.tools); + const toolsPrompt = injectTools ? generatedPrompt : ""; + + const isNewConversation = !conversationId; + const parts = [ + isNewConversation ? systemPrompt : "", + query, + isNewConversation ? toolsPrompt : "", + ].filter(Boolean); + const finalQuery = parts.join("\n\n"); + const inputs = this.settings.inputs || {}; + + const requestBody = { + inputs, + query: finalQuery, + response_mode: this.settings.responseMode, + conversation_id: conversationId, + user: userId, + ...(files.length ? { files } : {}), + }; + + return { + body: requestBody, + toolNames, + requestHeaders: cleanHeaders, + }; + } + + private findConversationId(headers: Record | undefined, messages: Message[]): string | undefined { + // Priority: Header > Last Message Provider Metadata + if (headers?.["chat-id"]) { + return headers["chat-id"]; } - // Extract the query from the latest user message - let query = ""; - if (typeof latestMessage.content === "string") { - query = latestMessage.content; - } else if (Array.isArray(latestMessage.content)) { - // Handle AI SDK v4 format with text objects in content array - query = latestMessage.content - .map((part) => { - if (typeof part === "string") { - return part; - } else if (typeof part === "object" && part !== null && "type" in part && part.type === "text" && "text" in part) { - return part.text; + // Iterate backwards to find the most recent conversation context + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (Array.isArray(msg.content)) { + for (const part of msg.content) { + const meta = part?.providerMetadata?.difyWorkflowData || part?.callProviderMetadata?.difyWorkflowData; + if (meta?.conversationId) { + return meta.conversationId as string; } + } + } + } + return undefined; + } + + private buildQuery(messages: Message[], conversationId: string | undefined): string { + // Dify expects a single query string. + // If conversation exists: only the new user input (+ tool results) + // If new conversation: combine all history as context + + let startIndex = 0; + + if (conversationId) { + // Existing Conversation: Extract latest interaction + // We walk backwards to find the start of the new interaction. + startIndex = messages.length; + + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (msg.role === "tool") { + startIndex = i; + continue; + } + if (msg.role === "user") { + startIndex = i; + break; // Found the primary user query + } + // Stop at assistant, system or other roles + break; + } + } + + return messages + .slice(startIndex) + .filter((msg) => msg.role === "user" || msg.role === "tool" || (!conversationId && msg.role === "assistant")) + .map((msg) => { + if (msg.role === "tool") return this.extractToolResult(msg); + return this.extractTextFromMessage(msg); + }) + .filter(Boolean) + .join("\n\n"); + } + + private extractSystemPrompt(messages: Message[]): string { + const systemMessage = messages.find(msg => msg.role === "system"); + const rawSystemPrompt = systemMessage ? this.extractTextFromMessage(systemMessage) : ""; + + return formatSystemPrompt(rawSystemPrompt); + } + + private extractTextFromMessage(msg: Message): string { + if (typeof msg.content === "string") { + return msg.content; + } + if (Array.isArray(msg.content)) { + return msg.content + .map((part: string | MessageContentPart) => { + if (typeof part === "string") return part; + if (part?.type === "text" && part.text) return part.text; return ""; }) .filter(Boolean) .join(" "); } + return ""; + } - const conversationId = options.headers?.["chat-id"]; - const userId = options.headers?.["user-id"] ?? "you_should_pass_user-id"; - const { - "chat-id": _, - "user-id": __, - ...cleanHeaders - } = options.headers || {}; - options.headers = cleanHeaders; - - return { - inputs: this.settings.inputs || {}, - query, - response_mode: this.settings.responseMode, - conversation_id: conversationId, - user: userId, - }; + private extractToolResult(msg: Message): string { + const content = msg.content; + if (!Array.isArray(content)) return ""; + + const results: string[] = []; + for (const part of content) { + if (part?.type === "tool-result") { + const name = part.toolName || "tool"; + const raw = part.output ?? part.result ?? part.content ?? part.text; + const output = + raw !== undefined && typeof raw === "object" && raw !== null && "value" in raw + ? (raw as { value: unknown }).value + : raw; + const result = output === undefined ? "[completed]" : + (typeof output === "string" ? output : JSON.stringify(output)); + results.push(`[${name} result]: ${result}`); + } + } + return results.join("\n"); } } diff --git a/src/dify-chat-schema.ts b/src/dify-chat-schema.ts index 432e8cc..4358779 100644 --- a/src/dify-chat-schema.ts +++ b/src/dify-chat-schema.ts @@ -48,12 +48,18 @@ export const workflowStartedSchema = difyStreamEventBase.extend({ export const workflowFinishedSchema = difyStreamEventBase.extend({ event: z.literal("workflow_finished"), workflow_run_id: z.string(), + task_id: z.string().optional(), data: z .object({ id: z.string(), workflow_id: z.string(), + outputs: z.record(z.unknown()).optional(), + status: z.string().optional(), + elapsed_time: z.number().optional(), total_tokens: z.number().optional(), - created_at: z.number(), + total_steps: z.string().optional(), + created_at: z.number().optional(), + finished_at: z.number().optional(), }) .passthrough(), }); @@ -91,19 +97,29 @@ export const messageSchema = difyStreamEventBase.extend({ export const messageEndSchema = difyStreamEventBase.extend({ event: z.literal("message_end"), - id: z.string(), + id: z.string().optional(), metadata: z .object({ usage: z .object({ - prompt_tokens: z.number(), - completion_tokens: z.number(), - total_tokens: z.number(), + prompt_tokens: z.number().optional(), + completion_tokens: z.number().optional(), + total_tokens: z.number().optional(), + prompt_unit_price: z.string().optional(), + prompt_price_unit: z.string().optional(), + prompt_price: z.string().optional(), + completion_unit_price: z.string().optional(), + completion_price_unit: z.string().optional(), + completion_price: z.string().optional(), + total_price: z.string().optional(), + currency: z.string().optional(), + latency: z.number().optional(), }) .passthrough(), + retriever_resources: z.array(z.unknown()).optional(), }) .passthrough(), - files: z.array(z.unknown()).optional(), + files: z.nullable(z.array(z.unknown())).optional(), }); export const ttsMessageSchema = difyStreamEventBase.extend({ @@ -192,14 +208,40 @@ export const agentMessageSchema = difyStreamEventBase.extend({ export const agentThoughtSchema = difyStreamEventBase.extend({ event: z.literal("agent_thought"), + id: z.string(), + position: z.number(), thought: z.string(), observation: z.string(), tool: z.string(), - tool_labels: z.record(z.string(), z.string()), + tool_labels: z.record(z.string(), z.string()).optional(), tool_input: z.string(), message_files: z.array(z.unknown()), }); +export const messageReplaceSchema = difyStreamEventBase.extend({ + event: z.literal("message_replace"), + answer: z.string(), +}); + +export const messageFileSchema = difyStreamEventBase.extend({ + event: z.literal("message_file"), + id: z.string(), + type: z.string(), + belongs_to: z.string(), + url: z.string(), +}); + +export const pingSchema = difyStreamEventBase.extend({ + event: z.literal("ping"), +}); + +export const errorSchema = difyStreamEventBase.extend({ + event: z.literal("error"), + status: z.number().optional(), + code: z.string().optional(), + message: z.string().optional(), +}); + // Combine all schemas with discriminatedUnion export const difyStreamEventSchema = z .discriminatedUnion("event", [ @@ -213,6 +255,10 @@ export const difyStreamEventSchema = z ttsMessageEndSchema, agentThoughtSchema, agentMessageSchema, + messageReplaceSchema, + messageFileSchema, + pingSchema, + errorSchema, ]) .or(difyStreamEventBase); // Fallback for any other event types diff --git a/src/dify-chat-settings.ts b/src/dify-chat-settings.ts index fc90d64..e70fe32 100644 --- a/src/dify-chat-settings.ts +++ b/src/dify-chat-settings.ts @@ -1,9 +1,20 @@ // Define model IDs for Dify export type DifyChatModelId = string; -/** - * Settings for the Dify chat API. - */ +export interface Logger { + debug: (message: string, extra?: Record) => void; + info: (message: string, extra?: Record) => void; + warn: (message: string, extra?: Record) => void; + error: (message: string, extra?: Record) => void; +} + +export const defaultLogger: Logger = { + debug: (message, extra) => console.debug(`[dify-ai-provider] ${message}`, extra ?? ""), + info: (message, extra) => console.log(`[dify-ai-provider] ${message}`, extra ?? ""), + warn: (message, extra) => console.warn(`[dify-ai-provider] ${message}`, extra ?? ""), + error: (message, extra) => console.error(`[dify-ai-provider] ${message}`, extra ?? ""), +}; + export interface DifyChatSettings { /** * Additional inputs to send with the request. @@ -20,4 +31,22 @@ export interface DifyChatSettings { * API key. */ apiKey?: string; + + /** + * Logger instance for debugging. If not provided, uses default console logger at info level. + * Set to `false` to disable logging. + */ + logger?: Logger | false; + + /** + * Inject tool definitions into the query prompt for LLMs that don't natively support tool calling. + * Defaults to `true` (appended to query). Set to `false` to disable. + */ + injectToolsPrompt?: boolean; + + /** + * Log request and response messages. Defaults to `false`. + */ + logMessages?: boolean; + } diff --git a/src/dify-provider.ts b/src/dify-provider.ts index 4a8e163..a6ae4b1 100644 --- a/src/dify-provider.ts +++ b/src/dify-provider.ts @@ -1,5 +1,5 @@ -import { FetchFunction, generateId, loadApiKey } from "@ai-sdk/provider-utils"; -import { DifyChatSettings, DifyChatModelId } from "./dify-chat-settings"; +import { FetchFunction, loadApiKey } from "@ai-sdk/provider-utils"; +import { DifyChatSettings, DifyChatModelId, Logger } from "./dify-chat-settings"; import { DifyChatLanguageModel } from "./dify-chat-language-model"; // model factory function with additional methods and properties @@ -32,6 +32,11 @@ export interface DifyProviderSettings { * or to provide a custom fetch implementation for e.g. testing. */ fetch?: FetchFunction; + + /** + * Logger instance for debugging. Set to `false` to disable logging. + */ + logger?: Logger | false; } export function createDifyProvider( diff --git a/src/stream-parser.ts b/src/stream-parser.ts new file mode 100644 index 0000000..754f04f --- /dev/null +++ b/src/stream-parser.ts @@ -0,0 +1,394 @@ +import type { LanguageModelV2CallOptions } from "@ai-sdk/provider"; +import type { Logger } from "./dify-chat-settings"; + +// ── Think Tag Parser ── + +export type ThinkSegment = { type: "reasoning" | "text"; content: string }; + +export class ThinkTagParser { + private state: "text" | "maybe-open" | "reasoning" | "maybe-close" = "text"; + private buffer = ""; + private accumulated = ""; + private hasStarted = false; + private closed = false; // Once is seen, treat all subsequent as plain text + private textAccumulated = ""; + private inString = false; + private escape = false; + + feed(chunk: string): ThinkSegment[] { + const raw: ThinkSegment[] = []; + this.accumulated += chunk; + + for (let i = 0; i < chunk.length; i++) { + const ch = chunk[i]; + + if (!this.hasStarted) { + if (ch.trim() === "") continue; + this.hasStarted = true; + } + + if (this.inString) { + if (this.escape) { + this.escape = false; + raw.push({ type: "text", content: ch }); + } else if (ch === "\\") { + this.escape = true; + raw.push({ type: "text", content: ch }); + } else if (ch === '"') { + this.inString = false; + raw.push({ type: "text", content: ch }); + } else { + raw.push({ type: "text", content: ch }); + } + continue; + } + + if (ch === '"') { + this.inString = true; + raw.push({ type: "text", content: ch }); + continue; + } + + switch (this.state) { + case "text": + if (ch === "<" && !this.closed) { + this.state = "maybe-open"; + this.buffer = "<"; + } else { + raw.push({ type: "text", content: ch }); + } + break; + + case "maybe-open": + this.buffer += ch; + if ("".startsWith(this.buffer)) { + if (this.buffer === "") { + this.state = "reasoning"; + this.buffer = ""; + } + } else { + raw.push({ type: "text", content: this.buffer }); + this.buffer = ""; + this.state = "text"; + } + break; + + case "reasoning": + if (ch === "<") { + this.state = "maybe-close"; + this.buffer = "<"; + } else { + raw.push({ type: "reasoning", content: ch }); + } + break; + + case "maybe-close": + this.buffer += ch; + if ("".startsWith(this.buffer)) { + if (this.buffer === "") { + this.state = "text"; + this.buffer = ""; + this.closed = true; + } + } else { + raw.push({ type: "reasoning", content: this.buffer }); + this.buffer = ""; + this.state = "reasoning"; + } + break; + } + } + + // Merge consecutive same-type segments + const merged: ThinkSegment[] = []; + for (const seg of raw) { + const last = merged[merged.length - 1]; + if (last && last.type === seg.type) { + last.content += seg.content; + } else { + merged.push({ ...seg }); + } + } + + // 实时累积解析后的纯文本 + for (const seg of merged) { + if (seg.type === "text") { + this.textAccumulated += seg.content; + } + } + + return merged; + } + + flush(): ThinkSegment[] { + if (!this.buffer) return []; + const type = this.state === "maybe-close" ? "reasoning" : "text"; + const seg = [{ type, content: this.buffer } as ThinkSegment]; + if (type === "text") { + this.textAccumulated += this.buffer; + } + this.buffer = ""; + return seg; + } + + reset(text: string): void { + this.accumulated = ""; + this.buffer = ""; + this.state = "text"; + this.closed = false; + this.hasStarted = false; + this.inString = false; + this.escape = false; + this.feed(text); + } + + getTextWithoutThink(): string { + return this.textAccumulated.trim(); + } +} + +// ── Tool Call Parser ── + +export interface ParsedToolCall { + id: string; + name: string; + input: string; +} + +export function parseToolCalls(text: string, toolNames: string[], genId: () => string): { toolCalls: ParsedToolCall[]; cleanText: string } { + const toolCalls: ParsedToolCall[] = []; + let cleanText = text; + + if (!toolNames.length) return { toolCalls, cleanText }; + + const lowerToolNames = toolNames.map((n) => n.toLowerCase()); + const matchedBlocks: string[] = []; + + let depth = 0; + let inString = false; + let escape = false; + let startIdx = -1; + + for (let i = 0; i < text.length; i++) { + const ch = text[i]; + if (escape) { escape = false; continue; } + if (ch === "\\") { escape = true; continue; } + if (ch === '"') { inString = !inString; continue; } + if (inString) continue; + + if (ch === "{") { + if (depth === 0) startIdx = i; + depth++; + } else if (ch === "}") { + depth--; + if (depth < 0) depth = 0; + if (depth === 0 && startIdx !== -1) { + const jsonCandidate = text.substring(startIdx, i + 1); + + try { + const parsed = JSON.parse(jsonCandidate); + if (parsed && typeof parsed === "object" && "name" in parsed && "arguments" in parsed) { + const toolName = String(parsed.name).toLowerCase(); + const originalIdx = lowerToolNames.indexOf(toolName); + if (originalIdx !== -1) { + toolCalls.push({ + id: genId(), + name: toolNames[originalIdx], + input: JSON.stringify(parsed.arguments) + }); + matchedBlocks.push(jsonCandidate); + } + } + } catch (e) { + // Not valid JSON, ignore + } + startIdx = -1; + } + } + } + + for (const block of matchedBlocks) { + cleanText = cleanText.replace(block, ""); + } + + return { toolCalls, cleanText: cleanText.trim() }; +} + +// ── Tool Prompt Formatting ── + +interface JsonSchema { + type?: string; + properties?: Record; + items?: JsonSchema; + required?: string[]; +} + +function placeholder(schema: JsonSchema | undefined): string { + if (!schema) return '"..."'; + const t = schema.type; + if (t === "array") { + const item = schema.items ? placeholder(schema.items) : '"..."'; + return `[${item}]`; + } + if (t === "object") { + const props = schema.properties; + if (!props) return "{}"; + const entries = Object.entries(props) + .slice(0, 3) + .map(([k, v]) => `"${k}": ${placeholder(v)}`); + return `{${entries.join(", ")}}`; + } + if (t === "number" || t === "integer") return "0"; + if (t === "boolean") return "true"; + return '"..."'; +} + +const TOOL_SHORT_DESCRIPTIONS: Record = { + question: "Ask the user a question", + bash: "Execute a bash command", + read: "Read a file", + glob: "Find files by pattern", + grep: "Search file contents with regex", + edit: "Edit a file with string replacement", + write: "Write/create a file", + task: "Launch a sub-agent for complex tasks", + webfetch: "Fetch content from a URL", + todowrite: "Create/manage a task list", + skill: "Load a specialized skill", +}; + +function getShortDescription(name: string, description?: string): string { + if (TOOL_SHORT_DESCRIPTIONS[name]) return TOOL_SHORT_DESCRIPTIONS[name]; + if (!description) return name; + return description.replace(/[\r\n]+/g, " ").split(/[.!?。]/)[0].trim().substring(0, 60); +} + +export function formatToolsPrompt( + tools: LanguageModelV2CallOptions["tools"], +): { prompt: string; toolNames: string[] } { + if (!tools?.length) { + return { prompt: "", toolNames: [] }; + } + + const defs = tools + .filter((t): t is { type: "function"; name: string; description?: string; inputSchema: unknown } => t.type === "function") + .map((t) => ({ name: t.name, description: t.description, parameters: t.inputSchema })); + + if (!defs.length) return { prompt: "", toolNames: [] }; + + const toolNames = defs.map((d) => d.name); + + const toolDescriptions = defs.map((t) => { + const params = t.parameters as JsonSchema | undefined; + const paramList = params?.properties + ? Object.entries(params.properties) + .map(([k, v]) => `"${k}": ${placeholder(v)}`) + .join(", ") + : ""; + const shortDesc = getShortDescription(t.name, t.description); + return `- ${t.name}: ${shortDesc}\n {"name": "${t.name}", "arguments": {${paramList}}}`; + }); + + const prompt = `\n\n# Tools\nCall tools with JSON format:\n\n${toolDescriptions.join("\n")}`; + + return { prompt, toolNames }; +} + +// ── File Upload ── + +export interface DifyFileAttachment { + type: string; + transfer_method: string; + upload_file_id: string; +} + +export async function uploadFileToDify( + data: string | Uint8Array | Buffer, + mimeType: string, + filename: string, + user: string, + baseURL: string, + headers: () => Record, + logger?: Logger, +): Promise { + try { + let blob: Blob; + if (typeof data === "string") { + const buffer = Buffer.from(data, "base64"); + blob = new Blob([buffer], { type: mimeType }); + } else { + blob = new Blob([data], { type: mimeType }); + } + const form = new FormData(); + form.append("file", blob, filename); + form.append("user", user); + + const { "content-type": _, "Content-Type": __, ...authHeaders } = headers(); + const res = await fetch(`${baseURL}/files/upload`, { + method: "POST", + headers: { ...authHeaders }, + body: form, + }); + + if (!res.ok) { + const errBody = await res.text().catch(() => ""); + logger?.error("Dify file upload failed", { status: res.status, statusText: res.statusText, body: errBody }); + return undefined; + } + + const json = await res.json() as { id?: string }; + return json.id; + } catch (e) { + logger?.error("Dify file upload error", { error: String(e) }); + return undefined; + } +} + +export async function extractFileAttachments( + messages: any[], + userId: string, + baseURL: string, + headers: () => Record, + logger?: Logger, +): Promise { + const uploadPromises: Promise[] = []; + + for (const msg of messages) { + if (!Array.isArray(msg.content)) continue; + for (const part of msg.content as any[]) { + if (!part || typeof part !== "object" || part.type !== "file") continue; + const mimeType: string = part.mediaType || part.mimeType || ""; + if (!mimeType) continue; + + let uploadData: string | Uint8Array | undefined; + const filename = part.filename || `file.${mimeType.split("/")[1] || "bin"}`; + + if (typeof part.data === "string") { + const m = part.data.match(/^data:[^;]+;base64,(.+)$/); + uploadData = m ? m[1] : part.data; + } else if (part.data instanceof Uint8Array) { + uploadData = part.data; + } + + if (uploadData) { + const task = uploadFileToDify(uploadData, mimeType, filename, userId, baseURL, headers, logger) + .then((uploadId): DifyFileAttachment | undefined => { + if (uploadId) { + return { + type: mimeType.startsWith("image/") ? "image" : "document", + transfer_method: "local_file", + upload_file_id: uploadId + }; + } else { + logger?.warn("Dify file upload skipped, attachment will be missing", { filename, mimeType }); + return undefined; + } + }); + uploadPromises.push(task); + } + } + } + + const results = await Promise.all(uploadPromises); + return results.filter((item): item is DifyFileAttachment => item !== undefined); +} diff --git a/test/debug-stream.ts b/test/debug-stream.ts new file mode 100644 index 0000000..eb8172f --- /dev/null +++ b/test/debug-stream.ts @@ -0,0 +1,94 @@ +/** + * Debug script: calls provider doStream() or doGenerate() and prints every part. + * + * Usage: + * DIFY_BASE_URL=https://api.dify.ai/v1 DIFY_API_KEY=app-xxx npx tsx test/debug-stream.ts [options] + * + * Options: + * --user User ID (default: debug-user) + * --query Query text (default: "Introduce yourself") + * --chat-id Conversation ID for multi-turn + * --blocking Use doGenerate (blocking mode) instead of streaming + */ +import { DifyChatLanguageModel } from "../src/dify-chat-language-model"; +import { loadApiKey } from "@ai-sdk/provider-utils"; + +function parseArgs() { + const args = process.argv.slice(2); + let appId = ""; + let user = "debug-user"; + let query = "Introduce yourself"; + let chatId = ""; + let blocking = false; + + for (let i = 0; i < args.length; i++) { + if (args[i] === "--blocking") blocking = true; + else if (args[i] === "--user" && i + 1 < args.length) user = args[++i]; + else if (args[i] === "--query" && i + 1 < args.length) query = args[++i]; + else if (args[i] === "--chat-id" && i + 1 < args.length) chatId = args[++i]; + else if (!args[i].startsWith("--")) appId = args[i]; + } + return { appId: appId || "default", user, query, chatId, blocking }; +} + +const { appId, user: userId, query, chatId, blocking } = parseArgs(); + +const logger = { + debug: (msg: string, extra?: any) => console.debug(`[DEBUG] ${msg}`, extra ?? ""), + info: (msg: string, extra?: any) => console.log(`[INFO] ${msg}`, extra ?? ""), + warn: (msg: string, extra?: any) => console.warn(`[WARN] ${msg}`, extra ?? ""), + error: (msg: string, extra?: any) => console.error(`[ERROR] ${msg}`, extra ?? ""), +}; + +const baseURL = process.env.DIFY_BASE_URL || "https://api.dify.ai/v1"; +const logMessages = true; + +const model = new DifyChatLanguageModel( + appId, + {responseMode: blocking ? "blocking" : "streaming", logger, logMessages}, + { + provider: "dify.chat", + baseURL, + headers: () => ({ + Authorization: `Bearer ${loadApiKey({ apiKey: undefined, environmentVariableName: "DIFY_API_KEY", description: "Dify API Key" })}`, + "Content-Type": "application/json", + }), + } +); + +const callOptions: any = { + prompt: [{ role: "user", content: [{ type: "text", text: query }] }], + headers: { + "user-id": userId, + ...(chatId ? { "chat-id": chatId } : {}), + }, + inputFormat: "messages", + mode: { type: "regular" }, +}; + +(async () => { + try { + if (blocking) { + console.log(`=== doGenerate (blocking) === query: "${query}"\n`); + const result = await model.doGenerate(callOptions); + console.log("content:", JSON.stringify(result.content, null, 2)); + console.log("finishReason:", result.finishReason); + console.log("usage:", result.usage); + console.log("providerMetadata:", JSON.stringify(result.providerMetadata, null, 2)); + } else { + console.log(`=== doStream === query: "${query}"\n`); + const { stream } = await model.doStream(callOptions); + const reader = stream.getReader(); + let idx = 0; + while (true) { + const { value, done } = await reader.read(); + if (done) break; + console.log(`[${idx++}]`, JSON.stringify(value)); + } + } + console.log("\n=== Done ==="); + } catch (e) { + console.error("Failed:", e); + process.exit(1); + } +})(); diff --git a/test/stream-parser.test.ts b/test/stream-parser.test.ts new file mode 100644 index 0000000..b065e44 --- /dev/null +++ b/test/stream-parser.test.ts @@ -0,0 +1,204 @@ +import { describe, it, expect } from "vitest"; +import type { LanguageModelV2CallOptions } from "@ai-sdk/provider"; +import { ThinkTagParser, parseToolCalls, formatToolsPrompt } from "../src/stream-parser"; + +// Mock genId for parseToolCalls +const mockGenId = () => "test-id"; + +// Mock tools for formatToolsPrompt +const mockTools: LanguageModelV2CallOptions["tools"] = [ + { + type: "function", + name: "test-tool", + description: "A test tool", + inputSchema: { + type: "object", + properties: { + param1: { type: "string" }, + param2: { type: "number" }, + nested: { + type: "object", + properties: { inner: { type: "boolean" } }, + }, + array: { + type: "array", + items: { type: "string" }, + }, + }, + }, + }, +]; + +describe("ThinkTagParser", () => { + it("handles plain text", () => { + const parser = new ThinkTagParser(); + const segments = parser.feed("Hello world"); + expect(segments).toEqual([{ type: "text", content: "Hello world" }]); + }); + + it("starts with whitespace ignored", () => { + const parser = new ThinkTagParser(); + let segments = parser.feed(" "); + expect(segments).toEqual([]); + segments = parser.feed("Hello"); + expect(segments).toEqual([{ type: "text", content: "Hello" }]); + }); + + it("consumes tag and starts reasoning", () => { + const parser = new ThinkTagParser(); + expect(parser.feed("")).toEqual([]); + expect(parser.feed("reason")).toEqual([{ type: "reasoning", content: "reason" }]); + }); + + it("parses complete think block", () => { + const parser = new ThinkTagParser(); + parser.feed("Text before reasoning here text after"); + expect(parser.getTextWithoutThink()).toBe("Text before text after"); + }); + + it("merges consecutive same-type segments within feed", () => { + const parser = new ThinkTagParser(); + const segments = parser.feed("textreasonmore"); + expect(segments).toEqual([ + { type: "text", content: "text" }, + { type: "reasoning", content: "reason" }, + { type: "text", content: "more" }, + ]); + }); + + it("treats second think as text after first close", () => { + const parser = new ThinkTagParser(); + parser.feed("first"); + const segments = parser.feed("second"); + expect(segments[0]).toHaveProperty("type", "text"); + expect(segments[0].content).toContain(""); + }); + + it("flush outputs remaining buffer", () => { + const parser = new ThinkTagParser(); + parser.feed("text { + const parser = new ThinkTagParser(); + parser.feed("foo"); + parser.reset("new text"); + expect(parser.getTextWithoutThink()).toBe("new text"); + }); + + it("getTextWithoutThink removes only first block", () => { + const parser = new ThinkTagParser(); + parser.reset("first middle second"); + expect(parser.getTextWithoutThink()).toBe("middle second"); + }); + + it("getTextWithoutThink handles unclosed think", () => { + const parser = new ThinkTagParser(); + parser.reset("unclosed"); + expect(parser.getTextWithoutThink()).toBe(""); + }); + + it("does not treat inside double-quoted string (e.g. tool args) as think tag", () => { + const parser = new ThinkTagParser(); + const jsonLike = "{\"name\":\"my_tool\",\"arguments\":{\"prompt\":\"Use tag in your reply\"}}"; + const segments = parser.feed(jsonLike); + expect(segments.every((s) => s.type === "text")).toBe(true); + expect(parser.getTextWithoutThink()).toBe(jsonLike); + }); + + it("does not treat inside JSON string as think close", () => { + const parser = new ThinkTagParser(); + parser.feed("real think text {\"note\":\" here\"} more"); + expect(parser.getTextWithoutThink()).toBe("text {\"note\":\" here\"} more"); + }); +}); + +describe("parseToolCalls", () => { + const toolNames = ["test-tool"]; + + it("no tools", () => { + const { toolCalls, cleanText } = parseToolCalls("hello", [], mockGenId); + expect(toolCalls).toEqual([]); + expect(cleanText).toBe("hello"); + }); + + it("parses valid tool call", () => { + const text = "Response {\"name\": \"test-tool\", \"arguments\": {\"param\": \"value\"}}"; + const { toolCalls, cleanText } = parseToolCalls(text, toolNames, mockGenId); + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("test-tool"); + expect(toolCalls[0].input).toBe('{"param":"value"}'); + expect(cleanText).toBe("Response"); + }); + + it("ignores JSON inside strings", () => { + const text = "\"{\"name\": \"test-tool\"}\" normal text"; + const { toolCalls } = parseToolCalls(text, toolNames, mockGenId); + expect(toolCalls).toEqual([]); + }); + + it("ignores invalid JSON", () => { + const text = "{ invalid json "; + const { toolCalls } = parseToolCalls(text, toolNames, mockGenId); + expect(toolCalls).toEqual([]); + }); + + it("handles escaped quotes", () => { + const text = "{\"name\": \"test-tool\", \"arguments\": {\"param\": \"val\\\"ue\"}}"; + const { toolCalls } = parseToolCalls(text, toolNames, mockGenId); + expect(toolCalls).toHaveLength(1); + }); + + it("multiple tool calls", () => { + const toolCall1 = JSON.stringify({name: "test-tool", arguments: {a:1}}); + const toolCall2 = JSON.stringify({name: "test-tool", arguments: {b:2}}); + const text = `Call1 ${toolCall1} Call2 ${toolCall2}`; + const { toolCalls } = parseToolCalls(text, toolNames, mockGenId); + expect(toolCalls).toHaveLength(2); + }); + + it("case insensitive tool name", () => { + const text = JSON.stringify({name: "TEST-TOOL", arguments: {p: "v"}}); + const { toolCalls } = parseToolCalls(text, ["test-tool"], mockGenId); + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].name).toBe("test-tool"); + }); + + it("parses tool call when arguments value contains (simulating ThinkTagParser output)", () => { + const toolJson = "{\"name\": \"test-tool\", \"arguments\": {\"prompt\": \"Use tag\"}}"; + const parser = new ThinkTagParser(); + parser.feed(toolJson); + const textWithoutThink = parser.getTextWithoutThink(); + const { toolCalls } = parseToolCalls(textWithoutThink, toolNames, mockGenId); + expect(toolCalls).toHaveLength(1); + expect(toolCalls[0].input).toBe('{"prompt":"Use tag"}'); + }); +}); + +describe("formatToolsPrompt", () => { + it("no tools", () => { + const { prompt, toolNames } = formatToolsPrompt([]); + expect(prompt).toBe(""); + expect(toolNames).toEqual([]); + }); + + it("formats tools with schema", () => { + const { prompt } = formatToolsPrompt(mockTools); + expect(prompt).toContain("test-tool: A test tool"); + expect(prompt).toContain('{"name": "test-tool", "arguments": {"param1": "...", "param2": 0, "nested": {"inner": true}'); + }); + + it("uses short descriptions from map", () => { + const toolsWithBash = [{ type: "function", name: "bash", description: "long desc", inputSchema: {} }] as any; + const { prompt } = formatToolsPrompt(toolsWithBash); + expect(prompt).toContain("bash: Execute a bash command"); + }); + + it("truncates long descriptions", () => { + const longDesc = "Very long description that should be truncated to 60 chars. "; + const tools = [{ type: "function", name: "tool", description: longDesc.repeat(2), inputSchema: {} }] as any; + const { prompt } = formatToolsPrompt(tools); + expect(prompt.length).toBeLessThan(200); // Rough check + }); +}); \ No newline at end of file