From d6ab2a712ca0d673ed74af244ee2c355ef04c7f7 Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 4 May 2026 15:36:49 -0700 Subject: [PATCH 1/3] feat(mistral): Instrument classification and moderation APIs --- .../mistral-v1-10-0.log-payloads.json | 119 +++ .../mistral-v1-10-0.span-events.json | 124 +++ .../mistral-v1-14-1.log-payloads.json | 119 +++ .../mistral-v1-14-1.span-events.json | 124 +++ .../mistral-v1-15-1.log-payloads.json | 119 +++ .../mistral-v1-15-1.span-events.json | 124 +++ .../mistral-v1.log-payloads.json | 119 +++ .../__snapshots__/mistral-v1.span-events.json | 124 +++ .../mistral-v2.log-payloads.json | 119 +++ .../__snapshots__/mistral-v2.span-events.json | 124 +++ .../mistral-instrumentation/assertions.ts | 139 +++ .../mistral-instrumentation/constants.mjs | 2 + .../mistral-instrumentation/scenario.impl.mjs | 802 +++++++++++------- .../scenario.mistral-v1-3-4.mjs | 2 + .../scenario.mistral-v1-3-4.ts | 2 + .../mistral-instrumentation/scenario.mjs | 6 +- .../mistral-instrumentation/scenario.test.ts | 12 + .../mistral-instrumentation/scenario.ts | 6 +- .../auto-instrumentations/configs/mistral.ts | 112 +++ .../plugins/mistral-channels.ts | 36 + .../instrumentation/plugins/mistral-plugin.ts | 63 ++ js/src/vendor-sdk-types/mistral.ts | 42 + js/src/wrappers/mistral.ts | 92 +- 23 files changed, 2217 insertions(+), 314 deletions(-) diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json index f91bf06c0..11e25c2cc 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json @@ -500,5 +500,124 @@ "type": "embedding" }, "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "item_count": 1, + "roles": [ + "user" + ], + "type": "messages" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "keys": [ + "messages" + ], + "type": "object" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json index 837183598..8bdafb6b2 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json @@ -449,5 +449,129 @@ "" ], "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderate", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderateChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classify", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classifyChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json index f91bf06c0..11e25c2cc 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json @@ -500,5 +500,124 @@ "type": "embedding" }, "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "item_count": 1, + "roles": [ + "user" + ], + "type": "messages" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "keys": [ + "messages" + ], + "type": "object" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json index 837183598..8bdafb6b2 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json @@ -449,5 +449,129 @@ "" ], "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderate", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderateChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classify", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classifyChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json index f91bf06c0..11e25c2cc 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json @@ -500,5 +500,124 @@ "type": "embedding" }, "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "item_count": 1, + "roles": [ + "user" + ], + "type": "messages" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "keys": [ + "messages" + ], + "type": "object" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json index 837183598..8bdafb6b2 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json @@ -449,5 +449,129 @@ "" ], "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderate", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderateChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classify", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classifyChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json index f91bf06c0..11e25c2cc 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json @@ -500,5 +500,124 @@ "type": "embedding" }, "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "item_count": 1, + "roles": [ + "user" + ], + "type": "messages" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "keys": [ + "messages" + ], + "type": "object" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json index 837183598..8bdafb6b2 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json @@ -449,5 +449,129 @@ "" ], "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderate", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderateChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classify", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classifyChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json index 8dd8304f6..5e2517dc6 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json @@ -512,5 +512,124 @@ "type": "embedding" }, "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "item_count": 1, + "roles": [ + "user" + ], + "type": "messages" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": "", + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" + }, + { + "has_input": false, + "has_output": false, + "input": null, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "output": null, + "span_id": "" + }, + { + "has_input": true, + "has_output": true, + "input": { + "keys": [ + "messages" + ], + "type": "object" + }, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "output": { + "choice_count": 1, + "finish_reason": null, + "type": "array" + }, + "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json index df04617bf..18b697ea2 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json @@ -461,5 +461,129 @@ "" ], "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderate", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-moderate-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-moderate-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.moderateChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classify", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" + }, + { + "has_input": false, + "has_output": false, + "metadata": { + "operation": "classifiers-classify-chat" + }, + "metric_keys": [], + "name": "mistral-classifiers-classify-chat-operation", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": null + }, + { + "has_input": true, + "has_output": true, + "metadata": { + "model": "mistral-moderation-2603", + "provider": "mistral" + }, + "metric_keys": [], + "name": "mistral.classifiers.classifyChat", + "root_span_id": "", + "span_id": "", + "span_parents": [ + "" + ], + "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/assertions.ts b/e2e/scenarios/mistral-instrumentation/assertions.ts index 131558574..784510550 100644 --- a/e2e/scenarios/mistral-instrumentation/assertions.ts +++ b/e2e/scenarios/mistral-instrumentation/assertions.ts @@ -20,6 +20,7 @@ import { ADJUSTABLE_REASONING_MODEL, FIM_MODEL, CHAT_MODEL, + CLASSIFIER_MODEL, AGENT_MODEL, EMBEDDING_MODEL, NATIVE_REASONING_MODEL, @@ -439,6 +440,22 @@ function buildSpanSummary( events, "mistral-embeddings-operation", ); + const classifiersModerateOperation = findLatestSpan( + events, + "mistral-classifiers-moderate-operation", + ); + const classifiersModerateChatOperation = findLatestSpan( + events, + "mistral-classifiers-moderate-chat-operation", + ); + const classifiersClassifyOperation = findLatestSpan( + events, + "mistral-classifiers-classify-operation", + ); + const classifiersClassifyChatOperation = findLatestSpan( + events, + "mistral-classifiers-classify-chat-operation", + ); return normalizeForSnapshot( [ @@ -483,6 +500,22 @@ function buildSpanSummary( findMistralSpan(events, embeddingsOperation?.span.id, [ "mistral.embeddings.create", ]), + classifiersModerateOperation, + findMistralSpan(events, classifiersModerateOperation?.span.id, [ + "mistral.classifiers.moderate", + ]), + classifiersModerateChatOperation, + findMistralSpan(events, classifiersModerateChatOperation?.span.id, [ + "mistral.classifiers.moderateChat", + ]), + classifiersClassifyOperation, + findMistralSpan(events, classifiersClassifyOperation?.span.id, [ + "mistral.classifiers.classify", + ]), + classifiersClassifyChatOperation, + findMistralSpan(events, classifiersClassifyChatOperation?.span.id, [ + "mistral.classifiers.classifyChat", + ]), ] .filter((event): event is CapturedLogEvent => event !== undefined) .map((event) => @@ -533,6 +566,10 @@ function buildPayloadSummary( "mistral.agents.complete", "mistral.agents.stream", "mistral.embeddings.create", + "mistral.classifiers.moderate", + "mistral.classifiers.moderateChat", + "mistral.classifiers.classify", + "mistral.classifiers.classifyChat", ]); const parentAndNameWithOutput = new Set(); @@ -590,6 +627,8 @@ export function defineMistralInstrumentationAssertions(options: { name: string; runScenario: RunMistralScenario; snapshotName: string; + supportsClassifiers?: boolean; + supportsClassify?: boolean; supportsThinkingStream?: boolean; testFileUrl: string; timeoutMs: number; @@ -603,6 +642,8 @@ export function defineMistralInstrumentationAssertions(options: { `${options.snapshotName}.log-payloads.json`, ); const supportsThinkingStream = options.supportsThinkingStream ?? true; + const supportsClassifiers = options.supportsClassifiers ?? true; + const supportsClassify = options.supportsClassify ?? true; const testConfig = { timeout: options.timeoutMs, }; @@ -1043,6 +1084,104 @@ export function defineMistralInstrumentationAssertions(options: { expect(output?.embedding_length).toBeGreaterThan(0); }); + if (supportsClassifiers) { + test("captures trace for classifiers.moderate()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "mistral-classifiers-moderate-operation", + ); + const span = findMistralSpan(events, operation?.span.id, [ + "mistral.classifiers.moderate", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + model: CLASSIFIER_MODEL, + provider: "mistral", + }); + expect(span?.input).toEqual(expect.any(String)); + expect(span?.output).toEqual(expect.any(Array)); + expect((span?.output as unknown[] | undefined)?.length).toBe(1); + }); + + test("captures trace for classifiers.moderateChat()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "mistral-classifiers-moderate-chat-operation", + ); + const span = findMistralSpan(events, operation?.span.id, [ + "mistral.classifiers.moderateChat", + ]); + const input = span?.input as Array<{ role?: string }> | undefined; + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + model: CLASSIFIER_MODEL, + provider: "mistral", + }); + expect(input).toEqual(expect.any(Array)); + expect(input?.[0]?.role).toBe("user"); + expect(span?.output).toEqual(expect.any(Array)); + expect((span?.output as unknown[] | undefined)?.length).toBe(1); + }); + } + + if (supportsClassifiers && supportsClassify) { + test("captures trace for classifiers.classify()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "mistral-classifiers-classify-operation", + ); + const span = findMistralSpan(events, operation?.span.id, [ + "mistral.classifiers.classify", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + model: CLASSIFIER_MODEL, + provider: "mistral", + }); + expect(span?.input).toEqual(expect.any(String)); + expect(span?.output).toEqual(expect.any(Array)); + expect((span?.output as unknown[] | undefined)?.length).toBe(1); + }); + + test("captures trace for classifiers.classifyChat()", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "mistral-classifiers-classify-chat-operation", + ); + const span = findMistralSpan(events, operation?.span.id, [ + "mistral.classifiers.classifyChat", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + model: CLASSIFIER_MODEL, + provider: "mistral", + }); + expect(span?.input).toEqual(expect.any(Object)); + expect(span?.output).toEqual(expect.any(Array)); + expect((span?.output as unknown[] | undefined)?.length).toBe(1); + }); + } + test("matches the shared span snapshot", testConfig, async () => { await expect( formatJsonFileSnapshot(buildSpanSummary(events, options.snapshotName)), diff --git a/e2e/scenarios/mistral-instrumentation/constants.mjs b/e2e/scenarios/mistral-instrumentation/constants.mjs index e57644f74..d3573e305 100644 --- a/e2e/scenarios/mistral-instrumentation/constants.mjs +++ b/e2e/scenarios/mistral-instrumentation/constants.mjs @@ -4,6 +4,7 @@ const NATIVE_REASONING_MODEL = "magistral-small-latest"; const EMBEDDING_MODEL = "mistral-embed"; const FIM_MODEL = "codestral-2508"; const AGENT_MODEL = CHAT_MODEL; +const CLASSIFIER_MODEL = "mistral-moderation-2603"; const ROOT_NAME = "mistral-root"; const SCENARIO_NAME = "mistral-instrumentation"; @@ -11,6 +12,7 @@ export { ADJUSTABLE_REASONING_MODEL, AGENT_MODEL, CHAT_MODEL, + CLASSIFIER_MODEL, EMBEDDING_MODEL, FIM_MODEL, NATIVE_REASONING_MODEL, diff --git a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs index d8b73adf8..2a95212c5 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs @@ -1,4 +1,5 @@ import { wrapMistral } from "braintrust"; +import { createServer } from "node:http"; import { collectAsync, runOperation, @@ -8,6 +9,7 @@ import { ADJUSTABLE_REASONING_MODEL, AGENT_MODEL, CHAT_MODEL, + CLASSIFIER_MODEL, EMBEDDING_MODEL, FIM_MODEL, NATIVE_REASONING_MODEL, @@ -24,6 +26,8 @@ const MISTRAL_REQUEST_RETRY_OPTIONS = { }; const MISTRAL_THINKING_STREAM_OPTOUTS = new Set(["mistral-sdk-v1-3-4"]); +const MISTRAL_CLASSIFIER_OPTOUTS = new Set(["mistral-sdk-v1-3-4"]); +const MISTRAL_CLASSIFY_OPTOUTS = new Set(["mistral-sdk-v1-3-4"]); function createMistralScenarioSpec(spec) { return { @@ -31,6 +35,12 @@ function createMistralScenarioSpec(spec) { ...(MISTRAL_THINKING_STREAM_OPTOUTS.has(spec.dependencyName) ? { supportsThinkingStream: false } : {}), + ...(MISTRAL_CLASSIFIER_OPTOUTS.has(spec.dependencyName) + ? { supportsClassifiers: false } + : {}), + ...(MISTRAL_CLASSIFY_OPTOUTS.has(spec.dependencyName) + ? { supportsClassify: false } + : {}), }; } @@ -354,7 +364,13 @@ async function resolveAgentRuntime(client) { async function runMistralInstrumentationScenario( Mistral, - { decorateClient, supportsThinkingStream = true } = {}, + { + classifyChatRequestInputKey = "inputs", + decorateClient, + supportsClassifiers = true, + supportsClassify = true, + supportsThinkingStream = true, + } = {}, ) { const baseClient = new Mistral({ apiKey: process.env.MISTRAL_API_KEY, @@ -363,96 +379,76 @@ async function runMistralInstrumentationScenario( const { agentId, cleanup } = await resolveAgentRuntime(baseClient); try { - await runTracedScenario({ - callback: async () => { - await runOperation( - "mistral-chat-complete-operation", - "chat-complete", - async () => { - await withRetry( - async () => - client.chat.complete({ + await withFakeMistralClassifierServer(async (classifierServerURL) => { + const classifierRequestOptions = { serverURL: classifierServerURL }; + + await runTracedScenario({ + callback: async () => { + await runOperation( + "mistral-chat-complete-operation", + "chat-complete", + async () => { + await withRetry( + async () => + client.chat.complete({ + model: CHAT_MODEL, + messages: [ + { + role: "system", + content: + "You are concise. Keep responses under five words.", + }, + { + role: "user", + content: "Reply with exactly: observability", + }, + ], + maxTokens: 24, + temperature: 0, + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-chat-stream-operation", + "chat-stream", + async () => { + await withRetry(async () => { + const stream = await client.chat.stream({ model: CHAT_MODEL, messages: [ - { - role: "system", - content: - "You are concise. Keep responses under five words.", - }, { role: "user", - content: "Reply with exactly: observability", + content: "Reply with exactly: streamed output", }, ], maxTokens: 24, + stream: true, temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); - }, - ); - - await runOperation( - "mistral-chat-stream-operation", - "chat-stream", - async () => { - await withRetry(async () => { - const stream = await client.chat.stream({ - model: CHAT_MODEL, - messages: [ - { - role: "user", - content: "Reply with exactly: streamed output", - }, - ], - maxTokens: 24, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); - - await runOperation( - "mistral-chat-reasoning-stream-operation", - "chat-stream-reasoning", - async () => { - await withRetry(async () => { - const stream = await client.chat.stream({ - model: ADJUSTABLE_REASONING_MODEL, - messages: [ - { - role: "user", - content: - "John is one of 4 children. The first sister is 4 years old. Next year, the second sister will be twice as old as the first sister. The third sister is two years older than the second sister. The third sister is half the age of her older brother. How old is John? Reply with just the number.", - }, - ], - maxTokens: 256, - reasoning_effort: "high", - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); - if (supportsThinkingStream) { await runOperation( - "mistral-chat-thinking-stream-operation", - "chat-stream-thinking", + "mistral-chat-reasoning-stream-operation", + "chat-stream-reasoning", async () => { await withRetry(async () => { const stream = await client.chat.stream({ - model: NATIVE_REASONING_MODEL, + model: ADJUSTABLE_REASONING_MODEL, messages: [ { role: "user", - content: "What is 2+2? Reply with just the number.", + content: + "John is one of 4 children. The first sister is 4 years old. Next year, the second sister will be twice as old as the first sister. The third sister is two years older than the second sister. The third sister is half the age of her older brother. How old is John? Reply with just the number.", }, ], - maxTokens: 1024, + maxTokens: 256, + reasoning_effort: "high", stream: true, temperature: 0, }); @@ -460,272 +456,368 @@ async function runMistralInstrumentationScenario( }, MISTRAL_REQUEST_RETRY_OPTIONS); }, ); - } - await runOperation( - "mistral-chat-tool-call-operation", - "chat-tool-call", - async () => { - await withRetry(async () => { - const request = { - model: CHAT_MODEL, - messages: [ - { - role: "user", - content: - "Call the get_weather tool for Vienna. Do not answer with plain text.", - }, - ], - toolChoice: "required", - maxTokens: 48, - temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [getWeatherToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } + if (supportsThinkingStream) { + await runOperation( + "mistral-chat-thinking-stream-operation", + "chat-stream-thinking", + async () => { + await withRetry(async () => { + const stream = await client.chat.stream({ + model: NATIVE_REASONING_MODEL, + messages: [ + { + role: "user", + content: "What is 2+2? Reply with just the number.", + }, + ], + maxTokens: 1024, + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + } - return await client.chat.complete({ - ...request, - tools: [getWeatherToolDefinition({ legacy: true })], - }); - } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - - await withRetry(async () => { - const request = { - model: CHAT_MODEL, - messages: [ - { - role: "user", - content: - "Call the get_exchange_rate tool for USD to EUR. Do not answer with plain text.", - }, - ], - toolChoice: "required", - maxTokens: 48, - temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [getExchangeRateToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; + await runOperation( + "mistral-chat-tool-call-operation", + "chat-tool-call", + async () => { + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "user", + content: + "Call the get_weather tool for Vienna. Do not answer with plain text.", + }, + ], + toolChoice: "required", + maxTokens: 48, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [getWeatherToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; + } + + return await client.chat.complete({ + ...request, + tools: [getWeatherToolDefinition({ legacy: true })], + }); } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); - return await client.chat.complete({ - ...request, - tools: [getExchangeRateToolDefinition({ legacy: true })], - }); - } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - - await withRetry(async () => { - const request = { - model: CHAT_MODEL, - messages: [ - { - role: "system", - content: - "You must return only tool calls and no plain text.", - }, - { - role: "user", - content: - "In a single assistant response, call exactly two tools: get_weather with location Vienna and get_exchange_rate with from_currency USD and to_currency EUR.", - }, - ], - toolChoice: "required", - maxTokens: 96, - temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [ - getWeatherToolDefinition(), - getExchangeRateToolDefinition(), + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "user", + content: + "Call the get_exchange_rate tool for USD to EUR. Do not answer with plain text.", + }, ], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; + toolChoice: "required", + maxTokens: 48, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [getExchangeRateToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; + } + + return await client.chat.complete({ + ...request, + tools: [getExchangeRateToolDefinition({ legacy: true })], + }); } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); - return await client.chat.complete({ - ...request, - tools: [ - getWeatherToolDefinition({ legacy: true }), - getExchangeRateToolDefinition({ legacy: true }), + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "system", + content: + "You must return only tool calls and no plain text.", + }, + { + role: "user", + content: + "In a single assistant response, call exactly two tools: get_weather with location Vienna and get_exchange_rate with from_currency USD and to_currency EUR.", + }, ], - }); - } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - }, - ); - - await runOperation( - "mistral-fim-complete-operation", - "fim-complete", - async () => { - await withRetry( - async () => - client.fim.complete({ + toolChoice: "required", + maxTokens: 96, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [ + getWeatherToolDefinition(), + getExchangeRateToolDefinition(), + ], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; + } + + return await client.chat.complete({ + ...request, + tools: [ + getWeatherToolDefinition({ legacy: true }), + getExchangeRateToolDefinition({ legacy: true }), + ], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + }, + ); + + await runOperation( + "mistral-fim-complete-operation", + "fim-complete", + async () => { + await withRetry( + async () => + client.fim.complete({ + model: FIM_MODEL, + prompt: "function add(a, b) {", + suffix: "}", + maxTokens: 24, + temperature: 0, + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-fim-stream-operation", + "fim-stream", + async () => { + await withRetry(async () => { + const stream = await client.fim.stream({ model: FIM_MODEL, - prompt: "function add(a, b) {", - suffix: "}", - maxTokens: 24, + prompt: "const project = ", + suffix: ";", + maxTokens: 16, + stream: true, temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); - }, - ); - - await runOperation( - "mistral-fim-stream-operation", - "fim-stream", - async () => { - await withRetry(async () => { - const stream = await client.fim.stream({ - model: FIM_MODEL, - prompt: "const project = ", - suffix: ";", - maxTokens: 16, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); - - await runOperation( - "mistral-agents-complete-operation", - "agents-complete", - async () => { - await withRetry( - async () => - client.agents.complete({ + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + + await runOperation( + "mistral-agents-complete-operation", + "agents-complete", + async () => { + await withRetry( + async () => + client.agents.complete({ + agentId, + messages: [ + { + role: "user", + content: "Reply with exactly: agent complete", + }, + ], + responseFormat: { + type: "text", + }, + maxTokens: 16, + temperature: 0, + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-agents-tool-call-operation", + "agents-tool-call", + async () => { + await withRetry(async () => { + const request = { agentId, messages: [ { role: "user", - content: "Reply with exactly: agent complete", + content: + "Call the get_time_in_city tool for Vienna. Do not answer with plain text.", + }, + ], + responseFormat: { + type: "text", + }, + toolChoice: "required", + maxTokens: 32, + temperature: 0, + }; + + try { + return await client.agents.complete({ + ...request, + tools: [getAgentTimeToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; + } + + return await client.agents.complete({ + ...request, + tools: [getAgentTimeToolDefinition({ legacy: true })], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + }, + ); + + await runOperation( + "mistral-agents-stream-operation", + "agents-stream", + async () => { + await withRetry(async () => { + const stream = await client.agents.stream({ + agentId, + messages: [ + { + role: "user", + content: "Reply with exactly: agent stream", }, ], responseFormat: { type: "text", }, maxTokens: 16, + stream: true, temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + + await runOperation( + "mistral-embeddings-operation", + "embeddings-create", + async () => { + await withRetry( + async () => + client.embeddings.create({ + model: EMBEDDING_MODEL, + inputs: "braintrust mistral instrumentation", + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + if (supportsClassifiers) { + await runOperation( + "mistral-classifiers-moderate-operation", + "classifiers-moderate", + async () => { + await client.classifiers.moderate( + { + model: CLASSIFIER_MODEL, + inputs: "A short and harmless moderation fixture.", + }, + classifierRequestOptions, + ); + }, ); - }, - ); - - await runOperation( - "mistral-agents-tool-call-operation", - "agents-tool-call", - async () => { - await withRetry(async () => { - const request = { - agentId, - messages: [ + + await runOperation( + "mistral-classifiers-moderate-chat-operation", + "classifiers-moderate-chat", + async () => { + await client.classifiers.moderateChat( { - role: "user", - content: - "Call the get_time_in_city tool for Vienna. Do not answer with plain text.", + model: CLASSIFIER_MODEL, + inputs: [ + { + role: "user", + content: "Please classify this harmless chat message.", + }, + ], }, - ], - responseFormat: { - type: "text", - }, - toolChoice: "required", - maxTokens: 32, - temperature: 0, - }; - - try { - return await client.agents.complete({ - ...request, - tools: [getAgentTimeToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } + classifierRequestOptions, + ); + }, + ); + } + + if (supportsClassifiers && supportsClassify) { + await runOperation( + "mistral-classifiers-classify-operation", + "classifiers-classify", + async () => { + await client.classifiers.classify( + { + model: CLASSIFIER_MODEL, + inputs: "A positive product review.", + }, + classifierRequestOptions, + ); + }, + ); - return await client.agents.complete({ - ...request, - tools: [getAgentTimeToolDefinition({ legacy: true })], - }); - } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - }, - ); - - await runOperation( - "mistral-agents-stream-operation", - "agents-stream", - async () => { - await withRetry(async () => { - const stream = await client.agents.stream({ - agentId, - messages: [ + await runOperation( + "mistral-classifiers-classify-chat-operation", + "classifiers-classify-chat", + async () => { + await client.classifiers.classifyChat( { - role: "user", - content: "Reply with exactly: agent stream", + model: CLASSIFIER_MODEL, + [classifyChatRequestInputKey]: { + messages: [ + { + role: "user", + content: "I need help with my account.", + }, + ], + }, }, - ], - responseFormat: { - type: "text", - }, - maxTokens: 16, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); - - await runOperation( - "mistral-embeddings-operation", - "embeddings-create", - async () => { - await withRetry( - async () => - client.embeddings.create({ - model: EMBEDDING_MODEL, - inputs: "braintrust mistral instrumentation", - }), - MISTRAL_REQUEST_RETRY_OPTIONS, + classifierRequestOptions, + ); + }, ); - }, - ); - }, - metadata: { - scenario: SCENARIO_NAME, - }, - projectNameBase: "e2e-mistral-instrumentation", - rootName: ROOT_NAME, + } + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-mistral-instrumentation", + rootName: ROOT_NAME, + }); }); } finally { await cleanup(); @@ -742,3 +834,91 @@ export async function runWrappedMistralInstrumentation(Mistral, options) { export async function runAutoMistralInstrumentation(Mistral, options) { await runMistralInstrumentationScenario(Mistral, options); } + +async function readRequestBody(request) { + const chunks = []; + for await (const chunk of request) { + chunks.push(chunk); + } + return Buffer.concat(chunks).toString("utf8"); +} + +function createClassifierResponseForPath(pathname) { + if (pathname.endsWith("/moderations")) { + return { + id: "mod-text", + model: CLASSIFIER_MODEL, + results: [ + { + categories: {}, + category_scores: {}, + }, + ], + }; + } + + if (pathname.endsWith("/chat/moderations")) { + return { + id: "mod-chat", + model: CLASSIFIER_MODEL, + results: [ + { + categories: {}, + category_scores: {}, + }, + ], + }; + } + + if (pathname.endsWith("/classifications")) { + return { + id: "classify-text", + model: CLASSIFIER_MODEL, + results: [{}], + }; + } + + if (pathname.endsWith("/chat/classifications")) { + return { + id: "classify-chat", + model: CLASSIFIER_MODEL, + results: [{}], + }; + } + + return null; +} + +async function withFakeMistralClassifierServer(callback) { + const server = createServer(async (request, response) => { + const url = new URL(request.url || "/", "http://127.0.0.1"); + await readRequestBody(request); + + const body = createClassifierResponseForPath(url.pathname); + if (!body) { + response.writeHead(404, { "Content-Type": "application/json" }); + response.end(JSON.stringify({ message: "not found" })); + return; + } + + response.writeHead(200, { "Content-Type": "application/json" }); + response.end(JSON.stringify(body)); + }); + + await new Promise((resolve) => { + server.listen(0, "127.0.0.1", resolve); + }); + + try { + const address = server.address(); + if (!address || typeof address === "string") { + throw new Error("Fake Mistral server did not bind to a TCP port."); + } + + await callback(`http://127.0.0.1:${address.port}`); + } finally { + await new Promise((resolve, reject) => { + server.close((error) => (error ? reject(error) : resolve())); + }); + } +} 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 3ecf3d47f..7cdf65960 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.mjs @@ -4,6 +4,8 @@ import { runAutoMistralInstrumentation } from "./scenario.impl.mjs"; runMain(async () => runAutoMistralInstrumentation(Mistral, { + supportsClassifiers: false, + supportsClassify: 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 349d45244..757c18306 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts +++ b/e2e/scenarios/mistral-instrumentation/scenario.mistral-v1-3-4.ts @@ -4,6 +4,8 @@ import { runWrappedMistralInstrumentation } from "./scenario.impl.mjs"; runMain(async () => runWrappedMistralInstrumentation(Mistral, { + supportsClassifiers: false, + supportsClassify: false, supportsThinkingStream: false, }), ); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.mjs b/e2e/scenarios/mistral-instrumentation/scenario.mjs index d4eb85870..5c9618e46 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.mjs @@ -2,4 +2,8 @@ import { Mistral } from "mistral-sdk-v2"; import { runMain } from "../../helpers/provider-runtime.mjs"; import { runAutoMistralInstrumentation } from "./scenario.impl.mjs"; -runMain(async () => runAutoMistralInstrumentation(Mistral)); +runMain(async () => + runAutoMistralInstrumentation(Mistral, { + classifyChatRequestInputKey: "input", + }), +); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.test.ts b/e2e/scenarios/mistral-instrumentation/scenario.test.ts index f12eed393..a0af4124b 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.test.ts +++ b/e2e/scenarios/mistral-instrumentation/scenario.test.ts @@ -40,6 +40,12 @@ for (const scenario of mistralScenarios) { ...(scenario.supportsThinkingStream === false ? { supportsThinkingStream: false } : {}), + ...(scenario.supportsClassifiers === false + ? { supportsClassifiers: false } + : {}), + ...(scenario.supportsClassify === false + ? { supportsClassify: false } + : {}), testFileUrl: import.meta.url, timeoutMs: MISTRAL_SCENARIO_TIMEOUT_MS, }); @@ -59,6 +65,12 @@ for (const scenario of mistralScenarios) { ...(scenario.supportsThinkingStream === false ? { supportsThinkingStream: false } : {}), + ...(scenario.supportsClassifiers === false + ? { supportsClassifiers: false } + : {}), + ...(scenario.supportsClassify === false + ? { supportsClassify: false } + : {}), testFileUrl: import.meta.url, timeoutMs: MISTRAL_SCENARIO_TIMEOUT_MS, }); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.ts b/e2e/scenarios/mistral-instrumentation/scenario.ts index 62f148ed9..cbec9bbe5 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.ts +++ b/e2e/scenarios/mistral-instrumentation/scenario.ts @@ -2,4 +2,8 @@ import { Mistral } from "mistral-sdk-v2"; import { runMain } from "../../helpers/scenario-runtime"; import { runWrappedMistralInstrumentation } from "./scenario.impl.mjs"; -runMain(async () => runWrappedMistralInstrumentation(Mistral)); +runMain(async () => + runWrappedMistralInstrumentation(Mistral, { + classifyChatRequestInputKey: "input", + }), +); diff --git a/js/src/auto-instrumentations/configs/mistral.ts b/js/src/auto-instrumentations/configs/mistral.ts index 896f54db1..a16f40832 100644 --- a/js/src/auto-instrumentations/configs/mistral.ts +++ b/js/src/auto-instrumentations/configs/mistral.ts @@ -86,6 +86,118 @@ export const mistralConfigs: InstrumentationConfig[] = [ }, }, + { + channelName: mistralChannels.classifiersModerate.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=1.0.0 <2.0.0", + filePath: "sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "moderate", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersModerate.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=2.0.0 <3.0.0", + filePath: "esm/sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "moderate", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersModerateChat.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=1.0.0 <2.0.0", + filePath: "sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "moderateChat", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersModerateChat.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=2.0.0 <3.0.0", + filePath: "esm/sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "moderateChat", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersClassify.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=1.10.0 <2.0.0", + filePath: "sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "classify", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersClassify.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=2.0.0 <3.0.0", + filePath: "esm/sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "classify", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersClassifyChat.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=1.10.0 <2.0.0", + filePath: "sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "classifyChat", + kind: "Async", + }, + }, + + { + channelName: mistralChannels.classifiersClassifyChat.channelName, + module: { + name: "@mistralai/mistralai", + versionRange: ">=2.0.0 <3.0.0", + filePath: "esm/sdk/classifiers.js", + }, + functionQuery: { + className: "Classifiers", + methodName: "classifyChat", + kind: "Async", + }, + }, + { channelName: mistralChannels.fimComplete.channelName, module: { diff --git a/js/src/instrumentation/plugins/mistral-channels.ts b/js/src/instrumentation/plugins/mistral-channels.ts index 6328d3872..64f115db9 100644 --- a/js/src/instrumentation/plugins/mistral-channels.ts +++ b/js/src/instrumentation/plugins/mistral-channels.ts @@ -4,16 +4,20 @@ import type { MistralAgentsCompletionResponse, MistralAgentsCreateParams, MistralAgentsResult, + MistralChatClassificationCreateParams, MistralChatCompletionEvent, MistralChatCompletionResponse, MistralChatCreateParams, MistralChatResult, + MistralClassificationCreateParams, + MistralClassificationResponse, MistralEmbeddingCreateParams, MistralEmbeddingResponse, MistralFimCompletionEvent, MistralFimCompletionResponse, MistralFimCreateParams, MistralFimResult, + MistralModerationResponse, } from "../../vendor-sdk-types/mistral"; export const mistralChannels = defineChannels("@mistralai/mistralai", { @@ -43,6 +47,38 @@ export const mistralChannels = defineChannels("@mistralai/mistralai", { kind: "async", }), + classifiersModerate: channel< + [MistralClassificationCreateParams], + MistralModerationResponse + >({ + channelName: "classifiers.moderate", + kind: "async", + }), + + classifiersModerateChat: channel< + [MistralChatClassificationCreateParams], + MistralModerationResponse + >({ + channelName: "classifiers.moderateChat", + kind: "async", + }), + + classifiersClassify: channel< + [MistralClassificationCreateParams], + MistralClassificationResponse + >({ + channelName: "classifiers.classify", + kind: "async", + }), + + classifiersClassifyChat: channel< + [MistralChatClassificationCreateParams], + MistralClassificationResponse + >({ + channelName: "classifiers.classifyChat", + kind: "async", + }), + fimComplete: channel<[MistralFimCreateParams], MistralFimCompletionResponse>({ channelName: "fim.complete", kind: "async", diff --git a/js/src/instrumentation/plugins/mistral-plugin.ts b/js/src/instrumentation/plugins/mistral-plugin.ts index da9f122b3..f33f82554 100644 --- a/js/src/instrumentation/plugins/mistral-plugin.ts +++ b/js/src/instrumentation/plugins/mistral-plugin.ts @@ -72,6 +72,50 @@ export class MistralPlugin extends BasePlugin { }), ); + this.unsubscribers.push( + traceAsyncChannel(mistralChannels.classifiersModerate, { + name: "mistral.classifiers.moderate", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }), + ); + + this.unsubscribers.push( + traceAsyncChannel(mistralChannels.classifiersModerateChat, { + name: "mistral.classifiers.moderateChat", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }), + ); + + this.unsubscribers.push( + traceAsyncChannel(mistralChannels.classifiersClassify, { + name: "mistral.classifiers.classify", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }), + ); + + this.unsubscribers.push( + traceAsyncChannel(mistralChannels.classifiersClassifyChat, { + name: "mistral.classifiers.classifyChat", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }), + ); + this.unsubscribers.push( traceStreamingChannel(mistralChannels.fimComplete, { name: "mistral.fim.complete", @@ -301,6 +345,21 @@ function extractEmbeddingInputWithMetadata(args: unknown[] | unknown): { }; } +function extractClassifierInputWithMetadata(args: unknown[] | unknown): { + input: unknown; + metadata: Record; +} { + const params = getMistralRequestArg(args); + const { input, inputs, ...rawMetadata } = params || {}; + + return { + input: processInputAttachments(inputs ?? input), + metadata: addMistralProviderMetadata( + extractMistralRequestMetadata(rawMetadata), + ), + }; +} + function extractPromptInputWithMetadata(args: unknown[] | unknown): { input: unknown; metadata: Record; @@ -347,6 +406,10 @@ function extractMistralStreamOutput(result: unknown): unknown { return isObject(result) ? result.choices : undefined; } +function extractClassifierOutput(result: unknown): unknown { + return isObject(result) ? result.results : undefined; +} + function extractMistralStreamingMetrics( result: unknown, startTime?: number, diff --git a/js/src/vendor-sdk-types/mistral.ts b/js/src/vendor-sdk-types/mistral.ts index 1a7edc8d8..10f0b00c0 100644 --- a/js/src/vendor-sdk-types/mistral.ts +++ b/js/src/vendor-sdk-types/mistral.ts @@ -98,6 +98,18 @@ export type MistralAgentsCreateParams = { [key: string]: unknown; }; +export type MistralClassificationCreateParams = { + input?: unknown; + inputs?: unknown; + [key: string]: unknown; +}; + +export type MistralChatClassificationCreateParams = { + input?: unknown; + inputs?: unknown; + [key: string]: unknown; +}; + export type MistralEmbeddingResponse = { id?: string; object?: string; @@ -128,6 +140,16 @@ export type MistralAgentsCompletionResponse = MistralChatCompletionResponse; export type MistralFimCompletionEvent = MistralChatCompletionEvent; export type MistralAgentsCompletionEvent = MistralChatCompletionEvent; +export type MistralClassificationResponse = { + id?: string; + model?: string; + results?: unknown; + usage?: unknown; + [key: string]: unknown; +}; + +export type MistralModerationResponse = MistralClassificationResponse; + export type MistralChatResult = | MistralChatCompletionResponse | MistralChatStreamingResult; @@ -178,10 +200,30 @@ export type MistralAgents = { ) => Promise; }; +export type MistralClassifiers = { + moderate: ( + request: MistralClassificationCreateParams, + options?: unknown, + ) => Promise; + moderateChat: ( + request: MistralChatClassificationCreateParams, + options?: unknown, + ) => Promise; + classify?: ( + request: MistralClassificationCreateParams, + options?: unknown, + ) => Promise; + classifyChat?: ( + request: MistralChatClassificationCreateParams, + options?: unknown, + ) => Promise; +}; + export type MistralClient = { chat?: MistralChat; fim?: MistralFim; agents?: MistralAgents; embeddings?: MistralEmbeddings; + classifiers?: MistralClassifiers; [key: string]: unknown; }; diff --git a/js/src/wrappers/mistral.ts b/js/src/wrappers/mistral.ts index b2ce06aff..c449997f7 100644 --- a/js/src/wrappers/mistral.ts +++ b/js/src/wrappers/mistral.ts @@ -5,9 +5,13 @@ import type { MistralAgentsCreateParams, MistralAgentsStreamingResult, MistralChat, + MistralChatClassificationCreateParams, MistralChatCompletionResponse, MistralChatCreateParams, MistralChatStreamingResult, + MistralClassificationCreateParams, + MistralClassificationResponse, + MistralClassifiers, MistralClient, MistralEmbeddingCreateParams, MistralEmbeddingResponse, @@ -16,6 +20,7 @@ import type { MistralFimCompletionResponse, MistralFimCreateParams, MistralFimStreamingResult, + MistralModerationResponse, } from "../vendor-sdk-types/mistral"; /** @@ -52,7 +57,8 @@ function isSupportedMistralClient(value: unknown): value is MistralClient { (value.chat !== undefined && hasChat(value.chat)) || (value.embeddings !== undefined && hasEmbeddings(value.embeddings)) || (value.fim !== undefined && hasFim(value.fim)) || - (value.agents !== undefined && hasAgents(value.agents)) + (value.agents !== undefined && hasAgents(value.agents)) || + (value.classifiers !== undefined && hasClassifiers(value.classifiers)) ); } @@ -72,6 +78,10 @@ function hasAgents(value: unknown): value is MistralAgents { return hasFunction(value, "complete") && hasFunction(value, "stream"); } +function hasClassifiers(value: unknown): value is MistralClassifiers { + return hasFunction(value, "moderate") && hasFunction(value, "moderateChat"); +} + function mistralProxy(mistral: MistralClient): MistralClient { return new Proxy(mistral, { get(target, prop, receiver) { @@ -86,6 +96,10 @@ function mistralProxy(mistral: MistralClient): MistralClient { return target.embeddings ? embeddingsProxy(target.embeddings) : target.embeddings; + case "classifiers": + return target.classifiers + ? classifiersProxy(target.classifiers) + : target.classifiers; default: return Reflect.get(target, prop, receiver); } @@ -153,6 +167,30 @@ function agentsProxy(agents: MistralAgents): MistralAgents { }); } +function classifiersProxy(classifiers: MistralClassifiers): MistralClassifiers { + return new Proxy(classifiers, { + get(target, prop, receiver) { + if (prop === "moderate") { + return wrapClassifiersModerate(target.moderate.bind(target)); + } + + if (prop === "moderateChat") { + return wrapClassifiersModerateChat(target.moderateChat.bind(target)); + } + + if (prop === "classify" && target.classify) { + return wrapClassifiersClassify(target.classify.bind(target)); + } + + if (prop === "classifyChat" && target.classifyChat) { + return wrapClassifiersClassifyChat(target.classifyChat.bind(target)); + } + + return Reflect.get(target, prop, receiver); + }, + }); +} + function wrapChatComplete( complete: ( request: MistralChatCreateParams, @@ -193,6 +231,58 @@ function wrapEmbeddingsCreate( ); } +function wrapClassifiersModerate( + moderate: ( + request: MistralClassificationCreateParams, + options?: unknown, + ) => Promise, +): MistralClassifiers["moderate"] { + return (request, options) => + mistralChannels.classifiersModerate.tracePromise( + () => moderate(request, options), + { arguments: [request] }, + ); +} + +function wrapClassifiersModerateChat( + moderateChat: ( + request: MistralChatClassificationCreateParams, + options?: unknown, + ) => Promise, +): MistralClassifiers["moderateChat"] { + return (request, options) => + mistralChannels.classifiersModerateChat.tracePromise( + () => moderateChat(request, options), + { arguments: [request] }, + ); +} + +function wrapClassifiersClassify( + classify: ( + request: MistralClassificationCreateParams, + options?: unknown, + ) => Promise, +): NonNullable { + return (request, options) => + mistralChannels.classifiersClassify.tracePromise( + () => classify(request, options), + { arguments: [request] }, + ); +} + +function wrapClassifiersClassifyChat( + classifyChat: ( + request: MistralChatClassificationCreateParams, + options?: unknown, + ) => Promise, +): NonNullable { + return (request, options) => + mistralChannels.classifiersClassifyChat.tracePromise( + () => classifyChat(request, options), + { arguments: [request] }, + ); +} + function wrapFimComplete( complete: ( request: MistralFimCreateParams, From 87911743d3a4ba4c85ef1e8fccd3026b2f3c2f1a Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 4 May 2026 15:40:21 -0700 Subject: [PATCH 2/3] cs --- .changeset/petite-pugs-relax.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/petite-pugs-relax.md diff --git a/.changeset/petite-pugs-relax.md b/.changeset/petite-pugs-relax.md new file mode 100644 index 000000000..cd2b5968b --- /dev/null +++ b/.changeset/petite-pugs-relax.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat(mistral): Instrument classification and moderation APIs From 8dd0003fe7c13513605fc6d08f257cebd5d4828d Mon Sep 17 00:00:00 2001 From: Luca Forstner Date: Mon, 4 May 2026 16:55:27 -0700 Subject: [PATCH 3/3] no fake stuff --- .../mistral-v1-10-0.log-payloads.json | 59 -- .../mistral-v1-10-0.span-events.json | 62 -- .../mistral-v1-14-1.log-payloads.json | 59 -- .../mistral-v1-14-1.span-events.json | 62 -- .../mistral-v1-15-1.log-payloads.json | 59 -- .../mistral-v1-15-1.span-events.json | 62 -- .../mistral-v1.log-payloads.json | 59 -- .../__snapshots__/mistral-v1.span-events.json | 62 -- .../mistral-v2.log-payloads.json | 59 -- .../__snapshots__/mistral-v2.span-events.json | 62 -- .../mistral-instrumentation/assertions.ts | 14 +- .../mistral-instrumentation/scenario.impl.mjs | 822 ++++++++---------- 12 files changed, 370 insertions(+), 1071 deletions(-) diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json index 11e25c2cc..d0624284c 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.log-payloads.json @@ -560,64 +560,5 @@ "type": "array" }, "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": "", - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": { - "keys": [ - "messages" - ], - "type": "object" - }, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json index 8bdafb6b2..5242f88b3 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-10-0.span-events.json @@ -511,67 +511,5 @@ "" ], "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classify", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-chat-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classifyChat", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json index 11e25c2cc..d0624284c 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.log-payloads.json @@ -560,64 +560,5 @@ "type": "array" }, "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": "", - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": { - "keys": [ - "messages" - ], - "type": "object" - }, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json index 8bdafb6b2..5242f88b3 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-14-1.span-events.json @@ -511,67 +511,5 @@ "" ], "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classify", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-chat-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classifyChat", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json index 11e25c2cc..d0624284c 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.log-payloads.json @@ -560,64 +560,5 @@ "type": "array" }, "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": "", - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": { - "keys": [ - "messages" - ], - "type": "object" - }, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json index 8bdafb6b2..5242f88b3 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1-15-1.span-events.json @@ -511,67 +511,5 @@ "" ], "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classify", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-chat-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classifyChat", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json index 11e25c2cc..d0624284c 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.log-payloads.json @@ -560,64 +560,5 @@ "type": "array" }, "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": "", - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": { - "keys": [ - "messages" - ], - "type": "object" - }, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json index 8bdafb6b2..5242f88b3 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v1.span-events.json @@ -511,67 +511,5 @@ "" ], "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classify", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-chat-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classifyChat", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json index 5e2517dc6..d57f77c09 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.log-payloads.json @@ -572,64 +572,5 @@ "type": "array" }, "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": "", - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" - }, - { - "has_input": false, - "has_output": false, - "input": null, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "output": null, - "span_id": "" - }, - { - "has_input": true, - "has_output": true, - "input": { - "keys": [ - "messages" - ], - "type": "object" - }, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "output": { - "choice_count": 1, - "finish_reason": null, - "type": "array" - }, - "span_id": "" } ] diff --git a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json index 18b697ea2..842dff64e 100644 --- a/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json +++ b/e2e/scenarios/mistral-instrumentation/__snapshots__/mistral-v2.span-events.json @@ -523,67 +523,5 @@ "" ], "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classify", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" - }, - { - "has_input": false, - "has_output": false, - "metadata": { - "operation": "classifiers-classify-chat" - }, - "metric_keys": [], - "name": "mistral-classifiers-classify-chat-operation", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": null - }, - { - "has_input": true, - "has_output": true, - "metadata": { - "model": "mistral-moderation-2603", - "provider": "mistral" - }, - "metric_keys": [], - "name": "mistral.classifiers.classifyChat", - "root_span_id": "", - "span_id": "", - "span_parents": [ - "" - ], - "type": "llm" } ] diff --git a/e2e/scenarios/mistral-instrumentation/assertions.ts b/e2e/scenarios/mistral-instrumentation/assertions.ts index 784510550..ea50d39c0 100644 --- a/e2e/scenarios/mistral-instrumentation/assertions.ts +++ b/e2e/scenarios/mistral-instrumentation/assertions.ts @@ -67,6 +67,12 @@ function isRecord(value: Json | undefined): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } +function nonEmptyString(value: string | undefined): string | null { + return typeof value === "string" && value.trim().length > 0 + ? value.trim() + : null; +} + function pickMetadata( metadata: Record | undefined, keys: string[], @@ -643,7 +649,9 @@ export function defineMistralInstrumentationAssertions(options: { ); const supportsThinkingStream = options.supportsThinkingStream ?? true; const supportsClassifiers = options.supportsClassifiers ?? true; - const supportsClassify = options.supportsClassify ?? true; + const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); + const supportsClassify = + (options.supportsClassify ?? true) && !!classifyModel; const testConfig = { timeout: options.timeoutMs, }; @@ -1150,7 +1158,7 @@ export function defineMistralInstrumentationAssertions(options: { expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); expect(span?.span.type).toBe("llm"); expect(span?.row.metadata).toMatchObject({ - model: CLASSIFIER_MODEL, + model: classifyModel, provider: "mistral", }); expect(span?.input).toEqual(expect.any(String)); @@ -1173,7 +1181,7 @@ export function defineMistralInstrumentationAssertions(options: { expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); expect(span?.span.type).toBe("llm"); expect(span?.row.metadata).toMatchObject({ - model: CLASSIFIER_MODEL, + model: classifyModel, provider: "mistral", }); expect(span?.input).toEqual(expect.any(Object)); diff --git a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs index 2a95212c5..a59934602 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs @@ -1,5 +1,4 @@ import { wrapMistral } from "braintrust"; -import { createServer } from "node:http"; import { collectAsync, runOperation, @@ -376,79 +375,100 @@ async function runMistralInstrumentationScenario( apiKey: process.env.MISTRAL_API_KEY, }); const client = decorateClient ? decorateClient(baseClient) : baseClient; + const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); const { agentId, cleanup } = await resolveAgentRuntime(baseClient); try { - await withFakeMistralClassifierServer(async (classifierServerURL) => { - const classifierRequestOptions = { serverURL: classifierServerURL }; - - await runTracedScenario({ - callback: async () => { - await runOperation( - "mistral-chat-complete-operation", - "chat-complete", - async () => { - await withRetry( - async () => - client.chat.complete({ - model: CHAT_MODEL, - messages: [ - { - role: "system", - content: - "You are concise. Keep responses under five words.", - }, - { - role: "user", - content: "Reply with exactly: observability", - }, - ], - maxTokens: 24, - temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); - }, - ); - - await runOperation( - "mistral-chat-stream-operation", - "chat-stream", - async () => { - await withRetry(async () => { - const stream = await client.chat.stream({ + await runTracedScenario({ + callback: async () => { + await runOperation( + "mistral-chat-complete-operation", + "chat-complete", + async () => { + await withRetry( + async () => + client.chat.complete({ model: CHAT_MODEL, messages: [ + { + role: "system", + content: + "You are concise. Keep responses under five words.", + }, { role: "user", - content: "Reply with exactly: streamed output", + content: "Reply with exactly: observability", }, ], maxTokens: 24, - stream: true, temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-chat-stream-operation", + "chat-stream", + async () => { + await withRetry(async () => { + const stream = await client.chat.stream({ + model: CHAT_MODEL, + messages: [ + { + role: "user", + content: "Reply with exactly: streamed output", + }, + ], + maxTokens: 24, + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + + await runOperation( + "mistral-chat-reasoning-stream-operation", + "chat-stream-reasoning", + async () => { + await withRetry(async () => { + const stream = await client.chat.stream({ + model: ADJUSTABLE_REASONING_MODEL, + messages: [ + { + role: "user", + content: + "John is one of 4 children. The first sister is 4 years old. Next year, the second sister will be twice as old as the first sister. The third sister is two years older than the second sister. The third sister is half the age of her older brother. How old is John? Reply with just the number.", + }, + ], + maxTokens: 256, + reasoning_effort: "high", + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + if (supportsThinkingStream) { await runOperation( - "mistral-chat-reasoning-stream-operation", - "chat-stream-reasoning", + "mistral-chat-thinking-stream-operation", + "chat-stream-thinking", async () => { await withRetry(async () => { const stream = await client.chat.stream({ - model: ADJUSTABLE_REASONING_MODEL, + model: NATIVE_REASONING_MODEL, messages: [ { role: "user", - content: - "John is one of 4 children. The first sister is 4 years old. Next year, the second sister will be twice as old as the first sister. The third sister is two years older than the second sister. The third sister is half the age of her older brother. How old is John? Reply with just the number.", + content: "What is 2+2? Reply with just the number.", }, ], - maxTokens: 256, - reasoning_effort: "high", + maxTokens: 1024, stream: true, temperature: 0, }); @@ -456,368 +476,332 @@ async function runMistralInstrumentationScenario( }, MISTRAL_REQUEST_RETRY_OPTIONS); }, ); + } - if (supportsThinkingStream) { - await runOperation( - "mistral-chat-thinking-stream-operation", - "chat-stream-thinking", - async () => { - await withRetry(async () => { - const stream = await client.chat.stream({ - model: NATIVE_REASONING_MODEL, - messages: [ - { - role: "user", - content: "What is 2+2? Reply with just the number.", - }, - ], - maxTokens: 1024, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); - } + await runOperation( + "mistral-chat-tool-call-operation", + "chat-tool-call", + async () => { + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "user", + content: + "Call the get_weather tool for Vienna. Do not answer with plain text.", + }, + ], + toolChoice: "required", + maxTokens: 48, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [getWeatherToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; + } - await runOperation( - "mistral-chat-tool-call-operation", - "chat-tool-call", - async () => { - await withRetry(async () => { - const request = { - model: CHAT_MODEL, - messages: [ - { - role: "user", - content: - "Call the get_weather tool for Vienna. Do not answer with plain text.", - }, - ], - toolChoice: "required", - maxTokens: 48, - temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [getWeatherToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } - - return await client.chat.complete({ - ...request, - tools: [getWeatherToolDefinition({ legacy: true })], - }); + return await client.chat.complete({ + ...request, + tools: [getWeatherToolDefinition({ legacy: true })], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "user", + content: + "Call the get_exchange_rate tool for USD to EUR. Do not answer with plain text.", + }, + ], + toolChoice: "required", + maxTokens: 48, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [getExchangeRateToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - await withRetry(async () => { - const request = { - model: CHAT_MODEL, - messages: [ - { - role: "user", - content: - "Call the get_exchange_rate tool for USD to EUR. Do not answer with plain text.", - }, + return await client.chat.complete({ + ...request, + tools: [getExchangeRateToolDefinition({ legacy: true })], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + + await withRetry(async () => { + const request = { + model: CHAT_MODEL, + messages: [ + { + role: "system", + content: + "You must return only tool calls and no plain text.", + }, + { + role: "user", + content: + "In a single assistant response, call exactly two tools: get_weather with location Vienna and get_exchange_rate with from_currency USD and to_currency EUR.", + }, + ], + toolChoice: "required", + maxTokens: 96, + temperature: 0, + }; + + try { + return await client.chat.complete({ + ...request, + tools: [ + getWeatherToolDefinition(), + getExchangeRateToolDefinition(), ], - toolChoice: "required", - maxTokens: 48, - temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [getExchangeRateToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } - - return await client.chat.complete({ - ...request, - tools: [getExchangeRateToolDefinition({ legacy: true })], - }); + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - await withRetry(async () => { - const request = { - model: CHAT_MODEL, + return await client.chat.complete({ + ...request, + tools: [ + getWeatherToolDefinition({ legacy: true }), + getExchangeRateToolDefinition({ legacy: true }), + ], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + }, + ); + + await runOperation( + "mistral-fim-complete-operation", + "fim-complete", + async () => { + await withRetry( + async () => + client.fim.complete({ + model: FIM_MODEL, + prompt: "function add(a, b) {", + suffix: "}", + maxTokens: 24, + temperature: 0, + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-fim-stream-operation", + "fim-stream", + async () => { + await withRetry(async () => { + const stream = await client.fim.stream({ + model: FIM_MODEL, + prompt: "const project = ", + suffix: ";", + maxTokens: 16, + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + + await runOperation( + "mistral-agents-complete-operation", + "agents-complete", + async () => { + await withRetry( + async () => + client.agents.complete({ + agentId, messages: [ - { - role: "system", - content: - "You must return only tool calls and no plain text.", - }, { role: "user", - content: - "In a single assistant response, call exactly two tools: get_weather with location Vienna and get_exchange_rate with from_currency USD and to_currency EUR.", + content: "Reply with exactly: agent complete", }, ], - toolChoice: "required", - maxTokens: 96, + responseFormat: { + type: "text", + }, + maxTokens: 16, temperature: 0, - }; - - try { - return await client.chat.complete({ - ...request, - tools: [ - getWeatherToolDefinition(), - getExchangeRateToolDefinition(), - ], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } - - return await client.chat.complete({ - ...request, - tools: [ - getWeatherToolDefinition({ legacy: true }), - getExchangeRateToolDefinition({ legacy: true }), - ], - }); + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + + await runOperation( + "mistral-agents-tool-call-operation", + "agents-tool-call", + async () => { + await withRetry(async () => { + const request = { + agentId, + messages: [ + { + role: "user", + content: + "Call the get_time_in_city tool for Vienna. Do not answer with plain text.", + }, + ], + responseFormat: { + type: "text", + }, + toolChoice: "required", + maxTokens: 32, + temperature: 0, + }; + + try { + return await client.agents.complete({ + ...request, + tools: [getAgentTimeToolDefinition()], + }); + } catch (error) { + if (!isMistralInputValidationError(error)) { + throw error; } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); - }, - ); - await runOperation( - "mistral-fim-complete-operation", - "fim-complete", - async () => { - await withRetry( - async () => - client.fim.complete({ - model: FIM_MODEL, - prompt: "function add(a, b) {", - suffix: "}", - maxTokens: 24, - temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); - }, - ); + return await client.agents.complete({ + ...request, + tools: [getAgentTimeToolDefinition({ legacy: true })], + }); + } + }, MISTRAL_REQUEST_RETRY_OPTIONS); + await simulateToolExecutionDelay(); + }, + ); + + await runOperation( + "mistral-agents-stream-operation", + "agents-stream", + async () => { + await withRetry(async () => { + const stream = await client.agents.stream({ + agentId, + messages: [ + { + role: "user", + content: "Reply with exactly: agent stream", + }, + ], + responseFormat: { + type: "text", + }, + maxTokens: 16, + stream: true, + temperature: 0, + }); + await collectAsync(stream); + }, MISTRAL_REQUEST_RETRY_OPTIONS); + }, + ); + + await runOperation( + "mistral-embeddings-operation", + "embeddings-create", + async () => { + await withRetry( + async () => + client.embeddings.create({ + model: EMBEDDING_MODEL, + inputs: "braintrust mistral instrumentation", + }), + MISTRAL_REQUEST_RETRY_OPTIONS, + ); + }, + ); + if (supportsClassifiers) { await runOperation( - "mistral-fim-stream-operation", - "fim-stream", + "mistral-classifiers-moderate-operation", + "classifiers-moderate", async () => { - await withRetry(async () => { - const stream = await client.fim.stream({ - model: FIM_MODEL, - prompt: "const project = ", - suffix: ";", - maxTokens: 16, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); + await client.classifiers.moderate({ + model: CLASSIFIER_MODEL, + inputs: "A short and harmless moderation fixture.", + }); }, ); await runOperation( - "mistral-agents-complete-operation", - "agents-complete", + "mistral-classifiers-moderate-chat-operation", + "classifiers-moderate-chat", async () => { - await withRetry( - async () => - client.agents.complete({ - agentId, - messages: [ - { - role: "user", - content: "Reply with exactly: agent complete", - }, - ], - responseFormat: { - type: "text", - }, - maxTokens: 16, - temperature: 0, - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); + await client.classifiers.moderateChat({ + model: CLASSIFIER_MODEL, + inputs: [ + { + role: "user", + content: "Please classify this harmless chat message.", + }, + ], + }); }, ); + } + if (supportsClassifiers && supportsClassify && classifyModel) { await runOperation( - "mistral-agents-tool-call-operation", - "agents-tool-call", + "mistral-classifiers-classify-operation", + "classifiers-classify", async () => { - await withRetry(async () => { - const request = { - agentId, - messages: [ - { - role: "user", - content: - "Call the get_time_in_city tool for Vienna. Do not answer with plain text.", - }, - ], - responseFormat: { - type: "text", - }, - toolChoice: "required", - maxTokens: 32, - temperature: 0, - }; - - try { - return await client.agents.complete({ - ...request, - tools: [getAgentTimeToolDefinition()], - }); - } catch (error) { - if (!isMistralInputValidationError(error)) { - throw error; - } - - return await client.agents.complete({ - ...request, - tools: [getAgentTimeToolDefinition({ legacy: true })], - }); - } - }, MISTRAL_REQUEST_RETRY_OPTIONS); - await simulateToolExecutionDelay(); + await client.classifiers.classify({ + model: classifyModel, + inputs: "A positive product review.", + }); }, ); await runOperation( - "mistral-agents-stream-operation", - "agents-stream", + "mistral-classifiers-classify-chat-operation", + "classifiers-classify-chat", async () => { - await withRetry(async () => { - const stream = await client.agents.stream({ - agentId, + await client.classifiers.classifyChat({ + model: classifyModel, + [classifyChatRequestInputKey]: { messages: [ { role: "user", - content: "Reply with exactly: agent stream", + content: "I need help with my account.", }, ], - responseFormat: { - type: "text", - }, - maxTokens: 16, - stream: true, - temperature: 0, - }); - await collectAsync(stream); - }, MISTRAL_REQUEST_RETRY_OPTIONS); - }, - ); - - await runOperation( - "mistral-embeddings-operation", - "embeddings-create", - async () => { - await withRetry( - async () => - client.embeddings.create({ - model: EMBEDDING_MODEL, - inputs: "braintrust mistral instrumentation", - }), - MISTRAL_REQUEST_RETRY_OPTIONS, - ); + }, + }); }, ); - - if (supportsClassifiers) { - await runOperation( - "mistral-classifiers-moderate-operation", - "classifiers-moderate", - async () => { - await client.classifiers.moderate( - { - model: CLASSIFIER_MODEL, - inputs: "A short and harmless moderation fixture.", - }, - classifierRequestOptions, - ); - }, - ); - - await runOperation( - "mistral-classifiers-moderate-chat-operation", - "classifiers-moderate-chat", - async () => { - await client.classifiers.moderateChat( - { - model: CLASSIFIER_MODEL, - inputs: [ - { - role: "user", - content: "Please classify this harmless chat message.", - }, - ], - }, - classifierRequestOptions, - ); - }, - ); - } - - if (supportsClassifiers && supportsClassify) { - await runOperation( - "mistral-classifiers-classify-operation", - "classifiers-classify", - async () => { - await client.classifiers.classify( - { - model: CLASSIFIER_MODEL, - inputs: "A positive product review.", - }, - classifierRequestOptions, - ); - }, - ); - - await runOperation( - "mistral-classifiers-classify-chat-operation", - "classifiers-classify-chat", - async () => { - await client.classifiers.classifyChat( - { - model: CLASSIFIER_MODEL, - [classifyChatRequestInputKey]: { - messages: [ - { - role: "user", - content: "I need help with my account.", - }, - ], - }, - }, - classifierRequestOptions, - ); - }, - ); - } - }, - metadata: { - scenario: SCENARIO_NAME, - }, - projectNameBase: "e2e-mistral-instrumentation", - rootName: ROOT_NAME, - }); + } + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-mistral-instrumentation", + rootName: ROOT_NAME, }); } finally { await cleanup(); @@ -834,91 +818,3 @@ export async function runWrappedMistralInstrumentation(Mistral, options) { export async function runAutoMistralInstrumentation(Mistral, options) { await runMistralInstrumentationScenario(Mistral, options); } - -async function readRequestBody(request) { - const chunks = []; - for await (const chunk of request) { - chunks.push(chunk); - } - return Buffer.concat(chunks).toString("utf8"); -} - -function createClassifierResponseForPath(pathname) { - if (pathname.endsWith("/moderations")) { - return { - id: "mod-text", - model: CLASSIFIER_MODEL, - results: [ - { - categories: {}, - category_scores: {}, - }, - ], - }; - } - - if (pathname.endsWith("/chat/moderations")) { - return { - id: "mod-chat", - model: CLASSIFIER_MODEL, - results: [ - { - categories: {}, - category_scores: {}, - }, - ], - }; - } - - if (pathname.endsWith("/classifications")) { - return { - id: "classify-text", - model: CLASSIFIER_MODEL, - results: [{}], - }; - } - - if (pathname.endsWith("/chat/classifications")) { - return { - id: "classify-chat", - model: CLASSIFIER_MODEL, - results: [{}], - }; - } - - return null; -} - -async function withFakeMistralClassifierServer(callback) { - const server = createServer(async (request, response) => { - const url = new URL(request.url || "/", "http://127.0.0.1"); - await readRequestBody(request); - - const body = createClassifierResponseForPath(url.pathname); - if (!body) { - response.writeHead(404, { "Content-Type": "application/json" }); - response.end(JSON.stringify({ message: "not found" })); - return; - } - - response.writeHead(200, { "Content-Type": "application/json" }); - response.end(JSON.stringify(body)); - }); - - await new Promise((resolve) => { - server.listen(0, "127.0.0.1", resolve); - }); - - try { - const address = server.address(); - if (!address || typeof address === "string") { - throw new Error("Fake Mistral server did not bind to a TCP port."); - } - - await callback(`http://127.0.0.1:${address.port}`); - } finally { - await new Promise((resolve, reject) => { - server.close((error) => (error ? reject(error) : resolve())); - }); - } -}