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 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..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 @@ -500,5 +500,65 @@ "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": "" } ] 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..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 @@ -449,5 +449,67 @@ "" ], "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" } ] 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..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 @@ -500,5 +500,65 @@ "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": "" } ] 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..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 @@ -449,5 +449,67 @@ "" ], "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" } ] 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..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 @@ -500,5 +500,65 @@ "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": "" } ] 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..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 @@ -449,5 +449,67 @@ "" ], "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" } ] 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..d0624284c 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,65 @@ "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": "" } ] 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..5242f88b3 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,67 @@ "" ], "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" } ] 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..d57f77c09 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,65 @@ "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": "" } ] 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..842dff64e 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,67 @@ "" ], "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" } ] diff --git a/e2e/scenarios/mistral-instrumentation/assertions.ts b/e2e/scenarios/mistral-instrumentation/assertions.ts index 131558574..ea50d39c0 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, @@ -66,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[], @@ -439,6 +446,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 +506,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 +572,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 +633,8 @@ export function defineMistralInstrumentationAssertions(options: { name: string; runScenario: RunMistralScenario; snapshotName: string; + supportsClassifiers?: boolean; + supportsClassify?: boolean; supportsThinkingStream?: boolean; testFileUrl: string; timeoutMs: number; @@ -603,6 +648,10 @@ export function defineMistralInstrumentationAssertions(options: { `${options.snapshotName}.log-payloads.json`, ); const supportsThinkingStream = options.supportsThinkingStream ?? true; + const supportsClassifiers = options.supportsClassifiers ?? true; + const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); + const supportsClassify = + (options.supportsClassify ?? true) && !!classifyModel; const testConfig = { timeout: options.timeoutMs, }; @@ -1043,6 +1092,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: classifyModel, + 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: classifyModel, + 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..a59934602 100644 --- a/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/mistral-instrumentation/scenario.impl.mjs @@ -8,6 +8,7 @@ import { ADJUSTABLE_REASONING_MODEL, AGENT_MODEL, CHAT_MODEL, + CLASSIFIER_MODEL, EMBEDDING_MODEL, FIM_MODEL, NATIVE_REASONING_MODEL, @@ -24,6 +25,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 +34,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,12 +363,19 @@ 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, }); const client = decorateClient ? decorateClient(baseClient) : baseClient; + const classifyModel = nonEmptyString(process.env.MISTRAL_CLASSIFIER_MODEL); const { agentId, cleanup } = await resolveAgentRuntime(baseClient); try { @@ -720,6 +736,66 @@ async function runMistralInstrumentationScenario( ); }, ); + + 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.", + }); + }, + ); + + 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.", + }, + ], + }); + }, + ); + } + + if (supportsClassifiers && supportsClassify && classifyModel) { + await runOperation( + "mistral-classifiers-classify-operation", + "classifiers-classify", + async () => { + await client.classifiers.classify({ + model: classifyModel, + inputs: "A positive product review.", + }); + }, + ); + + await runOperation( + "mistral-classifiers-classify-chat-operation", + "classifiers-classify-chat", + async () => { + await client.classifiers.classifyChat({ + model: classifyModel, + [classifyChatRequestInputKey]: { + messages: [ + { + role: "user", + content: "I need help with my account.", + }, + ], + }, + }); + }, + ); + } }, metadata: { scenario: SCENARIO_NAME, 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,