diff --git a/.changeset/voyage-instrumentation.md b/.changeset/voyage-instrumentation.md new file mode 100644 index 000000000..698f148e0 --- /dev/null +++ b/.changeset/voyage-instrumentation.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +feat: Add voyage instrumentation diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 084ebc487..f888b67bf 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -420,6 +420,21 @@ } ] }, + { + "scenarioDirName": "voyageai-instrumentation", + "label": "Voyage AI Instrumentation", + "metadataScenario": "voyageai-instrumentation", + "variants": [ + { + "variantKey": "voyageai-v0", + "label": "v0 pinned" + }, + { + "variantKey": "voyageai-v0-latest", + "label": "v0 latest" + } + ] + }, { "scenarioDirName": "openrouter-instrumentation", "label": "OpenRouter Instrumentation", diff --git a/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.json b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.json new file mode 100644 index 000000000..c92b29bda --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.json @@ -0,0 +1,159 @@ +{ + "span_tree": [ + { + "name": "voyageai-instrumentation-root", + "type": "task", + "children": [ + { + "name": "voyageai-embed-operation", + "children": [ + { + "name": "voyageai.embed", + "type": "llm", + "children": [], + "input": [ + "Braintrust tracing" + ], + "output": { + "embedding_count": 1, + "embedding_length": 3 + }, + "metadata": { + "inputType": "document", + "model": "voyage-4", + "outputDimension": 3, + "provider": "voyage" + }, + "metrics": { + "tokens": 4 + } + } + ], + "metadata": { + "operation": "embed", + "testRunId": "" + } + }, + { + "name": "voyageai-multimodal-embed-operation", + "children": [ + { + "name": "voyageai.multimodalEmbed", + "type": "llm", + "children": [], + "input": [ + { + "content": [ + { + "text": "A yellow banana.", + "type": "text" + }, + { + "imageUrl": "https://example.com/banana.jpg", + "type": "image_url" + } + ] + } + ], + "output": { + "embedding_count": 1, + "embedding_length": 2 + }, + "metadata": { + "inputType": "document", + "model": "voyage-multimodal-3.5", + "provider": "voyage" + }, + "metrics": { + "tokens": 6 + } + } + ], + "metadata": { + "operation": "multimodal-embed", + "testRunId": "" + } + }, + { + "name": "voyageai-rerank-operation", + "children": [ + { + "name": "voyageai.rerank", + "type": "llm", + "children": [], + "input": { + "documents": [ + "Vienna is in Austria.", + "Paris is in France." + ], + "query": "Which document is about France?" + }, + "output": [ + { + "index": 1, + "relevance_score": 0.97 + }, + { + "index": 0, + "relevance_score": 0.21 + } + ], + "metadata": { + "document_count": 2, + "model": "rerank-2.5", + "provider": "voyage", + "returnDocuments": true, + "topK": 2 + }, + "metrics": { + "tokens": 11 + } + } + ], + "metadata": { + "operation": "rerank", + "testRunId": "" + } + }, + { + "name": "voyageai-contextualized-embed-operation", + "children": [ + { + "name": "voyageai.contextualizedEmbed", + "type": "llm", + "children": [], + "input": [ + [ + "Paris is a city.", + "It is the capital of France." + ] + ], + "output": { + "document_count": 1, + "embedding_count": 2, + "embedding_length": 4 + }, + "metadata": { + "inputType": "document", + "model": "voyage-context-3", + "outputDimension": 4, + "provider": "voyage" + }, + "metrics": { + "tokens": 8 + } + } + ], + "metadata": { + "operation": "contextualized-embed", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "voyageai-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.txt b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.txt new file mode 100644 index 000000000..6300a04ee --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0-latest.span-tree.txt @@ -0,0 +1,119 @@ +span_tree: +└── voyageai-instrumentation-root [task] + metadata: { + "scenario": "voyageai-instrumentation", + "testRunId": "" + } + ├── voyageai-embed-operation + │ metadata: { + │ "operation": "embed", + │ "testRunId": "" + │ } + │ └── voyageai.embed [llm] + │ input: [ + │ "Braintrust tracing" + │ ] + │ output: { + │ "embedding_count": 1, + │ "embedding_length": 3 + │ } + │ metadata: { + │ "inputType": "document", + │ "model": "voyage-4", + │ "outputDimension": 3, + │ "provider": "voyage" + │ } + │ metrics: { + │ "tokens": 4 + │ } + ├── voyageai-multimodal-embed-operation + │ metadata: { + │ "operation": "multimodal-embed", + │ "testRunId": "" + │ } + │ └── voyageai.multimodalEmbed [llm] + │ input: [ + │ { + │ "content": [ + │ { + │ "text": "A yellow banana.", + │ "type": "text" + │ }, + │ { + │ "imageUrl": "https://example.com/banana.jpg", + │ "type": "image_url" + │ } + │ ] + │ } + │ ] + │ output: { + │ "embedding_count": 1, + │ "embedding_length": 2 + │ } + │ metadata: { + │ "inputType": "document", + │ "model": "voyage-multimodal-3.5", + │ "provider": "voyage" + │ } + │ metrics: { + │ "tokens": 6 + │ } + ├── voyageai-rerank-operation + │ metadata: { + │ "operation": "rerank", + │ "testRunId": "" + │ } + │ └── voyageai.rerank [llm] + │ input: { + │ "documents": [ + │ "Vienna is in Austria.", + │ "Paris is in France." + │ ], + │ "query": "Which document is about France?" + │ } + │ output: [ + │ { + │ "index": 1, + │ "relevance_score": 0.97 + │ }, + │ { + │ "index": 0, + │ "relevance_score": 0.21 + │ } + │ ] + │ metadata: { + │ "document_count": 2, + │ "model": "rerank-2.5", + │ "provider": "voyage", + │ "returnDocuments": true, + │ "topK": 2 + │ } + │ metrics: { + │ "tokens": 11 + │ } + └── voyageai-contextualized-embed-operation + metadata: { + "operation": "contextualized-embed", + "testRunId": "" + } + └── voyageai.contextualizedEmbed [llm] + input: [ + [ + "Paris is a city.", + "It is the capital of France." + ] + ] + output: { + "document_count": 1, + "embedding_count": 2, + "embedding_length": 4 + } + metadata: { + "inputType": "document", + "model": "voyage-context-3", + "outputDimension": 4, + "provider": "voyage" + } + metrics: { + "tokens": 8 + } diff --git a/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.json b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.json new file mode 100644 index 000000000..c92b29bda --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.json @@ -0,0 +1,159 @@ +{ + "span_tree": [ + { + "name": "voyageai-instrumentation-root", + "type": "task", + "children": [ + { + "name": "voyageai-embed-operation", + "children": [ + { + "name": "voyageai.embed", + "type": "llm", + "children": [], + "input": [ + "Braintrust tracing" + ], + "output": { + "embedding_count": 1, + "embedding_length": 3 + }, + "metadata": { + "inputType": "document", + "model": "voyage-4", + "outputDimension": 3, + "provider": "voyage" + }, + "metrics": { + "tokens": 4 + } + } + ], + "metadata": { + "operation": "embed", + "testRunId": "" + } + }, + { + "name": "voyageai-multimodal-embed-operation", + "children": [ + { + "name": "voyageai.multimodalEmbed", + "type": "llm", + "children": [], + "input": [ + { + "content": [ + { + "text": "A yellow banana.", + "type": "text" + }, + { + "imageUrl": "https://example.com/banana.jpg", + "type": "image_url" + } + ] + } + ], + "output": { + "embedding_count": 1, + "embedding_length": 2 + }, + "metadata": { + "inputType": "document", + "model": "voyage-multimodal-3.5", + "provider": "voyage" + }, + "metrics": { + "tokens": 6 + } + } + ], + "metadata": { + "operation": "multimodal-embed", + "testRunId": "" + } + }, + { + "name": "voyageai-rerank-operation", + "children": [ + { + "name": "voyageai.rerank", + "type": "llm", + "children": [], + "input": { + "documents": [ + "Vienna is in Austria.", + "Paris is in France." + ], + "query": "Which document is about France?" + }, + "output": [ + { + "index": 1, + "relevance_score": 0.97 + }, + { + "index": 0, + "relevance_score": 0.21 + } + ], + "metadata": { + "document_count": 2, + "model": "rerank-2.5", + "provider": "voyage", + "returnDocuments": true, + "topK": 2 + }, + "metrics": { + "tokens": 11 + } + } + ], + "metadata": { + "operation": "rerank", + "testRunId": "" + } + }, + { + "name": "voyageai-contextualized-embed-operation", + "children": [ + { + "name": "voyageai.contextualizedEmbed", + "type": "llm", + "children": [], + "input": [ + [ + "Paris is a city.", + "It is the capital of France." + ] + ], + "output": { + "document_count": 1, + "embedding_count": 2, + "embedding_length": 4 + }, + "metadata": { + "inputType": "document", + "model": "voyage-context-3", + "outputDimension": 4, + "provider": "voyage" + }, + "metrics": { + "tokens": 8 + } + } + ], + "metadata": { + "operation": "contextualized-embed", + "testRunId": "" + } + } + ], + "metadata": { + "scenario": "voyageai-instrumentation", + "testRunId": "" + } + } + ] +} diff --git a/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.txt b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.txt new file mode 100644 index 000000000..6300a04ee --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/__snapshots__/voyageai-v0.span-tree.txt @@ -0,0 +1,119 @@ +span_tree: +└── voyageai-instrumentation-root [task] + metadata: { + "scenario": "voyageai-instrumentation", + "testRunId": "" + } + ├── voyageai-embed-operation + │ metadata: { + │ "operation": "embed", + │ "testRunId": "" + │ } + │ └── voyageai.embed [llm] + │ input: [ + │ "Braintrust tracing" + │ ] + │ output: { + │ "embedding_count": 1, + │ "embedding_length": 3 + │ } + │ metadata: { + │ "inputType": "document", + │ "model": "voyage-4", + │ "outputDimension": 3, + │ "provider": "voyage" + │ } + │ metrics: { + │ "tokens": 4 + │ } + ├── voyageai-multimodal-embed-operation + │ metadata: { + │ "operation": "multimodal-embed", + │ "testRunId": "" + │ } + │ └── voyageai.multimodalEmbed [llm] + │ input: [ + │ { + │ "content": [ + │ { + │ "text": "A yellow banana.", + │ "type": "text" + │ }, + │ { + │ "imageUrl": "https://example.com/banana.jpg", + │ "type": "image_url" + │ } + │ ] + │ } + │ ] + │ output: { + │ "embedding_count": 1, + │ "embedding_length": 2 + │ } + │ metadata: { + │ "inputType": "document", + │ "model": "voyage-multimodal-3.5", + │ "provider": "voyage" + │ } + │ metrics: { + │ "tokens": 6 + │ } + ├── voyageai-rerank-operation + │ metadata: { + │ "operation": "rerank", + │ "testRunId": "" + │ } + │ └── voyageai.rerank [llm] + │ input: { + │ "documents": [ + │ "Vienna is in Austria.", + │ "Paris is in France." + │ ], + │ "query": "Which document is about France?" + │ } + │ output: [ + │ { + │ "index": 1, + │ "relevance_score": 0.97 + │ }, + │ { + │ "index": 0, + │ "relevance_score": 0.21 + │ } + │ ] + │ metadata: { + │ "document_count": 2, + │ "model": "rerank-2.5", + │ "provider": "voyage", + │ "returnDocuments": true, + │ "topK": 2 + │ } + │ metrics: { + │ "tokens": 11 + │ } + └── voyageai-contextualized-embed-operation + metadata: { + "operation": "contextualized-embed", + "testRunId": "" + } + └── voyageai.contextualizedEmbed [llm] + input: [ + [ + "Paris is a city.", + "It is the capital of France." + ] + ] + output: { + "document_count": 1, + "embedding_count": 2, + "embedding_length": 4 + } + metadata: { + "inputType": "document", + "model": "voyage-context-3", + "outputDimension": 4, + "provider": "voyage" + } + metrics: { + "tokens": 8 + } diff --git a/e2e/scenarios/voyageai-instrumentation/assertions.ts b/e2e/scenarios/voyageai-instrumentation/assertions.ts new file mode 100644 index 000000000..67e0fe63f --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/assertions.ts @@ -0,0 +1,133 @@ +import { beforeAll, describe, expect, test } from "vitest"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { resolveFileSnapshotPath } from "../../helpers/file-snapshot"; +import { + withScenarioHarness, + type ScenarioRunContext, +} from "../../helpers/scenario-harness"; +import { matchSpanTreeSnapshot } from "../../helpers/span-tree"; +import { + findLatestChildSpan, + findLatestSpan, +} from "../../helpers/trace-selectors"; +import { ROOT_NAME, SCENARIO_NAME } from "./scenario.impl.mjs"; + +type RunVoyageAIScenario = (harness: { + runNodeScenarioDir: (options: { + entry: string; + env: Record; + nodeArgs: string[]; + runContext?: ScenarioRunContext; + scenarioDir: string; + }) => Promise; + runScenarioDir: (options: { + entry: string; + env: Record; + runContext?: ScenarioRunContext; + scenarioDir: string; + }) => Promise; +}) => Promise; + +const OPERATIONS = [ + { + model: "voyage-4", + operationName: "voyageai-embed-operation", + output: { embedding_count: 1, embedding_length: 3 }, + spanName: "voyageai.embed", + tokens: 4, + }, + { + model: "voyage-multimodal-3.5", + operationName: "voyageai-multimodal-embed-operation", + output: { embedding_count: 1, embedding_length: 2 }, + spanName: "voyageai.multimodalEmbed", + tokens: 6, + }, + { + model: "rerank-2.5", + operationName: "voyageai-rerank-operation", + output: [ + { index: 1, relevance_score: 0.97 }, + { index: 0, relevance_score: 0.21 }, + ], + spanName: "voyageai.rerank", + tokens: 11, + }, + { + model: "voyage-context-3", + operationName: "voyageai-contextualized-embed-operation", + output: { + document_count: 1, + embedding_count: 2, + embedding_length: 4, + }, + spanName: "voyageai.contextualizedEmbed", + tokens: 8, + }, +] as const; + +function spanTreeEvents(events: CapturedLogEvent[]): CapturedLogEvent[] { + return [ + findLatestSpan(events, ROOT_NAME), + ...OPERATIONS.flatMap(({ operationName, spanName }) => { + const operation = findLatestSpan(events, operationName); + return [ + operation, + findLatestChildSpan(events, spanName, operation?.span.id), + ]; + }), + ].map((event) => event!); +} + +export function defineVoyageAIInstrumentationAssertions(options: { + name: string; + runScenario: RunVoyageAIScenario; + snapshotName: string; + testFileUrl: string; +}): void { + const spanSnapshotPath = resolveFileSnapshotPath( + options.testFileUrl, + `${options.snapshotName}.span-tree.json`, + ); + + describe(options.name, () => { + let events: CapturedLogEvent[] = []; + + beforeAll(async () => { + await withScenarioHarness(async (harness) => { + await options.runScenario(harness); + events = harness.events(); + }); + }); + + test("captures the scenario root span", () => { + const root = findLatestSpan(events, ROOT_NAME); + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ + scenario: SCENARIO_NAME, + }); + }); + + test.each(OPERATIONS)( + "captures $spanName", + ({ model, operationName, output, spanName, tokens }) => { + const operation = findLatestSpan(events, operationName); + const span = findLatestChildSpan(events, spanName, operation?.span.id); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + model, + provider: "voyage", + }); + expect(span?.metrics).toMatchObject({ tokens }); + expect(span?.output).toEqual(output); + }, + ); + + test("matches the span tree snapshot", async () => { + await matchSpanTreeSnapshot(spanTreeEvents(events), spanSnapshotPath); + }); + }); +} diff --git a/e2e/scenarios/voyageai-instrumentation/package.json b/e2e/scenarios/voyageai-instrumentation/package.json new file mode 100644 index 000000000..86ebcc9c8 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/package.json @@ -0,0 +1,18 @@ +{ + "name": "@braintrust/e2e-voyageai-instrumentation", + "private": true, + "braintrustScenario": { + "bump": { + "dependencies": { + "voyageai-sdk-v0-latest": { + "package": "voyageai", + "range": "0" + } + } + } + }, + "dependencies": { + "voyageai-sdk-v0": "npm:voyageai@0.2.1", + "voyageai-sdk-v0-latest": "npm:voyageai@0.4.0" + } +} diff --git a/e2e/scenarios/voyageai-instrumentation/pnpm-lock.yaml b/e2e/scenarios/voyageai-instrumentation/pnpm-lock.yaml new file mode 100644 index 000000000..617c696b4 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/pnpm-lock.yaml @@ -0,0 +1,87 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + voyageai-sdk-v0: + specifier: npm:voyageai@0.2.1 + version: voyageai@0.2.1 + voyageai-sdk-v0-latest: + specifier: npm:voyageai@0.4.0 + version: voyageai@0.4.0 + +packages: + + node-fetch@2.7.0: + resolution: {integrity: sha512-c4FRfUm/dbcWZ7U+1Wq0AwCyFL+3nt2bEw05wfxSz+DWpWsitgmSgYmy2dQdWyKC1694ELPqMs/YzUSNozLt8A==} + engines: {node: 4.x || >=6.0.0} + peerDependencies: + encoding: ^0.1.0 + peerDependenciesMeta: + encoding: + optional: true + + tr46@0.0.3: + resolution: {integrity: sha512-N3WMsuqV66lT30CrXNbEjx4GEwlow3v6rr4mCcv6prnfwhS01rkgyFdjPNBYd9br7LpXV1+Emh01fHnq2Gdgrw==} + + voyageai@0.2.1: + resolution: {integrity: sha512-ym7Dk6p8Si6lR9wDh58EzxwT0ziD/pqXjzzzceOSySO3Ic3uosHZLOTAsb3Gq+1OaKdEMnni/p8TohKUNvLTkg==} + engines: {node: '>=18.0.0'} + peerDependencies: + '@huggingface/transformers': ^3.8.0 + onnxruntime-node: '>=1.17.0' + peerDependenciesMeta: + '@huggingface/transformers': + optional: true + onnxruntime-node: + optional: true + + voyageai@0.4.0: + resolution: {integrity: sha512-RCWGknFqwbZORdRo7dHwBuZYeO64TnE+G30Qv4IGxq5jkMa2uK+eqMbOv/DhszcFUtE1iZjb4LqO9G/hj2EjxA==} + engines: {node: '>=18.0.0'} + peerDependencies: + '@huggingface/transformers': ^3.8.0 + onnxruntime-node: '>=1.17.0' + peerDependenciesMeta: + '@huggingface/transformers': + optional: true + onnxruntime-node: + optional: true + + webidl-conversions@3.0.1: + resolution: {integrity: sha512-2JAn3z8AR6rjK8Sm8orRC0h/bcl/DqL7tRPdGZ4I1CjdF+EaMLmYxBHyXuKL849eucPFhvBoxMsflfOb8kxaeQ==} + + whatwg-url@5.0.0: + resolution: {integrity: sha512-saE57nupxk6v3HY35+jzBwYa0rKSy0XR8JSxZPwgLr7ys0IBzhGviA1/TUGJLmSVqs8pb9AnvICXEuOHLprYTw==} + +snapshots: + + node-fetch@2.7.0: + dependencies: + whatwg-url: 5.0.0 + + tr46@0.0.3: {} + + voyageai@0.2.1: + dependencies: + node-fetch: 2.7.0 + transitivePeerDependencies: + - encoding + + voyageai@0.4.0: + dependencies: + node-fetch: 2.7.0 + transitivePeerDependencies: + - encoding + + webidl-conversions@3.0.1: {} + + whatwg-url@5.0.0: + dependencies: + tr46: 0.0.3 + webidl-conversions: 3.0.1 diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.mjs b/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.mjs new file mode 100644 index 000000000..a63daa348 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.mjs @@ -0,0 +1,10 @@ +import { createRequire } from "node:module"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoVoyageAIInstrumentation } from "./scenario.impl.mjs"; + +const require = createRequire(import.meta.url); +const voyageAIPackageName = + process.env.VOYAGEAI_PACKAGE_NAME ?? "voyageai-sdk-v0"; +const voyageAI = require(voyageAIPackageName); + +runMain(async () => runAutoVoyageAIInstrumentation(voyageAI)); diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.ts b/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.ts new file mode 100644 index 000000000..0a28f5c86 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.cjs-client.ts @@ -0,0 +1,12 @@ +import { createRequire } from "node:module"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedVoyageAIInstrumentation } from "./scenario.impl.mjs"; + +const require = createRequire(import.meta.url); +const voyageAIPackageName = + process.env.VOYAGEAI_PACKAGE_NAME ?? "voyageai-sdk-v0"; + +runMain(async () => { + const voyageAI = require(voyageAIPackageName); + await runWrappedVoyageAIInstrumentation(voyageAI); +}); diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.impl.mjs b/e2e/scenarios/voyageai-instrumentation/scenario.impl.mjs new file mode 100644 index 000000000..76225d22c --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.impl.mjs @@ -0,0 +1,160 @@ +import { wrapVoyageAI } from "braintrust"; +import { + runOperation, + runTracedScenario, +} from "../../helpers/provider-runtime.mjs"; + +export const ROOT_NAME = "voyageai-instrumentation-root"; +export const SCENARIO_NAME = "voyageai-instrumentation"; + +const RESPONSES_BY_PATH = { + "/embeddings": { + data: [{ embedding: [0.1, 0.2, 0.3], index: 0, object: "embedding" }], + model: "voyage-4", + object: "list", + usage: { total_tokens: 4 }, + }, + "/multimodalembeddings": { + data: [{ embedding: [0.1, 0.2], index: 0, object: "embedding" }], + model: "voyage-multimodal-3.5", + object: "list", + usage: { total_tokens: 6 }, + }, + "/rerank": { + data: [ + { document: "Paris is in France.", index: 1, relevance_score: 0.97 }, + { document: "Vienna is in Austria.", index: 0, relevance_score: 0.21 }, + ], + model: "rerank-2.5", + object: "list", + usage: { total_tokens: 11 }, + }, + "/contextualizedembeddings": { + data: [ + { + data: [ + { embedding: [0.1, 0.2, 0.3, 0.4], index: 0, object: "embedding" }, + { embedding: [0.5, 0.6, 0.7, 0.8], index: 1, object: "embedding" }, + ], + index: 0, + object: "list", + }, + ], + model: "voyage-context-3", + object: "list", + usage: { total_tokens: 8 }, + }, +}; + +function createVoyageAIFetch() { + return async (input) => { + const url = new URL( + typeof input === "string" + ? input + : input instanceof URL + ? input.href + : input.url, + ); + const body = RESPONSES_BY_PATH[url.pathname]; + if (!body) { + return new Response(JSON.stringify({ detail: "Not found" }), { + headers: { "content-type": "application/json" }, + status: 404, + }); + } + + return new Response(JSON.stringify(body), { + headers: { "content-type": "application/json" }, + status: 200, + }); + }; +} + +export async function runVoyageAIInstrumentationScenario( + voyageAI, + { decorateClient } = {}, +) { + const baseClient = new voyageAI.VoyageAIClient({ + apiKey: "voyage-test-key", + baseUrl: "https://api.voyage.test", + fetch: createVoyageAIFetch(), + }); + const client = decorateClient ? decorateClient(baseClient) : baseClient; + + await runTracedScenario({ + callback: async () => { + await runOperation("voyageai-embed-operation", "embed", async () => { + const resultPromise = client.embed({ + input: ["Braintrust tracing"], + inputType: "document", + model: "voyage-4", + outputDimension: 3, + }); + if (typeof resultPromise.withRawResponse !== "function") { + throw new Error("embed() did not preserve withRawResponse()"); + } + await resultPromise; + }); + + await runOperation( + "voyageai-multimodal-embed-operation", + "multimodal-embed", + async () => { + await client.multimodalEmbed({ + inputs: [ + { + content: [ + { text: "A yellow banana.", type: "text" }, + { + imageUrl: "https://example.com/banana.jpg", + type: "image_url", + }, + ], + }, + ], + inputType: "document", + model: "voyage-multimodal-3.5", + }); + }, + ); + + await runOperation("voyageai-rerank-operation", "rerank", async () => { + await client.rerank({ + documents: ["Vienna is in Austria.", "Paris is in France."], + model: "rerank-2.5", + query: "Which document is about France?", + returnDocuments: true, + topK: 2, + }); + }); + + await runOperation( + "voyageai-contextualized-embed-operation", + "contextualized-embed", + async () => { + await client.contextualizedEmbed({ + inputType: "document", + inputs: [["Paris is a city.", "It is the capital of France."]], + model: "voyage-context-3", + outputDimension: 4, + }); + }, + ); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-voyageai-instrumentation", + rootName: ROOT_NAME, + }); +} + +export async function runWrappedVoyageAIInstrumentation(voyageAI) { + await runVoyageAIInstrumentationScenario(voyageAI, { + decorateClient: wrapVoyageAI, + }); +} + +export async function runAutoVoyageAIInstrumentation(voyageAI) { + await runVoyageAIInstrumentationScenario(voyageAI); +} diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.mjs b/e2e/scenarios/voyageai-instrumentation/scenario.mjs new file mode 100644 index 000000000..08d1d1dd2 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.mjs @@ -0,0 +1,7 @@ +const voyageAIPackageName = + process.env.VOYAGEAI_PACKAGE_NAME ?? "voyageai-sdk-v0-latest"; +const voyageAI = await import(voyageAIPackageName); +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoVoyageAIInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runAutoVoyageAIInstrumentation(voyageAI)); diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.test.ts b/e2e/scenarios/voyageai-instrumentation/scenario.test.ts new file mode 100644 index 000000000..6431662a7 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.test.ts @@ -0,0 +1,78 @@ +import { describe } from "vitest"; +import { + prepareScenarioDir, + readInstalledPackageVersion, + resolveScenarioDir, +} from "../../helpers/scenario-harness"; +import { defineVoyageAIInstrumentationAssertions } from "./assertions"; + +const originalScenarioDir = resolveScenarioDir(import.meta.url); +const scenarioDir = await prepareScenarioDir({ + scenarioDir: originalScenarioDir, +}); + +const scenarios = await Promise.all( + [ + { + autoEntry: "scenario.cjs-client.mjs", + dependencyName: "voyageai-sdk-v0", + snapshotName: "voyageai-v0", + wrapperEntry: "scenario.cjs-client.ts", + }, + { + autoEntry: "scenario.mjs", + dependencyName: "voyageai-sdk-v0-latest", + snapshotName: "voyageai-v0-latest", + wrapperEntry: "scenario.ts", + }, + ].map(async (scenario) => ({ + ...scenario, + version: await readInstalledPackageVersion( + scenarioDir, + scenario.dependencyName, + ), + })), +); + +describe.concurrent("variants", () => { + for (const scenario of scenarios) { + describe.sequential(`voyageai sdk ${scenario.version}`, () => { + const env = { + VOYAGEAI_PACKAGE_NAME: scenario.dependencyName, + }; + const runContext = { + originalScenarioDir, + variantKey: scenario.snapshotName, + }; + + defineVoyageAIInstrumentationAssertions({ + name: "wrapped instrumentation", + runScenario: async ({ runScenarioDir }) => { + await runScenarioDir({ + entry: scenario.wrapperEntry, + env, + runContext, + scenarioDir, + }); + }, + snapshotName: scenario.snapshotName, + testFileUrl: import.meta.url, + }); + + defineVoyageAIInstrumentationAssertions({ + name: "auto-hook instrumentation", + runScenario: async ({ runNodeScenarioDir }) => { + await runNodeScenarioDir({ + entry: scenario.autoEntry, + env, + nodeArgs: ["--import", "braintrust/hook.mjs"], + runContext, + scenarioDir, + }); + }, + snapshotName: scenario.snapshotName, + testFileUrl: import.meta.url, + }); + }); + } +}); diff --git a/e2e/scenarios/voyageai-instrumentation/scenario.ts b/e2e/scenarios/voyageai-instrumentation/scenario.ts new file mode 100644 index 000000000..8499bedd5 --- /dev/null +++ b/e2e/scenarios/voyageai-instrumentation/scenario.ts @@ -0,0 +1,9 @@ +const voyageAIPackageName = + process.env.VOYAGEAI_PACKAGE_NAME ?? "voyageai-sdk-v0-latest"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedVoyageAIInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => { + const voyageAI = await import(voyageAIPackageName); + await runWrappedVoyageAIInstrumentation(voyageAI); +}); diff --git a/js/src/auto-instrumentations/configs/all.ts b/js/src/auto-instrumentations/configs/all.ts index 905715e43..13cf8690d 100644 --- a/js/src/auto-instrumentations/configs/all.ts +++ b/js/src/auto-instrumentations/configs/all.ts @@ -30,6 +30,7 @@ import { openRouterConfigs } from "./openrouter"; import { openRouterAgentConfigs } from "./openrouter-agent"; import { piCodingAgentConfigs } from "./pi-coding-agent"; import { strandsAgentSDKConfigs } from "./strands-agent-sdk"; +import { voyageAIConfigs } from "./voyageai"; interface InstrumentationConfigGroup { integrations: readonly (keyof InstrumentationIntegrationsConfig)[]; @@ -111,6 +112,10 @@ const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[] integrations: ["flue"], configs: flueConfigs, }, + { + integrations: ["voyageai"], + configs: voyageAIConfigs, + }, // Note: `@mastra/core` is not listed here because its instrumentation // doesn't go through the AST `code-transformer` matcher — Mastra's // content-hashed chunks make `filePath`-based matching too brittle. diff --git a/js/src/auto-instrumentations/configs/voyageai.ts b/js/src/auto-instrumentations/configs/voyageai.ts new file mode 100644 index 000000000..f54edf84e --- /dev/null +++ b/js/src/auto-instrumentations/configs/voyageai.ts @@ -0,0 +1,50 @@ +import type { InstrumentationConfig } from "../orchestrion-js"; +import { voyageAIChannels } from "../../instrumentation/plugins/voyageai-channels"; + +const GENERATED_CLIENT_MODULES = [ + "dist/cjs/Client.js", + "dist/esm/Client.mjs", +] as const; + +const generatedClientOperations = [ + { + channelName: voyageAIChannels.embed.channelName, + methodName: "embed", + versionRange: ">=0.2.0 <1.0.0", + }, + { + channelName: voyageAIChannels.multimodalEmbed.channelName, + methodName: "multimodalEmbed", + versionRange: ">=0.2.0 <1.0.0", + }, + { + channelName: voyageAIChannels.rerank.channelName, + methodName: "rerank", + versionRange: ">=0.2.0 <1.0.0", + }, + { + channelName: voyageAIChannels.contextualizedEmbed.channelName, + methodName: "contextualizedEmbed", + versionRange: ">=0.2.0 <1.0.0", + }, +] as const; + +export const voyageAIConfigs: InstrumentationConfig[] = [ + ...GENERATED_CLIENT_MODULES.flatMap((filePath) => + generatedClientOperations.map( + ({ channelName, methodName, versionRange }) => ({ + channelName, + module: { + name: "voyageai", + versionRange, + filePath, + }, + functionQuery: { + className: "VoyageAIClient", + methodName, + kind: "Async" as const, + }, + }), + ), + ), +]; diff --git a/js/src/exports.ts b/js/src/exports.ts index 02104f8c2..7e1a25563 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -222,6 +222,7 @@ export { wrapOpenRouterAgent } from "./wrappers/openrouter-agent"; export { wrapOpenRouter } from "./wrappers/openrouter"; export { wrapMistral } from "./wrappers/mistral"; export { wrapCohere } from "./wrappers/cohere"; +export { wrapVoyageAI } from "./wrappers/voyageai"; export { wrapGroq } from "./wrappers/groq"; export { wrapBedrockRuntime } from "./wrappers/bedrock-runtime"; export { wrapCopilotClient } from "./wrappers/github-copilot"; diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index 05c8d4838..c7ace5d4c 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -19,6 +19,7 @@ import { LangChainPlugin } from "./plugins/langchain-plugin"; import { LangSmithPlugin } from "./plugins/langsmith-plugin"; import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; +import { VoyageAIPlugin } from "./plugins/voyageai-plugin"; import { CloudflareAIChatPlugin } from "./plugins/cloudflare-ai-chat-plugin"; import { CloudflareAgentsPlugin } from "./plugins/cloudflare-agents-plugin"; @@ -118,6 +119,10 @@ vi.mock("./plugins/strands-agent-sdk-plugin", () => ({ StrandsAgentSDKPlugin: createPluginClassMock(), })); +vi.mock("./plugins/voyageai-plugin", () => ({ + VoyageAIPlugin: createPluginClassMock(), +})); + vi.mock("./plugins/cloudflare-ai-chat-plugin", () => ({ CloudflareAIChatPlugin: createPluginClassMock(), })); @@ -247,6 +252,15 @@ describe("BraintrustPlugin", () => { expect(mockInstance.enable).toHaveBeenCalledTimes(1); }); + it("should create and enable Voyage AI plugin by default", () => { + const plugin = new BraintrustPlugin(); + plugin.enable(); + + expect(VoyageAIPlugin).toHaveBeenCalledTimes(1); + const mockInstance = vi.mocked(VoyageAIPlugin).mock.results[0].value; + expect(mockInstance.enable).toHaveBeenCalledTimes(1); + }); + it("should create and enable Groq plugin by default", () => { const plugin = new BraintrustPlugin(); plugin.enable(); @@ -1091,6 +1105,15 @@ describe("BraintrustPlugin", () => { expect(MistralPlugin).not.toHaveBeenCalled(); expect(CoherePlugin).not.toHaveBeenCalled(); }); + + it("should not create Voyage AI plugin when voyageai: false", () => { + const plugin = new BraintrustPlugin({ + integrations: { voyageai: false }, + }); + plugin.enable(); + + expect(VoyageAIPlugin).not.toHaveBeenCalled(); + }); }); }); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index ec88ddec2..045b47df9 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -23,6 +23,7 @@ import { LangChainPlugin } from "./plugins/langchain-plugin"; import { LangSmithPlugin } from "./plugins/langsmith-plugin"; import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; +import { VoyageAIPlugin } from "./plugins/voyageai-plugin"; import { CloudflareAIChatPlugin } from "./plugins/cloudflare-ai-chat-plugin"; import { CloudflareAgentsPlugin } from "./plugins/cloudflare-agents-plugin"; import type { InstrumentationIntegrationsConfig } from "./config"; @@ -44,6 +45,7 @@ export interface BraintrustPluginConfig { * - LangChain.js and LangGraph * - Mistral SDK * - Cohere SDK + * - Voyage AI SDK * * The plugin is automatically enabled when the Braintrust library is loaded. * Individual integrations can be disabled via configuration. @@ -75,6 +77,7 @@ export class BraintrustPlugin extends BasePlugin { private langSmithPlugin: LangSmithPlugin | null = null; private piCodingAgentPlugin: PiCodingAgentPlugin | null = null; private strandsAgentSDKPlugin: StrandsAgentSDKPlugin | null = null; + private voyageAIPlugin: VoyageAIPlugin | null = null; private cloudflareAIChatPlugin: CloudflareAIChatPlugin | null = null; private cloudflareAgentsPlugin: CloudflareAgentsPlugin | null = null; @@ -167,6 +170,11 @@ export class BraintrustPlugin extends BasePlugin { this.coherePlugin.enable(); } + if (integrations.voyageai !== false) { + this.voyageAIPlugin = new VoyageAIPlugin(); + this.voyageAIPlugin.enable(); + } + if (integrations.groq !== false) { this.groqPlugin = new GroqPlugin(); this.groqPlugin.enable(); @@ -311,6 +319,11 @@ export class BraintrustPlugin extends BasePlugin { this.coherePlugin = null; } + if (this.voyageAIPlugin) { + this.voyageAIPlugin.disable(); + this.voyageAIPlugin = null; + } + if (this.groqPlugin) { this.groqPlugin.disable(); this.groqPlugin = null; diff --git a/js/src/instrumentation/config.ts b/js/src/instrumentation/config.ts index c51a5b48d..42383ba73 100644 --- a/js/src/instrumentation/config.ts +++ b/js/src/instrumentation/config.ts @@ -31,6 +31,7 @@ export interface InstrumentationIntegrationsConfig { langchain?: boolean; langgraph?: boolean; langsmith?: boolean; + voyageai?: boolean; } export interface InstrumentationConfig { @@ -114,6 +115,9 @@ const envIntegrationAliases: Record< "@langchain": "langchain", langgraph: "langgraph", langsmith: "langsmith", + voyage: "voyageai", + "voyage-ai": "voyageai", + voyageai: "voyageai", }; export function getDefaultInstrumentationIntegrations(): Record< @@ -150,6 +154,7 @@ export function getDefaultInstrumentationIntegrations(): Record< langchain: true, langgraph: true, langsmith: true, + voyageai: true, piCodingAgent: true, strandsAgentSDK: true, cloudflareAgents: true, diff --git a/js/src/instrumentation/plugins/instrumentation-names.test.ts b/js/src/instrumentation/plugins/instrumentation-names.test.ts index 2e342f9b2..b7eb092b0 100644 --- a/js/src/instrumentation/plugins/instrumentation-names.test.ts +++ b/js/src/instrumentation/plugins/instrumentation-names.test.ts @@ -30,6 +30,7 @@ import { openRouterAgentChannels } from "./openrouter-agent-channels"; import { openRouterChannels } from "./openrouter-channels"; import { piCodingAgentChannels } from "./pi-coding-agent-channels"; import { strandsAgentSDKChannels } from "./strands-agent-sdk-channels"; +import { voyageAIChannels } from "./voyageai-channels"; describe("built-in instrumentation provenance names", () => { it.each([ @@ -74,6 +75,7 @@ describe("built-in instrumentation provenance names", () => { strandsAgentSDKChannels.agentStream, INSTRUMENTATION_NAMES.STRANDS_AGENT_SDK, ], + [voyageAIChannels.embed, INSTRUMENTATION_NAMES.VOYAGEAI], ])("uses %s for its canonical channel group", (channel, expected) => { expect(channel.instrumentationName).toBe(expected); }); diff --git a/js/src/instrumentation/plugins/voyageai-channels.ts b/js/src/instrumentation/plugins/voyageai-channels.ts new file mode 100644 index 000000000..d93c01622 --- /dev/null +++ b/js/src/instrumentation/plugins/voyageai-channels.ts @@ -0,0 +1,43 @@ +import { channel, defineChannels } from "../core/channel-definitions"; +import { INSTRUMENTATION_NAMES } from "../../span-origin"; +import type { + VoyageAIContextualizedEmbedRequest, + VoyageAIContextualizedResult, + VoyageAIEmbeddingResponse, + VoyageAIEmbedRequest, + VoyageAIMultimodalEmbedRequest, + VoyageAIRerankRequest, + VoyageAIRerankResponse, +} from "../../vendor-sdk-types/voyageai"; + +export const voyageAIChannels = defineChannels( + "voyageai", + { + embed: channel<[VoyageAIEmbedRequest], VoyageAIEmbeddingResponse>({ + channelName: "embed", + kind: "async", + }), + + multimodalEmbed: channel< + [VoyageAIMultimodalEmbedRequest], + VoyageAIEmbeddingResponse + >({ + channelName: "multimodalEmbed", + kind: "async", + }), + + rerank: channel<[VoyageAIRerankRequest], VoyageAIRerankResponse>({ + channelName: "rerank", + kind: "async", + }), + + contextualizedEmbed: channel< + [VoyageAIContextualizedEmbedRequest], + VoyageAIContextualizedResult + >({ + channelName: "contextualizedEmbed", + kind: "async", + }), + }, + { instrumentationName: INSTRUMENTATION_NAMES.VOYAGEAI }, +); diff --git a/js/src/instrumentation/plugins/voyageai-plugin.test.ts b/js/src/instrumentation/plugins/voyageai-plugin.test.ts new file mode 100644 index 000000000..78ab9693d --- /dev/null +++ b/js/src/instrumentation/plugins/voyageai-plugin.test.ts @@ -0,0 +1,234 @@ +import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { Attachment, _exportsForTestingOnly, initLogger } from "../../logger"; +import { configureNode } from "../../node/config"; +import { voyageAIChannels } from "./voyageai-channels"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +describe("VoyageAIPlugin", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "voyageai-plugin.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + it("captures embedding and reranking operations", async () => { + await voyageAIChannels.embed.tracePromise( + async () => ({ + data: [{ embedding: [0.1, 0.2, 0.3], index: 0 }], + model: "voyage-4", + usage: { totalTokens: 4 }, + }), + { + arguments: [ + { + input: ["Braintrust tracing"], + inputType: "document", + model: "voyage-4", + outputDimension: 3, + }, + ], + }, + ); + + await voyageAIChannels.rerank.tracePromise( + async () => ({ + data: [ + { index: 1, relevanceScore: 0.95 }, + { index: 0, relevanceScore: 0.4 }, + ], + model: "rerank-2.5", + usage: { totalTokens: 9 }, + }), + { + arguments: [ + { + documents: ["Vienna", "Paris"], + model: "rerank-2.5", + query: "capital of France", + topK: 2, + }, + ], + }, + ); + + const spans = await backgroundLogger.drain(); + const embedSpan = spans.find( + (span: any) => span.span_attributes?.name === "voyageai.embed", + ) as Record | undefined; + const rerankSpan = spans.find( + (span: any) => span.span_attributes?.name === "voyageai.rerank", + ) as Record | undefined; + + expect(embedSpan).toMatchObject({ + input: ["Braintrust tracing"], + metadata: { + inputType: "document", + model: "voyage-4", + outputDimension: 3, + provider: "voyage", + }, + metrics: { tokens: 4 }, + output: { + embedding_count: 1, + embedding_length: 3, + }, + }); + expect(rerankSpan).toMatchObject({ + input: { + documents: ["Vienna", "Paris"], + query: "capital of France", + }, + metadata: { + document_count: 2, + model: "rerank-2.5", + provider: "voyage", + topK: 2, + }, + metrics: { tokens: 9 }, + output: [ + { index: 1, relevance_score: 0.95 }, + { index: 0, relevance_score: 0.4 }, + ], + }); + }); + + it("normalizes multimodal and contextualized embedding summaries", async () => { + await voyageAIChannels.multimodalEmbed.tracePromise( + async () => ({ + data: [{ embedding: [0.1, 0.2], index: 0 }], + model: "voyage-multimodal-3.5", + usage: { totalTokens: 5 }, + }), + { + arguments: [ + { + inputs: [ + { + content: [ + { type: "text", text: "A banana" }, + { + type: "image_url", + image_url: "https://example.com/banana.jpg", + }, + ], + }, + ], + model: "voyage-multimodal-3.5", + }, + ], + }, + ); + + await voyageAIChannels.contextualizedEmbed.tracePromise( + async () => ({ + results: [ + { + embeddings: [ + [0.1, 0.2, 0.3], + [0.4, 0.5, 0.6], + ], + index: 0, + }, + ], + totalTokens: 7, + rawResponse: { + model: "voyage-context-3", + }, + }), + { + arguments: [ + { + inputs: [["First chunk", "Second chunk"]], + model: "voyage-context-3", + }, + ], + }, + ); + + const spans = await backgroundLogger.drain(); + const multimodalSpan = spans.find( + (span: any) => span.span_attributes?.name === "voyageai.multimodalEmbed", + ) as Record | undefined; + const contextualizedSpan = spans.find( + (span: any) => + span.span_attributes?.name === "voyageai.contextualizedEmbed", + ) as Record | undefined; + + expect(multimodalSpan?.output).toEqual({ + embedding_count: 1, + embedding_length: 2, + }); + expect(contextualizedSpan).toMatchObject({ + metadata: { + model: "voyage-context-3", + provider: "voyage", + }, + metrics: { tokens: 7 }, + output: { + document_count: 1, + embedding_count: 2, + embedding_length: 3, + }, + }); + }); + + it("converts Voyage base64 media inputs to attachments", async () => { + await voyageAIChannels.multimodalEmbed.tracePromise( + async () => ({ + data: [{ embedding: [0.1, 0.2], index: 0 }], + model: "voyage-multimodal-3.5", + usage: { totalTokens: 5 }, + }), + { + arguments: [ + { + inputs: [ + { + content: [ + { + imageBase64: "data:image/png;base64,aGVsbG8=", + type: "image_base64", + }, + ], + }, + ], + model: "voyage-multimodal-3.5", + }, + ], + }, + ); + + const spans = await backgroundLogger.drain(); + const span = spans.find( + (candidate: any) => + candidate.span_attributes?.name === "voyageai.multimodalEmbed", + ) as Record | undefined; + const attachment = span?.input?.[0]?.content?.[0]?.imageBase64; + + expect(attachment).toBeInstanceOf(Attachment); + expect(attachment.reference).toMatchObject({ + content_type: "image/png", + filename: "image.png", + type: "braintrust_attachment", + }); + }); +}); diff --git a/js/src/instrumentation/plugins/voyageai-plugin.ts b/js/src/instrumentation/plugins/voyageai-plugin.ts new file mode 100644 index 000000000..2fe8ef6ae --- /dev/null +++ b/js/src/instrumentation/plugins/voyageai-plugin.ts @@ -0,0 +1,295 @@ +import { SpanTypeAttribute, isObject } from "../../../util/index"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import type { + VoyageAIContextualizedResult, + VoyageAIEmbeddingResponse, + VoyageAIRerankResponse, + VoyageAIUsage, +} from "../../vendor-sdk-types/voyageai"; +import { BasePlugin } from "../core"; +import { traceAsyncChannel, unsubscribeAll } from "../core/channel-tracing"; +import { voyageAIChannels } from "./voyageai-channels"; + +const EMBED_METADATA_ALLOWLIST = new Set([ + "encodingFormat", + "inputType", + "model", + "outputDimension", + "outputDtype", + "truncation", +]); + +const MULTIMODAL_EMBED_METADATA_ALLOWLIST = new Set([ + "encodingFormat", + "inputType", + "model", + "truncation", +]); + +const RERANK_METADATA_ALLOWLIST = new Set([ + "model", + "returnDocuments", + "topK", + "truncation", +]); + +const CONTEXTUALIZED_EMBED_METADATA_ALLOWLIST = new Set([ + "inputType", + "model", + "outputDimension", + "outputDtype", +]); + +export class VoyageAIPlugin extends BasePlugin { + protected onEnable(): void { + this.unsubscribers.push( + traceAsyncChannel(voyageAIChannels.embed, { + name: "voyageai.embed", + type: SpanTypeAttribute.LLM, + extractInput: (args) => + extractEmbeddingInput(args, "input", EMBED_METADATA_ALLOWLIST), + extractOutput: summarizeEmbeddingOutput, + extractMetadata: extractResponseMetadata, + extractMetrics: extractUsageMetrics, + }), + traceAsyncChannel(voyageAIChannels.multimodalEmbed, { + name: "voyageai.multimodalEmbed", + type: SpanTypeAttribute.LLM, + extractInput: (args) => + extractEmbeddingInput( + args, + "inputs", + MULTIMODAL_EMBED_METADATA_ALLOWLIST, + true, + ), + extractOutput: summarizeEmbeddingOutput, + extractMetadata: extractResponseMetadata, + extractMetrics: extractUsageMetrics, + }), + traceAsyncChannel(voyageAIChannels.rerank, { + name: "voyageai.rerank", + type: SpanTypeAttribute.LLM, + extractInput: extractRerankInput, + extractOutput: summarizeRerankOutput, + extractMetadata: extractResponseMetadata, + extractMetrics: extractUsageMetrics, + }), + traceAsyncChannel(voyageAIChannels.contextualizedEmbed, { + name: "voyageai.contextualizedEmbed", + type: SpanTypeAttribute.LLM, + extractInput: (args) => + extractEmbeddingInput( + args, + "inputs", + CONTEXTUALIZED_EMBED_METADATA_ALLOWLIST, + ), + extractOutput: summarizeContextualizedEmbeddingOutput, + extractMetadata: extractResponseMetadata, + extractMetrics: extractUsageMetrics, + }), + ); + } + + protected onDisable(): void { + this.unsubscribers = unsubscribeAll(this.unsubscribers); + } +} + +function getRequestArg(args: unknown): Record | undefined { + if (Array.isArray(args)) { + return isObject(args[0]) ? args[0] : undefined; + } + + if (!isObject(args)) { + return undefined; + } + + const firstArg = Reflect.get(args, "0"); + return isObject(firstArg) ? firstArg : undefined; +} + +function pickMetadata( + request: Record | undefined, + allowlist: ReadonlySet, +): Record { + const metadata: Record = {}; + if (request) { + for (const key of allowlist) { + if (!Object.hasOwn(request, key)) { + continue; + } + const value = request[key]; + if (value !== undefined) { + metadata[key] = value; + } + } + } + + return { + ...metadata, + provider: "voyage", + }; +} + +function extractEmbeddingInput( + args: unknown, + inputKey: "input" | "inputs", + metadataAllowlist: ReadonlySet, + processAttachments = false, +): { + input: unknown; + metadata: Record; +} { + const request = getRequestArg(args); + const input = request?.[inputKey]; + + return { + input: processAttachments ? processInputAttachments(input) : input, + metadata: pickMetadata(request, metadataAllowlist), + }; +} + +function extractRerankInput(args: unknown): { + input: unknown; + metadata: Record; +} { + const request = getRequestArg(args); + const documents = request?.documents; + + return { + input: { + documents, + query: request?.query, + }, + metadata: { + ...pickMetadata(request, RERANK_METADATA_ALLOWLIST), + ...(Array.isArray(documents) ? { document_count: documents.length } : {}), + }, + }; +} + +function extractResponseMetadata( + result: VoyageAIEmbeddingResponse | VoyageAIRerankResponse, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + const rawResponse = isObject(result.rawResponse) + ? result.rawResponse + : undefined; + const model = + typeof result.model === "string" + ? result.model + : typeof rawResponse?.model === "string" + ? rawResponse.model + : undefined; + + return model ? { model } : undefined; +} + +function summarizeEmbeddingOutput( + result: VoyageAIEmbeddingResponse, +): Record | undefined { + if (!isObject(result) || !Array.isArray(result.data)) { + return undefined; + } + + const firstEmbedding = + isObject(result.data[0]) && Array.isArray(result.data[0].embedding) + ? result.data[0].embedding + : undefined; + + return { + embedding_count: result.data.length, + ...(firstEmbedding ? { embedding_length: firstEmbedding.length } : {}), + }; +} + +function summarizeRerankOutput( + result: VoyageAIRerankResponse, +): unknown[] | undefined { + if (!isObject(result) || !Array.isArray(result.data)) { + return undefined; + } + + return result.data.slice(0, 100).map((item) => ({ + index: isObject(item) ? item.index : undefined, + relevance_score: isObject(item) + ? ((typeof item.relevanceScore === "number" + ? item.relevanceScore + : item.relevance_score) ?? null) + : null, + })); +} + +function summarizeContextualizedEmbeddingOutput( + result: VoyageAIContextualizedResult, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + if (Array.isArray(result.results)) { + const embeddings = result.results.flatMap((item) => + isObject(item) && Array.isArray(item.embeddings) ? item.embeddings : [], + ); + const firstEmbedding = Array.isArray(embeddings[0]) + ? embeddings[0] + : undefined; + + return { + document_count: result.results.length, + embedding_count: embeddings.length, + ...(firstEmbedding ? { embedding_length: firstEmbedding.length } : {}), + }; + } + + if (!Array.isArray(result.data)) { + return undefined; + } + + const embeddings = result.data.flatMap((item) => + isObject(item) && Array.isArray(item.data) ? item.data : [], + ); + const firstEmbedding = + isObject(embeddings[0]) && Array.isArray(embeddings[0].embedding) + ? embeddings[0].embedding + : undefined; + + return { + document_count: result.data.length, + embedding_count: embeddings.length, + ...(firstEmbedding ? { embedding_length: firstEmbedding.length } : {}), + }; +} + +function extractUsageMetrics( + result: + | VoyageAIEmbeddingResponse + | VoyageAIRerankResponse + | VoyageAIContextualizedResult, +): Record { + if (!isObject(result)) { + return {}; + } + + const rawResponse = isObject(result.rawResponse) + ? result.rawResponse + : undefined; + const usage = ( + isObject(result.usage) + ? result.usage + : isObject(rawResponse?.usage) + ? rawResponse.usage + : undefined + ) as VoyageAIUsage | undefined; + const tokens = + typeof result.totalTokens === "number" + ? result.totalTokens + : (usage?.totalTokens ?? usage?.total_tokens); + + return typeof tokens === "number" && Number.isFinite(tokens) && tokens >= 0 + ? { tokens } + : {}; +} diff --git a/js/src/span-origin.ts b/js/src/span-origin.ts index 0fd2f5fc8..a5f56076f 100644 --- a/js/src/span-origin.ts +++ b/js/src/span-origin.ts @@ -36,6 +36,7 @@ export const INSTRUMENTATION_NAMES = { OPENROUTER_AGENT: "openrouter-agent", PI_CODING_AGENT: "pi-coding-agent", STRANDS_AGENT_SDK: "strands-agent-sdk", + VOYAGEAI: "voyageai", } as const; export type SpanInstrumentationName = diff --git a/js/src/vendor-sdk-types/voyageai.ts b/js/src/vendor-sdk-types/voyageai.ts new file mode 100644 index 000000000..c535baa0b --- /dev/null +++ b/js/src/vendor-sdk-types/voyageai.ts @@ -0,0 +1,122 @@ +export type VoyageAIUsage = { + totalTokens?: number; + total_tokens?: number; + [key: string]: unknown; +}; + +export type VoyageAIEmbedRequest = { + input: unknown; + model: string; + inputType?: unknown; + truncation?: boolean; + encodingFormat?: unknown; + outputDimension?: number; + outputDtype?: unknown; + [key: string]: unknown; +}; + +export type VoyageAIMultimodalEmbedRequest = { + inputs: unknown; + model: string; + inputType?: unknown; + truncation?: boolean; + encodingFormat?: unknown; + [key: string]: unknown; +}; + +export type VoyageAIRerankRequest = { + query: string; + documents: unknown; + model: string; + topK?: number; + returnDocuments?: boolean; + truncation?: boolean; + [key: string]: unknown; +}; + +export type VoyageAIContextualizedEmbedRequest = { + inputs: unknown; + model: string; + inputType?: unknown; + outputDimension?: number; + outputDtype?: unknown; + [key: string]: unknown; +}; + +export type VoyageAIEmbeddingDataItem = { + embedding?: unknown; + index?: number; + [key: string]: unknown; +}; + +export type VoyageAIEmbeddingResponse = { + data?: VoyageAIEmbeddingDataItem[]; + model?: string; + usage?: VoyageAIUsage; + [key: string]: unknown; +}; + +export type VoyageAIRerankResponse = { + data?: Array<{ + index?: number; + relevanceScore?: number; + relevance_score?: number; + [key: string]: unknown; + }>; + model?: string; + usage?: VoyageAIUsage; + [key: string]: unknown; +}; + +export type VoyageAIContextualizedEmbeddingResponse = { + data?: Array<{ + data?: VoyageAIEmbeddingDataItem[]; + index?: number; + [key: string]: unknown; + }>; + model?: string; + usage?: VoyageAIUsage; + [key: string]: unknown; +}; + +export type VoyageAIContextualizedEmbedResult = { + results?: Array<{ + embeddings?: unknown; + index?: number; + [key: string]: unknown; + }>; + totalTokens?: number; + rawResponse?: VoyageAIContextualizedEmbeddingResponse; + [key: string]: unknown; +}; + +export type VoyageAIContextualizedResult = + | VoyageAIContextualizedEmbeddingResponse + | VoyageAIContextualizedEmbedResult; + +export type VoyageAIResponsePromise = PromiseLike & { + withRawResponse?: () => Promise<{ + data: T; + rawResponse: unknown; + }>; +}; + +export type VoyageAIClient = { + embed?: ( + request: VoyageAIEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise; + multimodalEmbed?: ( + request: VoyageAIMultimodalEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise; + rerank?: ( + request: VoyageAIRerankRequest, + options?: unknown, + ) => VoyageAIResponsePromise; + contextualizedEmbed?: ( + request: VoyageAIContextualizedEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise; + [key: string]: unknown; +}; diff --git a/js/src/wrappers/attachment-utils.ts b/js/src/wrappers/attachment-utils.ts index 7815f8907..7fb2b96f5 100644 --- a/js/src/wrappers/attachment-utils.ts +++ b/js/src/wrappers/attachment-utils.ts @@ -141,6 +141,40 @@ export function processInputAttachments(input: any): any { } } + // Voyage AI multimodal image_base64/video_base64 content format + const voyageBase64Key = + node.type === "image_base64" + ? Object.hasOwn(node, "imageBase64") + ? "imageBase64" + : "image_base64" + : node.type === "video_base64" + ? Object.hasOwn(node, "videoBase64") + ? "videoBase64" + : "video_base64" + : undefined; + const voyageBase64Value = voyageBase64Key + ? node[voyageBase64Key] + : undefined; + if ( + voyageBase64Key && + typeof voyageBase64Value === "string" && + voyageBase64Value.startsWith("data:") + ) { + const mediaType = inferMediaTypeFromDataUrl( + voyageBase64Value, + node.type === "video_base64" ? "video/mp4" : "image/png", + ); + const filename = `${node.type === "video_base64" ? "video" : "image"}.${getExtensionFromMediaType(mediaType)}`; + const attachment = toAttachment(voyageBase64Value, mediaType, filename); + + if (attachment) { + return { + ...node, + [voyageBase64Key]: attachment, + }; + } + } + // OpenAI chat file content format if ( node.type === "file" && diff --git a/js/src/wrappers/voyageai.test.ts b/js/src/wrappers/voyageai.test.ts new file mode 100644 index 000000000..196590605 --- /dev/null +++ b/js/src/wrappers/voyageai.test.ts @@ -0,0 +1,158 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { _exportsForTestingOnly, initLogger } from "../logger"; +import { configureNode } from "../node/config"; +import { wrapVoyageAI } from "./voyageai"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +describe("Voyage AI wrapper", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "voyageai.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.restoreAllMocks(); + }); + + it("returns unsupported clients unchanged", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const unsupported = { tokenize: vi.fn() }; + + expect(wrapVoyageAI(unsupported)).toBe(unsupported); + expect(warnSpy).toHaveBeenCalledWith( + "Unsupported Voyage AI library. Not wrapping.", + ); + }); + + it("wraps all Voyage execution methods", async () => { + const client = wrapVoyageAI({ + embed: vi.fn(async () => ({ + data: [{ embedding: [0.1, 0.2], index: 0 }], + model: "voyage-4", + usage: { totalTokens: 2 }, + })), + multimodalEmbed: vi.fn(async () => ({ + data: [{ embedding: [0.1, 0.2, 0.3], index: 0 }], + model: "voyage-multimodal-3.5", + usage: { totalTokens: 3 }, + })), + rerank: vi.fn(async () => ({ + data: [{ index: 0, relevanceScore: 0.9 }], + model: "rerank-2.5", + usage: { totalTokens: 4 }, + })), + contextualizedEmbed: vi.fn(async () => ({ + data: [ + { + data: [{ embedding: [0.1, 0.2, 0.3, 0.4], index: 0 }], + index: 0, + }, + ], + model: "voyage-context-3", + usage: { totalTokens: 5 }, + })), + } as any) as any; + + await client.embed({ input: "hello", model: "voyage-4" }); + await client.multimodalEmbed({ + inputs: [{ content: [{ type: "text", text: "hello" }] }], + model: "voyage-multimodal-3.5", + }); + await client.rerank({ + documents: ["Paris"], + model: "rerank-2.5", + query: "France", + }); + await client.contextualizedEmbed({ + inputs: [["hello"]], + model: "voyage-context-3", + }); + + const spans = await backgroundLogger.drain(); + expect(spans.map((span: any) => span.span_attributes?.name).sort()).toEqual( + [ + "voyageai.contextualizedEmbed", + "voyageai.embed", + "voyageai.multimodalEmbed", + "voyageai.rerank", + ], + ); + }); + + it("preserves Voyage HttpResponsePromise helpers", async () => { + class MockResponsePromise extends Promise { + withRawResponse() { + return Promise.resolve({ + data: "raw-data", + rawResponse: "raw-response", + }); + } + } + + const originalPromise = new MockResponsePromise((resolve) => { + resolve({ + data: [{ embedding: [0.1], index: 0 }], + model: "voyage-4", + usage: { totalTokens: 1 }, + }); + }); + const client = wrapVoyageAI({ + embed: vi.fn(() => originalPromise), + } as any) as any; + + const resultPromise = client.embed({ + input: "hello", + model: "voyage-4", + }); + + expect(resultPromise).toBe(originalPromise); + await expect(resultPromise.withRawResponse()).resolves.toEqual({ + data: "raw-data", + rawResponse: "raw-response", + }); + await resultPromise; + expect(await backgroundLogger.drain()).toHaveLength(1); + }); + + it("preserves provider errors", async () => { + const providerError = new Error("Voyage request failed"); + const client = wrapVoyageAI({ + rerank: vi.fn(async () => { + throw providerError; + }), + } as any) as any; + + await expect( + client.rerank({ + documents: ["Paris"], + model: "rerank-2.5", + query: "France", + }), + ).rejects.toBe(providerError); + }); +}); diff --git a/js/src/wrappers/voyageai.ts b/js/src/wrappers/voyageai.ts new file mode 100644 index 000000000..7bc13a046 --- /dev/null +++ b/js/src/wrappers/voyageai.ts @@ -0,0 +1,138 @@ +import { voyageAIChannels } from "../instrumentation/plugins/voyageai-channels"; +import type { + VoyageAIClient, + VoyageAIContextualizedEmbedRequest, + VoyageAIContextualizedResult, + VoyageAIEmbeddingResponse, + VoyageAIEmbedRequest, + VoyageAIMultimodalEmbedRequest, + VoyageAIRerankRequest, + VoyageAIRerankResponse, + VoyageAIResponsePromise, +} from "../vendor-sdk-types/voyageai"; + +/** + * Wrap a Voyage AI client so method calls emit diagnostics-channel events that + * Braintrust plugins can consume. + */ +export function wrapVoyageAI(client: T): T { + if (!isSupportedVoyageAIClient(client)) { + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.warn("Unsupported Voyage AI library. Not wrapping."); + return client; + } + + return voyageAIProxy(client) as T; +} + +const voyageAIProxyCache = new WeakMap(); + +function isObject(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function hasFunction(value: Record, name: string): boolean { + return typeof Reflect.get(value, name) === "function"; +} + +function isSupportedVoyageAIClient(value: unknown): value is VoyageAIClient { + if (!isObject(value)) { + return false; + } + + return ( + hasFunction(value, "embed") || + hasFunction(value, "multimodalEmbed") || + hasFunction(value, "rerank") || + hasFunction(value, "contextualizedEmbed") + ); +} + +function voyageAIProxy(client: VoyageAIClient): VoyageAIClient { + const cached = voyageAIProxyCache.get(client); + if (cached) { + return cached; + } + + const proxy = new Proxy(client, { + get(target, prop, receiver) { + switch (prop) { + case "embed": + return typeof target.embed === "function" + ? wrapEmbed(target.embed.bind(target)) + : target.embed; + case "multimodalEmbed": + return typeof target.multimodalEmbed === "function" + ? wrapMultimodalEmbed(target.multimodalEmbed.bind(target)) + : target.multimodalEmbed; + case "rerank": + return typeof target.rerank === "function" + ? wrapRerank(target.rerank.bind(target)) + : target.rerank; + case "contextualizedEmbed": + return typeof target.contextualizedEmbed === "function" + ? wrapContextualizedEmbed(target.contextualizedEmbed.bind(target)) + : target.contextualizedEmbed; + default: + return Reflect.get(target, prop, receiver); + } + }, + }); + + voyageAIProxyCache.set(client, proxy); + return proxy; +} + +function wrapEmbed( + embed: ( + request: VoyageAIEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise, +): NonNullable { + return (request, options) => + voyageAIChannels.embed.tracePromise(() => embed(request, options), { + arguments: [request], + }); +} + +function wrapMultimodalEmbed( + multimodalEmbed: ( + request: VoyageAIMultimodalEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise, +): NonNullable { + return (request, options) => + voyageAIChannels.multimodalEmbed.tracePromise( + () => multimodalEmbed(request, options), + { + arguments: [request], + }, + ); +} + +function wrapRerank( + rerank: ( + request: VoyageAIRerankRequest, + options?: unknown, + ) => VoyageAIResponsePromise, +): NonNullable { + return (request, options) => + voyageAIChannels.rerank.tracePromise(() => rerank(request, options), { + arguments: [request], + }); +} + +function wrapContextualizedEmbed( + contextualizedEmbed: ( + request: VoyageAIContextualizedEmbedRequest, + options?: unknown, + ) => VoyageAIResponsePromise, +): NonNullable { + return (request, options) => + voyageAIChannels.contextualizedEmbed.tracePromise( + () => contextualizedEmbed(request, options), + { + arguments: [request], + }, + ); +}