diff --git a/.changeset/wip-feat-add-instrumentation-api-for-mistral-batch-api.md b/.changeset/wip-feat-add-instrumentation-api-for-mistral-batch-api.md new file mode 100644 index 000000000..ccf0eda00 --- /dev/null +++ b/.changeset/wip-feat-add-instrumentation-api-for-mistral-batch-api.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +wip feat: Add instrumentation API for mistral batch API diff --git a/e2e/helpers/normalize.ts b/e2e/helpers/normalize.ts index c0c79b904..c5d368ae0 100644 --- a/e2e/helpers/normalize.ts +++ b/e2e/helpers/normalize.ts @@ -96,6 +96,7 @@ const DYNAMIC_HEADER_KEYS = new Set([ ]); const PROVIDER_ID_KEYS = new Set([ "agentId", + "agent_id", "claude_agent_sdk.task_id", "interaction_id", "itemId", diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.json index 84f243bfd..4c15557a9 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.json @@ -4,6 +4,201 @@ "name": "mistral-root", "type": "task", "children": [ + { + "name": "mistral-batch-file-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "parallel_tool_calls": false, + "provider": "mistral", + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "tool_configuration": { + "include": [ + "news" + ] + }, + "type": "web_search" + } + ] + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-file", + "testRunId": "" + } + }, + { + "name": "mistral-batch-inline-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-inline", + "testRunId": "" + } + }, { "name": "mistral-chat-complete-operation", "children": [ @@ -23,20 +218,18 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "observability", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", @@ -70,7 +263,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "streamed output", @@ -81,7 +274,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", @@ -116,7 +309,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -127,7 +320,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 256, + "max_tokens": 256, "model": "mistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -163,7 +356,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "4", @@ -174,7 +367,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 1024, + "max_tokens": 1024, "model": "magistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -209,20 +402,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" } ] @@ -232,12 +423,33 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 13, @@ -258,20 +470,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", "name": "get_exchange_rate" }, "id": "", - "index": 0, "type": "function" } ] @@ -281,12 +491,38 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 22, @@ -311,20 +547,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" }, { @@ -333,7 +567,6 @@ "name": "get_exchange_rate" }, "id": "", - "index": 1, "type": "function" } ] @@ -343,12 +576,57 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 96, + "max_tokens": 96, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 34, @@ -373,20 +651,18 @@ "input": "function add(a, b) {", "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "\n return a + b;\n", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "codestral-2508", "object": "chat.completion", "provider": "mistral", @@ -415,7 +691,7 @@ "input": "const project = ", "output": [ { - "finishReason": "length", + "finish_reason": "length", "index": 0, "message": { "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -426,7 +702,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "codestral-2508", "object": "chat.completion.chunk", "provider": "mistral", @@ -461,25 +737,23 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Agent complete.", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0 @@ -512,20 +786,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"city\": \"Vienna\"}", "name": "get_time_in_city" }, "id": "", - "index": 0, "type": "function" } ] @@ -533,18 +805,39 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 32, + "max_tokens": 32, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the local time in a city.", + "name": "get_time_in_city", + "parameters": { + "properties": { + "city": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 14, @@ -574,7 +867,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Understood.", @@ -583,14 +876,14 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.txt b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.txt index ad10dca0f..2f0fbd853 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.txt +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-latest.span-tree.txt @@ -4,6 +4,167 @@ span_tree: "scenario": "mistral-instrumentation", "testRunId": "" } + ├── mistral-batch-file-operation + │ metadata: { + │ "operation": "batch-file", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "file batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "file batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "parallel_tool_calls": false, + │ │ "provider": "mistral", + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ }, + │ │ { + │ │ "tool_configuration": { + │ │ "include": [ + │ │ "news" + │ │ ] + │ │ }, + │ │ "type": "web_search" + │ │ } + │ │ ] + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "file batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "file batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + ├── mistral-batch-inline-operation + │ metadata: { + │ "operation": "batch-inline", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "inline batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "inline batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "provider": "mistral" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "inline batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "inline batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } ├── mistral-chat-complete-operation │ metadata: { │ "operation": "chat-complete", @@ -22,20 +183,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "observability", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", @@ -61,7 +220,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "streamed output", @@ -72,7 +231,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -99,7 +258,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -110,7 +269,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 256, + │ "max_tokens": 256, │ "model": "mistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -138,7 +297,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "4", @@ -149,7 +308,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 1024, + │ "max_tokens": 1024, │ "model": "magistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -176,20 +335,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"location\": \"Vienna\"}", │ │ "name": "get_weather" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -199,12 +356,33 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 13, @@ -221,20 +399,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", │ │ "name": "get_exchange_rate" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -244,12 +420,38 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": { + │ │ "properties": { + │ │ "from_currency": { + │ │ "description": "Base currency code, e.g. USD.", + │ │ "type": "string" + │ │ }, + │ │ "to_currency": { + │ │ "description": "Target currency code, e.g. EUR.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "from_currency", + │ │ "to_currency" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 22, @@ -270,20 +472,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"location\": \"Vienna\"}", │ "name": "get_weather" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ }, │ { @@ -292,7 +492,6 @@ span_tree: │ "name": "get_exchange_rate" │ }, │ "id": "", - │ "index": 1, │ "type": "function" │ } │ ] @@ -302,12 +501,57 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 96, + │ "max_tokens": 96, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ }, + │ { + │ "function": { + │ "description": "Get currency exchange rate.", + │ "name": "get_exchange_rate", + │ "parameters": { + │ "properties": { + │ "from_currency": { + │ "description": "Base currency code, e.g. USD.", + │ "type": "string" + │ }, + │ "to_currency": { + │ "description": "Target currency code, e.g. EUR.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "from_currency", + │ "to_currency" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 34, @@ -324,20 +568,18 @@ span_tree: │ input: "function add(a, b) {" │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "\n return a + b;\n", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "codestral-2508", │ "object": "chat.completion", │ "provider": "mistral", @@ -358,7 +600,7 @@ span_tree: │ input: "const project = " │ output: [ │ { - │ "finishReason": "length", + │ "finish_reason": "length", │ "index": 0, │ "message": { │ "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -369,7 +611,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "codestral-2508", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -396,25 +638,23 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Agent complete.", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0 @@ -439,20 +679,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"city\": \"Vienna\"}", │ "name": "get_time_in_city" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ } │ ] @@ -460,18 +698,39 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 32, + │ "max_tokens": 32, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the local time in a city.", + │ "name": "get_time_in_city", + │ "parameters": { + │ "properties": { + │ "city": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "city" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 14, @@ -493,7 +752,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Understood.", @@ -502,14 +761,14 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.json index c16212b6d..66f8adbea 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.json @@ -4,6 +4,119 @@ "name": "mistral-root", "type": "task", "children": [ + { + "name": "mistral-batch-file-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "parallel_tool_calls": false, + "provider": "mistral", + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "tool_configuration": { + "include": [ + "news" + ] + }, + "type": "web_search" + } + ] + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-file", + "testRunId": "" + } + }, { "name": "mistral-chat-complete-operation", "children": [ @@ -23,20 +136,18 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "observability", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", @@ -70,7 +181,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "streamed output", @@ -81,7 +192,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", @@ -116,7 +227,7 @@ ], "output": [ { - "finishReason": null, + "finish_reason": null, "index": 0, "message": { "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister (S1):**\n - Current age: 4 years old.\n\n2. **Second Sister (S2):**\n - Next year, S2 will be twice as old as S1.\n - Next year, S1 will be 4 + 1 = 5 years old.\n - Therefore, S2's age next year: 2 × 5 = 10 years old.\n - Current age of S2: 10 - 1 = 9 years old.\n\n3. **Third Sister (S3):**\n - S3 is two years older than S2.\n - Current age of S3: 9 + 2 = 11 years old.\n\n4. **Older Brother (John):**\n - S3 is half the age of her older brother.\n - Therefore, John", @@ -127,7 +238,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 256, + "max_tokens": 256, "model": "mistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -159,11 +270,32 @@ } ], "metadata": { - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "error": "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" }, @@ -179,13 +311,12 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"city\": \"Vienna\"}", @@ -201,12 +332,22 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": {} + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 13, @@ -226,11 +367,37 @@ } ], "metadata": { - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "error": "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" }, @@ -246,13 +413,12 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"from\": \"USD\", \"to\": \"EUR\"}", @@ -268,12 +434,22 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": {} + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 20, @@ -297,11 +473,56 @@ } ], "metadata": { - "maxTokens": 96, + "max_tokens": 96, "model": "mistral-small-2506", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "error": "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" }, @@ -321,13 +542,12 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", @@ -351,12 +571,30 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 96, + "max_tokens": 96, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": {} + }, + "type": "function" + }, + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": {} + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 34, @@ -381,20 +619,18 @@ "input": "function add(a, b) {", "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "\n return a + b;\n", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "codestral-2508", "object": "chat.completion", "provider": "mistral", @@ -423,7 +659,7 @@ "input": "const project = ", "output": [ { - "finishReason": "length", + "finish_reason": "length", "index": 0, "message": { "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -434,7 +670,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "codestral-2508", "object": "chat.completion.chunk", "provider": "mistral", @@ -469,25 +705,23 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Agent complete.", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0 @@ -519,14 +753,35 @@ } ], "metadata": { - "agentId": "", - "maxTokens": 32, + "agent_id": "", + "max_tokens": 32, "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the local time in a city.", + "name": "get_time_in_city", + "parameters": { + "properties": { + "city": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "error": "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" }, @@ -542,13 +797,12 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{}", @@ -562,18 +816,28 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 32, + "max_tokens": 32, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the local time in a city.", + "name": "get_time_in_city", + "parameters": {} + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 8, @@ -603,7 +867,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Understood.", @@ -612,14 +876,14 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.txt b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.txt index fda964a3d..709bd194d 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.txt +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-tree.txt @@ -4,6 +4,102 @@ span_tree: "scenario": "mistral-instrumentation", "testRunId": "" } + ├── mistral-batch-file-operation + │ metadata: { + │ "operation": "batch-file", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "file batch two", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "file batch 2", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "provider": "mistral" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "file batch one", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "file batch 1", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "parallel_tool_calls": false, + │ "provider": "mistral", + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ }, + │ { + │ "tool_configuration": { + │ "include": [ + │ "news" + │ ] + │ }, + │ "type": "web_search" + │ } + │ ] + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } ├── mistral-chat-complete-operation │ metadata: { │ "operation": "chat-complete", @@ -22,20 +118,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "observability", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", @@ -61,7 +155,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "streamed output", @@ -72,7 +166,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -99,7 +193,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": null, + │ "finish_reason": null, │ "index": 0, │ "message": { │ "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister (S1):**\n - Current age: 4 years old.\n\n2. **Second Sister (S2):**\n - Next year, S2 will be twice as old as S1.\n - Next year, S1 will be 4 + 1 = 5 years old.\n - Therefore, S2's age next year: 2 × 5 = 10 years old.\n - Current age of S2: 10 - 1 = 9 years old.\n\n3. **Third Sister (S3):**\n - S3 is two years older than S2.\n - Current age of S3: 9 + 2 = 11 years old.\n\n4. **Older Brother (John):**\n - S3 is half the age of her older brother.\n - Therefore, John", @@ -110,7 +204,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 256, + │ "max_tokens": 256, │ "model": "mistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -134,11 +228,32 @@ span_tree: │ │ } │ │ ] │ │ metadata: { - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ error: "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" │ ├── mistral.chat.complete [llm] @@ -150,13 +265,12 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"city\": \"Vienna\"}", @@ -172,12 +286,22 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": {} + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 13, @@ -193,11 +317,37 @@ span_tree: │ │ } │ │ ] │ │ metadata: { - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": { + │ │ "properties": { + │ │ "from_currency": { + │ │ "description": "Base currency code, e.g. USD.", + │ │ "type": "string" + │ │ }, + │ │ "to_currency": { + │ │ "description": "Target currency code, e.g. EUR.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "from_currency", + │ │ "to_currency" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ error: "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" │ ├── mistral.chat.complete [llm] @@ -209,13 +359,12 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"from\": \"USD\", \"to\": \"EUR\"}", @@ -231,12 +380,22 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": {} + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 20, @@ -256,11 +415,56 @@ span_tree: │ │ } │ │ ] │ │ metadata: { - │ │ "maxTokens": 96, + │ │ "max_tokens": 96, │ │ "model": "mistral-small-2506", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ }, + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": { + │ │ "properties": { + │ │ "from_currency": { + │ │ "description": "Base currency code, e.g. USD.", + │ │ "type": "string" + │ │ }, + │ │ "to_currency": { + │ │ "description": "Target currency code, e.g. EUR.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "from_currency", + │ │ "to_currency" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ error: "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" │ └── mistral.chat.complete [llm] @@ -276,13 +480,12 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"location\": \"Vienna\"}", @@ -306,12 +509,30 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 96, + │ "max_tokens": 96, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get weather for a city.", + │ "name": "get_weather", + │ "parameters": {} + │ }, + │ "type": "function" + │ }, + │ { + │ "function": { + │ "description": "Get currency exchange rate.", + │ "name": "get_exchange_rate", + │ "parameters": {} + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 34, @@ -328,20 +549,18 @@ span_tree: │ input: "function add(a, b) {" │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "\n return a + b;\n", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "codestral-2508", │ "object": "chat.completion", │ "provider": "mistral", @@ -362,7 +581,7 @@ span_tree: │ input: "const project = " │ output: [ │ { - │ "finishReason": "length", + │ "finish_reason": "length", │ "index": 0, │ "message": { │ "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -373,7 +592,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "codestral-2508", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -400,25 +619,23 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Agent complete.", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0 @@ -442,14 +659,35 @@ span_tree: │ │ } │ │ ] │ │ metadata: { - │ │ "agentId": "", - │ │ "maxTokens": 32, + │ │ "agent_id": "", + │ │ "max_tokens": 32, │ │ "provider": "mistral", - │ │ "responseFormat": { + │ │ "response_format": { │ │ "type": "text" │ │ }, │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get the local time in a city.", + │ │ "name": "get_time_in_city", + │ │ "parameters": { + │ │ "properties": { + │ │ "city": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "city" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ error: "Input validation failed: TypeError: Cannot read properties of undefined (reading '_zod')" │ └── mistral.agents.complete [llm] @@ -461,13 +699,12 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{}", @@ -481,18 +718,28 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 32, + │ "max_tokens": 32, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the local time in a city.", + │ "name": "get_time_in_city", + │ "parameters": {} + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 8, @@ -514,7 +761,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Understood.", @@ -523,14 +770,14 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.json index 49db6f2f6..10c6f2213 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.json @@ -4,6 +4,201 @@ "name": "mistral-root", "type": "task", "children": [ + { + "name": "mistral-batch-file-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "parallel_tool_calls": false, + "provider": "mistral", + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "tool_configuration": { + "include": [ + "news" + ] + }, + "type": "web_search" + } + ] + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-file", + "testRunId": "" + } + }, + { + "name": "mistral-batch-inline-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-inline", + "testRunId": "" + } + }, { "name": "mistral-chat-complete-operation", "children": [ @@ -23,20 +218,18 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "observability", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", @@ -71,7 +264,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "streamed output", @@ -82,7 +275,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", @@ -118,7 +311,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -129,7 +322,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 256, + "max_tokens": 256, "model": "mistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -166,7 +359,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "4", @@ -177,7 +370,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 1024, + "max_tokens": 1024, "model": "magistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -213,20 +406,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" } ] @@ -236,12 +427,33 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 13, @@ -263,20 +475,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", "name": "get_exchange_rate" }, "id": "", - "index": 0, "type": "function" } ] @@ -286,12 +496,38 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 22, @@ -317,20 +553,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" }, { @@ -339,7 +573,6 @@ "name": "get_exchange_rate" }, "id": "", - "index": 1, "type": "function" } ] @@ -349,12 +582,57 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 96, + "max_tokens": 96, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 34, @@ -380,20 +658,18 @@ "input": "function add(a, b) {", "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "\n return a + b;\n", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "codestral-2508", "object": "chat.completion", "provider": "mistral", @@ -423,7 +699,7 @@ "input": "const project = ", "output": [ { - "finishReason": "length", + "finish_reason": "length", "index": 0, "message": { "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -434,7 +710,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "codestral-2508", "object": "chat.completion.chunk", "provider": "mistral", @@ -470,25 +746,23 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Agent complete.", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0 @@ -522,20 +796,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"city\": \"Vienna\"}", "name": "get_time_in_city" }, "id": "", - "index": 0, "type": "function" } ] @@ -543,18 +815,39 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 32, + "max_tokens": 32, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the local time in a city.", + "name": "get_time_in_city", + "parameters": { + "properties": { + "city": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 14, @@ -585,7 +878,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Understood.", @@ -594,14 +887,14 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.txt b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.txt index 334930b65..9bfb24662 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.txt +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2-latest.span-tree.txt @@ -4,6 +4,167 @@ span_tree: "scenario": "mistral-instrumentation", "testRunId": "" } + ├── mistral-batch-file-operation + │ metadata: { + │ "operation": "batch-file", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "file batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "file batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "parallel_tool_calls": false, + │ │ "provider": "mistral", + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ }, + │ │ { + │ │ "tool_configuration": { + │ │ "include": [ + │ │ "news" + │ │ ] + │ │ }, + │ │ "type": "web_search" + │ │ } + │ │ ] + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "file batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "file batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + ├── mistral-batch-inline-operation + │ metadata: { + │ "operation": "batch-inline", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "inline batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "inline batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "provider": "mistral" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "inline batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "inline batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } ├── mistral-chat-complete-operation │ metadata: { │ "operation": "chat-complete", @@ -22,20 +183,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "observability", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", @@ -62,7 +221,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "streamed output", @@ -73,7 +232,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -101,7 +260,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -112,7 +271,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 256, + │ "max_tokens": 256, │ "model": "mistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -141,7 +300,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "4", @@ -152,7 +311,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 1024, + │ "max_tokens": 1024, │ "model": "magistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -180,20 +339,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"location\": \"Vienna\"}", │ │ "name": "get_weather" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -203,12 +360,33 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 13, @@ -226,20 +404,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", │ │ "name": "get_exchange_rate" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -249,12 +425,38 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": { + │ │ "properties": { + │ │ "from_currency": { + │ │ "description": "Base currency code, e.g. USD.", + │ │ "type": "string" + │ │ }, + │ │ "to_currency": { + │ │ "description": "Target currency code, e.g. EUR.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "from_currency", + │ │ "to_currency" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 22, @@ -276,20 +478,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"location\": \"Vienna\"}", │ "name": "get_weather" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ }, │ { @@ -298,7 +498,6 @@ span_tree: │ "name": "get_exchange_rate" │ }, │ "id": "", - │ "index": 1, │ "type": "function" │ } │ ] @@ -308,12 +507,57 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 96, + │ "max_tokens": 96, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ }, + │ { + │ "function": { + │ "description": "Get currency exchange rate.", + │ "name": "get_exchange_rate", + │ "parameters": { + │ "properties": { + │ "from_currency": { + │ "description": "Base currency code, e.g. USD.", + │ "type": "string" + │ }, + │ "to_currency": { + │ "description": "Target currency code, e.g. EUR.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "from_currency", + │ "to_currency" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 34, @@ -331,20 +575,18 @@ span_tree: │ input: "function add(a, b) {" │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "\n return a + b;\n", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "codestral-2508", │ "object": "chat.completion", │ "provider": "mistral", @@ -366,7 +608,7 @@ span_tree: │ input: "const project = " │ output: [ │ { - │ "finishReason": "length", + │ "finish_reason": "length", │ "index": 0, │ "message": { │ "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -377,7 +619,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "codestral-2508", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -405,25 +647,23 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Agent complete.", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0 @@ -449,20 +689,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"city\": \"Vienna\"}", │ "name": "get_time_in_city" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ } │ ] @@ -470,18 +708,39 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 32, + │ "max_tokens": 32, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the local time in a city.", + │ "name": "get_time_in_city", + │ "parameters": { + │ "properties": { + │ "city": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "city" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 14, @@ -504,7 +763,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Understood.", @@ -513,14 +772,14 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.json index c843b05d8..94fa3ba47 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.json @@ -4,6 +4,201 @@ "name": "mistral-root", "type": "task", "children": [ + { + "name": "mistral-batch-file-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "parallel_tool_calls": false, + "provider": "mistral", + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "tool_configuration": { + "include": [ + "news" + ] + }, + "type": "web_search" + } + ] + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "file batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "file batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-file", + "testRunId": "" + } + }, + { + "name": "mistral-batch-inline-operation", + "children": [ + { + "name": "mistral.batch", + "type": "task", + "children": [ + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch one", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 1", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + }, + { + "name": "mistral.chat.complete", + "type": "llm", + "children": [], + "input": [ + { + "content": "inline batch two", + "role": "user" + } + ], + "output": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "inline batch 2", + "role": "assistant" + } + } + ], + "metadata": { + "id": "", + "model": "mistral-small-2506", + "object": "chat.completion", + "provider": "mistral" + }, + "metrics": { + "completion_tokens": 1, + "prompt_tokens": 4, + "tokens": 5 + } + } + ], + "metadata": { + "provider": "mistral" + } + } + ], + "metadata": { + "operation": "batch-inline", + "testRunId": "" + } + }, { "name": "mistral-chat-complete-operation", "children": [ @@ -23,20 +218,18 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "observability", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", @@ -71,7 +264,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "streamed output", @@ -82,7 +275,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", @@ -118,7 +311,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -129,7 +322,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 256, + "max_tokens": 256, "model": "mistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -166,7 +359,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "4", @@ -177,7 +370,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 1024, + "max_tokens": 1024, "model": "magistral-small-latest", "object": "chat.completion.chunk", "provider": "mistral", @@ -213,20 +406,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" } ] @@ -236,12 +427,33 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 13, @@ -263,20 +475,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", "name": "get_exchange_rate" }, "id": "", - "index": 0, "type": "function" } ] @@ -286,12 +496,38 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 48, + "max_tokens": 48, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 22, @@ -317,20 +553,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"location\": \"Vienna\"}", "name": "get_weather" }, "id": "", - "index": 0, "type": "function" }, { @@ -339,7 +573,6 @@ "name": "get_exchange_rate" }, "id": "", - "index": 1, "type": "function" } ] @@ -349,12 +582,57 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 96, + "max_tokens": 96, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get weather for a city.", + "name": "get_weather", + "parameters": { + "properties": { + "location": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "location" + ], + "type": "object" + } + }, + "type": "function" + }, + { + "function": { + "description": "Get currency exchange rate.", + "name": "get_exchange_rate", + "parameters": { + "properties": { + "from_currency": { + "description": "Base currency code, e.g. USD.", + "type": "string" + }, + "to_currency": { + "description": "Target currency code, e.g. EUR.", + "type": "string" + } + }, + "required": [ + "from_currency", + "to_currency" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 34, @@ -380,20 +658,18 @@ "input": "function add(a, b) {", "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "\n return a + b;\n", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { "created": 0, "id": "", - "maxTokens": 24, + "max_tokens": 24, "model": "codestral-2508", "object": "chat.completion", "provider": "mistral", @@ -423,7 +699,7 @@ "input": "const project = ", "output": [ { - "finishReason": "length", + "finish_reason": "length", "index": 0, "message": { "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -434,7 +710,7 @@ "metadata": { "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "codestral-2508", "object": "chat.completion.chunk", "provider": "mistral", @@ -470,25 +746,23 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Agent complete.", - "prefix": false, - "role": "assistant", - "toolCalls": null + "role": "assistant" } } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0 @@ -522,20 +796,18 @@ ], "output": [ { - "finishReason": "tool_calls", + "finish_reason": "tool_calls", "index": 0, "message": { "content": "", - "prefix": false, "role": "assistant", - "toolCalls": [ + "tool_calls": [ { "function": { "arguments": "{\"city\": \"Vienna\"}", "name": "get_time_in_city" }, "id": "", - "index": 0, "type": "function" } ] @@ -543,18 +815,39 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 32, + "max_tokens": 32, "model": "mistral-small-2506", "object": "chat.completion", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "temperature": 0, - "toolChoice": "required" + "tool_choice": "required", + "tools": [ + { + "function": { + "description": "Get the local time in a city.", + "name": "get_time_in_city", + "parameters": { + "properties": { + "city": { + "description": "City name, e.g. Vienna.", + "type": "string" + } + }, + "required": [ + "city" + ], + "type": "object" + } + }, + "type": "function" + } + ] }, "metrics": { "completion_tokens": 14, @@ -585,7 +878,7 @@ ], "output": [ { - "finishReason": "stop", + "finish_reason": "stop", "index": 0, "message": { "content": "Understood.", @@ -594,14 +887,14 @@ } ], "metadata": { - "agentId": "", + "agent_id": "", "created": 0, "id": "", - "maxTokens": 16, + "max_tokens": 16, "model": "mistral-small-2506", "object": "chat.completion.chunk", "provider": "mistral", - "responseFormat": { + "response_format": { "type": "text" }, "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.txt b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.txt index 441d2c832..5eadc2cd0 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.txt +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-tree.txt @@ -4,6 +4,167 @@ span_tree: "scenario": "mistral-instrumentation", "testRunId": "" } + ├── mistral-batch-file-operation + │ metadata: { + │ "operation": "batch-file", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "file batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "file batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "parallel_tool_calls": false, + │ │ "provider": "mistral", + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ }, + │ │ { + │ │ "tool_configuration": { + │ │ "include": [ + │ │ "news" + │ │ ] + │ │ }, + │ │ "type": "web_search" + │ │ } + │ │ ] + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "file batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "file batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } + ├── mistral-batch-inline-operation + │ metadata: { + │ "operation": "batch-inline", + │ "testRunId": "" + │ } + │ └── mistral.batch [task] + │ metadata: { + │ "provider": "mistral" + │ } + │ ├── mistral.chat.complete [llm] + │ │ input: [ + │ │ { + │ │ "content": "inline batch one", + │ │ "role": "user" + │ │ } + │ │ ] + │ │ output: [ + │ │ { + │ │ "finish_reason": "stop", + │ │ "index": 0, + │ │ "message": { + │ │ "content": "inline batch 1", + │ │ "role": "assistant" + │ │ } + │ │ } + │ │ ] + │ │ metadata: { + │ │ "id": "", + │ │ "model": "mistral-small-2506", + │ │ "object": "chat.completion", + │ │ "provider": "mistral" + │ │ } + │ │ metrics: { + │ │ "completion_tokens": 1, + │ │ "prompt_tokens": 4, + │ │ "tokens": 5 + │ │ } + │ └── mistral.chat.complete [llm] + │ input: [ + │ { + │ "content": "inline batch two", + │ "role": "user" + │ } + │ ] + │ output: [ + │ { + │ "finish_reason": "stop", + │ "index": 0, + │ "message": { + │ "content": "inline batch 2", + │ "role": "assistant" + │ } + │ } + │ ] + │ metadata: { + │ "id": "", + │ "model": "mistral-small-2506", + │ "object": "chat.completion", + │ "provider": "mistral" + │ } + │ metrics: { + │ "completion_tokens": 1, + │ "prompt_tokens": 4, + │ "tokens": 5 + │ } ├── mistral-chat-complete-operation │ metadata: { │ "operation": "chat-complete", @@ -22,20 +183,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "observability", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", @@ -62,7 +221,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "streamed output", @@ -73,7 +232,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -101,7 +260,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Let's break down the problem step by step to find John's age.\n\n1. **First Sister's Age:**\n - Given: The first sister is **4 years old**.\n\n2. **Second Sister's Age Next Year:**\n - Next year, the second sister will be twice as old as the first sister.\n - First sister's age next year: \\(4 + 1 = 5\\) years.\n - Therefore, second sister's age next year: \\(2 \\times 5 = 10\\) years.\n - **Current age of the second sister:** \\(10 - 1 = 9\\) years.\n\n3. **Third Sister's Age:**\n - The third sister is two years older than the second sister.\n - **Current age of the third sister:** \\(9 + 2 = 11\\) years.\n\n4. **John's Age:**\n - The third sister is half the age of her older brother (John).\n - Therefore, John's age: \\(11 \\times 2 = 22\\) years.\n\n**Final Answer:**\n\\[\n\\boxed{22}\n\\]", @@ -112,7 +271,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 256, + │ "max_tokens": 256, │ "model": "mistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -141,7 +300,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "4", @@ -152,7 +311,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 1024, + │ "max_tokens": 1024, │ "model": "magistral-small-latest", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -180,20 +339,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"location\": \"Vienna\"}", │ │ "name": "get_weather" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -203,12 +360,33 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get weather for a city.", + │ │ "name": "get_weather", + │ │ "parameters": { + │ │ "properties": { + │ │ "location": { + │ │ "description": "City name, e.g. Vienna.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "location" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 13, @@ -226,20 +404,18 @@ span_tree: │ │ ] │ │ output: [ │ │ { - │ │ "finishReason": "tool_calls", + │ │ "finish_reason": "tool_calls", │ │ "index": 0, │ │ "message": { │ │ "content": "", - │ │ "prefix": false, │ │ "role": "assistant", - │ │ "toolCalls": [ + │ │ "tool_calls": [ │ │ { │ │ "function": { │ │ "arguments": "{\"from_currency\": \"USD\", \"to_currency\": \"EUR\"}", │ │ "name": "get_exchange_rate" │ │ }, │ │ "id": "", - │ │ "index": 0, │ │ "type": "function" │ │ } │ │ ] @@ -249,12 +425,38 @@ span_tree: │ │ metadata: { │ │ "created": 0, │ │ "id": "", - │ │ "maxTokens": 48, + │ │ "max_tokens": 48, │ │ "model": "mistral-small-2506", │ │ "object": "chat.completion", │ │ "provider": "mistral", │ │ "temperature": 0, - │ │ "toolChoice": "required" + │ │ "tool_choice": "required", + │ │ "tools": [ + │ │ { + │ │ "function": { + │ │ "description": "Get currency exchange rate.", + │ │ "name": "get_exchange_rate", + │ │ "parameters": { + │ │ "properties": { + │ │ "from_currency": { + │ │ "description": "Base currency code, e.g. USD.", + │ │ "type": "string" + │ │ }, + │ │ "to_currency": { + │ │ "description": "Target currency code, e.g. EUR.", + │ │ "type": "string" + │ │ } + │ │ }, + │ │ "required": [ + │ │ "from_currency", + │ │ "to_currency" + │ │ ], + │ │ "type": "object" + │ │ } + │ │ }, + │ │ "type": "function" + │ │ } + │ │ ] │ │ } │ │ metrics: { │ │ "completion_tokens": 22, @@ -276,20 +478,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"location\": \"Vienna\"}", │ "name": "get_weather" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ }, │ { @@ -298,7 +498,6 @@ span_tree: │ "name": "get_exchange_rate" │ }, │ "id": "", - │ "index": 1, │ "type": "function" │ } │ ] @@ -308,12 +507,57 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 96, + │ "max_tokens": 96, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get weather for a city.", + │ "name": "get_weather", + │ "parameters": { + │ "properties": { + │ "location": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "location" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ }, + │ { + │ "function": { + │ "description": "Get currency exchange rate.", + │ "name": "get_exchange_rate", + │ "parameters": { + │ "properties": { + │ "from_currency": { + │ "description": "Base currency code, e.g. USD.", + │ "type": "string" + │ }, + │ "to_currency": { + │ "description": "Target currency code, e.g. EUR.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "from_currency", + │ "to_currency" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 34, @@ -331,20 +575,18 @@ span_tree: │ input: "function add(a, b) {" │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "\n return a + b;\n", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 24, + │ "max_tokens": 24, │ "model": "codestral-2508", │ "object": "chat.completion", │ "provider": "mistral", @@ -366,7 +608,7 @@ span_tree: │ input: "const project = " │ output: [ │ { - │ "finishReason": "length", + │ "finish_reason": "length", │ "index": 0, │ "message": { │ "content": "{\n name: 'project',\n title: 'Projects',\n type:", @@ -377,7 +619,7 @@ span_tree: │ metadata: { │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "codestral-2508", │ "object": "chat.completion.chunk", │ "provider": "mistral", @@ -405,25 +647,23 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Agent complete.", - │ "prefix": false, - │ "role": "assistant", - │ "toolCalls": null + │ "role": "assistant" │ } │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0 @@ -449,20 +689,18 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "tool_calls", + │ "finish_reason": "tool_calls", │ "index": 0, │ "message": { │ "content": "", - │ "prefix": false, │ "role": "assistant", - │ "toolCalls": [ + │ "tool_calls": [ │ { │ "function": { │ "arguments": "{\"city\": \"Vienna\"}", │ "name": "get_time_in_city" │ }, │ "id": "", - │ "index": 0, │ "type": "function" │ } │ ] @@ -470,18 +708,39 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 32, + │ "max_tokens": 32, │ "model": "mistral-small-2506", │ "object": "chat.completion", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "temperature": 0, - │ "toolChoice": "required" + │ "tool_choice": "required", + │ "tools": [ + │ { + │ "function": { + │ "description": "Get the local time in a city.", + │ "name": "get_time_in_city", + │ "parameters": { + │ "properties": { + │ "city": { + │ "description": "City name, e.g. Vienna.", + │ "type": "string" + │ } + │ }, + │ "required": [ + │ "city" + │ ], + │ "type": "object" + │ } + │ }, + │ "type": "function" + │ } + │ ] │ } │ metrics: { │ "completion_tokens": 14, @@ -504,7 +763,7 @@ span_tree: │ ] │ output: [ │ { - │ "finishReason": "stop", + │ "finish_reason": "stop", │ "index": 0, │ "message": { │ "content": "Understood.", @@ -513,14 +772,14 @@ span_tree: │ } │ ] │ metadata: { - │ "agentId": "", + │ "agent_id": "", │ "created": 0, │ "id": "", - │ "maxTokens": 16, + │ "max_tokens": 16, │ "model": "mistral-small-2506", │ "object": "chat.completion.chunk", │ "provider": "mistral", - │ "responseFormat": { + │ "response_format": { │ "type": "text" │ }, │ "stream": true, diff --git a/e2e/scenarios/mistral-instrumentation/assertions.ts b/e2e/scenarios/mistral-instrumentation/assertions.ts index b6479d91d..bbc9ea571 100644 --- a/e2e/scenarios/mistral-instrumentation/assertions.ts +++ b/e2e/scenarios/mistral-instrumentation/assertions.ts @@ -143,17 +143,13 @@ function summarizeOutput(output: unknown): Json { const message = firstChoice.message; const finishReason = - typeof firstChoice.finishReason === "string" || - firstChoice.finishReason === null - ? firstChoice.finishReason - : typeof firstChoice.finish_reason === "string" || - firstChoice.finish_reason === null - ? firstChoice.finish_reason - : null; - const toolCalls = - (Array.isArray(message.tool_calls) && message.tool_calls) || - (Array.isArray(message.toolCalls) && message.toolCalls) || - []; + typeof firstChoice.finish_reason === "string" || + firstChoice.finish_reason === null + ? firstChoice.finish_reason + : null; + const toolCalls = Array.isArray(message.tool_calls) + ? message.tool_calls + : []; const contentParts = Array.isArray(message.content) ? message.content : []; const contentPartTypes = contentParts .map((part) => @@ -166,9 +162,7 @@ function summarizeOutput(output: unknown): Json { has_content: typeof message.content === "string" ? message.content.length > 0 - : contentPartTypes.includes("text") - ? true - : false, + : contentPartTypes.includes("text"), role: typeof message.role === "string" ? message.role : null, tool_call_count: toolCalls.length, type: "choices", @@ -308,6 +302,24 @@ function pickOutputSpans(spans: T[]): T[] { } function snapshotEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { + const batchFileOperation = findLatestSpan( + events, + "mistral-batch-file-operation", + ); + const batchFileTask = findChildSpans( + events, + "mistral.batch", + batchFileOperation?.span.id, + )[0]; + const batchInlineOperation = findLatestSpan( + events, + "mistral-batch-inline-operation", + ); + const batchInlineTask = findChildSpans( + events, + "mistral.batch", + batchInlineOperation?.span.id, + )[0]; const chatCompleteOperation = findLatestSpan( events, "mistral-chat-complete-operation", @@ -383,6 +395,16 @@ function snapshotEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { return [ findLatestSpan(events, ROOT_NAME), + batchFileOperation, + batchFileTask, + ...findChildSpans(events, "mistral.chat.complete", batchFileTask?.span.id), + batchInlineOperation, + batchInlineTask, + ...findChildSpans( + events, + "mistral.chat.complete", + batchInlineTask?.span.id, + ), chatCompleteOperation, findMistralSpan(events, chatCompleteOperation?.span.id, [ "mistral.chat.complete", @@ -479,6 +501,7 @@ export function defineMistralInstrumentationAssertions(options: { snapshotName: string; supportsClassifiers?: boolean; supportsClassify?: boolean; + supportsInlineBatch?: boolean; supportsThinkingStream?: boolean; testFileUrl: string; timeoutMs: number; @@ -492,6 +515,7 @@ export function defineMistralInstrumentationAssertions(options: { const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); const supportsClassify = (options.supportsClassify ?? true) && !!classifyModel; + const supportsInlineBatch = options.supportsInlineBatch ?? true; const timeoutMs = effectiveScenarioTimeoutMs(options.timeoutMs); const testConfig = { timeout: timeoutMs, @@ -516,6 +540,90 @@ export function defineMistralInstrumentationAssertions(options: { }); }); + test("captures file batch.jobs.create() lifecycle", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan(events, "mistral-batch-file-operation"); + const task = findChildSpans( + events, + "mistral.batch", + operation?.span.id, + )[0]; + const children = findChildSpans( + events, + "mistral.chat.complete", + task?.span.id, + ); + + expect(operation).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(task?.span.type).toBe("task"); + expect(task?.row.metadata).toMatchObject({ provider: "mistral" }); + expect(children).toHaveLength(2); + for (const child of children) { + expect(child.span.type).toBe("llm"); + expect(child.row.metadata).toMatchObject({ + model: CHAT_MODEL, + provider: "mistral", + }); + expect(child.input).toBeDefined(); + expect(child.output).toBeDefined(); + expect(child.metrics).toMatchObject({ + prompt_tokens: 4, + completion_tokens: 1, + tokens: 5, + }); + } + expect( + children.find((child) => Array.isArray(child.row.metadata?.tools))?.row + .metadata, + ).toMatchObject({ + tool_choice: "required", + parallel_tool_calls: false, + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get weather for a city.", + parameters: expect.any(Object), + }, + }, + { + type: "web_search", + tool_configuration: { include: ["news"] }, + }, + ], + }); + }); + + if (supportsInlineBatch) { + test("captures inline batch.jobs.create() lifecycle", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "mistral-batch-inline-operation", + ); + const task = findChildSpans( + events, + "mistral.batch", + operation?.span.id, + )[0]; + const children = findChildSpans( + events, + "mistral.chat.complete", + task?.span.id, + ); + + expect(operation).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(task?.span.type).toBe("task"); + expect(children).toHaveLength(2); + expect(children.every((child) => child.output !== undefined)).toBe( + true, + ); + }); + } + test("captures trace for chat.complete()", testConfig, () => { const root = findLatestSpan(events, ROOT_NAME); const operation = findLatestSpan( @@ -667,18 +775,14 @@ export function defineMistralInstrumentationAssertions(options: { const output = span.output as | Array<{ message?: { - toolCalls?: unknown; tool_calls?: unknown; }; }> | undefined; const firstChoice = Array.isArray(output) ? output[0] : undefined; - const toolCalls = - (Array.isArray(firstChoice?.message?.tool_calls) && - firstChoice.message.tool_calls) || - (Array.isArray(firstChoice?.message?.toolCalls) && - firstChoice.message.toolCalls) || - []; + const toolCalls = Array.isArray(firstChoice?.message?.tool_calls) + ? firstChoice.message.tool_calls + : []; return [span.span.id, toolCalls.length]; }), ); @@ -689,14 +793,10 @@ export function defineMistralInstrumentationAssertions(options: { .map((span) => { const output = span.output as | Array<{ - finishReason?: unknown; finish_reason?: unknown; }> | undefined; const firstChoice = Array.isArray(output) ? output[0] : undefined; - if (typeof firstChoice?.finishReason === "string") { - return firstChoice.finishReason; - } if (typeof firstChoice?.finish_reason === "string") { return firstChoice.finish_reason; } @@ -708,18 +808,14 @@ export function defineMistralInstrumentationAssertions(options: { const output = span.output as | Array<{ message?: { - toolCalls?: unknown; tool_calls?: unknown; }; }> | undefined; const firstChoice = Array.isArray(output) ? output[0] : undefined; - const toolCalls = - (Array.isArray(firstChoice?.message?.tool_calls) && - firstChoice.message.tool_calls) || - (Array.isArray(firstChoice?.message?.toolCalls) && - firstChoice.message.toolCalls) || - []; + const toolCalls = Array.isArray(firstChoice?.message?.tool_calls) + ? firstChoice.message.tool_calls + : []; return toolCalls .map((toolCall) => { @@ -825,7 +921,7 @@ export function defineMistralInstrumentationAssertions(options: { expect(span?.span.type).toBe("llm"); expect(metadata).toMatchObject({ provider: "mistral", - agentId: expect.any(String), + agent_id: expect.any(String), }); if (typeof metadata?.model === "string") { expect(metadata.model).toBe(AGENT_MODEL); @@ -855,18 +951,14 @@ export function defineMistralInstrumentationAssertions(options: { const output = span.output as | Array<{ message?: { - toolCalls?: unknown; tool_calls?: unknown; }; }> | undefined; const firstChoice = Array.isArray(output) ? output[0] : undefined; - const toolCalls = - (Array.isArray(firstChoice?.message?.tool_calls) && - firstChoice.message.tool_calls) || - (Array.isArray(firstChoice?.message?.toolCalls) && - firstChoice.message.toolCalls) || - []; + const toolCalls = Array.isArray(firstChoice?.message?.tool_calls) + ? firstChoice.message.tool_calls + : []; return toolCalls.length > 0; }); @@ -904,7 +996,7 @@ export function defineMistralInstrumentationAssertions(options: { expect(span?.span.type).toBe("llm"); expect(metadata).toMatchObject({ provider: "mistral", - agentId: expect.any(String), + agent_id: expect.any(String), }); if (typeof metadata?.model === "string") { expect(metadata.model).toBe(AGENT_MODEL); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs index 2506dbbd9..8b6e64872 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs @@ -1,4 +1,8 @@ -import { wrapMistral } from "braintrust"; +import { + completeMistralBatchTrace, + startMistralBatchTrace, + wrapMistral, +} from "braintrust"; import { collectAsync, runOperation, @@ -27,6 +31,7 @@ const MISTRAL_REQUEST_RETRY_OPTIONS = { const MISTRAL_THINKING_STREAM_OPTOUTS = new Set(["mistral-sdk-v1"]); const MISTRAL_CLASSIFIER_OPTOUTS = new Set(["mistral-sdk-v1"]); const MISTRAL_CLASSIFY_OPTOUTS = new Set(["mistral-sdk-v1"]); +const MISTRAL_INLINE_BATCH_OPTOUTS = new Set(["mistral-sdk-v1"]); function createMistralScenarioSpec(spec) { return { @@ -40,6 +45,9 @@ function createMistralScenarioSpec(spec) { ...(MISTRAL_CLASSIFY_OPTOUTS.has(spec.dependencyName) ? { supportsClassify: false } : {}), + ...(MISTRAL_INLINE_BATCH_OPTOUTS.has(spec.dependencyName) + ? { supportsInlineBatch: false } + : {}), }; } @@ -230,6 +238,80 @@ async function simulateToolExecutionDelay() { await new Promise((resolve) => setTimeout(resolve, TEST_TOOL_DELAY_MS)); } +function batchResult(customId, content) { + return { + custom_id: customId, + response: { + status_code: 200, + body: { + id: `mistral-batch-${customId}`, + object: "chat.completion", + model: CHAT_MODEL, + choices: [ + { + index: 0, + message: { role: "assistant", content }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 4, + completion_tokens: 1, + total_tokens: 5, + }, + }, + }, + }; +} + +function createBatchHttpClient() { + return { + async request(request) { + const requestBody = JSON.parse(await request.text()); + const requests = Array.isArray(requestBody.requests) + ? requestBody.requests + : []; + const requestCount = requests.length || 2; + const now = Math.ceil(Date.now() / 1000); + return new Response( + JSON.stringify({ + id: requests.length > 0 ? "batch-inline-e2e" : "batch-file-e2e", + object: "batch", + input_files: requestBody.input_files || [], + metadata: requestBody.metadata || null, + endpoint: requestBody.endpoint, + model: requestBody.model, + output_file: null, + error_file: null, + errors: [], + ...(requests.length > 0 + ? { + outputs: requests.map((batchRequest, index) => + batchResult( + batchRequest.custom_id, + `inline batch ${index + 1}`, + ), + ), + } + : {}), + status: "SUCCESS", + created_at: now, + total_requests: requestCount, + completed_requests: requestCount, + succeeded_requests: requestCount, + failed_requests: 0, + started_at: now, + completed_at: now, + }), + { + headers: { "Content-Type": "application/json" }, + status: 200, + }, + ); + }, + }; +} + async function createAgentViaHttp(client, apiKey) { const baseUrl = getMistralApiBaseUrl(client); const response = await withRetry( @@ -356,6 +438,8 @@ async function runMistralInstrumentationScenario( decorateClient, supportsClassifiers = true, supportsClassify = true, + supportsInlineBatch = true, + supportsSignedBatch = true, supportsThinkingStream = true, } = {}, ) { @@ -364,12 +448,116 @@ async function runMistralInstrumentationScenario( serverURL: process.env.MISTRAL_BASE_URL || process.env.MISTRAL_API_URL, }); const client = decorateClient ? decorateClient(baseClient) : baseClient; + const baseBatchClient = new Mistral({ + apiKey: "mistral-batch-e2e-key", + httpClient: createBatchHttpClient(), + }); + const batchClient = decorateClient + ? decorateClient(baseBatchClient) + : baseBatchClient; const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); const { agentId, cleanup } = await resolveAgentRuntime(baseClient); try { await runTracedScenario({ callback: async () => { + await runOperation( + "mistral-batch-file-operation", + "batch-file", + async () => { + const records = [ + { + custom_id: "file-one", + body: { + messages: [{ role: "user", content: "file batch one" }], + toolChoice: "required", + parallelToolCalls: false, + tools: [ + getWeatherToolDefinition(), + { + type: "web_search", + toolConfiguration: { include: ["news"] }, + }, + ], + }, + }, + { + custom_id: "file-two", + body: { + messages: [{ role: "user", content: "file batch two" }], + }, + }, + ]; + const input = { + files: [ + { + file: { id: "mistral-batch-input-file" }, + content: records, + }, + ], + }; + const params = supportsSignedBatch + ? await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: CHAT_MODEL, + }, + }) + : { + inputFiles: ["mistral-batch-input-file"], + endpoint: "/v1/chat/completions", + model: CHAT_MODEL, + }; + const batch = await batchClient.batch.jobs.create(params); + await completeMistralBatchTrace({ + batch, + input, + outputContent: [ + batchResult("file-two", "file batch 2"), + batchResult("file-one", "file batch 1"), + ], + }); + }, + ); + + if (supportsInlineBatch) { + await runOperation( + "mistral-batch-inline-operation", + "batch-inline", + async () => { + const input = { + requests: [ + { + customId: "inline-one", + body: { + messages: [{ role: "user", content: "inline batch one" }], + }, + }, + { + customId: "inline-two", + body: { + messages: [{ role: "user", content: "inline batch two" }], + }, + }, + ], + }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: CHAT_MODEL, + }, + }); + const batch = await batchClient.batch.jobs.create(params); + await completeMistralBatchTrace({ + batch, + input, + }); + }, + ); + } + await runOperation( "mistral-chat-complete-operation", "chat-complete", diff --git a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs index c39fdd2dd..8a56a5a27 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs @@ -7,6 +7,8 @@ runMain(async () => runAutoMistralInstrumentation(Mistral, { supportsClassifiers: false, supportsClassify: false, + supportsInlineBatch: false, + supportsSignedBatch: false, supportsThinkingStream: false, }), ); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts index 90ed108d3..1e7b1bf8f 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts +++ b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts @@ -7,6 +7,8 @@ runMain(async () => { await runWrappedMistralInstrumentation(Mistral, { supportsClassifiers: false, supportsClassify: false, + supportsInlineBatch: false, + supportsSignedBatch: false, supportsThinkingStream: false, }); }); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.test.ts b/e2e/scenarios/mistral-instrumentation/scenario.test.ts index 6bd3264a2..4e1fa432a 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.test.ts +++ b/e2e/scenarios/mistral-instrumentation/scenario.test.ts @@ -52,6 +52,9 @@ describe.concurrent("variants", () => { ...(scenario.supportsClassify === false ? { supportsClassify: false } : {}), + ...(scenario.supportsInlineBatch === false + ? { supportsInlineBatch: false } + : {}), testFileUrl: import.meta.url, timeoutMs: MISTRAL_SCENARIO_TIMEOUT_MS, }); @@ -81,6 +84,9 @@ describe.concurrent("variants", () => { ...(scenario.supportsClassify === false ? { supportsClassify: false } : {}), + ...(scenario.supportsInlineBatch === false + ? { supportsInlineBatch: false } + : {}), testFileUrl: import.meta.url, timeoutMs: MISTRAL_SCENARIO_TIMEOUT_MS, }); diff --git a/js/src/exports.ts b/js/src/exports.ts index 4072fcc54..8f8e7f372 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -129,6 +129,27 @@ export { export { registerSandbox } from "./sandbox"; +export { + completeMistralBatchTrace, + failMistralBatchTrace, + startMistralBatchTrace, +} from "./mistral-batch"; +export type { + CompleteMistralBatchTraceArgs, + FailMistralBatchTraceArgs, + MistralBatchCreateInputParams, + MistralBatchCreateParams, + MistralBatchCollectionContext, + MistralBatchFileLike, + MistralBatchInput, + MistralBatchInputFileContent, + MistralBatchJSONLSource, + MistralBatchLike, + MistralBatchTraceCollection, + MistralBatchRequestLike, + StartMistralBatchTraceArgs, +} from "./mistral-batch-types"; + // Internal isomorph layer for platform-specific implementations import _internalIso from "./isomorph"; export { _internalIso }; diff --git a/js/src/instrumentation/plugins/mistral-batch-constants.ts b/js/src/instrumentation/plugins/mistral-batch-constants.ts new file mode 100644 index 000000000..1c3381866 --- /dev/null +++ b/js/src/instrumentation/plugins/mistral-batch-constants.ts @@ -0,0 +1 @@ +export const BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY = "braintrust.batch_context"; diff --git a/js/src/instrumentation/plugins/mistral-batch-instrumentation.test.ts b/js/src/instrumentation/plugins/mistral-batch-instrumentation.test.ts new file mode 100644 index 000000000..e262e8903 --- /dev/null +++ b/js/src/instrumentation/plugins/mistral-batch-instrumentation.test.ts @@ -0,0 +1,1389 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { + _exportsForTestingOnly, + _internalGetGlobalState, + Attachment, + initLogger, + startSpan, + withCurrent, +} from "../../logger"; +import { configureNode } from "../../node/config"; +import { + completeMistralBatchTrace, + failMistralBatchTrace, + startMistralBatchTrace, +} from "../../mistral-batch"; +import { BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY } from "./mistral-batch-constants"; +import { MistralPlugin } from "./mistral-plugin"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +const requests = [ + { + customId: "one", + body: { messages: [{ role: "user", content: "first" }] }, + }, + { + customId: "two", + body: { messages: [{ role: "user", content: "second" }] }, + }, +]; + +function success(customId: string, content: string) { + return { + custom_id: customId, + response: { + status_code: 200, + body: { + id: `response-${customId}`, + model: "mistral-small-latest", + choices: [{ message: { role: "assistant", content } }], + usage: { + prompt_tokens: 2, + completion_tokens: 1, + total_tokens: 3, + }, + }, + }, + }; +} + +describe("Mistral Batch instrumentation", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + let plugin: MistralPlugin; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + _internalGetGlobalState()._internalSetTraceContextSigningSecret( + "mistral-batch-test-secret", + ); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "mistral-batch-instrumentation.test.ts", + projectId: "test-project-id", + }); + plugin = new MistralPlugin(); + plugin.enable(); + }); + + afterEach(() => { + plugin.disable(); + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + it("prepares inline parameters and starts pending spans", async () => { + const sourceParams = { + endpoint: "/v1/chat/completions" as const, + model: "mistral-small-latest", + metadata: { owner: "sdk" }, + }; + const params = await startMistralBatchTrace({ + input: { requests }, + params: sourceParams, + }); + + expect(params).toMatchObject({ + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + requests, + metadata: { + owner: "sdk", + [BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY]: expect.any(String), + }, + }); + expect(sourceParams.metadata).toEqual({ owner: "sdk" }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + const task = rows.find( + (row) => row.span_attributes?.name === "mistral.batch", + ); + const children = rows.filter( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + expect(task).toMatchObject({ + metadata: { provider: "mistral" }, + span_attributes: { type: "task" }, + span_parents: [], + }); + expect(task?.input).toBeUndefined(); + expect(task?.metrics).not.toHaveProperty("end"); + expect(children).toHaveLength(2); + expect(children.map((row) => row.input[0].content).sort()).toEqual([ + "first", + "second", + ]); + expect(children[0].metadata).toMatchObject({ + model: "mistral-small-latest", + provider: "mistral", + }); + expect(children[0].metrics).not.toHaveProperty("end"); + }); + + it("traces agent batches and records the resolved response model", async () => { + const input = { requests: [requests[0]] }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + agentId: "agent-123", + }, + }); + + expect(params).toMatchObject({ + agentId: "agent-123", + metadata: { + [BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY]: expect.any(String), + }, + }); + const serializedContext = JSON.parse( + params.metadata?.[BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY] ?? "{}", + ); + expect(serializedContext).toMatchObject({ agentId: "agent-123" }); + expect(serializedContext).not.toHaveProperty("model"); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + const startedChild = startedRows.find( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + expect(startedChild?.metadata).toMatchObject({ + agent_id: "agent-123", + provider: "mistral", + }); + expect(startedChild?.metadata).not.toHaveProperty("model"); + + await completeMistralBatchTrace({ + batch: { + id: "batch-agent", + endpoint: params.endpoint, + agentId: params.agentId, + model: null, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + outputs: [success("one", "A")], + }, + input, + }); + + const completedRows = (await backgroundLogger.drain()) as Array< + Record + >; + expect(completedRows).toHaveLength(2); + expect( + completedRows.find((row) => row.id === startedChild?.id)?.metadata, + ).toMatchObject({ model: "mistral-small-latest" }); + }); + + it("preserves propagated events on signed and collect-only batch spans", async () => { + const assertPropagation = (rows: Array>): void => { + const batchRows = rows.filter((row) => + ["mistral.batch", "mistral.chat.complete"].includes( + row.span_attributes?.name, + ), + ); + expect(batchRows).toHaveLength(2); + expect( + batchRows.every( + (row) => row.metadata?.propagated === "batch-descendant", + ), + ).toBe(true); + expect(batchRows.every((row) => row.tags?.includes("batch-tag"))).toBe( + true, + ); + }; + const input = { requests: [requests[0]] }; + + const signedParent = startSpan({ + name: "signed-parent", + propagatedEvent: { + metadata: { propagated: "batch-descendant" }, + tags: ["batch-tag"], + }, + }); + await withCurrent(signedParent, () => + startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }), + ); + signedParent.end(); + assertPropagation( + (await backgroundLogger.drain()) as Array>, + ); + + const collectParent = startSpan({ + name: "collect-parent", + propagatedEvent: { + metadata: { propagated: "batch-descendant" }, + tags: ["batch-tag"], + }, + }); + await withCurrent(collectParent, () => + completeMistralBatchTrace({ + batch: { + id: "batch-propagated-collect-only", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 1, + outputs: [success("one", "A")], + }, + input, + }), + ); + collectParent.end(); + assertPropagation( + (await backgroundLogger.drain()) as Array>, + ); + }); + + it("removes signed context and closes initialized spans when startup fails", async () => { + const recursiveMessages: any[] = []; + recursiveMessages.push(recursiveMessages); + Object.defineProperty(recursiveMessages, "toJSON", { + value: () => [{ role: "user", content: "serializable" }], + }); + const traceErrors: Error[] = []; + + const params = await startMistralBatchTrace({ + input: { + requests: [ + requests[0], + { + customId: "startup-failure", + body: { messages: recursiveMessages }, + }, + ], + }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(params.metadata).toBeUndefined(); + expect(traceErrors).toHaveLength(1); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect(rows.every((row) => typeof row.error === "string")).toBe(true); + }); + + it("uses provider JSON semantics when correlating inline inputs", async () => { + const params = await startMistralBatchTrace({ + input: { + requests: [ + { + customId: "json-semantics", + body: { + messages: [{ role: "user", content: "hello" }], + optional: undefined, + stop: [undefined, "done"], + }, + }, + ], + }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await completeMistralBatchTrace({ + batch: { + id: "batch-json-semantics", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }, + outputContent: [success("json-semantics", "Hi")], + input: { + requests: [ + { + customId: "json-semantics", + body: { + messages: [{ role: "user", content: "hello" }], + stop: [null, "done"], + }, + }, + ], + }, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("uses locale-independent key ordering for input correlation", async () => { + const startLocaleCompare = vi + .spyOn(String.prototype, "localeCompare") + .mockImplementation(() => { + throw new Error("locale-dependent comparison was used"); + }); + let params: Awaited>; + try { + params = await startMistralBatchTrace({ + input: { + requests: [ + { + customId: "unicode-keys", + body: { + messages: [{ role: "user", content: "hello" }], + z: 1, + ä: 2, + }, + }, + ], + }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + } finally { + startLocaleCompare.mockRestore(); + } + await backgroundLogger.drain(); + + const completeLocaleCompare = vi + .spyOn(String.prototype, "localeCompare") + .mockImplementation(() => { + throw new Error("locale-dependent comparison was used"); + }); + try { + await completeMistralBatchTrace({ + batch: { + id: "batch-unicode-keys", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }, + outputContent: [success("unicode-keys", "Hi")], + input: { + requests: [ + { + customId: "unicode-keys", + body: { + ä: 2, + z: 1, + messages: [{ role: "user", content: "hello" }], + }, + }, + ], + }, + }); + } finally { + completeLocaleCompare.mockRestore(); + } + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("accepts nullable fields from Mistral BatchJob values", async () => { + interface ProviderBatchJob { + id: string; + inputFiles: string[]; + metadata?: Record | null; + endpoint: string; + model?: string | null; + outputs?: Array> | null; + outputFile?: string | null; + errorFile?: string | null; + status: string; + createdAt: number; + totalRequests: number; + completedRequests: number; + succeededRequests: number; + failedRequests: number; + } + + const input = { requests: [requests[0]] }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + const batch: ProviderBatchJob = { + id: "batch-nullable-provider-fields", + inputFiles: [], + metadata: params.metadata, + endpoint: params.endpoint, + model: null, + outputs: null, + outputFile: "provider-output-file-id", + errorFile: null, + status: "SUCCESS", + createdAt: 1, + totalRequests: 1, + completedRequests: 1, + succeededRequests: 1, + failedRequests: 0, + }; + const completed = await completeMistralBatchTrace({ + batch, + input, + outputContent: [success("one", "A")], + }); + + expect(completed.batch).toBe(batch); + expect(completed.collectionContext).toEqual({ + version: 1, + value: expect.any(String), + }); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("keeps explicit batch tracing enabled when the Mistral plugin is disabled", async () => { + plugin.disable(); + + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + + expect(params.metadata).toHaveProperty( + BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY, + ); + expect(await backgroundLogger.drain()).toHaveLength(3); + }); + + it("supports multiple input files and preserves provider-ready ids", async () => { + const params = await startMistralBatchTrace({ + input: { + files: [ + { + file: { id: "file-one" }, + content: `${JSON.stringify({ custom_id: "one", body: requests[0].body })}\n`, + }, + { + file: { id: "file-two" }, + content: [{ custom_id: "two", body: requests[1].body }], + }, + ], + }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + + expect(params.inputFiles).toEqual(["file-one", "file-two"]); + expect(params.requests).toBeUndefined(); + expect(await backgroundLogger.drain()).toHaveLength(3); + }); + + it("reopens one-shot input sources across the batch lifecycle", async () => { + let sourceReads = 0; + const input = { + files: [ + { + file: { id: "response-file" }, + content: () => { + sourceReads++; + return new Response( + requests + .map(({ customId, body }) => + JSON.stringify({ custom_id: customId, body }), + ) + .join("\n"), + ); + }, + }, + ], + }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await completeMistralBatchTrace({ + batch: { + id: "batch-response-input", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 2, + inputFiles: ["response-file"], + metadata: params.metadata, + }, + outputContent: [success("one", "A"), success("two", "B")], + input, + }); + + expect(sourceReads).toBeGreaterThanOrEqual(2); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("captures canonical tool configuration on batch children", async () => { + await startMistralBatchTrace({ + input: { + requests: [ + { + customId: "tool-request", + body: { + messages: [{ role: "user", content: "look this up" }], + toolChoice: "any", + parallelToolCalls: false, + tools: [ + { + type: "function", + function: { + name: "lookup", + description: "Look something up", + parameters: { type: "object" }, + }, + }, + ], + }, + }, + ], + }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + const child = rows.find( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + expect(child?.metadata).toMatchObject({ + tool_choice: "required", + parallel_tool_calls: false, + tools: [ + { + type: "function", + function: { + name: "lookup", + description: "Look something up", + parameters: { type: "object" }, + }, + }, + ], + }); + }); + + it("completes signed spans out of order and returns the exact batch", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + const startedChildren = startedRows.filter( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + const childIds = new Set(startedChildren.map((row) => row.id)); + const taskId = startedRows.find( + (row) => row.span_attributes?.name === "mistral.batch", + )?.id; + const batch = { + id: "batch-signed", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 2, + metadata: params.metadata, + outputs: [success("two", "B"), success("one", "A")], + }; + + const returned = await completeMistralBatchTrace({ + batch, + input: { requests }, + }); + expect(returned.batch).toBe(batch); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + const children = rows.filter((row) => childIds.has(row.id)); + expect(children).toHaveLength(2); + expect(children.every((row) => row.input === undefined)).toBe(true); + expect(children.every((row) => row.span_attributes === undefined)).toBe( + true, + ); + for (const child of children) { + const started = startedChildren.find((row) => row.id === child.id); + expect(child.root_span_id).toBe(started?.root_span_id); + expect(child.span_parents).toEqual(started?.span_parents); + } + expect(children.map((row) => row.output[0].message.content).sort()).toEqual( + ["A", "B"], + ); + expect(children[0].metrics).toMatchObject({ + prompt_tokens: 2, + completion_tokens: 1, + tokens: 3, + end: expect.any(Number), + }); + expect(rows.find((row) => row.id === taskId)?.metrics).toHaveProperty( + "end", + ); + }); + + it("resumes signed attachment spans without rewriting their input", async () => { + const dataUrl = "data:image/png;base64,aGVsbG8gd29ybGQ="; + const input = { + requests: [ + { + customId: "attachment-input", + body: { + messages: [ + { + role: "user", + content: [{ type: "image_url", image_url: { url: dataUrl } }], + }, + ], + }, + }, + ], + }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + const startedChild = startedRows.find( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + const attachment = startedChild?.input?.[0]?.content?.[0]?.image_url?.url; + expect(attachment).toBeInstanceOf(Attachment); + + const batch = { + id: "batch-signed-attachment-retry", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }; + const outputContent = [success("attachment-input", "A")]; + for (let retry = 0; retry < 2; retry++) { + await completeMistralBatchTrace({ batch, input, outputContent }); + const rows = (await backgroundLogger.drain()) as Array< + Record + >; + const child = rows.find((row) => row.id === startedChild?.id); + expect(child?.input).toBeUndefined(); + expect(child?.span_attributes).toBeUndefined(); + expect(child?.metrics?.end).toEqual(expect.any(Number)); + } + }); + + it("ends pending spans when submission fails", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await failMistralBatchTrace({ + params, + input: { requests }, + error: new Error("submission rejected"), + }); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect(rows.every((row) => row.error === "submission rejected")).toBe(true); + }); + + it("reports mismatched file ids when submission fails", async () => { + const input = { + files: [ + { + file: { id: "expected-file" }, + content: [ + { custom_id: requests[0].customId, body: requests[0].body }, + ], + }, + ], + }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + const traceErrors: Error[] = []; + + await failMistralBatchTrace({ + params: { ...params, inputFiles: ["different-file"] }, + input, + error: new Error("submission rejected"), + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(traceErrors).toHaveLength(1); + expect(traceErrors[0].message).toContain("file ids do not match"); + expect(await backgroundLogger.drain()).toHaveLength(0); + }); + + it("closes failed batches despite a provider request-count mismatch", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + const idsByInput = new Map( + startedRows + .filter((row) => row.input?.[0]?.content) + .map((row) => [row.input[0].content, row.id]), + ); + const taskId = startedRows.find( + (row) => row.span_attributes?.name === "mistral.batch", + )?.id; + + await completeMistralBatchTrace({ + batch: { + id: "batch-failed-count-mismatch", + endpoint: params.endpoint, + model: params.model, + status: "FAILED", + totalRequests: 1, + metadata: params.metadata, + }, + outputContent: [success("one", "A")], + input: { requests }, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect(rows.find((row) => row.id === taskId)?.error).toBe( + "Mistral Batch FAILED", + ); + expect(rows.find((row) => row.id === idsByInput.get("second"))?.error).toBe( + "Mistral Batch FAILED", + ); + expect( + rows.find((row) => row.id === idsByInput.get("first"))?.output?.[0] + ?.message?.content, + ).toBe("A"); + }); + + it("closes every signed child when failed input reconstruction is partial", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + const startedRows = (await backgroundLogger.drain()) as Array< + Record + >; + const taskId = startedRows.find( + (row) => row.span_attributes?.name === "mistral.batch", + )?.id; + + await completeMistralBatchTrace({ + batch: { + id: "batch-failed-partial-input", + endpoint: params.endpoint, + model: params.model, + status: "FAILED", + totalRequests: 2, + metadata: params.metadata, + }, + outputContent: [success("one", "A")], + input: { + requests: [requests[0], { customId: "", body: requests[1].body }], + }, + }); + + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect(rows.find((row) => row.id === taskId)?.error).toBe( + "Mistral Batch FAILED", + ); + expect( + rows.some( + (row) => row.output?.[0]?.message?.content === "A" && !row.error, + ), + ).toBe(true); + }); + + it("collects a terminal batch without prior signed metadata", async () => { + const batch = { + id: "batch-collect-only", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 2, + createdAt: 1000, + startedAt: 1500, + completedAt: 2000, + }; + + await completeMistralBatchTrace({ + batch, + input: { requests }, + outputContent: [success("one", "A"), success("two", "B")], + }); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect( + rows.filter( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ), + ).toHaveLength(2); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + expect( + rows.find((row) => row.span_attributes?.name === "mistral.batch")?.metrics + ?.start, + ).toBe(1000); + expect( + rows + .filter((row) => row.span_attributes?.name === "mistral.chat.complete") + .every((row) => row.metrics?.start === 1500), + ).toBe(true); + expect(rows.every((row) => row.metrics?.end === 2000)).toBe(true); + }); + + it("rejects mismatched collect-only terminal batches", async () => { + const countErrors: Error[] = []; + await completeMistralBatchTrace({ + batch: { + id: "batch-collect-only-failed-count-mismatch", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "FAILED", + totalRequests: 1, + }, + input: { requests }, + onTraceError: (error) => { + countErrors.push(error); + }, + }); + + expect(countErrors.at(-1)?.message).toContain("request count"); + expect(await backgroundLogger.drain()).toHaveLength(0); + + const fileErrors: Error[] = []; + await completeMistralBatchTrace({ + batch: { + id: "batch-collect-only-cancelled-file-mismatch", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "CANCELLED", + totalRequests: 1, + inputFiles: ["different-file"], + }, + input: { + files: [ + { + file: { id: "expected-file" }, + content: [ + { custom_id: requests[0].customId, body: requests[0].body }, + ], + }, + ], + }, + onTraceError: (error) => { + fileErrors.push(error); + }, + }); + + expect(fileErrors.at(-1)?.message).toContain("input files"); + expect(await backgroundLogger.drain()).toHaveLength(0); + }); + + it("falls back to collect-only for invalid tracing metadata", async () => { + await completeMistralBatchTrace({ + batch: { + id: "batch-invalid", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 2, + metadata: { [BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY]: "tampered" }, + }, + outputContent: [success("one", "A"), success("two", "B")], + input: { requests }, + }); + expect(await backgroundLogger.drain()).toHaveLength(3); + }); + + it("returns context that keeps collect-only retries on the same trace", async () => { + const batch = { + id: "batch-collect-only-retry", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 2, + }; + const outputContent = [success("one", "A"), success("two", "B")]; + + let collectionContext: + | Awaited< + ReturnType + >["collectionContext"] + | undefined; + const collectUnderParent = async (name: string) => { + const parent = startSpan({ name }); + const result = await withCurrent(parent, () => + completeMistralBatchTrace({ + batch, + input: { requests }, + ...(collectionContext ? { collectionContext } : {}), + outputContent, + }), + ); + collectionContext ??= result.collectionContext; + parent.end(); + const rows = (await backgroundLogger.drain()) as Array< + Record + >; + return rows + .filter((row) => + ["mistral.batch", "mistral.chat.complete"].includes( + row.span_attributes?.name, + ), + ) + .map((row) => ({ + id: row.id, + rootSpanId: row.root_span_id, + spanParents: row.span_parents, + })) + .sort((left, right) => left.id.localeCompare(right.id)); + }; + + expect(await collectUnderParent("first-collection-parent")).toEqual( + await collectUnderParent("second-collection-parent"), + ); + expect(collectionContext).toEqual({ + version: 1, + value: expect.any(String), + }); + }); + + it("uses stable collect-only attachment identities for inputs and outputs", async () => { + const dataUrl = "data:image/png;base64,aGVsbG8gd29ybGQ="; + const input = { + requests: [ + { + customId: "collect-attachment", + body: { + messages: [ + { + role: "user", + content: [{ type: "image_url", image_url: { url: dataUrl } }], + }, + ], + }, + }, + ], + }; + const batch = { + id: "batch-collect-only-attachments", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 1, + }; + const outputContent = [ + { + custom_id: "collect-attachment", + response: { + status_code: 200, + body: { + choices: [ + { + message: { + role: "assistant", + content: [{ type: "image_url", image_url: { url: dataUrl } }], + }, + }, + ], + }, + }, + }, + ]; + let collectionContext: + | Awaited< + ReturnType + >["collectionContext"] + | undefined; + const collectKeys = async () => { + const result = await completeMistralBatchTrace({ + batch, + input, + ...(collectionContext ? { collectionContext } : {}), + outputContent, + }); + collectionContext ??= result.collectionContext; + const rows = (await backgroundLogger.drain()) as Array< + Record + >; + const child = rows.find( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + const inputAttachment = child?.input?.[0]?.content?.[0]?.image_url?.url; + const outputAttachment = + child?.output?.[0]?.message?.content?.[0]?.image_url?.url; + expect(inputAttachment).toBeInstanceOf(Attachment); + expect(outputAttachment).toBeInstanceOf(Attachment); + return { + input: inputAttachment.reference.key, + output: outputAttachment.reference.key, + }; + }; + + const first = await collectKeys(); + const second = await collectKeys(); + expect(second).toEqual(first); + expect(first.output).not.toBe(first.input); + expect(first.input).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/, + ); + expect(first.output).toMatch( + /^[0-9a-f]{8}-[0-9a-f]{4}-4[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/, + ); + }); + + it("uses a new trace for independent collections without retry context", async () => { + const batch = { + id: "batch-collect-only-parent-scope", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 2, + }; + const outputContent = [success("one", "A"), success("two", "B")]; + const collectIds = async (name: string) => { + const parent = startSpan({ name }); + await withCurrent(parent, () => + completeMistralBatchTrace({ + batch, + input: { requests }, + outputContent, + }), + ); + parent.end(); + return ((await backgroundLogger.drain()) as Array>) + .filter((row) => row.span_attributes?.name === "mistral.batch") + .map((row) => row.id); + }; + + expect(await collectIds("first-parent")).not.toEqual( + await collectIds("second-parent"), + ); + }); + + it("leaves a successful batch pending when choices are missing or malformed", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await completeMistralBatchTrace({ + batch: { + id: "batch-malformed-result", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 2, + metadata: params.metadata, + }, + outputContent: [ + { + custom_id: "one", + response: { status_code: 200, body: { id: "response-one" } }, + }, + { + custom_id: "two", + response: { status_code: 200, body: { choices: {} } }, + }, + ], + input: { requests }, + }); + + expect(await backgroundLogger.drain()).toHaveLength(0); + }); + + it("completes a fully correlated batch despite unrelated result records", async () => { + const params = await startMistralBatchTrace({ + input: { requests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + const traceErrors: Error[] = []; + + await completeMistralBatchTrace({ + batch: { + id: "batch-extra-results", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 2, + metadata: params.metadata, + }, + input: { requests }, + outputContent: [ + success("unknown", "ignored"), + { malformed: true }, + success("one", "A"), + success("one", "duplicate"), + success("two", "B"), + ], + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(traceErrors).toHaveLength(1); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("creates pending collect-only children for missing success results", async () => { + const traceErrors: Error[] = []; + await completeMistralBatchTrace({ + batch: { + id: "batch-collect-only-incomplete", + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + status: "SUCCESS", + totalRequests: 2, + }, + input: { requests }, + outputContent: [success("one", "A")], + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(traceErrors.at(-1)?.message).toContain("results are incomplete"); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows).toHaveLength(3); + const task = rows.find( + (row) => row.span_attributes?.name === "mistral.batch", + ); + const children = rows.filter( + (row) => row.span_attributes?.name === "mistral.chat.complete", + ); + expect(task?.metrics?.end).toBeUndefined(); + expect(children).toHaveLength(2); + expect( + children.filter((row) => row.metrics?.end === undefined), + ).toHaveLength(1); + }); + + it("stops reading results once every input is correlated", async () => { + let reads = 0; + async function* outputs() { + reads++; + yield success("one", "A"); + reads++; + throw new Error("result source was read after correlation completed"); + } + + const input = { requests: [requests[0]] }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await completeMistralBatchTrace({ + batch: { + id: "batch-stop-after-correlation", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }, + outputContent: outputs(), + input, + }); + + expect(reads).toBe(1); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows.every((row) => row.metrics?.end !== undefined)).toBe(true); + }); + + it("cancels a result stream after every input is correlated", async () => { + let cancelled = false; + const encodedResult = new TextEncoder().encode( + `${JSON.stringify(success("one", "A"))}\n`, + ); + const outputFile = new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(encodedResult); + }, + cancel() { + cancelled = true; + }, + }), + ); + const input = { requests: [requests[0]] }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + await completeMistralBatchTrace({ + batch: { + id: "batch-cancel-result-stream", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }, + input, + outputContent: outputFile, + }); + + expect(cancelled).toBe(true); + }); + + it("bounds result records and reports truncated correlation", async () => { + const input = { requests: [requests[0]] }; + const params = await startMistralBatchTrace({ + input, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + }); + await backgroundLogger.drain(); + + const traceErrors: Error[] = []; + await completeMistralBatchTrace({ + batch: { + id: "batch-result-record-limit", + endpoint: params.endpoint, + model: params.model, + status: "SUCCESS", + totalRequests: 1, + metadata: params.metadata, + }, + outputContent: [ + ...Array.from({ length: 18 }, (_, index) => + success(`unknown-${index}`, "ignored"), + ), + success("one", "A"), + ], + input, + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(traceErrors.at(-1)?.message).toContain("tracing record limit"); + const rows = (await backgroundLogger.drain()) as Array>; + expect(rows.every((row) => row.metrics?.end === undefined)).toBe(true); + }); + + it("declines batches above the provider request limit without partial spans", async () => { + const oversizedRequests = Array.from({ length: 100_001 }, (_, index) => ({ + customId: `request-${index}`, + body: { messages: [{ role: "user", content: "hello" }] }, + })); + + const traceErrors: Error[] = []; + const params = await startMistralBatchTrace({ + input: { requests: oversizedRequests }, + params: { + endpoint: "/v1/chat/completions", + model: "mistral-small-latest", + }, + onTraceError: (error) => { + traceErrors.push(error); + }, + }); + + expect(params.metadata).toBeUndefined(); + expect(traceErrors[0]?.message).toContain("correlation record limit"); + expect(await backgroundLogger.drain()).toHaveLength(0); + }); +}); diff --git a/js/src/instrumentation/plugins/mistral-batch-instrumentation.ts b/js/src/instrumentation/plugins/mistral-batch-instrumentation.ts new file mode 100644 index 000000000..f8e4f9782 --- /dev/null +++ b/js/src/instrumentation/plugins/mistral-batch-instrumentation.ts @@ -0,0 +1,1922 @@ +import { debugLogger } from "../../debug-logger"; +import { + _internalGetGlobalState, + _internalResumeSpan, + _internalStartSpanWithInitialMerge, + _internalStartSpanWithInitialMergeAndParentSpanIds, + getSpanParentObject, + NOOP_SPAN, + type Span, + utf8ByteLength, + withCurrent, +} from "../../logger"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "../../span-origin"; +import { getCurrentUnixTimestamp } from "../../util"; +import { + deterministicReplacer, + SpanTypeAttribute, + isObject, +} from "../../../util/index"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { + SpanComponentsV4, + type SpanComponentsV4Data, +} from "../../../util/span_identifier_v4"; +import type { + CompleteMistralBatchTraceArgs, + FailMistralBatchTraceArgs, + MistralBatchCreateParams, + MistralBatchInput, + MistralBatchLike, + MistralBatchTraceCollection, + StartMistralBatchTraceArgs, +} from "../../mistral-batch-types"; +import { BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY } from "./mistral-batch-constants"; +import { + extractMistralChatInput, + extractMistralResponseMetadata, + normalizeMistralChatChoices, + parseMistralMetricsFromUsage, +} from "./mistral-span-data"; + +const SUPPORTED_ENDPOINT = "/v1/chat/completions"; +const TERMINAL_STATUSES = new Set([ + "SUCCESS", + "FAILED", + "TIMEOUT_EXCEEDED", + "CANCELLED", +]); +const INPUT_DIGEST_CHUNK_SIZE = 1024 * 1024; +const MAX_JSONL_LINE_LENGTH = 16 * 1024 * 1024; +const MAX_CORRELATION_REQUESTS = 100_000; +const MAX_CORRELATION_BYTES = 32 * 1024 * 1024; +const MAX_RESULT_BYTES = 512 * 1024 * 1024; +const MAX_CONTEXT_LENGTH = 2048; +const MAX_RECORDED_ISSUES = 16; + +type BatchTarget = + | { model: string; agentId?: never } + | { agentId: string; model?: never }; + +type BatchContextBase = { + version: 1; + parent: string; + inputDigest: string; + inputKind: "files" | "requests"; + inputFileIds: string[]; + endpoint: string; + requestCount: number; + startTime: number; + operationNonce: string; +}; + +type BatchContext = BatchContextBase & BatchTarget; + +type BatchBindings = Pick< + BatchContextBase, + "inputKind" | "inputFileIds" | "endpoint" +> & + BatchTarget; + +type SerializedBatchContext = Pick< + BatchContextBase, + | "version" + | "parent" + | "inputDigest" + | "requestCount" + | "startTime" + | "operationNonce" +> & + BatchTarget & { signature: string }; + +type BatchInputRecord = { + customId: string; + ordinal: number; + spanData?: ReturnType; +}; + +type InputData = { + inputDigest: string; + inputs: Map; + inputKind: BatchContext["inputKind"]; + inputFileIds: string[]; + issues: Error[]; + requestCount: number; +}; + +type PreparedBatch = { + context: BatchContext; + inputs: Map; + params: MistralBatchCreateParams; + startParent: string; +}; + +type CollectOnlyBatch = { + childStartTime: number; + context: BatchContext; + startParent: string; +}; + +type BatchResultRecord = { + value: Record; + source: "output" | "error"; +}; + +type ResultConsumption = { + bytes: number; + limitIssue?: Error; + records: number; + maxRecords: number; +}; + +function read(value: unknown, key: PropertyKey): unknown { + if (!isObject(value)) { + return undefined; + } + try { + return Reflect.get(value, key); + } catch { + return undefined; + } +} + +function isBatchRecordIterable( + value: unknown, +): value is Iterable | AsyncIterable { + if ( + value === null || + (typeof value !== "object" && typeof value !== "function") + ) { + return false; + } + try { + return ( + typeof Reflect.get(value, Symbol.iterator) === "function" || + typeof Reflect.get(value, Symbol.asyncIterator) === "function" + ); + } catch { + return false; + } +} + +function validCustomId(value: unknown): value is string { + return typeof value === "string" && value.length > 0 && value.length <= 64; +} + +function logBatchInstrumentationError(context: string, error: unknown): void { + debugLogger.debug(`Mistral Batch instrumentation ${context}:`, error); +} + +function reportBatchTraceError( + onTraceError: ((error: Error) => void | Promise) | undefined, + context: string, + error: Error, +): void { + logBatchInstrumentationError(context, error); + if (!onTraceError) { + return; + } + try { + void Promise.resolve(onTraceError(error)).catch((callbackError) => + logBatchInstrumentationError( + "trace error callback rejected", + callbackError, + ), + ); + } catch (callbackError) { + logBatchInstrumentationError("trace error callback threw", callbackError); + } +} + +function recordIssue(issues: Error[], issue: Error): void { + if (issues.length < MAX_RECORDED_ISSUES) { + issues.push(issue); + } +} + +function batchContextPayload(context: BatchContext): string { + return JSON.stringify([BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY, context]); +} + +async function importBatchSigningKey(secret: string) { + return await globalThis.crypto.subtle.importKey( + "raw", + new TextEncoder().encode(secret), + { name: "HMAC", hash: "SHA-256" }, + false, + ["sign", "verify"], + ); +} + +async function signBatchContext( + context: BatchContext, + secret: string, +): Promise { + try { + const signature = await globalThis.crypto.subtle.sign( + "HMAC", + await importBatchSigningKey(secret), + new TextEncoder().encode(batchContextPayload(context)), + ); + return digestHex(new Uint8Array(signature), 32); + } catch { + return undefined; + } +} + +function signatureBytes(value: string): Uint8Array | undefined { + if (!/^[0-9a-f]{64}$/.test(value)) { + return undefined; + } + return Uint8Array.from(value.match(/.{2}/g) ?? [], (byte) => + Number.parseInt(byte, 16), + ); +} + +async function verifyBatchContext( + context: BatchContext, + signature: string, +): Promise { + const secret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + const bytes = signatureBytes(signature); + if (!secret || !bytes) { + return false; + } + try { + return await globalThis.crypto.subtle.verify( + "HMAC", + await importBatchSigningKey(secret), + bytes, + new TextEncoder().encode(batchContextPayload(context)), + ); + } catch { + return false; + } +} + +function parseSerializedContext( + value: unknown, +): SerializedBatchContext | undefined { + if (typeof value !== "string" || value.length > MAX_CONTEXT_LENGTH) { + return undefined; + } + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + return undefined; + } + if ( + !isObject(parsed) || + parsed.version !== 1 || + typeof parsed.parent !== "string" || + typeof parsed.inputDigest !== "string" || + !/^[0-9a-f]{64}$/.test(parsed.inputDigest) || + !( + (typeof parsed.model === "string" && + parsed.model.length > 0 && + parsed.agentId === undefined) || + (typeof parsed.agentId === "string" && + parsed.agentId.length > 0 && + parsed.model === undefined) + ) || + typeof parsed.startTime !== "number" || + !Number.isFinite(parsed.startTime) || + typeof parsed.requestCount !== "number" || + !Number.isInteger(parsed.requestCount) || + parsed.requestCount < 0 || + parsed.requestCount > MAX_CORRELATION_REQUESTS || + typeof parsed.operationNonce !== "string" || + !/^[0-9a-f]{32}$/.test(parsed.operationNonce) || + typeof parsed.signature !== "string" + ) { + return undefined; + } + let target: BatchTarget; + if (typeof parsed.agentId === "string") { + target = { agentId: parsed.agentId }; + } else if (typeof parsed.model === "string") { + target = { model: parsed.model }; + } else { + return undefined; + } + return { + version: 1, + parent: parsed.parent, + inputDigest: parsed.inputDigest, + ...target, + requestCount: parsed.requestCount, + startTime: parsed.startTime, + operationNonce: parsed.operationNonce, + signature: parsed.signature, + }; +} + +async function parseBatchContext( + value: unknown, + bindings: BatchBindings, +): Promise { + const parsed = parseSerializedContext(value); + if (!parsed) { + return undefined; + } + const context: BatchContext = { + version: 1, + parent: parsed.parent, + inputDigest: parsed.inputDigest, + ...bindings, + requestCount: parsed.requestCount, + startTime: parsed.startTime, + operationNonce: parsed.operationNonce, + }; + if (parsed.model !== context.model || parsed.agentId !== context.agentId) { + return undefined; + } + if (!(await verifyBatchContext(context, parsed.signature))) { + return undefined; + } + try { + SpanComponentsV4.fromStr(context.parent); + } catch { + return undefined; + } + return context; +} + +async function exportParent(parent: ReturnType) { + if ("toStr" in parent && typeof parent.toStr === "function") { + return parent.toStr(); + } + return await parent.export(); +} + +function exportMinimalParent(exported: string): string | undefined { + if (!exported) { + return undefined; + } + const parsed = SpanComponentsV4.fromStr(exported).data; + const projectId = read(parsed.compute_object_metadata_args, "project_id"); + const projectName = read(parsed.compute_object_metadata_args, "project_name"); + const computeObjectMetadataArgs = + typeof projectId === "string" || typeof projectName === "string" + ? { + ...(typeof projectId === "string" ? { project_id: projectId } : {}), + ...(typeof projectName === "string" + ? { project_name: projectName } + : {}), + } + : undefined; + if (!parsed.object_id && !computeObjectMetadataArgs) { + return undefined; + } + + const routing = parsed.object_id + ? { object_id: parsed.object_id } + : { compute_object_metadata_args: computeObjectMetadataArgs ?? {} }; + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + return new SpanComponentsV4({ + object_type: parsed.object_type, + ...routing, + ...(parsed.row_id && parsed.span_id && parsed.root_span_id + ? { + row_id: parsed.row_id, + span_id: parsed.span_id, + root_span_id: parsed.root_span_id, + } + : {}), + } as SpanComponentsV4Data).toStr(); +} + +function batchAttachmentKey( + rowId: string, + kind: "input" | "output", + index: number, +): string { + const [timeLow, timeMid, timeHigh, clockSeq, node] = rowId.split("-"); + const suffix = BigInt(`0x${node}`); + const domain = kind === "output" ? 1n << 47n : 0n; + const stableSuffix = (suffix ^ domain ^ BigInt(index + 1)) + .toString(16) + .padStart(12, "0"); + const variant = ((Number.parseInt(clockSeq[0], 16) & 0x3) | 0x8).toString(16); + return `${timeLow}-${timeMid}-4${timeHigh.slice(1)}-${variant}${clockSeq.slice(1)}-${stableSuffix}`; +} + +async function deterministicDigest(namespace: string, ...parts: string[]) { + const encoded = new TextEncoder().encode( + [namespace, ...parts].map((part) => `${part.length}:${part}`).join("\0"), + ); + return new Uint8Array( + await globalThis.crypto.subtle.digest("SHA-256", encoded), + ); +} + +function digestHex(bytes: Uint8Array, length: number): string { + return Array.from(bytes.slice(0, length)) + .map((byte) => byte.toString(16).padStart(2, "0")) + .join(""); +} + +function digestUuid(bytes: Uint8Array): string { + const hex = digestHex(bytes, 16); + return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}`; +} + +function randomBatchNonce(): string { + return digestHex(globalThis.crypto.getRandomValues(new Uint8Array(16)), 16); +} + +async function batchSpanIds(operationNonce: string): Promise<{ + rowId: string; + spanId: string; + rootSpanId: string; +}> { + const [row, span, root] = await Promise.all([ + deterministicDigest("mistral:batch:row", operationNonce), + deterministicDigest("mistral:batch:span", operationNonce), + deterministicDigest("mistral:batch:root", operationNonce), + ]); + return { + rowId: digestUuid(row), + spanId: digestHex(span, 8), + rootSpanId: digestHex(root, 16), + }; +} + +async function childSpanIds(operationNonce: string, ordinal: number) { + const [row, span] = await Promise.all([ + deterministicDigest( + "mistral:batch:child:row", + operationNonce, + String(ordinal), + ), + deterministicDigest( + "mistral:batch:child:span", + operationNonce, + String(ordinal), + ), + ]); + return { rowId: digestUuid(row), spanId: digestHex(span, 8) }; +} + +async function* jsonlRecords( + file: unknown, + onIssue: (error: Error) => boolean | void = () => {}, + onBytes?: (bytes: number) => boolean, +): AsyncGenerator { + const reportIssue = (error: Error) => onIssue(error) !== false; + const resolvedFile = await file; + if (resolvedFile === undefined) { + return; + } + if (typeof resolvedFile === "string") { + let offset = 0; + while (offset < resolvedFile.length) { + const newline = resolvedFile.indexOf("\n", offset); + const end = newline === -1 ? resolvedFile.length : newline; + const rawLine = resolvedFile.slice(offset, end); + offset = newline === -1 ? resolvedFile.length : newline + 1; + if ( + onBytes && + !onBytes(utf8ByteLength(rawLine) + (newline === -1 ? 0 : 1)) + ) { + return; + } + const line = rawLine.replace(/\r$/, ""); + if (!line.trim()) { + continue; + } + if (line.length > MAX_JSONL_LINE_LENGTH) { + if ( + !reportIssue( + new Error("Mistral Batch JSONL line exceeds the tracing limit"), + ) + ) { + return; + } + continue; + } + try { + yield JSON.parse(line); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + if ( + !reportIssue(new Error("Mistral Batch file contains malformed JSONL")) + ) { + return; + } + } + } + return; + } + + const body = read(resolvedFile, "body"); + const stream = isObject(body) ? body : resolvedFile; + const getReader = read(stream, "getReader"); + if (isObject(stream) && typeof getReader === "function") { + let reader: unknown; + let reachedEof = false; + try { + reader = Reflect.apply(getReader, stream, []); + if (!isObject(reader)) { + throw new Error("Batch stream returned an invalid reader"); + } + + const decoder = new TextDecoder(); + let pending = ""; + let overLimit = false; + while (true) { + const readChunk = read(reader, "read"); + if (typeof readChunk !== "function") { + throw new Error("Batch stream reader has no read method"); + } + const chunk = await Reflect.apply(readChunk, reader, []); + if (!isObject(chunk)) { + throw new Error("Batch stream returned an invalid chunk"); + } + if (chunk.done === true) { + reachedEof = true; + pending += decoder.decode(); + break; + } + if (!(chunk.value instanceof Uint8Array)) { + throw new Error("Batch stream returned a non-byte chunk"); + } + if (onBytes && !onBytes(chunk.value.byteLength)) { + return; + } + + pending += decoder.decode(chunk.value, { stream: true }); + if (pending.length > MAX_JSONL_LINE_LENGTH && !pending.includes("\n")) { + pending = ""; + overLimit = true; + if ( + !reportIssue( + new Error("Mistral Batch JSONL line exceeds the tracing limit"), + ) + ) { + return; + } + } + let newline = pending.indexOf("\n"); + while (newline !== -1) { + const line = pending.slice(0, newline).replace(/\r$/, ""); + pending = pending.slice(newline + 1); + if (overLimit) { + overLimit = false; + } else if (line.trim()) { + if (line.length > MAX_JSONL_LINE_LENGTH) { + if ( + !reportIssue( + new Error( + "Mistral Batch JSONL line exceeds the tracing limit", + ), + ) + ) { + return; + } + } else { + try { + yield JSON.parse(line); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + if ( + !reportIssue( + new Error( + "Mistral Batch response contains malformed JSONL", + ), + ) + ) { + return; + } + } + } + } + newline = pending.indexOf("\n"); + } + } + + if (!overLimit && pending.trim()) { + if (pending.length > MAX_JSONL_LINE_LENGTH) { + if ( + !reportIssue( + new Error("Mistral Batch JSONL line exceeds the tracing limit"), + ) + ) { + return; + } + } else { + try { + yield JSON.parse(pending.replace(/\r$/, "")); + } catch (error) { + logBatchInstrumentationError("skipped malformed JSONL", error); + if ( + !reportIssue( + new Error("Mistral Batch response contains malformed JSONL"), + ) + ) { + return; + } + } + } + } + } catch (error) { + logBatchInstrumentationError("could not read JSONL stream", error); + reportIssue(new Error("Mistral Batch stream could not be read")); + } finally { + const cancel = read(reader, "cancel"); + if (!reachedEof && isObject(reader) && typeof cancel === "function") { + try { + await Reflect.apply(cancel, reader, []); + } catch (error) { + logBatchInstrumentationError("could not cancel JSONL stream", error); + } + } + const releaseLock = read(reader, "releaseLock"); + if (isObject(reader) && typeof releaseLock === "function") { + try { + Reflect.apply(releaseLock, reader, []); + } catch (error) { + logBatchInstrumentationError( + "could not release stream reader", + error, + ); + } + } + } + return; + } + + if (isBatchRecordIterable(resolvedFile)) { + for await (const record of resolvedFile) { + if (onBytes) { + let serialized: string | undefined; + try { + serialized = JSON.stringify(record); + } catch (error) { + logBatchInstrumentationError("could not measure JSONL record", error); + } + if (serialized === undefined) { + if ( + !reportIssue( + new Error("Mistral Batch result contains a non-JSON record"), + ) + ) { + return; + } + continue; + } + if (!onBytes(utf8ByteLength(serialized))) { + return; + } + } + yield record; + } + } else { + logBatchInstrumentationError("skipped invalid JSONL source", resolvedFile); + reportIssue(new Error("Mistral Batch file source is invalid")); + } +} + +async function readBatchInputs(input: MistralBatchInput): Promise { + const inputs = new Map(); + const issues: Error[] = []; + const inputKind = "files" in input ? "files" : "requests"; + const inputFileIds = + "files" in input ? input.files.map(({ file }) => file.id) : []; + const chunkDigests: string[] = []; + let correlationBytes = 0; + let digestChunk = ""; + + const addRecord = async (value: unknown): Promise => { + const customId = + inputKind === "files" + ? read(value, "custom_id") + : read(value, "customId"); + const body = read(value, "body"); + if (!validCustomId(customId) || inputs.has(customId) || !isObject(body)) { + recordIssue( + issues, + new Error("Mistral Batch input contains an invalid record"), + ); + return false; + } + if (inputs.size >= MAX_CORRELATION_REQUESTS) { + recordIssue( + issues, + new Error("Mistral Batch input exceeds the correlation record limit"), + ); + return false; + } + + let canonicalRecord: string; + try { + const serializedRecord = JSON.stringify( + { body, custom_id: customId }, + deterministicReplacer, + ); + if (serializedRecord === undefined) { + throw new Error("Mistral Batch input could not be serialized"); + } + canonicalRecord = serializedRecord; + } catch (error) { + recordIssue( + issues, + error instanceof Error + ? error + : new Error("Mistral Batch input could not be canonicalized"), + ); + return false; + } + const canonicalBytes = utf8ByteLength(canonicalRecord); + if ( + canonicalBytes > MAX_JSONL_LINE_LENGTH || + correlationBytes + canonicalBytes > MAX_CORRELATION_BYTES + ) { + recordIssue( + issues, + new Error("Mistral Batch input exceeds the correlation byte limit"), + ); + return false; + } + correlationBytes += canonicalBytes; + const framedRecord = `${canonicalRecord.length}:${canonicalRecord}`; + if ( + digestChunk.length > 0 && + digestChunk.length + framedRecord.length > INPUT_DIGEST_CHUNK_SIZE + ) { + chunkDigests.push( + digestHex( + await deterministicDigest("mistral:batch:input:chunk", digestChunk), + 32, + ), + ); + digestChunk = ""; + } + digestChunk += framedRecord; + + let spanData: BatchInputRecord["spanData"]; + try { + spanData = extractMistralChatInput(body, { processAttachments: false }); + } catch (error) { + logBatchInstrumentationError("could not extract batch input", error); + } + inputs.set(customId, { customId, ordinal: inputs.size, spanData }); + return true; + }; + + try { + if ("files" in input) { + if ( + input.files.length === 0 || + inputFileIds.some((id) => typeof id !== "string" || id.length === 0) || + new Set(inputFileIds).size !== inputFileIds.length + ) { + recordIssue( + issues, + new Error("Mistral Batch input contains invalid files"), + ); + } + if (issues.length === 0) { + fileLoop: for (const { content } of input.files) { + const source = typeof content === "function" ? content() : content; + for await (const value of jsonlRecords(source, (issue) => { + recordIssue(issues, issue); + return false; + })) { + if (!(await addRecord(value))) { + break fileLoop; + } + } + } + } + } else { + for (const value of input.requests) { + if (!(await addRecord(value))) { + break; + } + } + } + } catch (error) { + logBatchInstrumentationError("could not process batch input", error); + recordIssue( + issues, + new Error("Mistral Batch input could not be processed"), + ); + } + + chunkDigests.push( + digestHex( + await deterministicDigest("mistral:batch:input:chunk", digestChunk), + 32, + ), + ); + return { + inputDigest: digestHex( + await deterministicDigest("mistral:batch:input", chunkDigests.join("")), + 32, + ), + inputs, + inputKind, + inputFileIds, + issues, + requestCount: inputs.size, + }; +} + +function providerReadyParams( + args: StartMistralBatchTraceArgs, +): MistralBatchCreateParams { + return "files" in args.input + ? { + ...args.params, + inputFiles: args.input.files.map(({ file }) => file.id), + } + : { ...args.params, requests: [...args.input.requests] }; +} + +async function prepareCreateParams( + args: StartMistralBatchTraceArgs, + parent: ReturnType, + startTime: number, +): Promise { + const params = providerReadyParams(args); + const endpoint = read(params, "endpoint"); + const model = read(params, "model"); + const agentId = read(params, "agentId"); + const validModel = typeof model === "string" && model.length > 0; + const validAgentId = typeof agentId === "string" && agentId.length > 0; + if (endpoint !== SUPPORTED_ENDPOINT || validModel === validAgentId) { + return undefined; + } + let target: BatchTarget; + if (validAgentId && typeof agentId === "string") { + target = { agentId }; + } else if (validModel && typeof model === "string") { + target = { model }; + } else { + return undefined; + } + + const metadataValue = read(params, "metadata"); + if ( + metadataValue !== undefined && + metadataValue !== null && + (!isObject(metadataValue) || Array.isArray(metadataValue)) + ) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch metadata is invalid"), + ); + return undefined; + } + const metadata = metadataValue ?? {}; + if ( + Object.prototype.hasOwnProperty.call( + metadata, + BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY, + ) + ) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch metadata contains a reserved tracing key"), + ); + return undefined; + } + + const startParent = await exportParent(parent); + const serializedParent = exportMinimalParent(startParent); + if (!serializedParent) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch tracing has no routable Braintrust parent"), + ); + return undefined; + } + const signingSecret = + _internalGetGlobalState()._internalGetTraceContextSigningSecret(); + if (!signingSecret) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch tracing has no signing secret"), + ); + return undefined; + } + + const inputData = await readBatchInputs(args.input); + if (inputData.issues.length > 0) { + reportBatchTraceError( + args.onTraceError, + "skipped invalid batch input", + inputData.issues[0], + ); + return undefined; + } + const context: BatchContext = { + version: 1, + parent: serializedParent, + inputDigest: inputData.inputDigest, + inputKind: inputData.inputKind, + inputFileIds: inputData.inputFileIds, + endpoint, + ...target, + requestCount: inputData.requestCount, + startTime, + operationNonce: randomBatchNonce(), + }; + const signature = await signBatchContext(context, signingSecret); + if (!signature) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch tracing context could not be signed"), + ); + return undefined; + } + const serialized = JSON.stringify({ + version: context.version, + parent: context.parent, + inputDigest: context.inputDigest, + ...target, + requestCount: context.requestCount, + startTime: context.startTime, + operationNonce: context.operationNonce, + signature, + } satisfies SerializedBatchContext); + if (serialized.length > MAX_CONTEXT_LENGTH) { + reportBatchTraceError( + args.onTraceError, + "skipped context injection", + new Error("Mistral Batch tracing context exceeds the tracing limit"), + ); + return undefined; + } + + return { + context, + inputs: inputData.inputs, + startParent, + params: { + ...params, + metadata: { + ...metadata, + [BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY]: serialized, + }, + }, + }; +} + +async function startBatchSpan( + context: BatchContext, + startParent = context.parent, +): Promise { + const ids = await batchSpanIds(context.operationNonce); + const parent = SpanComponentsV4.fromStr(startParent); + const hasParentSpan = Boolean( + parent.data.row_id && parent.data.span_id && parent.data.root_span_id, + ); + return withCurrent(NOOP_SPAN, () => + _internalStartSpanWithInitialMergeAndParentSpanIds( + withSpanInstrumentationName( + { + name: "mistral.batch", + type: SpanTypeAttribute.TASK, + parent: parent.toStr(), + parentSpanIds: hasParentSpan + ? { + spanId: parent.data.span_id ?? "", + rootSpanId: parent.data.root_span_id ?? "", + } + : { + parentSpanIds: [], + rootSpanId: ids.rootSpanId, + }, + spanId: ids.spanId, + startTime: context.startTime, + event: { + id: ids.rowId, + metadata: { provider: "mistral" }, + }, + }, + INSTRUMENTATION_NAMES.MISTRAL, + ), + ), + ); +} + +async function resumeBatchSpan(context: BatchContext): Promise { + const ids = await batchSpanIds(context.operationNonce); + const parent = SpanComponentsV4.fromStr(context.parent).data; + const routing = parent.object_id + ? { object_id: parent.object_id } + : { + compute_object_metadata_args: parent.compute_object_metadata_args ?? {}, + }; + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + const exported = new SpanComponentsV4({ + object_type: parent.object_type, + ...routing, + row_id: ids.rowId, + span_id: ids.spanId, + root_span_id: parent.root_span_id ?? ids.rootSpanId, + } as SpanComponentsV4Data).toStr(); + return _internalResumeSpan({ + exported, + spanParents: parent.span_id ? [parent.span_id] : [], + }); +} + +async function startBatchChild( + context: BatchContext, + taskParent: string, + input: BatchInputRecord, + startTime = context.startTime, +): Promise { + const ids = await childSpanIds(context.operationNonce, input.ordinal); + const spanInput = + input.spanData?.input === undefined + ? undefined + : processInputAttachments(input.spanData.input, { + attachmentKey: (index) => + batchAttachmentKey(ids.rowId, "input", index), + }); + return withCurrent(NOOP_SPAN, () => + _internalStartSpanWithInitialMerge( + withSpanInstrumentationName( + { + name: "mistral.chat.complete", + type: SpanTypeAttribute.LLM, + parent: taskParent, + spanId: ids.spanId, + startTime, + event: { + id: ids.rowId, + ...(spanInput !== undefined ? { input: spanInput } : {}), + metadata: { + ...input.spanData?.metadata, + ...(context.model ? { model: context.model } : {}), + ...(context.agentId ? { agent_id: context.agentId } : {}), + provider: "mistral", + }, + }, + }, + INSTRUMENTATION_NAMES.MISTRAL, + ), + ), + ); +} + +async function resumeBatchChild( + context: BatchContext, + taskParent: string, + input: BatchInputRecord, +): Promise { + const ids = await childSpanIds(context.operationNonce, input.ordinal); + const task = SpanComponentsV4.fromStr(taskParent).data; + if (!task.span_id || !task.root_span_id) { + throw new Error("Mistral Batch task span cannot be resumed"); + } + const routing = task.object_id + ? { object_id: task.object_id } + : { + compute_object_metadata_args: task.compute_object_metadata_args ?? {}, + }; + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + const exported = new SpanComponentsV4({ + object_type: task.object_type, + ...routing, + row_id: ids.rowId, + span_id: ids.spanId, + root_span_id: task.root_span_id, + } as SpanComponentsV4Data).toStr(); + return _internalResumeSpan({ exported, spanParents: [task.span_id] }); +} + +async function startPreparedBatch( + inputs: Map, + context: BatchContext, + startParent: string, +): Promise { + let task: Span | undefined; + let taskParent: string | undefined; + let startedChildCount = 0; + try { + task = await startBatchSpan(context, startParent); + taskParent = await task.export(); + for (const input of inputs.values()) { + await startBatchChild(context, taskParent, input); + startedChildCount++; + } + } catch (error) { + const initializationError = + error instanceof Error + ? error + : new Error("Mistral Batch tracing could not start spans"); + const endTime = Math.max(getCurrentUnixTimestamp(), context.startTime); + for (let ordinal = startedChildCount - 1; ordinal >= 0; ordinal--) { + try { + if (!taskParent) { + break; + } + const child = await resumeBatchChild(context, taskParent, { + customId: `startup-rollback-${ordinal}`, + ordinal, + }); + child.log({ error: initializationError }); + child.end({ endTime }); + } catch (cleanupError) { + logBatchInstrumentationError( + "could not roll back an initialized child", + cleanupError, + ); + } + } + if (task) { + try { + task.log({ error: initializationError }); + task.end({ endTime }); + } catch (cleanupError) { + logBatchInstrumentationError( + "could not roll back the initialized task", + cleanupError, + ); + } + } + throw initializationError; + } +} + +export async function startMistralBatchTraceImpl( + args: StartMistralBatchTraceArgs, +): Promise { + const parent = getSpanParentObject(); + const startTime = getCurrentUnixTimestamp(); + let prepared: PreparedBatch | undefined; + try { + prepared = await prepareCreateParams(args, parent, startTime); + } catch (error) { + reportBatchTraceError( + args.onTraceError, + "could not prepare batch", + error instanceof Error + ? error + : new Error("Mistral Batch tracing could not prepare the batch"), + ); + } + + if (!prepared) { + return providerReadyParams(args); + } + try { + await startPreparedBatch( + prepared.inputs, + prepared.context, + prepared.startParent, + ); + return prepared.params; + } catch (error) { + reportBatchTraceError( + args.onTraceError, + "could not start batch spans", + error instanceof Error + ? error + : new Error("Mistral Batch tracing could not start spans"), + ); + return providerReadyParams(args); + } +} + +function errorFromResult(result: BatchResultRecord): Error | undefined { + const error = read(result.value, "error"); + if (isObject(error)) { + const message = read(error, "message"); + return new Error( + typeof message === "string" ? message : "Mistral Batch request failed", + ); + } + const response = read(result.value, "response"); + const statusCode = + read(response, "status_code") ?? read(response, "statusCode"); + if ( + result.source === "error" || + (typeof statusCode === "number" && (statusCode < 200 || statusCode >= 300)) + ) { + const body = read(response, "body"); + const bodyError = read(body, "error"); + const message = read(bodyError, "message") ?? read(body, "message"); + return new Error( + typeof message === "string" ? message : "Mistral Batch request failed", + ); + } + return undefined; +} + +async function completeBatchResult( + context: BatchContext, + taskParent: string, + childStartTime: number, + endTime: number, + input: BatchInputRecord, + result: BatchResultRecord, + resume: boolean, + initializedOrdinals: Set, +): Promise { + const child = + resume || initializedOrdinals.has(input.ordinal) + ? await resumeBatchChild(context, taskParent, input) + : await startBatchChild(context, taskParent, input, childStartTime); + if (!resume && !initializedOrdinals.has(input.ordinal)) { + initializedOrdinals.add(input.ordinal); + } + try { + const resultError = errorFromResult(result); + const response = read(result.value, "response"); + const responseBody = read(response, "body"); + if (resultError) { + child.log({ error: resultError }); + } else if (isObject(responseBody)) { + const choices = normalizeMistralChatChoices( + read(responseBody, "choices"), + ); + if (!choices) { + return false; + } + child.log({ + output: processInputAttachments(choices, { + attachmentKey: (index) => + batchAttachmentKey(child.id, "output", index), + }), + metadata: extractMistralResponseMetadata(responseBody), + metrics: parseMistralMetricsFromUsage(read(responseBody, "usage")), + }); + } else { + return false; + } + child.end({ endTime }); + return true; + } catch (error) { + logBatchInstrumentationError("could not complete batch child", error); + return false; + } +} + +function batchBindings( + inputData: InputData, + endpoint: unknown, + model: unknown, + agentId: unknown, +): BatchBindings | undefined { + if (endpoint !== SUPPORTED_ENDPOINT) { + return undefined; + } + let target: BatchTarget; + if (typeof agentId === "string" && agentId.length > 0) { + target = { agentId }; + } else if (typeof model === "string" && model.length > 0) { + target = { model }; + } else { + return undefined; + } + return { + inputKind: inputData.inputKind, + inputFileIds: inputData.inputFileIds, + endpoint, + ...target, + }; +} + +function hasReservedContext(metadata: unknown): boolean { + return ( + isObject(metadata) && + Object.prototype.hasOwnProperty.call( + metadata, + BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY, + ) + ); +} + +async function signedContextFromBatch( + batch: MistralBatchLike, + inputData: InputData, +): Promise { + const metadata = read(batch, "metadata"); + const serialized = parseSerializedContext( + read(metadata, BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY), + ); + const providerModel = read(batch, "model"); + const providerAgentId = read(batch, "agentId"); + if ( + !serialized || + (serialized.agentId && + typeof providerAgentId === "string" && + providerAgentId !== serialized.agentId) || + (serialized.model && + typeof providerModel === "string" && + providerModel !== serialized.model) + ) { + return undefined; + } + const bindings = batchBindings( + inputData, + read(batch, "endpoint"), + serialized.model, + serialized.agentId, + ); + if (!bindings || !hasReservedContext(metadata)) { + return undefined; + } + return await parseBatchContext( + read(metadata, BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY), + bindings, + ); +} + +function validProviderTime(value: unknown, minimum = 0): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= minimum + ? value + : undefined; +} + +async function collectOnlyContext( + batch: MistralBatchLike, + inputData: InputData, + parentObject: ReturnType | string, + collectionTime: number, +): Promise { + const bindings = batchBindings( + inputData, + read(batch, "endpoint"), + read(batch, "model"), + read(batch, "agentId"), + ); + const batchId = read(batch, "id"); + if (!bindings || typeof batchId !== "string" || !batchId) { + return undefined; + } + const startParent = + typeof parentObject === "string" + ? parentObject + : await exportParent(parentObject); + const serializedParent = exportMinimalParent(startParent); + if (!serializedParent) { + return undefined; + } + const createdAt = validProviderTime(read(batch, "createdAt")); + const startTime = + createdAt !== undefined && createdAt <= collectionTime + ? createdAt + : collectionTime; + const startedAt = validProviderTime(read(batch, "startedAt"), startTime); + const childStartTime = + startedAt !== undefined && startedAt <= collectionTime + ? startedAt + : startTime; + const nonce = await deterministicDigest( + "mistral:batch:collect-only", + batchId, + bindings.endpoint, + bindings.agentId ? "agent" : "model", + bindings.agentId ?? bindings.model, + inputData.inputDigest, + serializedParent, + ); + return { + childStartTime, + startParent, + context: { + version: 1, + parent: serializedParent, + inputDigest: inputData.inputDigest, + requestCount: inputData.requestCount, + ...bindings, + startTime, + operationNonce: digestHex(nonce, 16), + }, + }; +} + +function terminalEndTime(batch: MistralBatchLike, startTime: number): number { + return ( + validProviderTime(read(batch, "completedAt"), startTime) ?? + Math.max(getCurrentUnixTimestamp(), startTime) + ); +} + +async function completeResultFile({ + childStartTime, + context, + endTime, + file, + completedOrdinals, + initializedOrdinals, + consumption, + inputs, + issues, + resume, + source, + taskParent, +}: { + childStartTime: number; + context: BatchContext; + endTime: number; + file: unknown; + completedOrdinals: Set; + initializedOrdinals: Set; + consumption: ResultConsumption; + inputs: Map; + issues: Error[]; + resume: boolean; + source: BatchResultRecord["source"]; + taskParent: string; +}): Promise { + if (inputs.size === 0) { + return; + } + try { + for await (const value of jsonlRecords( + file, + (issue) => recordIssue(issues, issue), + (bytes) => { + if (consumption.bytes + bytes > MAX_RESULT_BYTES) { + consumption.limitIssue = new Error( + "Mistral Batch results exceed the tracing byte limit", + ); + recordIssue(issues, consumption.limitIssue); + return false; + } + consumption.bytes += bytes; + return true; + }, + )) { + if (consumption.records >= consumption.maxRecords) { + consumption.limitIssue = new Error( + "Mistral Batch results exceed the tracing record limit", + ); + recordIssue(issues, consumption.limitIssue); + return; + } + consumption.records++; + try { + const customId = read(value, "custom_id") ?? read(value, "customId"); + if (!validCustomId(customId) || !isObject(value)) { + recordIssue( + issues, + new Error("Mistral Batch result is missing a valid custom id"), + ); + continue; + } + const input = inputs.get(customId); + if (!input) { + recordIssue( + issues, + new Error("Mistral Batch result does not match an input record"), + ); + continue; + } + const completed = await completeBatchResult( + context, + taskParent, + childStartTime, + endTime, + input, + { + value, + source, + }, + resume, + initializedOrdinals, + ); + if (!completed) { + recordIssue( + issues, + new Error("Mistral Batch result contains an invalid response"), + ); + continue; + } + completedOrdinals.add(input.ordinal); + inputs.delete(customId); + if (inputs.size === 0) { + return; + } + } catch (error) { + logBatchInstrumentationError("skipped invalid result", error); + recordIssue( + issues, + new Error("Mistral Batch result could not be processed"), + ); + } + } + } catch (error) { + logBatchInstrumentationError(`could not process ${source} file`, error); + recordIssue( + issues, + new Error(`Mistral Batch ${source} file could not be processed`), + ); + } +} + +function sameStrings(left: unknown, right: string[]): boolean { + return ( + Array.isArray(left) && + left.length === right.length && + left.every((value, index) => value === right[index]) + ); +} + +async function completeBatch( + args: CompleteMistralBatchTraceArgs, + parent: ReturnType, + collectionTime: number, +): Promise { + const { batch } = args; + const batchId = read(batch, "id"); + const status = read(batch, "status"); + if ( + typeof batchId !== "string" || + typeof status !== "string" || + !TERMINAL_STATUSES.has(status) + ) { + return; + } + + const inputData = await readBatchInputs(args.input); + const metadata = read(batch, "metadata"); + let context: BatchContext | undefined; + let childStartTime: number | undefined; + let startParent: string | undefined; + let resumedSignedContext = false; + if (hasReservedContext(metadata)) { + context = await signedContextFromBatch(batch, inputData); + resumedSignedContext = context !== undefined; + if (!resumedSignedContext) { + reportBatchTraceError( + args.onTraceError, + "ignored invalid signed context", + new Error("Mistral Batch tracing context verification failed"), + ); + } + } + if (!context && inputData.issues.length === 0) { + const collectOnly = await collectOnlyContext( + batch, + inputData, + args.collectionContext?.value ?? parent, + collectionTime, + ); + context = collectOnly?.context; + childStartTime = collectOnly?.childStartTime; + startParent = collectOnly?.startParent; + } + if (!context) { + if (inputData.issues.length > 0) { + reportBatchTraceError( + args.onTraceError, + "could not reconstruct batch input", + inputData.issues[0], + ); + } + return; + } + + const requiresCompleteIntegrity = + status === "SUCCESS" || !resumedSignedContext; + const integrityIssues = [...inputData.issues]; + if (resumedSignedContext && inputData.inputDigest !== context.inputDigest) { + integrityIssues.push( + new Error("Mistral Batch input does not match the signed input"), + ); + } + const batchInputFiles = read(batch, "inputFiles"); + if ( + inputData.inputKind === "files" && + !sameStrings(batchInputFiles, inputData.inputFileIds) + ) { + integrityIssues.push( + new Error("Mistral Batch input files do not match the batch input"), + ); + } + const totalRequests = read(batch, "totalRequests"); + const requestCountMismatch = resumedSignedContext + ? typeof totalRequests === "number" && + totalRequests !== inputData.requestCount + : totalRequests !== inputData.requestCount; + if (requestCountMismatch) { + const issue = new Error( + "Mistral Batch input does not match its request count", + ); + if (requiresCompleteIntegrity) { + integrityIssues.push(issue); + } else { + logBatchInstrumentationError( + "completed terminal batch despite request count mismatch", + issue, + ); + } + } + if (integrityIssues.length > 0) { + let diagnosticContext: string; + if (!resumedSignedContext) { + diagnosticContext = + "skipped collect-only batch with incomplete correlation"; + } else if (status === "SUCCESS") { + diagnosticContext = "left batch spans pending"; + } else { + diagnosticContext = + "completed terminal batch with incomplete correlation"; + } + reportBatchTraceError( + args.onTraceError, + diagnosticContext, + integrityIssues[0], + ); + if (requiresCompleteIntegrity) { + return; + } + } + + const resolvedChildStartTime = childStartTime ?? context.startTime; + const endTime = terminalEndTime( + batch, + Math.max(context.startTime, resolvedChildStartTime), + ); + const task = resumedSignedContext + ? await resumeBatchSpan(context) + : await startBatchSpan(context, startParent); + const taskParent = await task.export(); + const issues: Error[] = []; + const completedOrdinals = new Set(); + const initializedOrdinals = new Set(); + const consumption: ResultConsumption = { + bytes: 0, + records: 0, + maxRecords: Math.min( + MAX_CORRELATION_REQUESTS, + context.requestCount + MAX_RECORDED_ISSUES, + ), + }; + const canProcessResults = + integrityIssues.length === 0 || + (status !== "SUCCESS" && + resumedSignedContext && + (inputData.issues.length > 0 || + inputData.inputDigest === context.inputDigest)); + if (canProcessResults) { + await completeResultFile({ + childStartTime: resolvedChildStartTime, + context, + endTime, + file: args.outputContent ?? read(batch, "outputs"), + completedOrdinals, + initializedOrdinals, + consumption, + inputs: inputData.inputs, + issues, + resume: resumedSignedContext, + source: "output", + taskParent, + }); + await completeResultFile({ + childStartTime: resolvedChildStartTime, + context, + endTime, + file: args.errorContent, + completedOrdinals, + initializedOrdinals, + consumption, + inputs: inputData.inputs, + issues, + resume: resumedSignedContext, + source: "error", + taskParent, + }); + } + + if (status === "SUCCESS" && inputData.inputs.size > 0) { + if (!resumedSignedContext) { + for (const input of inputData.inputs.values()) { + if (initializedOrdinals.has(input.ordinal)) { + continue; + } + try { + await startBatchChild( + context, + taskParent, + input, + resolvedChildStartTime, + ); + initializedOrdinals.add(input.ordinal); + } catch (error) { + logBatchInstrumentationError( + "could not start pending collect-only child", + error, + ); + recordIssue( + issues, + new Error("Mistral Batch pending child could not be initialized"), + ); + } + } + } + reportBatchTraceError( + args.onTraceError, + "left batch spans pending", + consumption.limitIssue ?? + issues[0] ?? + new Error("Mistral Batch results are incomplete"), + ); + return; + } + + if (issues.length > 0) { + reportBatchTraceError( + args.onTraceError, + "ignored invalid batch result records", + issues[0], + ); + } + + const statusError = new Error(`Mistral Batch ${status}`); + const remainingByOrdinal = new Map( + [...inputData.inputs.values()].map((input) => [input.ordinal, input]), + ); + const closureCount = + status === "SUCCESS" || !resumedSignedContext + ? inputData.requestCount + : context.requestCount; + for (let ordinal = 0; ordinal < closureCount; ordinal++) { + if (completedOrdinals.has(ordinal)) { + continue; + } + const input = remainingByOrdinal.get(ordinal) ?? { + customId: `unresolved-${ordinal}`, + ordinal, + }; + try { + const child = + resumedSignedContext || initializedOrdinals.has(input.ordinal) + ? await resumeBatchChild(context, taskParent, input) + : await startBatchChild( + context, + taskParent, + input, + resolvedChildStartTime, + ); + child.log({ error: statusError }); + child.end({ endTime }); + } catch (error) { + logBatchInstrumentationError("could not complete missing result", error); + } + } + + if (status !== "SUCCESS") { + task.log({ error: statusError }); + } + task.end({ endTime }); +} + +export async function completeMistralBatchTraceImpl< + TBatch extends MistralBatchLike, +>( + args: CompleteMistralBatchTraceArgs, + parent: ReturnType, +): Promise> { + const collectionTime = getCurrentUnixTimestamp(); + const suppliedCollectionContext = args.collectionContext; + let collectionContext = + suppliedCollectionContext && + read(suppliedCollectionContext, "version") === 1 && + typeof read(suppliedCollectionContext, "value") === "string" + ? suppliedCollectionContext + : undefined; + if (suppliedCollectionContext && !collectionContext) { + reportBatchTraceError( + args.onTraceError, + "ignored invalid collection context", + new Error("Mistral Batch collection context is invalid"), + ); + } + if (!collectionContext) { + try { + collectionContext = { version: 1, value: await exportParent(parent) }; + } catch (error) { + reportBatchTraceError( + args.onTraceError, + "could not create collection context", + error instanceof Error + ? error + : new Error("Mistral Batch collection context could not be created"), + ); + } + } + try { + await completeBatch( + { ...args, ...(collectionContext ? { collectionContext } : {}) }, + parent, + collectionTime, + ); + } catch (error) { + reportBatchTraceError( + args.onTraceError, + "could not complete batch", + error instanceof Error + ? error + : new Error("Mistral Batch tracing could not complete the batch"), + ); + } + return { + batch: args.batch, + ...(collectionContext ? { collectionContext } : {}), + }; +} + +function normalizeError(error: unknown): Error { + if (error instanceof Error) { + return error; + } + return new Error(typeof error === "string" ? error : "Mistral Batch failed"); +} + +async function failBatch(args: FailMistralBatchTraceArgs): Promise { + const inputData = await readBatchInputs(args.input); + if (inputData.issues.length > 0) { + reportBatchTraceError( + args.onTraceError, + "could not reconstruct failed batch input", + inputData.issues[0], + ); + return; + } + const metadata = read(args.params, "metadata"); + const bindings = batchBindings( + inputData, + read(args.params, "endpoint"), + read(args.params, "model"), + read(args.params, "agentId"), + ); + if (!bindings || !hasReservedContext(metadata)) { + return; + } + const context = await parseBatchContext( + read(metadata, BRAINTRUST_MISTRAL_BATCH_CONTEXT_KEY), + bindings, + ); + if (!context || context.inputDigest !== inputData.inputDigest) { + reportBatchTraceError( + args.onTraceError, + "could not verify failed batch context", + new Error("Mistral Batch tracing context does not match the input"), + ); + return; + } + if ( + inputData.inputKind === "files" && + !sameStrings(read(args.params, "inputFiles"), inputData.inputFileIds) + ) { + reportBatchTraceError( + args.onTraceError, + "could not verify failed batch files", + new Error("Mistral Batch tracing file ids do not match the input"), + ); + return; + } + + const endTime = Math.max(getCurrentUnixTimestamp(), context.startTime); + const error = normalizeError(args.error); + const task = await resumeBatchSpan(context); + const taskParent = await task.export(); + for (const input of inputData.inputs.values()) { + try { + const child = await resumeBatchChild(context, taskParent, input); + child.log({ error }); + child.end({ endTime }); + } catch (spanError) { + logBatchInstrumentationError("could not fail batch child", spanError); + } + } + task.log({ error }); + task.end({ endTime }); +} + +export async function failMistralBatchTraceImpl( + args: FailMistralBatchTraceArgs, +): Promise { + try { + await failBatch(args); + } catch (error) { + reportBatchTraceError( + args.onTraceError, + "could not fail batch", + error instanceof Error + ? error + : new Error("Mistral Batch tracing could not fail the batch"), + ); + } +} diff --git a/js/src/instrumentation/plugins/mistral-plugin.test.ts b/js/src/instrumentation/plugins/mistral-plugin.test.ts index 96d410440..02d5ecdf2 100644 --- a/js/src/instrumentation/plugins/mistral-plugin.test.ts +++ b/js/src/instrumentation/plugins/mistral-plugin.test.ts @@ -1,10 +1,12 @@ import { describe, expect, it } from "vitest"; +import { aggregateMistralStreamChunks } from "./mistral-plugin"; import { - aggregateMistralStreamChunks, + extractMistralChatInput, extractMistralRequestMetadata, extractMistralResponseMetadata, + normalizeMistralChatChoices, parseMistralMetricsFromUsage, -} from "./mistral-plugin"; +} from "./mistral-span-data"; describe("extractMistralRequestMetadata", () => { it("keeps only allowlisted request metadata", () => { @@ -17,19 +19,100 @@ describe("extractMistralRequestMetadata", () => { n: 2, safe_prompt: true, toolChoice: "auto", + parallelToolCalls: false, + max_tool_calls: 2, messages: [{ role: "user", content: "hi" }], - tools: [{ type: "function" }], + tools: [ + { + type: "function", + function: { + name: "lookup", + description: "Look something up", + parameters: { type: "object" }, + strict: true, + }, + }, + { + type: "web_search", + toolConfiguration: { + include: ["news"], + requiresConfirmation: ["open_url"], + }, + }, + { + type: "connector", + connectorId: "connector-123", + authorization: { type: "api-key", value: "secret" }, + }, + ], suffix: "ignored", arbitrary: "ignored", }), ).toEqual({ model: "mistral-large-latest", - maxTokens: 128, + max_tokens: 128, reasoning_effort: "high", temperature: 0.4, n: 2, safe_prompt: true, - toolChoice: "auto", + tool_choice: "auto", + parallel_tool_calls: false, + max_tool_calls: 2, + tools: [ + { + type: "function", + function: { + name: "lookup", + description: "Look something up", + parameters: { type: "object" }, + strict: true, + }, + }, + { + type: "web_search", + tool_configuration: { + include: ["news"], + requires_confirmation: ["open_url"], + }, + }, + { + type: "connector", + connector_id: "connector-123", + }, + ], + }); + }); + + it("normalizes Mistral's required tool choice", () => { + expect(extractMistralRequestMetadata({ toolChoice: "any" })).toEqual({ + tool_choice: "required", + }); + }); + + it("canonicalizes camel-case and snake-case request metadata", () => { + const camelCase = extractMistralRequestMetadata({ + agentId: "agent-123", + maxTokens: 128, + reasoningEffort: "high", + responseFormat: { type: "json_object" }, + topP: 0.8, + }); + + expect(camelCase).toEqual( + extractMistralRequestMetadata({ + agent_id: "agent-123", + max_tokens: 128, + reasoning_effort: "high", + response_format: { type: "json_object" }, + top_p: 0.8, + }), + ); + expect(camelCase).toEqual({ + agent_id: "agent-123", + max_tokens: 128, + reasoning_effort: "high", + response_format: { type: "json_object" }, + top_p: 0.8, }); }); @@ -46,7 +129,8 @@ describe("extractMistralResponseMetadata", () => { created: 1234, object: "chat.completion", model: "mistral-large-latest", - agentId: "agent_123", + agentId: "legacy_agent", + agent_id: "agent_123", usage: { total_tokens: 20 }, choices: [{ index: 0 }], data: [{ embedding: [0.1] }], @@ -57,7 +141,7 @@ describe("extractMistralResponseMetadata", () => { created: 1234, object: "chat.completion", model: "mistral-large-latest", - agentId: "agent_123", + agent_id: "agent_123", }); }); @@ -71,6 +155,85 @@ describe("extractMistralResponseMetadata", () => { }); }); +describe("Mistral chat payload normalization", () => { + it("normalizes input messages through the shared chat extractor", () => { + expect( + extractMistralChatInput({ + messages: [ + { + role: "assistant", + content: null, + toolCalls: [ + { + id: "call-1", + type: "function", + function: { name: "lookup", arguments: '{"q":"hi"}' }, + }, + ], + }, + { role: "tool", toolCallId: "call-1", content: "result" }, + ], + }).input, + ).toEqual([ + { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call-1", + type: "function", + function: { name: "lookup", arguments: '{"q":"hi"}' }, + }, + ], + }, + { role: "tool", tool_call_id: "call-1", content: "result" }, + ]); + }); + + it("normalizes camel-case and snake-case choices identically", () => { + const expected = [ + { + index: 0, + finish_reason: "tool_calls", + message: { + role: "assistant", + content: null, + tool_calls: [ + { + id: "call-1", + type: "function", + function: { name: "lookup", arguments: '{"q":"hi"}' }, + }, + ], + }, + }, + ]; + const camelCase = [ + { + index: 0, + finishReason: "tool_calls", + message: { + role: "assistant", + content: null, + prefix: false, + toolCalls: expected[0].message.tool_calls, + }, + }, + ]; + + expect(normalizeMistralChatChoices(camelCase)).toEqual(expected); + expect( + normalizeMistralChatChoices([ + { + index: 0, + finish_reason: "tool_calls", + message: expected[0].message, + }, + ]), + ).toEqual(expected); + }); +}); + describe("parseMistralMetricsFromUsage", () => { it("returns empty metrics for missing usage", () => { expect(parseMistralMetricsFromUsage(undefined)).toEqual({}); @@ -89,7 +252,20 @@ describe("parseMistralMetricsFromUsage", () => { prompt_tokens: 10, completion_tokens: 6, tokens: 16, - prompt_audio_seconds: 3, + }); + }); + + it("derives total tokens when the provider total is missing or invalid", () => { + expect( + parseMistralMetricsFromUsage({ + prompt_tokens: 10, + completionTokens: 6, + total_tokens: Number.POSITIVE_INFINITY, + }), + ).toEqual({ + prompt_tokens: 10, + completion_tokens: 6, + tokens: 16, }); }); @@ -108,6 +284,35 @@ describe("parseMistralMetricsFromUsage", () => { completion_reasoning_tokens: 2, }); }); + + it("omits unsupported and invalid metric values", () => { + expect( + parseMistralMetricsFromUsage({ + promptTokens: 10, + completionTokens: -1, + totalTokens: Number.POSITIVE_INFINITY, + promptAudioSeconds: 3, + arbitraryCounter: 99, + toString: 42, + inputTokensDetails: { + cachedTokens: 7, + reasoningTokens: 4, + audioTokens: 1.5, + arbitraryTokens: 8, + toString: 11, + }, + outputTokensDetails: { + reasoningTokens: 2, + cachedTokens: 5, + imageTokens: Number.NaN, + }, + }), + ).toEqual({ + prompt_tokens: 10, + prompt_cached_tokens: 7, + completion_reasoning_tokens: 2, + }); + }); }); describe("aggregateMistralStreamChunks", () => { @@ -164,7 +369,7 @@ describe("aggregateMistralStreamChunks", () => { role: "assistant", content: "Hello world", }, - finishReason: "stop", + finish_reason: "stop", }); expect(aggregated.metadata).toMatchObject({ @@ -220,7 +425,7 @@ describe("aggregateMistralStreamChunks", () => { expect(aggregated.output?.[0]).toMatchObject({ message: { content: null, - toolCalls: [ + tool_calls: [ { id: "tool_1", function: { @@ -230,7 +435,7 @@ describe("aggregateMistralStreamChunks", () => { }, ], }, - finishReason: "tool_calls", + finish_reason: "tool_calls", }); }); @@ -289,9 +494,8 @@ describe("aggregateMistralStreamChunks", () => { expect(aggregated.output?.[0]).toMatchObject({ message: { content: null, - toolCalls: [ + tool_calls: [ { - index: 0, id: "tool_0", function: { name: "first_tool", @@ -299,7 +503,6 @@ describe("aggregateMistralStreamChunks", () => { }, }, { - index: 1, id: "tool_1", function: { name: "second_tool", @@ -308,7 +511,7 @@ describe("aggregateMistralStreamChunks", () => { }, ], }, - finishReason: "tool_calls", + finish_reason: "tool_calls", }); }); @@ -357,7 +560,7 @@ describe("aggregateMistralStreamChunks", () => { role: "assistant", content: "a", }, - finishReason: "stop", + finish_reason: "stop", }, { index: 1, @@ -365,7 +568,7 @@ describe("aggregateMistralStreamChunks", () => { role: "assistant", content: "b", }, - finishReason: "length", + finish_reason: "length", }, ]); }); @@ -433,7 +636,7 @@ describe("aggregateMistralStreamChunks", () => { }, ], }, - finishReason: "stop", + finish_reason: "stop", }, ]); }); diff --git a/js/src/instrumentation/plugins/mistral-plugin.ts b/js/src/instrumentation/plugins/mistral-plugin.ts index f33f82554..167ed0ba4 100644 --- a/js/src/instrumentation/plugins/mistral-plugin.ts +++ b/js/src/instrumentation/plugins/mistral-plugin.ts @@ -8,11 +8,17 @@ import { SpanTypeAttribute, isObject } from "../../../util/index"; import { processInputAttachments } from "../../wrappers/attachment-utils"; import { getCurrentUnixTimestamp } from "../../util"; import { mistralChannels } from "./mistral-channels"; +import { + extractMistralChatInput, + extractMistralRequestMetadata, + extractMistralResponseMetadata, + normalizeMistralChatChoices, + parseMistralMetricsFromUsage, +} from "./mistral-span-data"; import type { MistralChatCompletionChunk, MistralChatCompletionChunkChoice, MistralChatCompletionEvent, - MistralChatCompletionResponse, MistralContentPart, MistralTextContentPart, MistralThinkingContentPart, @@ -33,10 +39,8 @@ export class MistralPlugin extends BasePlugin { traceStreamingChannel(mistralChannels.chatComplete, { name: "mistral.chat.complete", type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, + extractInput: extractMistralChatInputFromArgs, + extractOutput: (result) => normalizeMistralChatChoices(result?.choices), extractMetadata: (result) => extractMistralResponseMetadata(result), extractMetrics: (result, startTime) => extractMistralMetrics(result?.usage, startTime), @@ -47,7 +51,7 @@ export class MistralPlugin extends BasePlugin { traceStreamingChannel(mistralChannels.chatStream, { name: "mistral.chat.stream", type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, + extractInput: extractMistralChatInputFromArgs, extractOutput: extractMistralStreamOutput, extractMetadata: (result) => extractMistralResponseMetadata(result), extractMetrics: (result, startTime) => @@ -121,9 +125,7 @@ export class MistralPlugin extends BasePlugin { name: "mistral.fim.complete", type: SpanTypeAttribute.LLM, extractInput: extractPromptInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, + extractOutput: (result) => normalizeMistralChatChoices(result?.choices), extractMetadata: (result) => extractMistralResponseMetadata(result), extractMetrics: (result, startTime) => extractMistralMetrics(result?.usage, startTime), @@ -147,10 +149,8 @@ export class MistralPlugin extends BasePlugin { traceStreamingChannel(mistralChannels.agentsComplete, { name: "mistral.agents.complete", type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, + extractInput: extractMistralChatInputFromArgs, + extractOutput: (result) => normalizeMistralChatChoices(result?.choices), extractMetadata: (result) => extractMistralResponseMetadata(result), extractMetrics: (result, startTime) => extractMistralMetrics(result?.usage, startTime), @@ -161,7 +161,7 @@ export class MistralPlugin extends BasePlugin { traceStreamingChannel(mistralChannels.agentsStream, { name: "mistral.agents.stream", type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, + extractInput: extractMistralChatInputFromArgs, extractOutput: extractMistralStreamOutput, extractMetadata: (result) => extractMistralResponseMetadata(result), extractMetrics: (result, startTime) => @@ -172,75 +172,6 @@ export class MistralPlugin extends BasePlugin { } } -const TOKEN_NAME_MAP: Record = { - promptTokens: "prompt_tokens", - inputTokens: "prompt_tokens", - completionTokens: "completion_tokens", - outputTokens: "completion_tokens", - totalTokens: "tokens", - prompt_tokens: "prompt_tokens", - input_tokens: "prompt_tokens", - completion_tokens: "completion_tokens", - output_tokens: "completion_tokens", - total_tokens: "tokens", - promptAudioSeconds: "prompt_audio_seconds", - prompt_audio_seconds: "prompt_audio_seconds", -}; - -const TOKEN_DETAIL_PREFIX_MAP: Record = { - promptTokensDetails: "prompt", - inputTokensDetails: "prompt", - completionTokensDetails: "completion", - outputTokensDetails: "completion", - prompt_tokens_details: "prompt", - input_tokens_details: "prompt", - completion_tokens_details: "completion", - output_tokens_details: "completion", -}; - -const MISTRAL_REQUEST_METADATA_ALLOWLIST = new Set([ - "agentId", - "agent_id", - "encodingFormat", - "encoding_format", - "frequencyPenalty", - "frequency_penalty", - "maxTokens", - "max_tokens", - "model", - "n", - "presencePenalty", - "presence_penalty", - "randomSeed", - "random_seed", - "reasoningEffort", - "reasoning_effort", - "responseFormat", - "response_format", - "safePrompt", - "safe_prompt", - "stream", - "stop", - "temperature", - "toolChoice", - "tool_choice", - "topP", - "top_p", -]); - -const MISTRAL_RESPONSE_METADATA_ALLOWLIST = new Set([ - "agentId", - "agent_id", - "created", - "id", - "model", - "object", -]); - -function camelToSnake(value: string): string { - return value.replace(/[A-Z]/g, (match) => `_${match.toLowerCase()}`); -} - function normalizeArgs(args: unknown[] | unknown): unknown[] { if (Array.isArray(args)) { return args; @@ -279,30 +210,6 @@ function addMistralProviderMetadata( }; } -function pickAllowedMetadata( - metadata: Record | undefined, - allowlist: ReadonlySet, -): Record { - if (!metadata) { - return {}; - } - - const picked: Record = {}; - for (const key of allowlist) { - const value = metadata[key]; - if (value !== undefined) { - picked[key] = value; - } - } - return picked; -} - -export function extractMistralRequestMetadata( - metadata: Record | undefined, -): Record { - return pickAllowedMetadata(metadata, MISTRAL_REQUEST_METADATA_ALLOWLIST); -} - function isMistralChatCompletionChunk( value: unknown, ): value is MistralChatCompletionChunk { @@ -315,19 +222,8 @@ function isMistralChunkChoice( return isObject(value); } -function extractMessagesInputWithMetadata(args: unknown[] | unknown): { - input: unknown; - metadata: Record; -} { - const params = getMistralRequestArg(args); - const { messages, ...rawMetadata } = params || {}; - - return { - input: processInputAttachments(messages), - metadata: addMistralProviderMetadata( - extractMistralRequestMetadata(rawMetadata), - ), - }; +function extractMistralChatInputFromArgs(args: unknown[] | unknown) { + return extractMistralChatInput(getMistralRequestArg(args) ?? {}); } function extractEmbeddingInputWithMetadata(args: unknown[] | unknown): { @@ -375,22 +271,6 @@ function extractPromptInputWithMetadata(args: unknown[] | unknown): { }; } -export function extractMistralResponseMetadata( - result: unknown, -): Record | undefined { - if (!isObject(result)) { - return undefined; - } - - const { choices: _choices, usage: _usage, data: _data, ...metadata } = result; - const picked = pickAllowedMetadata( - metadata, - MISTRAL_RESPONSE_METADATA_ALLOWLIST, - ); - - return Object.keys(picked).length > 0 ? picked : undefined; -} - function extractMistralMetrics( usage: unknown, startTime?: number, @@ -403,7 +283,9 @@ function extractMistralMetrics( } function extractMistralStreamOutput(result: unknown): unknown { - return isObject(result) ? result.choices : undefined; + return isObject(result) + ? normalizeMistralChatChoices(result.choices) + : undefined; } function extractClassifierOutput(result: unknown): unknown { @@ -414,13 +296,10 @@ function extractMistralStreamingMetrics( result: unknown, startTime?: number, ): Record { - const metrics = isObject(result) - ? parseMistralMetricsFromUsage(result.usage) - : {}; - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; + return extractMistralMetrics( + isObject(result) ? result.usage : undefined, + startTime, + ); } function extractDeltaText(content: unknown): string | undefined { @@ -757,46 +636,10 @@ type MistralChoiceAccumulator = { toolCalls?: MistralToolCallDelta[]; }; -export function parseMistralMetricsFromUsage( - usage: unknown, -): Record { - if (!isObject(usage)) { - return {}; - } - - const metrics: Record = {}; - - for (const [name, value] of Object.entries(usage)) { - if (typeof value === "number") { - metrics[TOKEN_NAME_MAP[name] || camelToSnake(name)] = value; - continue; - } - - if (!isObject(value)) { - continue; - } - - const prefix = TOKEN_DETAIL_PREFIX_MAP[name]; - if (!prefix) { - continue; - } - - for (const [nestedName, nestedValue] of Object.entries(value)) { - if (typeof nestedValue !== "number") { - continue; - } - - metrics[`${prefix}_${camelToSnake(nestedName)}`] = nestedValue; - } - } - - return metrics; -} - export function aggregateMistralStreamChunks( chunks: MistralChatCompletionEvent[], ): { - output: MistralChatCompletionResponse["choices"]; + output: Record[] | undefined; metrics: Record; metadata?: Record; } { @@ -885,23 +728,25 @@ export function aggregateMistralStreamChunks( } } - const output = Array.from(choiceAccumulators.values()) - .sort((left, right) => - left.index === right.index - ? left.order - right.order - : left.index - right.index, - ) - .map((choice) => ({ - index: choice.index, - message: { - ...(choice.role ? { role: choice.role } : {}), - content: choice.contentParts ?? choice.content ?? null, - ...(choice.toolCalls ? { toolCalls: choice.toolCalls } : {}), - }, - ...(choice.finishReason !== undefined - ? { finishReason: choice.finishReason } - : {}), - })); + const output = normalizeMistralChatChoices( + Array.from(choiceAccumulators.values()) + .sort((left, right) => + left.index === right.index + ? left.order - right.order + : left.index - right.index, + ) + .map((choice) => ({ + index: choice.index, + message: { + ...(choice.role ? { role: choice.role } : {}), + content: choice.contentParts ?? choice.content ?? null, + ...(choice.toolCalls ? { toolCalls: choice.toolCalls } : {}), + }, + ...(choice.finishReason !== undefined + ? { finishReason: choice.finishReason } + : {}), + })), + ); return { output, diff --git a/js/src/instrumentation/plugins/mistral-span-data.ts b/js/src/instrumentation/plugins/mistral-span-data.ts new file mode 100644 index 000000000..dac37d1ba --- /dev/null +++ b/js/src/instrumentation/plugins/mistral-span-data.ts @@ -0,0 +1,431 @@ +import { isObject } from "../../../util/index"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; + +const TOKEN_NAME_MAP: Record = { + promptTokens: "prompt_tokens", + inputTokens: "prompt_tokens", + completionTokens: "completion_tokens", + outputTokens: "completion_tokens", + totalTokens: "tokens", + prompt_tokens: "prompt_tokens", + input_tokens: "prompt_tokens", + completion_tokens: "completion_tokens", + output_tokens: "completion_tokens", + total_tokens: "tokens", +}; + +const PROMPT_TOKEN_DETAIL_NAME_MAP: Record = { + cachedTokens: "prompt_cached_tokens", + cached_tokens: "prompt_cached_tokens", + cacheCreationTokens: "prompt_cache_creation_tokens", + cache_creation_tokens: "prompt_cache_creation_tokens", + cacheCreation5mTokens: "prompt_cache_creation_5m_tokens", + cache_creation_5m_tokens: "prompt_cache_creation_5m_tokens", + cacheCreation1hTokens: "prompt_cache_creation_1h_tokens", + cache_creation_1h_tokens: "prompt_cache_creation_1h_tokens", + audioTokens: "prompt_audio_tokens", + audio_tokens: "prompt_audio_tokens", +}; + +const COMPLETION_TOKEN_DETAIL_NAME_MAP: Record = { + reasoningTokens: "completion_reasoning_tokens", + reasoning_tokens: "completion_reasoning_tokens", + audioTokens: "completion_audio_tokens", + audio_tokens: "completion_audio_tokens", + imageTokens: "completion_image_tokens", + image_tokens: "completion_image_tokens", +}; + +const TOKEN_DETAIL_NAME_MAP: Record> = { + promptTokensDetails: PROMPT_TOKEN_DETAIL_NAME_MAP, + inputTokensDetails: PROMPT_TOKEN_DETAIL_NAME_MAP, + prompt_tokens_details: PROMPT_TOKEN_DETAIL_NAME_MAP, + input_tokens_details: PROMPT_TOKEN_DETAIL_NAME_MAP, + completionTokensDetails: COMPLETION_TOKEN_DETAIL_NAME_MAP, + outputTokensDetails: COMPLETION_TOKEN_DETAIL_NAME_MAP, + completion_tokens_details: COMPLETION_TOKEN_DETAIL_NAME_MAP, + output_tokens_details: COMPLETION_TOKEN_DETAIL_NAME_MAP, +}; + +const MISTRAL_REQUEST_METADATA_MAP: Record = { + agentId: "agent_id", + agent_id: "agent_id", + encodingFormat: "encoding_format", + encoding_format: "encoding_format", + frequencyPenalty: "frequency_penalty", + frequency_penalty: "frequency_penalty", + maxTokens: "max_tokens", + max_tokens: "max_tokens", + model: "model", + n: "n", + presencePenalty: "presence_penalty", + presence_penalty: "presence_penalty", + randomSeed: "random_seed", + random_seed: "random_seed", + reasoningEffort: "reasoning_effort", + reasoning_effort: "reasoning_effort", + responseFormat: "response_format", + response_format: "response_format", + safePrompt: "safe_prompt", + safe_prompt: "safe_prompt", + stream: "stream", + stop: "stop", + temperature: "temperature", + topP: "top_p", + top_p: "top_p", +}; + +const MISTRAL_RESPONSE_METADATA_MAP: Record = { + agentId: "agent_id", + agent_id: "agent_id", + created: "created", + id: "id", + model: "model", + object: "object", +}; + +function pickMappedMetadata( + metadata: Record | undefined, + mapping: Readonly>, +): Record { + if (!metadata) { + return {}; + } + + const picked: Record = {}; + for (const [source, canonical] of Object.entries(mapping)) { + const value = metadata[source]; + if (value !== undefined) { + picked[canonical] = value; + } + } + return picked; +} + +export function extractMistralRequestMetadata( + metadata: Record | undefined, +): Record { + const picked = pickMappedMetadata(metadata, MISTRAL_REQUEST_METADATA_MAP); + if (!metadata) { + return picked; + } + + if (Array.isArray(metadata.tools)) { + const tools = metadata.tools.flatMap((tool) => { + if (!isObject(tool)) { + return []; + } + if (isObject(tool.function)) { + const fn = tool.function; + if (typeof fn.name !== "string" || fn.name.length === 0) { + return []; + } + return [ + { + type: "function", + function: { + name: fn.name, + ...(typeof fn.description === "string" + ? { description: fn.description } + : {}), + ...(isObject(fn.parameters) ? { parameters: fn.parameters } : {}), + ...(typeof fn.strict === "boolean" ? { strict: fn.strict } : {}), + }, + }, + ]; + } + + const type = tool.type; + if ( + type !== "web_search" && + type !== "web_search_premium" && + type !== "code_interpreter" && + type !== "image_generation" && + type !== "document_library" && + type !== "connector" + ) { + return []; + } + + const normalized: Record = { type }; + const rawConfiguration = + tool.tool_configuration ?? tool.toolConfiguration; + if (isObject(rawConfiguration)) { + const configuration: Record = {}; + for (const key of ["exclude", "include"] as const) { + const value = rawConfiguration[key]; + if ( + Array.isArray(value) && + value.every((entry) => typeof entry === "string") + ) { + configuration[key] = value; + } + } + const requiresConfirmation = + rawConfiguration.requires_confirmation ?? + rawConfiguration.requiresConfirmation; + if ( + Array.isArray(requiresConfirmation) && + requiresConfirmation.every((entry) => typeof entry === "string") + ) { + configuration.requires_confirmation = requiresConfirmation; + } + if (Object.keys(configuration).length > 0) { + normalized.tool_configuration = configuration; + } + } + const libraryIds = tool.library_ids ?? tool.libraryIds; + if ( + type === "document_library" && + Array.isArray(libraryIds) && + libraryIds.every((entry) => typeof entry === "string") + ) { + normalized.library_ids = libraryIds; + } + const connectorId = tool.connector_id ?? tool.connectorId; + if (type === "connector" && typeof connectorId === "string") { + normalized.connector_id = connectorId; + } + return [normalized]; + }); + if (tools.length > 0) { + picked.tools = tools; + } + } + + const rawToolChoice = metadata.tool_choice ?? metadata.toolChoice; + if ( + rawToolChoice === "auto" || + rawToolChoice === "none" || + rawToolChoice === "required" + ) { + picked.tool_choice = rawToolChoice; + } else if (rawToolChoice === "any") { + picked.tool_choice = "required"; + } else if ( + isObject(rawToolChoice) && + isObject(rawToolChoice.function) && + typeof rawToolChoice.function.name === "string" && + rawToolChoice.function.name.length > 0 + ) { + picked.tool_choice = { + type: "function", + function: { name: rawToolChoice.function.name }, + }; + } + + const parallelToolCalls = + metadata.parallel_tool_calls ?? metadata.parallelToolCalls; + if (typeof parallelToolCalls === "boolean") { + picked.parallel_tool_calls = parallelToolCalls; + } + const maxToolCalls = metadata.max_tool_calls ?? metadata.maxToolCalls; + if (typeof maxToolCalls === "number" && Number.isFinite(maxToolCalls)) { + picked.max_tool_calls = maxToolCalls; + } + + return picked; +} + +function normalizeMistralToolCall( + value: unknown, +): Record | undefined { + if ( + !isObject(value) || + typeof value.id !== "string" || + !isObject(value.function) || + typeof value.function.name !== "string" || + typeof value.function.arguments !== "string" + ) { + return undefined; + } + + return { + id: value.id, + type: typeof value.type === "string" ? value.type : "function", + function: { + name: value.function.name, + arguments: value.function.arguments, + }, + }; +} + +function normalizeMistralChatMessage( + value: unknown, + defaultRole?: "assistant", +): Record | undefined { + if (!isObject(value)) { + return undefined; + } + + const role = typeof value.role === "string" ? value.role : defaultRole; + const toolCallId = value.tool_call_id ?? value.toolCallId; + const rawToolCalls = value.tool_calls ?? value.toolCalls; + const toolCalls = Array.isArray(rawToolCalls) + ? rawToolCalls + .map(normalizeMistralToolCall) + .filter( + (toolCall): toolCall is Record => + toolCall !== undefined, + ) + : []; + + const message: Record = {}; + if (role) { + message.role = role; + } + if (Object.hasOwn(value, "content")) { + message.content = value.content; + } else if (defaultRole) { + message.content = null; + } + if (typeof value.name === "string") { + message.name = value.name; + } + if (typeof toolCallId === "string") { + message.tool_call_id = toolCallId; + } + if (toolCalls.length > 0) { + message.tool_calls = toolCalls; + } + return message; +} + +function normalizeMistralChatMessages(messages: unknown): unknown { + if (!Array.isArray(messages)) { + return messages; + } + + return messages.map( + (message) => normalizeMistralChatMessage(message) ?? message, + ); +} + +export function normalizeMistralChatChoices( + choices: unknown, +): Record[] | undefined { + if (!Array.isArray(choices) || choices.length === 0) { + return undefined; + } + + const normalized: Record[] = []; + for (const [position, choice] of choices.entries()) { + if (!isObject(choice)) { + return undefined; + } + const message = normalizeMistralChatMessage(choice.message, "assistant"); + if (!message) { + return undefined; + } + const finishReason = choice.finish_reason ?? choice.finishReason; + normalized.push({ + index: + typeof choice.index === "number" && choice.index >= 0 + ? choice.index + : position, + ...(typeof finishReason === "string" || finishReason === null + ? { finish_reason: finishReason } + : {}), + message, + ...(Object.hasOwn(choice, "logprobs") + ? { logprobs: choice.logprobs } + : {}), + }); + } + + return normalized; +} + +export function extractMistralChatInput( + params: Record, + options: { processAttachments?: boolean } = {}, +): { + input: unknown; + metadata: Record; +} { + const { messages, ...rawMetadata } = params; + const normalizedMessages = normalizeMistralChatMessages(messages); + return { + input: + options.processAttachments === false + ? normalizedMessages + : processInputAttachments(normalizedMessages), + metadata: { + ...extractMistralRequestMetadata(rawMetadata), + provider: "mistral", + }, + }; +} + +export function extractMistralResponseMetadata( + result: unknown, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + const { choices: _choices, usage: _usage, data: _data, ...metadata } = result; + const picked = pickMappedMetadata(metadata, MISTRAL_RESPONSE_METADATA_MAP); + return Object.keys(picked).length > 0 ? picked : undefined; +} + +export function parseMistralMetricsFromUsage( + usage: unknown, +): Record { + if (!isObject(usage)) { + return {}; + } + + const metrics: Record = {}; + for (const [name, value] of Object.entries(usage)) { + const metricName = Object.hasOwn(TOKEN_NAME_MAP, name) + ? TOKEN_NAME_MAP[name] + : undefined; + if (metricName) { + if ( + typeof value === "number" && + Number.isFinite(value) && + Number.isInteger(value) && + value >= 0 + ) { + metrics[metricName] = value; + } + continue; + } + if (!isObject(value)) { + continue; + } + + const detailNames = Object.hasOwn(TOKEN_DETAIL_NAME_MAP, name) + ? TOKEN_DETAIL_NAME_MAP[name] + : undefined; + if (!detailNames) { + continue; + } + for (const [nestedName, nestedValue] of Object.entries(value)) { + const metricName = Object.hasOwn(detailNames, nestedName) + ? detailNames[nestedName] + : undefined; + if ( + metricName && + typeof nestedValue === "number" && + Number.isFinite(nestedValue) && + Number.isInteger(nestedValue) && + nestedValue >= 0 + ) { + metrics[metricName] = nestedValue; + } + } + } + if ( + metrics.tokens === undefined && + Number.isSafeInteger(metrics.prompt_tokens) && + Number.isSafeInteger(metrics.completion_tokens) + ) { + const tokens = metrics.prompt_tokens + metrics.completion_tokens; + if (Number.isSafeInteger(tokens)) { + metrics.tokens = tokens; + } + } + return metrics; +} diff --git a/js/src/logger.ts b/js/src/logger.ts index ad316c481..76253abe1 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -2484,19 +2484,21 @@ export function updateSpan({ }); } -/** @internal Rehydrate an exported root span so work can continue in another process. */ +/** @internal Rehydrate an exported span without emitting another initial row. */ export function _internalResumeSpan({ exported, + spanParents = [], state, }: { exported: string; + spanParents?: string[]; state?: BraintrustState; }): Span { const resolvedState = state ?? _globalState; const components = SpanComponentsV4.fromStr(exported); const { row_id, root_span_id, span_id } = components.data; if (!row_id || !root_span_id || !span_id) { - throw new Error("Only exported root spans can be resumed"); + throw new Error("Only exported spans can be resumed"); } return new SpanImpl({ state: resolvedState, @@ -2505,7 +2507,7 @@ export function _internalResumeSpan({ spanComponentsToObjectIdLambda(resolvedState, components), ), parentComputeObjectMetadataArgs: undefined, - parentSpanIds: { parentSpanIds: [], rootSpanId: root_span_id }, + parentSpanIds: { parentSpanIds: spanParents, rootSpanId: root_span_id }, spanId: span_id, event: { id: row_id }, propagatedEvent: (components.data.propagated_event ?? undefined) as @@ -6732,6 +6734,22 @@ export function _internalStartSpanWithInitialMerge< InitialSpanWriteAsMergeArg).span; } +/** @internal Start a deterministic span under an exported object-only parent. */ +export function _internalStartSpanWithInitialMergeAndParentSpanIds< + IsAsyncFlush extends boolean = true, +>(args: StartSpanArgs & AsyncFlushArg & OptionalStateArg): Span { + return startSpanAndIsLogger( + { + ...args, + [INITIAL_SPAN_WRITE_AS_MERGE]: true, + } as StartSpanArgs & + AsyncFlushArg & + OptionalStateArg & + InitialSpanWriteAsMergeArg, + { useParentSpanIdsForObjectParent: true }, + ).span; +} + /** @internal Start a span with SDK-controlled context fields. */ export function _internalStartSpanWithContext< IsAsyncFlush extends boolean = true, @@ -6768,6 +6786,7 @@ function startSpanAndIsLogger( AsyncFlushArg & OptionalStateArg & InternalSpanContextArg, + internalOptions?: { useParentSpanIdsForObjectParent?: boolean }, ): { span: Span; isSyncFlushLogger: boolean } { const state = args?.state ?? _globalState; @@ -6785,17 +6804,21 @@ function startSpanAndIsLogger( parentObject instanceof SpanComponentsV3 || parentObject instanceof SpanComponentsV4 ) { - const parentSpanIds: ParentSpanIds | undefined = + let parentSpanIds: ParentSpanIds | MultiParentSpanIds | undefined; + if ( parentObject.data.row_id && parentSpanIdsUsable( parentObject.data.span_id, parentObject.data.root_span_id, ) - ? { - spanId: parentObject.data.span_id, - rootSpanId: parentObject.data.root_span_id, - } - : undefined; + ) { + parentSpanIds = { + spanId: parentObject.data.span_id, + rootSpanId: parentObject.data.root_span_id, + }; + } else if (internalOptions?.useParentSpanIdsForObjectParent) { + parentSpanIds = args?.parentSpanIds; + } // The parent object/state are already resolved from `parent` above; drop // the raw `parent` so it isn't re-normalized. const { parent: _ignoredParent, ...spanArgs } = args ?? {}; diff --git a/js/src/mistral-batch-types.ts b/js/src/mistral-batch-types.ts new file mode 100644 index 000000000..aa114fef0 --- /dev/null +++ b/js/src/mistral-batch-types.ts @@ -0,0 +1,111 @@ +type MistralBatchJSONLContent = + | string + | Iterable + | AsyncIterable + | ReadableStream + | Response; + +export type MistralBatchJSONLSource = + | MistralBatchJSONLContent + | PromiseLike; + +/** + * Repeatable input content. Use a factory for one-shot sources such as a + * Response, ReadableStream, generator, or async iterable because start and + * completion may read the input in different processes or at different times. + */ +export type MistralBatchInputFileContent = + | string + | ReadonlyArray + | (() => MistralBatchJSONLSource); + +export interface MistralBatchFileLike { + id: string; + [key: string]: unknown; +} + +export interface MistralBatchRequestLike { + customId: string; + body: Record; + [key: string]: unknown; +} + +export type MistralBatchInput = + | { + files: ReadonlyArray<{ + file: MistralBatchFileLike; + content: MistralBatchInputFileContent; + }>; + } + | { requests: ReadonlyArray }; + +interface MistralBatchCreateInputParamsBase { + endpoint: "/v1/chat/completions"; + metadata?: Record | null; + timeoutHours?: number; + [key: string]: unknown; +} + +export type MistralBatchCreateInputParams = MistralBatchCreateInputParamsBase & + ({ model: string; agentId?: never } | { agentId: string; model?: never }); + +export type MistralBatchCreateParams = MistralBatchCreateInputParams & { + inputFiles?: string[]; + requests?: MistralBatchRequestLike[]; +}; + +export interface MistralBatchLike { + id: string; + endpoint: string; + status: string; + model?: string | null; + agentId?: string | null; + inputFiles?: string[]; + metadata?: Record | null; + outputFile?: string | null; + errorFile?: string | null; + outputs?: unknown | null; + createdAt?: number; + startedAt?: number | null; + completedAt?: number | null; + totalRequests?: number; + completedRequests?: number; + succeededRequests?: number; + failedRequests?: number; +} + +export interface StartMistralBatchTraceArgs { + input: MistralBatchInput; + params: MistralBatchCreateInputParams; + onTraceError?: (error: Error) => void | Promise; +} + +/** Opaque context that keeps collect-only retries on the same trace. */ +export interface MistralBatchCollectionContext { + readonly version: 1; + readonly value: string; +} + +export interface MistralBatchTraceCollection { + batch: TBatch; + collectionContext?: MistralBatchCollectionContext; +} + +export interface CompleteMistralBatchTraceArgs< + TBatch extends MistralBatchLike, +> { + batch: TBatch; + input: MistralBatchInput; + outputContent?: MistralBatchJSONLSource; + errorContent?: MistralBatchJSONLSource; + /** Required when retrying a collect-only completion. */ + collectionContext?: MistralBatchCollectionContext; + onTraceError?: (error: Error) => void | Promise; +} + +export interface FailMistralBatchTraceArgs { + params: MistralBatchCreateParams; + input: MistralBatchInput; + error: unknown; + onTraceError?: (error: Error) => void | Promise; +} diff --git a/js/src/mistral-batch.ts b/js/src/mistral-batch.ts new file mode 100644 index 000000000..89ed62b3b --- /dev/null +++ b/js/src/mistral-batch.ts @@ -0,0 +1,64 @@ +import { + completeMistralBatchTraceImpl, + failMistralBatchTraceImpl, + startMistralBatchTraceImpl, +} from "./instrumentation/plugins/mistral-batch-instrumentation"; +import { getSpanParentObject } from "./logger"; +import type { + CompleteMistralBatchTraceArgs, + FailMistralBatchTraceArgs, + MistralBatchCreateParams, + MistralBatchLike, + MistralBatchTraceCollection, + StartMistralBatchTraceArgs, +} from "./mistral-batch-types"; + +/** + * Start a trace for a new Mistral Batch, with pending spans for its requests, + * and return the parameters to pass to `mistral.batch.jobs.create()`. + */ +export async function startMistralBatchTrace( + args: StartMistralBatchTraceArgs, +): Promise { + return await startMistralBatchTraceImpl(args); +} + +/** + * Complete a Mistral Batch trace from caller-supplied Batch API data. + * + * This helper performs no Mistral API requests. Supplied Response and stream + * bodies are consumed directly. Inline `batch.outputs` are used when + * `outputContent` is omitted. Provider file IDs must be downloaded before + * their contents are supplied as `outputContent` or `errorContent`. For a + * collect-only trace, persist the returned `collectionContext` and pass it to + * every retry so that all attempts update the same spans. + */ +export async function completeMistralBatchTrace< + TBatch extends MistralBatchLike, +>( + args: CompleteMistralBatchTraceArgs, +): Promise> { + const outputContent = + args.outputContent === undefined + ? undefined + : Promise.resolve(args.outputContent); + const errorContent = + args.errorContent === undefined + ? undefined + : Promise.resolve(args.errorContent); + + void outputContent?.catch(() => undefined); + void errorContent?.catch(() => undefined); + + return await completeMistralBatchTraceImpl( + { ...args, outputContent, errorContent }, + getSpanParentObject(), + ); +} + +/** End pending Mistral Batch spans after `batch.jobs.create()` fails. */ +export async function failMistralBatchTrace( + args: FailMistralBatchTraceArgs, +): Promise { + await failMistralBatchTraceImpl(args); +} diff --git a/js/src/wrappers/attachment-utils.ts b/js/src/wrappers/attachment-utils.ts index 7fb2b96f5..a95d78895 100644 --- a/js/src/wrappers/attachment-utils.ts +++ b/js/src/wrappers/attachment-utils.ts @@ -71,10 +71,12 @@ export function convertDataToBlob(data: any, mediaType: string): Blob | null { } /** - * Process input to extract and convert image/file content parts to Attachments - * Similar to processImagesInOutput in oai_responses.ts - replaces data in-place + * Return a payload copy with inline image/file content converted to Attachments. */ -export function processInputAttachments(input: any): any { +export function processInputAttachments( + input: any, + options?: { attachmentKey?: (index: number) => string }, +): any { if (!input) { return input; } @@ -99,11 +101,17 @@ export function processInputAttachments(input: any): any { return null; } - return new Attachment({ + const attachment = new Attachment({ data: blob, filename, contentType: mediaType, }); + const key = options?.attachmentKey?.(attachmentIndex); + attachmentIndex++; + if (key) { + attachment.reference.key = key; + } + return attachment; }; const processNode = (node: any): any => { @@ -217,7 +225,6 @@ export function processInputAttachments(input: any): any { const attachment = toAttachment(node.image, mediaType, filename); if (attachment) { - attachmentIndex++; return { ...node, image: attachment, @@ -234,7 +241,6 @@ export function processInputAttachments(input: any): any { const attachment = toAttachment(node.data, mediaType, filename); if (attachment) { - attachmentIndex++; return { ...node, data: attachment, diff --git a/js/util/json_util.test.ts b/js/util/json_util.test.ts index e85470f88..5bc754d4d 100644 --- a/js/util/json_util.test.ts +++ b/js/util/json_util.test.ts @@ -19,3 +19,12 @@ test("deterministicReplacer basic", () => { JSON.stringify(obj2, deterministicReplacer), ); }); + +test("deterministicReplacer handles prototype keys safely", () => { + const value = JSON.parse('{"z":1,"__proto__":{"polluted":true},"a":2}'); + + expect(JSON.stringify(value, deterministicReplacer)).toBe( + '{"__proto__":{"polluted":true},"a":2,"z":1}', + ); + expect(Object.prototype).not.toHaveProperty("polluted"); +}); diff --git a/js/util/json_util.ts b/js/util/json_util.ts index dd81d6e73..e8a5983e9 100644 --- a/js/util/json_util.ts +++ b/js/util/json_util.ts @@ -2,15 +2,21 @@ // serialized in a deterministic key order. // // https://gist.github.com/davidfurlong/463a83a33b70a3b6618e97ec9679e490 -export function deterministicReplacer(_key: string, value: any) { - return value instanceof Object && !(value instanceof Array) - ? Object.keys(value) - .sort() - .reduce((sorted: { [key: string]: unknown }, key) => { - sorted[key] = value[key]; - return sorted; - }, {}) - : value; +export function deterministicReplacer(_key: string, value: any): any { + if (value && typeof value === "object" && !Array.isArray(value)) { + return Object.fromEntries( + Object.entries(value).sort(([left], [right]) => { + if (left < right) { + return -1; + } + if (left > right) { + return 1; + } + return 0; + }), + ); + } + return value; } export function constructJsonArray(items: string[]) {