diff --git a/.speakeasy/gen.yaml b/.speakeasy/gen.yaml index 9e9c64be..c4fefac0 100644 --- a/.speakeasy/gen.yaml +++ b/.speakeasy/gen.yaml @@ -38,15 +38,21 @@ typescript: acceptHeaderEnum: true additionalDependencies: dependencies: + '@opentelemetry/semantic-conventions': ^1.40.0 ws: ^8.18.0 zod-to-json-schema: ^3.25.0 devDependencies: + '@opentelemetry/api': ^1.9.0 '@types/node': ^22.0.0 '@types/ws': ^8.5.13 - peerDependencies: {} + peerDependencies: + '@opentelemetry/api': ^1.9.0 additionalPackageJSON: description: TypeScript client library for the Mistral AI API license: Apache-2.0 + peerDependenciesMeta: + '@opentelemetry/api': + optional: true additionalScripts: {} alwaysIncludeInboundAndOutbound: false author: Mistral AI diff --git a/package-lock.json b/package-lock.json index e8fd08aa..aad3cf74 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,16 +9,45 @@ "version": "2.2.5", "license": "Apache-2.0", "dependencies": { + "@opentelemetry/semantic-conventions": "^1.40.0", "ws": "^8.18.0", "zod": "^3.25.0 || ^4.0.0", "zod-to-json-schema": "^3.25.0" }, "devDependencies": { + "@opentelemetry/api": "^1.9.0", "@types/node": "^22.0.0", "@types/ws": "^8.5.13", "@typescript/native-preview": "7.0.0-dev.20260302.1", "oxlint": "^1.60.0", "typescript": "~5.8.3" + }, + "peerDependencies": { + "@opentelemetry/api": "^1.9.0" + }, + "peerDependenciesMeta": { + "@opentelemetry/api": { + "optional": true + } + } + }, + "node_modules/@opentelemetry/api": { + "version": "1.9.1", + "resolved": "https://registry.npmjs.org/@opentelemetry/api/-/api-1.9.1.tgz", + "integrity": "sha512-gLyJlPHPZYdAk1JENA9LeHejZe1Ti77/pTeFm/nMXmQH/HFZlcS/O2XJB+L8fkbrNSqhdtlvjBVjxwUYanNH5Q==", + "dev": true, + "license": "Apache-2.0", + "engines": { + "node": ">=8.0.0" + } + }, + "node_modules/@opentelemetry/semantic-conventions": { + "version": "1.40.0", + "resolved": "https://registry.npmjs.org/@opentelemetry/semantic-conventions/-/semantic-conventions-1.40.0.tgz", + "integrity": "sha512-cifvXDhcqMwwTlTK04GBNeIe7yyo28Mfby85QXFe1Yk8nmi36Ab/5UQwptOx84SsoGNRg+EVSjwzfSZMy6pmlw==", + "license": "Apache-2.0", + "engines": { + "node": ">=14" } }, "node_modules/@oxlint/binding-android-arm-eabi": { diff --git a/package.json b/package.json index 93d07bf9..556be87f 100644 --- a/package.json +++ b/package.json @@ -38,6 +38,16 @@ "types": "./esm/extra/realtime/index.d.ts", "default": "./esm/extra/realtime/index.js" }, + "./extra/observability": { + "source": "./src/extra/observability/index.ts", + "types": "./esm/extra/observability/index.d.ts", + "default": "./esm/extra/observability/index.js" + }, + "./hooks/tracing": { + "source": "./src/hooks/tracing.ts", + "types": "./esm/hooks/tracing.d.ts", + "default": "./esm/hooks/tracing.js" + }, "./*.js": { "source": "./src/*.ts", "types": "./esm/*.d.ts", @@ -60,6 +70,7 @@ "prepublishOnly": "npm run build" }, "devDependencies": { + "@opentelemetry/api": "^1.9.0", "@types/node": "^22.0.0", "@types/ws": "^8.5.13", "@typescript/native-preview": "7.0.0-dev.20260302.1", @@ -67,8 +78,17 @@ "typescript": "~5.8.3" }, "dependencies": { + "@opentelemetry/semantic-conventions": "^1.40.0", "ws": "^8.18.0", "zod": "^3.25.0 || ^4.0.0", "zod-to-json-schema": "^3.25.0" + }, + "peerDependencies": { + "@opentelemetry/api": "^1.9.0" + }, + "peerDependenciesMeta": { + "@opentelemetry/api": { + "optional": true + } } } diff --git a/src/extra/observability/formatting.ts b/src/extra/observability/formatting.ts new file mode 100644 index 00000000..43a97d85 --- /dev/null +++ b/src/extra/observability/formatting.ts @@ -0,0 +1,194 @@ +/** + * Formatting helpers for converting Mistral API payloads to OTEL GenAI convention formats. + * + * These are pure functions with no OTEL dependencies - they transform objects to objects + * matching the GenAI semantic convention schemas for input/output messages and tool definitions. + * The caller is responsible for the final JSON serialization (single JSON.stringify on the whole + * collection) before setting span attributes. + * + * Schemas: + * - Input messages: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-input-messages.json + * - Output messages: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-output-messages.json + * - Tool definitions: https://github.com/Cirilla-zmh/semantic-conventions/blob/cc4d07e7e56b80e9aa5904a3d524c134699da37f/docs/gen-ai/gen-ai-tool-definitions.json + */ + +/** + * Convert Mistral message content to OTEL parts array. + * + * Mistral content is either a string or an array of content chunks. + */ +function contentToParts(content: unknown): Array> { + if (content == null) { + return []; + } + if (typeof content === "string") { + return [{ type: "text", content }]; + } + if (!Array.isArray(content)) { + return []; + } + // Content chunks array - map known Mistral types to OTEL part types + const parts: Array> = []; + for (const chunk of content) { + if (typeof chunk === "string") { + parts.push({ type: "text", content: chunk }); + } else if (typeof chunk === "object" && chunk !== null) { + const c = chunk as Record; + const chunkType = (c["type"] as string) || ""; + if (chunkType === "text") { + parts.push({ type: "text", content: c["text"] || "" }); + } else if (chunkType === "thinking") { + const thinking = c["thinking"]; + let contentStr: string; + if (Array.isArray(thinking)) { + const textParts = thinking + .filter( + (sub): sub is Record => + typeof sub === "object" && sub !== null && (sub as Record)["type"] === "text" + ) + .map((sub) => (sub["text"] as string) || ""); + contentStr = textParts.join("\n"); + } else { + contentStr = String(thinking || ""); + } + parts.push({ type: "reasoning", content: contentStr }); + } else if (chunkType === "image_url") { + const url = c["image_url"]; + const uri = + typeof url === "object" && url !== null + ? ((url as Record)["url"] as string) || "" + : String(url || ""); + parts.push({ type: "uri", modality: "image", uri }); + } else { + // Catch-all for other content chunk types + parts.push({ type: chunkType }); + } + } + } + return parts; +} + +/** + * Convert Mistral tool_calls to OTEL ToolCallRequestPart entries. + */ +function toolCallsToParts(toolCalls: unknown): Array> { + if (!toolCalls || !Array.isArray(toolCalls)) { + return []; + } + const parts: Array> = []; + for (const tc of toolCalls) { + if (typeof tc !== "object" || tc === null) continue; + const tcObj = tc as Record; + const func = (tcObj["function"] as Record) || {}; + const part: Record = { + type: "tool_call", + name: func["name"] || "", + }; + const tcId = tcObj["id"]; + if (tcId != null) { + part["id"] = tcId; + } + const args = func["arguments"]; + if (args != null) { + part["arguments"] = args; + } + parts.push(part); + } + return parts; +} + +/** + * Format a single input message per the OTEL GenAI convention. + * + * Schema: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-input-messages.json + * ChatMessage: {role (required), parts (required), name?} + * + * Conversation entry objects (e.g. function.result) don't carry a "role" + * field - they are detected via their "type" and mapped to the closest OTEL role. + */ +export function formatInputMessage(message: Record): Record { + const entryType = message["type"]; + + // Conversation entry: function.result to OTEL tool role + if (entryType === "function.result") { + const part: Record = { type: "tool_call_response", response: message["result"] }; + const toolCallId = message["tool_call_id"]; + if (toolCallId != null) { + part["id"] = toolCallId; + } + return { role: "tool", parts: [part] }; + } + + // TODO: may need to handle other types for conversations (e.g. agent handoff) + + const role = (message["role"] as string) || "unknown"; + const parts: Array> = []; + + if (role === "tool") { + // Tool messages are responses to tool calls + const toolPart: Record = { + type: "tool_call_response", + response: message["content"], + }; + const toolCallId = message["tool_call_id"]; + if (toolCallId != null) { + toolPart["id"] = toolCallId; + } + parts.push(toolPart); + } else { + parts.push(...contentToParts(message["content"])); + parts.push(...toolCallsToParts(message["tool_calls"])); + } + + return { role, parts }; +} + +/** + * Format a single output choice/message per the OTEL GenAI convention. + * + * Schema: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-output-messages.json + * OutputMessage: {role (required), parts (required), finish_reason (required), name?} + */ +export function formatOutputMessage(choice: Record): Record { + const message = (choice["message"] as Record) || {}; + const parts: Array> = []; + parts.push(...contentToParts(message["content"])); + parts.push(...toolCallsToParts(message["tool_calls"])); + + return { + role: message["role"] || "assistant", + parts, + finish_reason: choice["finish_reason"] || "", + }; +} + +/** + * Flatten a Mistral tool definition to the OTEL GenAI convention schema. + * + * Mistral format: {"type": "function", "function": {"name": ..., "description": ..., "parameters": ...}} + * OTEL format: {"type": "function", "name": ..., "description": ..., "parameters": ...} + * + * Schema, still under review: https://github.com/open-telemetry/semantic-conventions + */ +export function formatToolDefinition(tool: Record): Record | null { + // Early exit conditions: only functions supported for now, and name is required + const type = (tool["type"] as string) || "function"; + const func = tool["function"] as Record | undefined; + if (!func) { + return null; + } + const name = func["name"]; + if (!name) { + return null; + } + const formatted: Record = { type, name }; + const description = func["description"]; + if (description != null) { + formatted["description"] = description; + } + const parameters = func["parameters"]; + if (parameters != null) { + formatted["parameters"] = parameters; + } + return formatted; +} diff --git a/src/extra/observability/index.ts b/src/extra/observability/index.ts new file mode 100644 index 00000000..db3ec4ce --- /dev/null +++ b/src/extra/observability/index.ts @@ -0,0 +1,2 @@ +export { registerTracerProvider, traceAsync } from "./otel.js"; +export { MISTRAL_SDK_OTEL_TRACER_NAME } from "./otel.js"; diff --git a/src/extra/observability/otel.ts b/src/extra/observability/otel.ts new file mode 100644 index 00000000..2a30fd38 --- /dev/null +++ b/src/extra/observability/otel.ts @@ -0,0 +1,750 @@ +/** + * OTEL conventions for gen AI may be found at: + * + * https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/ + * https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-agent-spans/ + */ + +import { + context as contextApi, + propagation, + trace, + type Context, + type Span, + type Tracer, + type TracerProvider, + SpanStatusCode, +} from "@opentelemetry/api"; +import * as semConvAttributes from "@opentelemetry/semantic-conventions/incubating"; + +import { + formatInputMessage, + formatOutputMessage, + formatToolDefinition, +} from "./formatting.js"; +import { accumulateChunksToResponseDict, parseSseChunks } from "./streaming.js"; + +export type { Context, Span, Tracer }; +export { semConvAttributes }; + +export const OTEL_SERVICE_NAME = "mistralai_sdk"; +export const MISTRAL_SDK_OTEL_TRACER_NAME = `${OTEL_SERVICE_NAME}_tracer`; +const MISTRAL_AGENT_TRACE_PUBLIC_ATTRIBUTE = "agent.trace.public"; +let registeredTracerProvider: TracerProvider | undefined; + +/** + * Route SDK tracing spans to a specific OpenTelemetry tracer provider. + */ +export function registerTracerProvider(tracerProvider?: TracerProvider): void { + registeredTracerProvider = tracerProvider; +} + +// Safe check for environment variable in both Node.js and browser +const getDebugTracing = (): boolean => { + try { + if (typeof globalThis !== "undefined" && "process" in globalThis) { + const proc = (globalThis as { process?: { env?: Record } }).process; + return proc?.env?.["MISTRAL_SDK_DEBUG_TRACING"]?.toLowerCase() === "true"; + } + } catch { + // Ignore + } + return false; +}; +const MISTRAL_SDK_DEBUG_TRACING = getDebugTracing(); +const DEBUG_HINT = "To see detailed tracing logs, set MISTRAL_SDK_DEBUG_TRACING=true."; + +export const MistralAIAttributes = { + MISTRAL_AI_OCR_USAGE_PAGES_PROCESSED: "mistral_ai.ocr.usage.pages_processed", + MISTRAL_AI_OCR_USAGE_DOC_SIZE_BYTES: "mistral_ai.ocr.usage.doc_size_bytes", + MISTRAL_AI_ERROR_CODE: "mistral_ai.error.code", +} as const; + +export const TracingErrors = { + FAILED_TO_CREATE_SPAN_FOR_REQUEST: "Failed to create span for request.", + FAILED_TO_ENRICH_SPAN_WITH_RESPONSE: "Failed to enrich span with response.", + FAILED_TO_HANDLE_ERROR_IN_SPAN: "Failed to handle error in span.", + FAILED_TO_END_SPAN: "Failed to end span.", +} as const; + +function parseTimeToMillis(ts: string): number { + return new Date(ts).getTime(); +} + +function inferGenAiOperationName(operationId: string): string | null { + if (operationId.includes("chat_completion") || operationId === "stream_chat") { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CHAT; + } + if ( + (operationId.includes("agents_create") || operationId.includes("agents_update")) && + !operationId.includes("alias") + ) { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CREATE_AGENT; + } + if (operationId.includes("agents_completion") || operationId === "stream_agents") { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_INVOKE_AGENT; + } + if ( + operationId.includes("conversations") && + (operationId.includes("start") || + operationId.includes("append") || + operationId.includes("restart")) + ) { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_INVOKE_AGENT; + } + if (operationId.includes("fim")) { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_TEXT_COMPLETION; + } + if (operationId.includes("embeddings")) { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_EMBEDDINGS; + } + if (operationId.includes("ocr_post")) { + return semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_GENERATE_CONTENT; + } + return null; +} + +function isKnownGenAiOperation(operationId: string): boolean { + return inferGenAiOperationName(operationId) !== null; +} + +function isJsonContentType(contentType: string | null): boolean { + const mediaType = contentType?.split(";")[0]?.trim().toLowerCase(); + return mediaType === "application/json" || + mediaType === "text/json" || + mediaType?.endsWith("+json") === true; +} + +function buildGenaiSpanName( + genAiOp: string, + body: Record +): string { + if ( + genAiOp === semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CREATE_AGENT || + genAiOp === semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_INVOKE_AGENT + ) { + const agentName = body["name"] as string | undefined; + return agentName ? `${genAiOp} ${agentName}` : genAiOp; + } + if (genAiOp === semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_EXECUTE_TOOL) { + const toolName = body["name"] as string | undefined; + return toolName ? `${genAiOp} ${toolName}` : genAiOp; + } + const model = body["model"] as string | undefined; + return model ? `${genAiOp} ${model}` : genAiOp; +} + +function setAvailableAttributes(span: Span, attributes: Record): void { + for (const [key, value] of Object.entries(attributes)) { + if (value != null && value !== "") { + span.setAttribute(key, value as string | number | boolean); + } + } +} + +function setMistralAgentTraceMarker(span: Span): void { + // Mistral-internal marker consumed by downstream trace processors. + span.setAttribute(MISTRAL_AGENT_TRACE_PUBLIC_ATTRIBUTE, ""); +} + +function setHttpAttributes( + span: Span, + url: URL, + method: string, + host: string +): void { + let port = url.port ? parseInt(url.port, 10) : -1; + if (port === -1) { + if (url.protocol === "https:") { + port = 443; + } else if (url.protocol === "http:") { + port = 80; + } + } + + span.setAttributes({ + [semConvAttributes.ATTR_HTTP_REQUEST_METHOD]: method, + [semConvAttributes.ATTR_URL_FULL]: url.toString(), + [semConvAttributes.ATTR_SERVER_ADDRESS]: host, + [semConvAttributes.ATTR_SERVER_PORT]: port, + }); +} + +function enrichRequestGenaiAttrs( + span: Span, + genAiOp: string, + requestBody: Record +): void { + span.updateName(buildGenaiSpanName(genAiOp, requestBody)); + + const attributes: Record = { + [semConvAttributes.ATTR_GEN_AI_REQUEST_CHOICE_COUNT]: requestBody["n"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_ENCODING_FORMATS]: requestBody["encoding_formats"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_FREQUENCY_PENALTY]: requestBody["frequency_penalty"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_MAX_TOKENS]: requestBody["max_tokens"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_MODEL]: requestBody["model"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_PRESENCE_PENALTY]: requestBody["presence_penalty"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_SEED]: requestBody["random_seed"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_STOP_SEQUENCES]: requestBody["stop"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_TEMPERATURE]: requestBody["temperature"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_TOP_P]: requestBody["top_p"], + [semConvAttributes.ATTR_GEN_AI_REQUEST_TOP_K]: requestBody["top_k"], + }; + + const inputMessages = (requestBody["messages"] || requestBody["inputs"]) as + | string + | Array> + | undefined; + if (typeof inputMessages === "string") { + attributes[semConvAttributes.ATTR_GEN_AI_INPUT_MESSAGES] = JSON.stringify([ + formatInputMessage({ role: "user", content: inputMessages }), + ]); + } else if (Array.isArray(inputMessages)) { + attributes[semConvAttributes.ATTR_GEN_AI_INPUT_MESSAGES] = JSON.stringify( + inputMessages.map((msg) => formatInputMessage(msg)) + ); + } + + const tools = requestBody["tools"] as Array> | undefined; + if (tools) { + const formattedTools = tools + .map((tool) => formatToolDefinition(tool)) + .filter((t): t is Record => t !== null); + if (formattedTools.length > 0) { + attributes[semConvAttributes.ATTR_GEN_AI_TOOL_DEFINITIONS] = JSON.stringify(formattedTools); + } + } + + // TODO: For agent start conversation, add agent id and version attributes here ? + + setAvailableAttributes(span, attributes); +} + +export function enrichSpanFromRequest( + span: Span, + operationId: string, + url: URL, + method: string, + host: string, + body: string | null +): Span { + setHttpAttributes(span, url, method, host); + + const genAiOp = inferGenAiOperationName(operationId); + if (genAiOp === null) { + return span; + } + + span.setAttributes({ + [semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]: genAiOp, + [semConvAttributes.ATTR_GEN_AI_PROVIDER_NAME]: semConvAttributes.GEN_AI_PROVIDER_NAME_VALUE_MISTRAL_AI, + }); + + if (body) { + try { + const requestBody = JSON.parse(body) as Record; + enrichRequestGenaiAttrs(span, genAiOp, requestBody); + } catch { + // Ignore parse errors + } + } + + return span; +} + +function enrichResponseGenaiAttrs( + span: Span, + genAiOp: string, + responseData: Record +): void { + const attributes: Record = {}; + + if (genAiOp !== semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CREATE_AGENT) { + attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_ID] = responseData["id"]; + } + attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_MODEL] = responseData["model"]; + + const choices = (responseData["choices"] || []) as Array>; + const finishReasons = choices + .map((c) => c["finish_reason"] as string | undefined) + .filter((r): r is string => !!r); + if (finishReasons.length > 0) { + attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_FINISH_REASONS] = finishReasons; + } + if (choices.length > 0) { + attributes[semConvAttributes.ATTR_GEN_AI_OUTPUT_MESSAGES] = JSON.stringify( + choices.map((choice) => formatOutputMessage(choice)) + ); + } + + const usage = responseData["usage"] as Record | undefined; + if (usage) { + attributes[semConvAttributes.ATTR_GEN_AI_USAGE_INPUT_TOKENS] = usage["prompt_tokens"] || 0; + attributes[semConvAttributes.ATTR_GEN_AI_USAGE_OUTPUT_TOKENS] = usage["completion_tokens"] || 0; + } + + setAvailableAttributes(span, attributes); +} + +function enrichCreateAgent(span: Span, responseData: Record): void { + const agentAttributes: Record = { + [semConvAttributes.ATTR_GEN_AI_AGENT_DESCRIPTION]: responseData["description"], + [semConvAttributes.ATTR_GEN_AI_AGENT_ID]: responseData["id"], + [semConvAttributes.ATTR_GEN_AI_AGENT_NAME]: responseData["name"], + "gen_ai.agent.version": String(responseData["version"]), + [semConvAttributes.ATTR_GEN_AI_REQUEST_MODEL]: responseData["model"], + [semConvAttributes.ATTR_GEN_AI_SYSTEM_INSTRUCTIONS]: responseData["instructions"], + }; + setAvailableAttributes(span, agentAttributes); +} + +function stringifyToolAttribute(value: unknown): string | null { + if (typeof value === "string") { + return value; + } + if (value) { + return JSON.stringify(value); + } + return null; +} + +function createToolExecutionChildSpan( + tracer: Tracer, + parentContext: ReturnType, + output: Record +): void { + const startMs = parseTimeToMillis(output["created_at"] as string); + const endMs = parseTimeToMillis(output["completed_at"] as string); + const opName = semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_EXECUTE_TOOL; + const spanName = buildGenaiSpanName(opName, output); + const childSpan = tracer.startSpan(spanName, { startTime: startMs }, parentContext); + setMistralAgentTraceMarker(childSpan); + + const toolArguments = output["arguments"]; + const toolResult = output["info"]; + const toolAttributes: Record = { + [semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]: opName, + [semConvAttributes.ATTR_GEN_AI_PROVIDER_NAME]: semConvAttributes.GEN_AI_PROVIDER_NAME_VALUE_MISTRAL_AI, + [semConvAttributes.ATTR_GEN_AI_TOOL_CALL_ID]: output["id"], + [semConvAttributes.ATTR_GEN_AI_TOOL_CALL_ARGUMENTS]: stringifyToolAttribute(toolArguments), + [semConvAttributes.ATTR_GEN_AI_TOOL_CALL_RESULT]: stringifyToolAttribute(toolResult), + [semConvAttributes.ATTR_GEN_AI_TOOL_NAME]: output["name"], + [semConvAttributes.ATTR_GEN_AI_TOOL_TYPE]: "extension", + }; + setAvailableAttributes(childSpan, toolAttributes); + childSpan.end(endMs); +} + +function createMessageOutputChildSpan( + tracer: Tracer, + parentContext: ReturnType, + output: Record +): void { + const startMs = parseTimeToMillis(output["created_at"] as string); + const endMs = parseTimeToMillis(output["completed_at"] as string); + const opName = semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CHAT; + const spanName = buildGenaiSpanName(opName, output); + const childSpan = tracer.startSpan(spanName, { startTime: startMs }, parentContext); + setMistralAgentTraceMarker(childSpan); + + const choiceWrapper = { + message: output, + finish_reason: (output["finish_reason"] as string) || "", + }; + const messageAttributes: Record = { + [semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]: opName, + [semConvAttributes.ATTR_GEN_AI_PROVIDER_NAME]: semConvAttributes.GEN_AI_PROVIDER_NAME_VALUE_MISTRAL_AI, + [semConvAttributes.ATTR_GEN_AI_RESPONSE_ID]: output["id"], + [semConvAttributes.ATTR_GEN_AI_AGENT_ID]: output["agent_id"], + [semConvAttributes.ATTR_GEN_AI_RESPONSE_MODEL]: output["model"], + [semConvAttributes.ATTR_GEN_AI_OUTPUT_MESSAGES]: JSON.stringify([formatOutputMessage(choiceWrapper)]), + }; + setAvailableAttributes(childSpan, messageAttributes); + childSpan.end(endMs); +} + +function enrichInvokeAgent( + tracer: Tracer, + span: Span, + responseData: Record +): void { + const conversationAttributes: Record = { + [semConvAttributes.ATTR_GEN_AI_CONVERSATION_ID]: responseData["conversation_id"], + }; + setAvailableAttributes(span, conversationAttributes); + + const outputs = (responseData["outputs"] || []) as Array>; + const parentContext = trace.setSpan(contextApi.active(), span); + for (const output of outputs) { + const outputType = output["type"] as string | undefined; + if (!outputType) continue; + if (outputType === "function.call") { + continue; + } + if (outputType === "tool.execution") { + createToolExecutionChildSpan(tracer, parentContext, output); + } else if (outputType === "message.output") { + createMessageOutputChildSpan(tracer, parentContext, output); + } + } +} + +function enrichOcr(span: Span, responseData: Record): void { + const usageInfo = responseData["usage_info"] as Record | undefined; + if (!usageInfo) return; + const ocrAttributes: Record = { + [MistralAIAttributes.MISTRAL_AI_OCR_USAGE_PAGES_PROCESSED]: usageInfo["pages_processed"], + [MistralAIAttributes.MISTRAL_AI_OCR_USAGE_DOC_SIZE_BYTES]: usageInfo["doc_size_bytes"], + }; + setAvailableAttributes(span, ocrAttributes); +} + +export function enrichSpanFromResponse( + tracer: Tracer, + span: Span, + operationId: string, + responseData: Record +): void { + const genAiOp = inferGenAiOperationName(operationId); + if (genAiOp === null) { + return; + } + + enrichResponseGenaiAttrs(span, genAiOp, responseData); + + if (genAiOp === semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CREATE_AGENT) { + enrichCreateAgent(span, responseData); + } else if (genAiOp === semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_INVOKE_AGENT) { + enrichInvokeAgent(tracer, span, responseData); + } + + if (operationId === "ocr_v1_ocr_post") { + enrichOcr(span, responseData); + } +} + +/** + * Get a tracer from the registered or global TracerProvider. + * + * The SDK does not set up its own TracerProvider. It relies on the application + * to register one explicitly or configure OpenTelemetry's global provider. + * + * If no TracerProvider is configured, the ProxyTracerProvider (default) will + * return a NoOp tracer, effectively disabling tracing. Once the application + * sets up a real TracerProvider, subsequent spans will be recorded. + */ +export function getOrCreateOtelTracer(): Tracer { + const tracerProvider = registeredTracerProvider ?? trace.getTracerProvider(); + return tracerProvider.getTracer(MISTRAL_SDK_OTEL_TRACER_NAME); +} + +export function getSpanContext(span: Span): Context { + return trace.setSpan(contextApi.active(), span); +} + +export function runWithContext(context: Context, fn: () => T): T { + return contextApi.with(context, fn); +} + +export async function recordRequestError( + context: Context, + error: unknown +): Promise { + const span = trace.getSpan(context); + if (span) { + await getResponseAndError(span, null, error); + } +} + +function warn(error: string, details: unknown): void { + if (MISTRAL_SDK_DEBUG_TRACING) { + console.warn(error, details); + } else { + console.warn(error, DEBUG_HINT); + } +} + +export async function getTracedRequestAndSpan( + tracer: Tracer, + operationId: string, + request: Request +): Promise<{ request: Request; span: Span; body: string | null }> { + const span = tracer.startSpan(operationId); + + if (!span.isRecording()) { + return { request, span, body: null }; + } + + try { + setMistralAgentTraceMarker(span); + + // Propagate gen_ai.conversation.id from OTEL baggage if present + const baggage = propagation.getBaggage(contextApi.active()); + const conversationId = baggage?.getEntry(semConvAttributes.ATTR_GEN_AI_CONVERSATION_ID)?.value; + if (conversationId) { + span.setAttribute(semConvAttributes.ATTR_GEN_AI_CONVERSATION_ID, conversationId); + } + + let body: string | null = null; + if ( + request.body && + isKnownGenAiOperation(operationId) && + isJsonContentType(request.headers.get("content-type")) + ) { + try { + const cloned = request.clone(); + body = await cloned.text(); + } catch { + // Ignore - body might not be cloneable + } + } + + const headers = new Headers(request.headers); + const headerCarrier: Record = {}; + propagation.inject(trace.setSpan(contextApi.active(), span), headerCarrier); + for (const [key, value] of Object.entries(headerCarrier)) { + headers.set(key, value); + } + + const newRequest = new Request(request, { headers }); + + const url = new URL(request.url); + enrichSpanFromRequest( + span, + operationId, + url, + request.method, + headers.get("host") || url.host, + body + ); + + return { request: newRequest, span, body }; + } catch (err) { + warn(TracingErrors.FAILED_TO_CREATE_SPAN_FOR_REQUEST, err); + try { + span.end(); + } catch { + // Ignore + } + return { request, span, body: null }; + } +} + +export async function getTracedResponse( + tracer: Tracer, + span: Span, + operationId: string, + response: Response, +): Promise { + if (!span.isRecording()) { + return response; + } + + try { + span.setStatus({ code: SpanStatusCode.OK }); + span.setAttribute(semConvAttributes.ATTR_HTTP_RESPONSE_STATUS_CODE, response.status); + + const knownGenAiOperation = isKnownGenAiOperation(operationId); + const responseContentType = response.headers.get("content-type"); + const isStreamResponse = + knownGenAiOperation && + !response.bodyUsed && + responseContentType?.toLowerCase().includes("text/event-stream") === true; + + if (isStreamResponse && response.body) { + return createTracedStreamResponse(response, span, tracer, operationId); + } + + if ( + knownGenAiOperation && + !response.bodyUsed && + isJsonContentType(responseContentType) + ) { + const clonedResponse = response.clone(); + try { + const responseData = await clonedResponse.json(); + enrichSpanFromResponse( + tracer, + span, + operationId, + responseData as Record + ); + } catch { + // Ignore parse errors + } + } + endSpan(span); + + return response; + } catch (err) { + warn(TracingErrors.FAILED_TO_ENRICH_SPAN_WITH_RESPONSE, err); + endSpan(span); + return response; + } +} + +function createTracedStreamResponse( + response: Response, + span: Span, + tracer: Tracer, + operationId: string +): Response { + const originalBody = response.body!; + const chunks: string[] = []; + const decoder = new TextDecoder(); + const reader = originalBody.getReader(); + let finalized = false; + + const finalize = (): void => { + if (finalized) { + return; + } + const remaining = decoder.decode(); + if (remaining) { + chunks.push(remaining); + } + finalized = true; + finalizeStreamSpan(chunks, span, tracer, operationId); + }; + + const tracedStream = new ReadableStream( + { + async pull(controller) { + try { + const { done, value } = await reader.read(); + if (done) { + finalize(); + controller.close(); + return; + } + + chunks.push(decoder.decode(value, { stream: true })); + controller.enqueue(value); + } catch (err) { + finalize(); + controller.error(err); + try { + await reader.cancel(err); + } catch { + // Ignore cancellation errors from already-failed streams. + } + } + }, + async cancel(reason) { + finalize(); + await reader.cancel(reason); + }, + }, + { highWaterMark: 0 } + ); + + return new Response(tracedStream, { + status: response.status, + statusText: response.statusText, + headers: response.headers, + }); +} + +function finalizeStreamSpan( + chunks: string[], + span: Span, + tracer: Tracer, + operationId: string +): void { + try { + const fullText = chunks.join(""); + const parsedChunks = parseSseChunks(fullText); + if (parsedChunks.length > 0) { + const responseData = accumulateChunksToResponseDict(parsedChunks); + enrichSpanFromResponse(tracer, span, operationId, responseData); + } + } catch (err) { + warn("Failed to enrich span with streaming response.", err); + } + endSpan(span); +} + +export async function getResponseAndError( + span: Span, + response: Response | null, + error: unknown +): Promise<{ response: Response | null; error: unknown }> { + if (!span.isRecording()) { + return { response, error }; + } + + try { + if (error) { + span.recordException(error as Error); + span.setStatus({ code: SpanStatusCode.ERROR, message: String(error) }); + } + + if (!response) { + endSpan(span); + return { response, error }; + } + + try { + const responseBody = await response.clone().json(); + const body = responseBody as Record; + if (body["object"] === "error") { + const errorMsg = (body["message"] as string) || ""; + const errorType = (body["type"] as string) || ""; + if (errorMsg) { + span.setStatus({ code: SpanStatusCode.ERROR, message: errorMsg }); + span.addEvent("exception", { + "exception.type": errorType || "api_error", + "exception.message": errorMsg, + }); + const attrs: Record = { + [semConvAttributes.ATTR_HTTP_RESPONSE_STATUS_CODE]: response.status, + }; + if (errorType) { + attrs[semConvAttributes.ATTR_ERROR_TYPE] = errorType; + } + if (body["code"]) { + attrs[MistralAIAttributes.MISTRAL_AI_ERROR_CODE] = body["code"]; + } + setAvailableAttributes(span, attrs); + } + } + } catch { + // Non-JSON error bodies are valid API responses; leave the span as-is. + } + endSpan(span); + } catch (err) { + warn(TracingErrors.FAILED_TO_HANDLE_ERROR_IN_SPAN, err); + try { + span.end(); + } catch { + // Ignore + } + } + + return { response, error }; +} + +function endSpan(span: Span): void { + try { + span.end(); + } catch (err) { + warn(TracingErrors.FAILED_TO_END_SPAN, err); + } +} + +/** + * Create a traced span using a callback pattern. + * This is the TypeScript equivalent of Python's context manager. + */ +export async function traceAsync( + name: string, + fn: (span: Span) => Promise +): Promise { + const tracer = getOrCreateOtelTracer(); + const span = tracer.startSpan(name); + try { + return await fn(span); + } finally { + span.end(); + } +} diff --git a/src/extra/observability/streaming.ts b/src/extra/observability/streaming.ts new file mode 100644 index 00000000..d9ec7dc3 --- /dev/null +++ b/src/extra/observability/streaming.ts @@ -0,0 +1,174 @@ +/** + * Streaming response helpers for OTEL tracing. + * + * Pure functions that parse SSE byte streams and accumulate CompletionChunk + * deltas into a ChatCompletionResponse-shaped dict suitable for span enrichment. + * + * NOTE: The SSE bytes are re-parsed here even though EventStream already + * parsed them during iteration. + * TracedResponse sits below EventStream and can only accumulate raw bytes; it + * has no access to the decoded events. Hooking into EventStream could eliminate + * this double-parse, but EventStream is Speakeasy-generated code. + */ + +import { + CompletionChunk, + completionChunkFromJSON, +} from "../../models/components/completionchunk.js"; +import { + UsageInfo, + UsageInfo$outboundSchema, +} from "../../models/components/usageinfo.js"; + +type AccumulatedChoice = { + message: { + role: string; + content: string; + tool_calls?: Array<{ + id: string | null; + function: { + name: string; + arguments: string; + }; + }>; + }; + finish_reason: string; +}; + +function normalizeCompletionChunkPayload(payload: string): string { + const parsed = JSON.parse(payload) as unknown; + if (typeof parsed !== "object" || parsed === null || Array.isArray(parsed)) { + return payload; + } + + const chunk = parsed as Record; + + // Python accepts null for these optional fields; the TS schema expects omission. + for (const field of ["object", "created", "usage"]) { + if (chunk[field] === null) { + delete chunk[field]; + } + } + + const choices = chunk["choices"]; + if (Array.isArray(choices)) { + for (const choice of choices) { + if (typeof choice === "object" && choice !== null && !Array.isArray(choice)) { + const choiceObj = choice as Record; + if (!("finish_reason" in choiceObj)) { + choiceObj["finish_reason"] = null; + } + } + } + } + + return JSON.stringify(chunk); +} + +/** + * Parse raw SSE text into a list of typed CompletionChunk models. + * + * Only CompletionChunk is handled. If new SSE-streamed response types + * are added, parsing and typing here will need updating. + */ +export function parseSseChunks(rawSseText: string): CompletionChunk[] { + const chunks: CompletionChunk[] = []; + const lines = rawSseText.split("\n"); + for (const line of lines) { + const trimmed = line.trim(); + if (!trimmed.startsWith("data: ")) { + continue; + } + const payload = trimmed.slice(6); + if (payload === "[DONE]") { + continue; + } + try { + const result = completionChunkFromJSON(normalizeCompletionChunkPayload(payload)); + if (result.ok) { + chunks.push(result.value); + } + } catch { + continue; + } + } + return chunks; +} + +/** + * Accumulate streaming CompletionChunk deltas into a ChatCompletionResponse-shaped dict. + */ +export function accumulateChunksToResponseDict( + chunks: CompletionChunk[] +): Record { + let responseId: string | undefined; + let model: string | undefined; + let usage: UsageInfo | undefined; + const choices: Map = new Map(); + + for (const chunk of chunks) { + responseId = responseId || chunk.id; + model = model || chunk.model; + usage = usage || chunk.usage; + + for (const choice of chunk.choices) { + let accumulated = choices.get(choice.index); + if (!accumulated) { + accumulated = { + message: { role: "assistant", content: "" }, + finish_reason: "", + }; + choices.set(choice.index, accumulated); + } + const msg = accumulated.message; + const delta = choice.delta; + if (typeof delta.role === "string") { + msg.role = delta.role; + } + if (typeof delta.content === "string" && delta.content) { + msg.content += delta.content; + } + if (typeof choice.finishReason === "string") { + accumulated.finish_reason = choice.finishReason; + } + if (Array.isArray(delta.toolCalls)) { + if (!msg.tool_calls) { + msg.tool_calls = []; + } + const toolCalls = msg.tool_calls; + for (const tc of delta.toolCalls) { + const tcIdx = tc.index != null ? tc.index : toolCalls.length; + while (toolCalls.length <= tcIdx) { + toolCalls.push({ + id: null, + function: { name: "", arguments: "" }, + }); + } + const existing = toolCalls[tcIdx]!; + // ToolCall.id defaults to "null" string in some cases (Speakeasy codegen quirk) + if (tc.id != null && tc.id !== "null") { + existing.id = tc.id; + } + if (tc.function.name) { + existing.function.name += tc.function.name; + } + if (typeof tc.function.arguments === "string" && tc.function.arguments) { + existing.function.arguments += tc.function.arguments; + } + } + } + } + } + + const sortedIndices = Array.from(choices.keys()).sort((a, b) => a - b); + const result: Record = { + id: responseId, + model, + choices: sortedIndices.map((idx) => choices.get(idx)), + }; + if (usage != null) { + // Convert to snake_case JSON format for OTEL attributes + result["usage"] = UsageInfo$outboundSchema.parse(usage); + } + return result; +} diff --git a/src/hooks/registration.ts b/src/hooks/registration.ts index 5acdc376..42e783f0 100644 --- a/src/hooks/registration.ts +++ b/src/hooks/registration.ts @@ -1,6 +1,7 @@ import { Hooks } from "./types.js"; import { CustomUserAgentHook } from "./custom_user_agent.js"; import { DeprecationWarningHook } from "./deprecation_warning.js"; +import { TracingHook } from "./tracing.js"; /* * This file is only ever generated once on the first generation and then is free to be modified. @@ -19,4 +20,9 @@ export function initHooks(hooks: Hooks) { const deprecationWarningHook = new DeprecationWarningHook(); hooks.registerAfterSuccessHook(deprecationWarningHook) + const tracingHook = new TracingHook(); + hooks.registerBeforeRequestHook(tracingHook) + hooks.registerAfterSuccessHook(tracingHook) + hooks.registerAfterErrorHook(tracingHook) + hooks.registerSDKInitHook(tracingHook) } diff --git a/src/hooks/tracing.ts b/src/hooks/tracing.ts new file mode 100644 index 00000000..ce6b7c01 --- /dev/null +++ b/src/hooks/tracing.ts @@ -0,0 +1,173 @@ +import type { Context, Span, Tracer } from "@opentelemetry/api"; +import { HTTPClient } from "../lib/http.js"; +import { + AfterErrorContext, + AfterErrorHook, + AfterSuccessContext, + AfterSuccessHook, + BeforeRequestContext, + BeforeRequestHook, + HookContext, + SDKInitHook, + SDKInitOptions, +} from "./types.js"; + +type ObservabilityModule = typeof import("../extra/observability/otel.js"); + +let observabilityModule: ObservabilityModule | null | undefined; + +async function getObservabilityModule(): Promise { + if (observabilityModule !== undefined) { + return observabilityModule; + } + + try { + observabilityModule = await import("../extra/observability/otel.js"); + } catch { + // OpenTelemetry is an optional peer; without it, tracing is a no-op. + observabilityModule = null; + } + + return observabilityModule; +} + +export const TRACING_SPAN_KEY = "_tracingSpan"; +export const TRACING_BODY_KEY = "_tracingBody"; +export const TRACING_TRACER_KEY = "_tracingTracer"; + +export type TracingContext = HookContext & { + [TRACING_SPAN_KEY]?: Span; + [TRACING_BODY_KEY]?: string | null; + [TRACING_TRACER_KEY]?: Tracer; +}; + +// Runs the actual HTTP send inside the GenAI span context so lower-level +// fetch/http auto-instrumentation parents its spans correctly. +class TracingHTTPClient extends HTTPClient { + constructor( + private readonly wrappedClient: HTTPClient, + private readonly requestContexts: WeakMap + ) { + super(); + } + + override async request(request: Request): Promise { + const activeContext = this.requestContexts.get(request); + if (!activeContext) { + return this.wrappedClient.request(request); + } + + let observability: ObservabilityModule | null = null; + try { + observability = await getObservabilityModule(); + if (!observability) { + return await this.wrappedClient.request(request); + } + + return await observability.runWithContext( + activeContext, + () => this.wrappedClient.request(request) + ); + } catch (error) { + if (observability) { + await observability.recordRequestError(activeContext, error); + } + throw error; + } finally { + this.requestContexts.delete(request); + } + } + + override clone(): HTTPClient { + return new TracingHTTPClient( + this.wrappedClient.clone(), + this.requestContexts + ); + } +} + +export class TracingHook implements SDKInitHook, BeforeRequestHook, AfterSuccessHook, AfterErrorHook { + readonly #requestContexts = new WeakMap(); + + sdkInit(opts: SDKInitOptions): SDKInitOptions { + return { + ...opts, + client: new TracingHTTPClient(opts.client, this.#requestContexts), + }; + } + + async beforeRequest( + hookCtx: BeforeRequestContext, + request: Request + ): Promise { + const ctx = hookCtx as TracingContext; + const observability = await getObservabilityModule(); + if (!observability) { + return request; + } + + const tracer = observability.getOrCreateOtelTracer(); + + const { request: tracedRequest, span, body } = await observability.getTracedRequestAndSpan( + tracer, + hookCtx.operationID, + request + ); + + ctx[TRACING_TRACER_KEY] = tracer; + ctx[TRACING_SPAN_KEY] = span; + ctx[TRACING_BODY_KEY] = body; + this.#requestContexts.set(tracedRequest, observability.getSpanContext(span)); + + return tracedRequest; + } + + async afterSuccess( + hookCtx: AfterSuccessContext, + response: Response + ): Promise { + const ctx = hookCtx as TracingContext; + const span = ctx[TRACING_SPAN_KEY]; + const tracer = ctx[TRACING_TRACER_KEY]; + + if (!span || !tracer) { + return response; + } + + const observability = await getObservabilityModule(); + if (!observability) { + return response; + } + + return observability.getTracedResponse( + tracer, + span, + hookCtx.operationID, + response + ); + } + + async afterError( + hookCtx: AfterErrorContext, + response: Response | null, + error: unknown + ): Promise<{ response: Response | null; error: unknown }> { + const ctx = hookCtx as TracingContext; + const span = ctx[TRACING_SPAN_KEY]; + + if (!span) { + return { response, error }; + } + + const observability = await getObservabilityModule(); + if (!observability) { + return { response, error }; + } + + return observability.getResponseAndError( + span, + response, + error + ); + } +} diff --git a/tests/extra/observability/customProvider.test.ts b/tests/extra/observability/customProvider.test.ts new file mode 100644 index 00000000..b2e5a5bb --- /dev/null +++ b/tests/extra/observability/customProvider.test.ts @@ -0,0 +1,98 @@ +import type { Tracer, TracerProvider } from "@opentelemetry/api"; +import { trace } from "@opentelemetry/api"; +import { + getOrCreateOtelTracer, + registerTracerProvider, +} from "../../../src/extra/observability/otel.js"; +import { + TracingHook, + type TracingContext, + TRACING_TRACER_KEY, +} from "../../../src/hooks/tracing.js"; +import { createMockTracer } from "./helpers.js"; + +function createMockTracerProvider(tracer: Tracer): TracerProvider { + return { getTracer: () => tracer } as TracerProvider; +} + +function createChatRequest(): Request { + return new Request("https://api.mistral.ai/v1/chat/completions", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ + model: "mistral-large", + messages: [{ role: "user", content: "Hello" }], + }), + }); +} + +function createHookContext(): TracingContext { + return { + operationID: "chat_completion_v1", + baseURL: "https://api.mistral.ai", + oAuth2Scopes: null, + retryConfig: { strategy: "none" as const }, + resolvedSecurity: null, + options: {}, + } as TracingContext; +} + +afterEach(() => { + registerTracerProvider(); +}); + +describe("getOrCreateOtelTracer", () => { + test("uses global provider when no provider is registered", () => { + const globalTracer = createMockTracer("global"); + const spy = vi + .spyOn(trace, "getTracerProvider") + .mockReturnValue(createMockTracerProvider(globalTracer)); + + const tracer = getOrCreateOtelTracer(); + + expect(spy).toHaveBeenCalled(); + expect((tracer as any).label).toBe("global"); + spy.mockRestore(); + }); + + test("uses registered provider before the global provider", () => { + const registeredTracer = createMockTracer("registered"); + const spy = vi.spyOn(trace, "getTracerProvider"); + + registerTracerProvider(createMockTracerProvider(registeredTracer)); + const tracer = getOrCreateOtelTracer(); + + expect(spy).not.toHaveBeenCalled(); + expect((tracer as any).label).toBe("registered"); + spy.mockRestore(); + }); +}); + +describe("TracingHook with registered TracerProvider", () => { + test("beforeRequest uses registered tracerProvider", async () => { + const hook = new TracingHook(); + const registeredTracer = createMockTracer("registered"); + const ctx = createHookContext(); + + registerTracerProvider(createMockTracerProvider(registeredTracer)); + + await hook.beforeRequest(ctx, createChatRequest()); + + expect((ctx[TRACING_TRACER_KEY] as any).label).toBe("registered"); + }); + + test("beforeRequest falls back to global provider when no provider is registered", async () => { + const hook = new TracingHook(); + const globalTracer = createMockTracer("global"); + const spy = vi + .spyOn(trace, "getTracerProvider") + .mockReturnValue(createMockTracerProvider(globalTracer)); + + const ctx = createHookContext(); + await hook.beforeRequest(ctx, createChatRequest()); + + expect(spy).toHaveBeenCalled(); + expect((ctx[TRACING_TRACER_KEY] as any).label).toBe("global"); + spy.mockRestore(); + }); +}); diff --git a/tests/extra/observability/formatting.test.ts b/tests/extra/observability/formatting.test.ts new file mode 100644 index 00000000..d3377db1 --- /dev/null +++ b/tests/extra/observability/formatting.test.ts @@ -0,0 +1,111 @@ +/** + * Unit tests for OTEL formatting helpers. + */ + +import { + formatInputMessage, + formatOutputMessage, + formatToolDefinition, +} from "../../../src/extra/observability/formatting.js"; + +describe("formatInputMessage", () => { + test("user message with text content", () => { + const result = formatInputMessage({ role: "user", content: "hello" }); + expect(result).toEqual({ + role: "user", + parts: [{ type: "text", content: "hello" }], + }); + }); + + test("assistant message with tool calls", () => { + const msg = { + role: "assistant", + content: "", + tool_calls: [{ id: "tc1", function: { name: "get_weather", arguments: '{"city":"Paris"}' } }], + }; + const result = formatInputMessage(msg); + expect(result.role).toBe("assistant"); + expect(result.parts).toContainEqual({ + type: "tool_call", + name: "get_weather", + id: "tc1", + arguments: '{"city":"Paris"}', + }); + }); + + test("tool message", () => { + const result = formatInputMessage({ role: "tool", content: "sunny", tool_call_id: "tc1" }); + expect(result).toEqual({ + role: "tool", + parts: [{ type: "tool_call_response", response: "sunny", id: "tc1" }], + }); + }); + + test("function.result entry (conversations API)", () => { + const msg = { type: "function.result", result: '{"status":"ok"}', tool_call_id: "tc1" }; + const result = formatInputMessage(msg); + expect(result).toEqual({ + role: "tool", + parts: [{ type: "tool_call_response", response: '{"status":"ok"}', id: "tc1" }], + }); + }); + + test("thinking chunk converts to reasoning", () => { + const content = [{ type: "thinking", thinking: [{ type: "text", text: "step 1" }] }]; + const result = formatInputMessage({ role: "assistant", content }); + expect(result.parts).toEqual([{ type: "reasoning", content: "step 1" }]); + }); + + test("image_url chunk", () => { + const content = [{ type: "image_url", image_url: { url: "https://example.com/img.png" } }]; + const result = formatInputMessage({ role: "user", content }); + expect(result.parts).toEqual([{ type: "uri", modality: "image", uri: "https://example.com/img.png" }]); + }); +}); + +describe("formatOutputMessage", () => { + test("simple response", () => { + const choice = { message: { role: "assistant", content: "hello" }, finish_reason: "stop" }; + const result = formatOutputMessage(choice); + expect(result).toEqual({ + role: "assistant", + parts: [{ type: "text", content: "hello" }], + finish_reason: "stop", + }); + }); + + test("tool calls response", () => { + const choice = { + message: { + role: "assistant", + content: null, + tool_calls: [{ id: "tc1", function: { name: "f", arguments: "{}" } }], + }, + finish_reason: "tool_calls", + }; + const result = formatOutputMessage(choice); + expect(result.finish_reason).toBe("tool_calls"); + expect((result.parts as Array>)[0].type).toBe("tool_call"); + }); +}); + +describe("formatToolDefinition", () => { + test("full definition", () => { + const tool = { + type: "function", + function: { name: "get_weather", description: "Get weather", parameters: { type: "object" } }, + }; + const result = formatToolDefinition(tool); + expect(result).toEqual({ + type: "function", + name: "get_weather", + description: "Get weather", + parameters: { type: "object" }, + }); + }); + + test("returns null when no function name", () => { + expect(formatToolDefinition({ type: "function" })).toBeNull(); + expect(formatToolDefinition({ function: {} })).toBeNull(); + }); +}); diff --git a/tests/extra/observability/helpers.ts b/tests/extra/observability/helpers.ts new file mode 100644 index 00000000..bde9ae38 --- /dev/null +++ b/tests/extra/observability/helpers.ts @@ -0,0 +1,62 @@ +import type { Span, Tracer } from "@opentelemetry/api"; + +export type MockSpan = Span & { + attributes: Record; + ended: boolean; + events: Array<{ name: string; attributes?: Record }>; + name: string; + status?: unknown; +}; + +export type MockTracer = Tracer & { + label?: string; + spans: Array<{ name: string }>; +}; + +export function createMockSpan(name = "test"): MockSpan { + const attributes: Record = {}; + const events: Array<{ name: string; attributes?: Record }> = []; + let spanName = name; + const span: MockSpan = { + attributes, + ended: false, + events, + get name() { return spanName; }, + set name(n: string) { spanName = n; }, + setAttributes(attrs) { Object.assign(attributes, attrs); return span; }, + setAttribute(k, v) { attributes[k] = v; return span; }, + updateName(n) { spanName = n; return span; }, + spanContext: () => ({ traceId: "", spanId: "", traceFlags: 0 }), + setStatus(status) { span.status = status; return span; }, + recordException: () => {}, + addEvent(name, attrs) { + const event: { name: string; attributes?: Record } = { name }; + if (attrs) { + event.attributes = attrs as Record; + } + events.push(event); + return span; + }, + end: () => { span.ended = true; }, + isRecording: () => true, + addLink: () => span, + addLinks: () => span, + } as MockSpan; + return span; +} + +export function createMockTracer(label?: string): MockTracer { + const spans: Array<{ name: string }> = []; + const tracer = { + spans, + startSpan(name: string) { + spans.push({ name }); + return createMockSpan(name); + }, + startActiveSpan: () => undefined as never, + } as MockTracer; + if (label) { + tracer.label = label; + } + return tracer; +} diff --git a/tests/extra/observability/otel.test.ts b/tests/extra/observability/otel.test.ts new file mode 100644 index 00000000..1f5a7921 --- /dev/null +++ b/tests/extra/observability/otel.test.ts @@ -0,0 +1,453 @@ +/** + * Tests for OTEL span enrichment helpers. + */ + +import { + context as contextApi, + ROOT_CONTEXT, + SpanStatusCode, + trace, + type Context, + type ContextManager, + type Span, + type Tracer, + type TracerProvider, +} from "@opentelemetry/api"; +import { HTTPClient, Mistral } from "../../../src/index.js"; +import { TracingHook, TracingContext, TRACING_SPAN_KEY, TRACING_TRACER_KEY } from "../../../src/hooks/tracing.js"; +import { + enrichSpanFromRequest, + enrichSpanFromResponse, + getTracedRequestAndSpan, + getTracedResponse, + MistralAIAttributes, + registerTracerProvider, + semConvAttributes, +} from "../../../src/extra/observability/otel.js"; +import { createMockSpan, createMockTracer } from "./helpers.js"; + +class TestContextManager implements ContextManager { + #activeContext: Context = ROOT_CONTEXT; + + active(): Context { + return this.#activeContext; + } + + with ReturnType>( + context: Context, + fn: F, + thisArg?: ThisParameterType, + ...args: A + ): ReturnType { + const previousContext = this.#activeContext; + this.#activeContext = context; + + const restore = () => { + this.#activeContext = previousContext; + }; + + try { + const result = fn.call(thisArg, ...args); + if (result instanceof Promise) { + return result.finally(restore) as ReturnType; + } + restore(); + return result; + } catch (error) { + restore(); + throw error; + } + } + + bind(_: Context, target: T): T { + return target; + } + + enable(): this { + return this; + } + + disable(): this { + this.#activeContext = ROOT_CONTEXT; + return this; + } +} + +describe("enrichSpanFromRequest", () => { + test("chat completion with model and messages", () => { + const span = createMockSpan(); + const body = JSON.stringify({ + model: "mistral-large", + messages: [{ role: "user", content: "Hi" }], + temperature: 0.7, + }); + enrichSpanFromRequest(span, "chat_completion_v1", new URL("https://api.mistral.ai/v1/chat"), "POST", "api.mistral.ai", body); + + expect(span.name).toBe("chat mistral-large"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]).toBe(semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_CHAT); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_REQUEST_MODEL]).toBe("mistral-large"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_REQUEST_TEMPERATURE]).toBe(0.7); + expect(span.attributes[semConvAttributes.ATTR_HTTP_REQUEST_METHOD]).toBe("POST"); + expect(span.attributes[semConvAttributes.ATTR_SERVER_PORT]).toBe(443); + }); + + test("conversations API with string inputs", () => { + const span = createMockSpan(); + const body = JSON.stringify({ model: "mistral-large", inputs: "Hello" }); + enrichSpanFromRequest(span, "conversations_start_v1", new URL("https://api.mistral.ai/v1/conversations/start"), "POST", "api.mistral.ai", body); + + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]).toBe(semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_INVOKE_AGENT); + const msgs = JSON.parse(span.attributes[semConvAttributes.ATTR_GEN_AI_INPUT_MESSAGES] as string); + expect(msgs[0].role).toBe("user"); + }); + + test("embeddings operation", () => { + const span = createMockSpan(); + const body = JSON.stringify({ model: "mistral-embed", input: ["text"] }); + enrichSpanFromRequest(span, "embeddings_v1", new URL("https://api.mistral.ai/v1/embeddings"), "POST", "api.mistral.ai", body); + + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]).toBe(semConvAttributes.GEN_AI_OPERATION_NAME_VALUE_EMBEDDINGS); + }); + + test("non-GenAI operation sets only HTTP attributes", () => { + const span = createMockSpan(); + enrichSpanFromRequest(span, "list_models", new URL("https://api.mistral.ai/v1/models"), "GET", "api.mistral.ai", null); + + expect(span.attributes[semConvAttributes.ATTR_HTTP_REQUEST_METHOD]).toBe("GET"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_OPERATION_NAME]).toBeUndefined(); + }); +}); + +describe("enrichSpanFromResponse", () => { + test("chat completion response", () => { + const span = createMockSpan(); + const response = { + id: "cmpl-1", + model: "mistral-large", + choices: [{ message: { role: "assistant", content: "Hi!" }, finish_reason: "stop" }], + usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }, + }; + enrichSpanFromResponse(createMockTracer(), span, "chat_completion_v1", response); + + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_ID]).toBe("cmpl-1"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_FINISH_REASONS]).toEqual(["stop"]); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_USAGE_INPUT_TOKENS]).toBe(10); + }); + + test("create agent response", () => { + const span = createMockSpan(); + const response = { id: "agent-1", name: "MyAgent", model: "mistral-large", instructions: "Be helpful" }; + enrichSpanFromResponse(createMockTracer(), span, "agents_create_v1", response); + + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_AGENT_ID]).toBe("agent-1"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_AGENT_NAME]).toBe("MyAgent"); + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_ID]).toBeUndefined(); + }); + + test("conversation response creates child spans for tool executions", () => { + const span = createMockSpan(); + const tracer = createMockTracer(); + const response = { + conversation_id: "conv-1", + outputs: [ + { type: "tool.execution", name: "search", created_at: "2024-01-01T00:00:00Z", completed_at: "2024-01-01T00:00:01Z" }, + { type: "message.output", model: "mistral-large", created_at: "2024-01-01T00:00:01Z", completed_at: "2024-01-01T00:00:02Z" }, + ], + }; + enrichSpanFromResponse(tracer, span, "conversations_start_v1", response); + + expect(span.attributes[semConvAttributes.ATTR_GEN_AI_CONVERSATION_ID]).toBe("conv-1"); + expect(tracer.spans).toHaveLength(2); + expect(tracer.spans[0].name).toBe("execute_tool search"); + expect(tracer.spans[1].name).toBe("chat mistral-large"); + }); + + test("OCR response", () => { + const span = createMockSpan(); + const response = { model: "pixtral", usage_info: { pages_processed: 3, doc_size_bytes: 1000 } }; + enrichSpanFromResponse(createMockTracer(), span, "ocr_v1_ocr_post", response); + + expect(span.attributes[MistralAIAttributes.MISTRAL_AI_OCR_USAGE_PAGES_PROCESSED]).toBe(3); + }); +}); + +describe("TracingHook concurrency", () => { + test("interleaved requests enrich correct spans", async () => { + const hook = new TracingHook(); + + const createCtx = (operationID: string) => ({ + operationID, + baseURL: "https://api.mistral.ai", + oAuth2Scopes: null, + retryConfig: { strategy: "none" as const }, + resolvedSecurity: null, + options: {}, + }); + + const createResp = (id: string) => + new Response(JSON.stringify({ + id, + model: "test", + choices: [{ message: { role: "assistant", content: "Hi" }, finish_reason: "stop" }], + }), { headers: { "Content-Type": "application/json" } }); + + const createMockSpanWithEndPromise = (name: string) => { + let resolveEnd!: () => void; + const endPromise = new Promise(r => { resolveEnd = r; }); + const span = createMockSpan(name); + span.end = () => { resolveEnd(); }; + return { span, endPromise }; + }; + + const ctxA = createCtx("chat_completion_v1") as TracingContext; + const ctxB = createCtx("chat_completion_v1") as TracingContext; + + const { span: mockSpanA, endPromise: endA } = createMockSpanWithEndPromise("spanA"); + const { span: mockSpanB, endPromise: endB } = createMockSpanWithEndPromise("spanB"); + const mockTracer = createMockTracer(); + + ctxA[TRACING_SPAN_KEY] = mockSpanA; + ctxB[TRACING_SPAN_KEY] = mockSpanB; + ctxA[TRACING_TRACER_KEY] = mockTracer; + ctxB[TRACING_TRACER_KEY] = mockTracer; + + await hook.afterSuccess(ctxA, createResp("resp-A")); + await hook.afterSuccess(ctxB, createResp("resp-B")); + + await Promise.all([endA, endB]); + + expect(mockSpanA.attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_ID]).toBe("resp-A"); + expect(mockSpanB.attributes[semConvAttributes.ATTR_GEN_AI_RESPONSE_ID]).toBe("resp-B"); + }); +}); + +describe("SDK HTTP span parenting", () => { + afterEach(() => { + contextApi.disable(); + registerTracerProvider(); + }); + + test("ends the GenAI span when the HTTP client rejects", async () => { + const genAiSpan = createMockSpan("chat"); + const tracer = { + startSpan: () => genAiSpan, + startActiveSpan: () => undefined as never, + } as Tracer; + const tracerProvider: TracerProvider = { + getTracer: () => tracer, + }; + registerTracerProvider(tracerProvider); + + const httpClient = new HTTPClient({ + async fetcher() { + throw new TypeError("fetch failed"); + }, + }); + const client = new Mistral({ + apiKey: "test-api-key", + httpClient, + }); + + await expect(client.chat.complete({ + model: "mistral-small-latest", + messages: [{ role: "user", content: "hello" }], + })).rejects.toThrow("Unable to make request"); + + expect(genAiSpan.status).toEqual({ + code: SpanStatusCode.ERROR, + message: "TypeError: fetch failed", + }); + expect(genAiSpan.ended).toBe(true); + }); + + test("runs the HTTP send with the GenAI span active", async () => { + contextApi.disable(); + contextApi.setGlobalContextManager(new TestContextManager()); + + const genAiSpan = createMockSpan("chat"); + const tracer = { + startSpan: () => genAiSpan, + startActiveSpan: () => undefined as never, + } as Tracer; + const tracerProvider: TracerProvider = { + getTracer: () => tracer, + }; + registerTracerProvider(tracerProvider); + + let activeSpanDuringFetch: Span | undefined; + const httpClient = new HTTPClient({ + async fetcher() { + activeSpanDuringFetch = trace.getSpan(contextApi.active()); + return new Response(JSON.stringify({ + id: "chatcmpl-parenting-test", + object: "chat.completion", + created: 1700000000, + model: "mistral-small-latest", + choices: [{ + index: 0, + message: { role: "assistant", content: "ok" }, + finish_reason: "stop", + }], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + }, + }); + + const client = new Mistral({ + apiKey: "test-api-key", + httpClient, + }); + + await client.chat.complete({ + model: "mistral-small-latest", + messages: [{ role: "user", content: "hello" }], + }); + + expect(activeSpanDuringFetch).toBe(genAiSpan); + }); +}); + +describe("TracingHook body and stream handling", () => { + test("does not read non-GenAI request bodies", async () => { + const body = new ReadableStream({ + pull(controller) { + controller.enqueue(new Uint8Array([1])); + controller.close(); + }, + }); + const request = new Request("https://api.mistral.ai/v1/files", { + method: "POST", + headers: { "content-type": "multipart/form-data; boundary=test" }, + body, + duplex: "half", + } as RequestInit & { duplex: "half" }); + const cloneSpy = vi.spyOn(request, "clone"); + + await getTracedRequestAndSpan( + createMockTracer(), + "files_api_routes_upload_file", + request + ); + + expect(cloneSpy).not.toHaveBeenCalled(); + }); + + test("does not read non-GenAI response bodies", async () => { + const body = new ReadableStream({ + pull(controller) { + controller.enqueue(new TextEncoder().encode("download bytes")); + controller.close(); + }, + }); + const response = new Response(body, { + headers: { "content-type": "application/octet-stream" }, + }); + const span = createMockSpan(); + const cloneSpy = vi.spyOn(response, "clone"); + + const tracedResponse = await getTracedResponse( + createMockTracer(), + span, + "files_api_routes_download_file", + response + ); + + expect(cloneSpy).not.toHaveBeenCalled(); + expect(span.ended).toBe(true); + await expect(tracedResponse.text()).resolves.toBe("download bytes"); + }); + + test("reads traced SSE responses lazily and cancels upstream", async () => { + const encoder = new TextEncoder(); + let reads = 0; + let cancelReason: unknown; + const body = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode( + 'data: {"id":"id-1","model":"m","choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":null}]}\n\n' + )); + }, + }); + const response = new Response(body, { + headers: { "content-type": "text/event-stream" }, + }); + const responseBody = response.body!; + const upstreamReader = responseBody.getReader(); + vi.spyOn(responseBody, "getReader").mockReturnValue({ + read() { + reads += 1; + return upstreamReader.read(); + }, + cancel(reason?: unknown) { + cancelReason = reason; + return upstreamReader.cancel(reason); + }, + releaseLock() { + return upstreamReader.releaseLock(); + }, + closed: upstreamReader.closed, + } as ReadableStreamDefaultReader); + const span = createMockSpan(); + + const tracedResponse = await getTracedResponse( + createMockTracer(), + span, + "stream_chat", + response + ); + + expect(reads).toBe(0); + + const reader = tracedResponse.body!.getReader(); + const firstChunk = await reader.read(); + expect(firstChunk.done).toBe(false); + expect(reads).toBe(1); + + await reader.cancel("stop"); + expect(cancelReason).toBe("stop"); + expect(span.ended).toBe(true); + }); +}); + +describe("TracingHook error handling", () => { + test("awaits error response enrichment before returning", async () => { + const hook = new TracingHook(); + const span = createMockSpan(); + const ctx = { + operationID: "chat_completion_v1", + baseURL: "https://api.mistral.ai", + oAuth2Scopes: null, + retryConfig: { strategy: "none" as const }, + resolvedSecurity: null, + options: {}, + [TRACING_SPAN_KEY]: span, + } as TracingContext; + const response = new Response(JSON.stringify({ + object: "error", + message: "Invalid request", + type: "invalid_request_error", + code: "bad_request", + }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + + await hook.afterError(ctx, response, null); + + expect(span.status).toEqual({ code: SpanStatusCode.ERROR, message: "Invalid request" }); + expect(span.events).toContainEqual({ + name: "exception", + attributes: { + "exception.type": "invalid_request_error", + "exception.message": "Invalid request", + }, + }); + expect(span.attributes[semConvAttributes.ATTR_HTTP_RESPONSE_STATUS_CODE]).toBe(400); + expect(span.attributes[semConvAttributes.ATTR_ERROR_TYPE]).toBe("invalid_request_error"); + expect(span.attributes[MistralAIAttributes.MISTRAL_AI_ERROR_CODE]).toBe("bad_request"); + expect(span.ended).toBe(true); + }); +}); diff --git a/tests/extra/observability/streaming.test.ts b/tests/extra/observability/streaming.test.ts new file mode 100644 index 00000000..f97d1e29 --- /dev/null +++ b/tests/extra/observability/streaming.test.ts @@ -0,0 +1,94 @@ +/** + * Tests for streaming SSE parsing and chunk accumulation. + */ + +import { + parseSseChunks, + accumulateChunksToResponseDict, +} from "../../../src/extra/observability/streaming.js"; + +function chunk(opts: { + content?: string; + role?: string; + toolCalls?: Array<{ id?: string; index?: number; function: { name?: string; arguments?: string } }>; + finishReason?: string | null; + usage?: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; +}): Record { + const delta: Record = {}; + if (opts.role) delta.role = opts.role; + if (opts.content !== undefined) delta.content = opts.content; + if (opts.toolCalls) { + delta.tool_calls = opts.toolCalls.map((tc) => ({ + id: tc.id ?? "null", + index: tc.index ?? 0, + function: { name: tc.function.name ?? "", arguments: tc.function.arguments ?? "" }, + })); + } + return { + id: "id-1", + model: "m", + choices: [{ index: 0, delta, finish_reason: opts.finishReason ?? null }], + usage: opts.usage, + }; +} + +function toSse(chunks: Record[]): string { + return chunks.map((c) => `data: ${JSON.stringify(c)}`).join("\n\n") + "\n\ndata: [DONE]\n\n"; +} + +describe("parseSseChunks", () => { + test("parses valid chunks and skips [DONE]", () => { + const result = parseSseChunks(toSse([chunk({ content: "hello" }), chunk({ content: " world" })])); + expect(result.length).toBe(2); + expect(result[0].choices[0].delta.content).toBe("hello"); + }); + + test("normalizes nullable optional fields without dropping first content chunk", () => { + const firstChunk = chunk({ role: "assistant", content: "Bonjour" }); + firstChunk["usage"] = null; + firstChunk["object"] = null; + const choice = (firstChunk["choices"] as Array>)[0]!; + delete choice["finish_reason"]; + + const result = parseSseChunks(toSse([firstChunk, chunk({ content: " monde" })])); + + expect(result.map((c) => c.choices[0].delta.content)).toEqual(["Bonjour", " monde"]); + }); + + test("skips invalid json", () => { + expect(parseSseChunks("data: {invalid}\n\ndata: [DONE]\n\n")).toEqual([]); + }); + + test("empty input", () => { + expect(parseSseChunks("")).toEqual([]); + }); +}); + +describe("accumulateChunksToResponseDict", () => { + test("concatenates content across chunks", () => { + const chunks = parseSseChunks(toSse([ + chunk({ role: "assistant", content: "Hello" }), + chunk({ content: " world", finishReason: "stop", usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 } }), + ])); + const result = accumulateChunksToResponseDict(chunks); + expect(result.choices).toEqual([{ message: { role: "assistant", content: "Hello world" }, finish_reason: "stop" }]); + expect(result.usage).toEqual({ prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 }); + }); + + test("accumulates fragmented tool call arguments", () => { + const chunks = parseSseChunks(toSse([ + chunk({ role: "assistant", content: "", toolCalls: [{ id: "tc-1", index: 0, function: { name: "f", arguments: '{"a"' } }] }), + chunk({ toolCalls: [{ index: 0, function: { arguments: ': 1}' } }], finishReason: "tool_calls" }), + ])); + const result = accumulateChunksToResponseDict(chunks) as { + choices: Array<{ message: { tool_calls?: Array<{ id: string; function: { name: string; arguments: string } }> } }>; + }; + expect(result.choices[0].message.tool_calls).toEqual([ + { id: "tc-1", function: { name: "f", arguments: '{"a": 1}' } }, + ]); + }); + + test("handles empty chunks", () => { + expect(accumulateChunksToResponseDict([])).toEqual({ id: undefined, model: undefined, choices: [] }); + }); +});