diff --git a/.github/workflows/cohere-typescript-tests.yml b/.github/workflows/cohere-typescript-tests.yml new file mode 100644 index 000000000..53a672f6f --- /dev/null +++ b/.github/workflows/cohere-typescript-tests.yml @@ -0,0 +1,48 @@ +name: Cohere TypeScript tests + +on: + pull_request: + branches: [dev, main] + paths: + - 'typescript/packages/traceai_cohere/**' + - 'typescript/packages/fi-core/**' + - 'typescript/packages/fi-semantic-conventions/**' + - 'typescript/pnpm-lock.yaml' + - 'typescript/pnpm-workspace.yaml' + - 'typescript/package.json' + - '.github/workflows/cohere-typescript-tests.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + cohere: + name: Cohere SDK floor and current + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: typescript + env: + FI_BASE_URL: http://127.0.0.1:9 + FI_GRPC_URL: http://127.0.0.1:9 + steps: + - uses: actions/checkout@v4 + with: + persist-credentials: false + - uses: actions/setup-node@v4 + with: + node-version: '22' + - run: corepack enable + - name: Install pinned SDK fixtures + run: pnpm install --frozen-lockfile --ignore-scripts --filter '@traceai/cohere...' + - name: Build package and shared dependencies + run: pnpm --filter '@traceai/cohere...' --workspace-concurrency=1 run build + - name: Real SDK, lifecycle, ESM and OTLP tests with network isolated + run: | + sudo unshare --net -- bash -c 'ip link set lo up; exec "$@"' _ \ + "$(command -v node)" node_modules/jest/bin/jest.js \ + --config packages/traceai_cohere/jest.config.js --runInBand + - name: Package tarball + run: pnpm --filter '@traceai/cohere' pack --pack-destination "$RUNNER_TEMP/cohere-package" diff --git a/typescript/packages/traceai_cohere/README.md b/typescript/packages/traceai_cohere/README.md index 7aa631782..d27614c53 100644 --- a/typescript/packages/traceai_cohere/README.md +++ b/typescript/packages/traceai_cohere/README.md @@ -1,436 +1,126 @@ # @traceai/cohere -OpenTelemetry instrumentation for the [Cohere](https://cohere.com/) TypeScript/JavaScript SDK. +Trace Cohere chat, streaming chat, embeddings and reranking in Future AGI with the maintained Cohere SDK instrumentor. -## Installation +## Install ```bash -npm install @traceai/cohere -# or -pnpm add @traceai/cohere -# or -yarn add @traceai/cohere +npm install @traceai/cohere @traceai/fi-core cohere-ai ``` -## Features +Use Node.js 18 or later and `cohere-ai` 7.x or 8.x. Use `cohere-ai` 8.3.0 for the v2 examples below. The 7.0.0 SDK only exposes the older v1 operations; streaming requires a version with `chatStream`. -- Automatic instrumentation of Cohere API calls -- Supports chat, embed, and rerank APIs -- Captures request parameters, model information, and responses -- Configurable input/output hiding for privacy -- Compatible with OpenTelemetry ecosystem +## Configure tracing -## Usage +Set `FI_API_KEY`, `FI_SECRET_KEY` and `COHERE_API_KEY` in your environment. Get your Future AGI credentials from your organization's settings. -### Basic Setup +Initialize tracing **before constructing any Cohere clients**. Explicit module instrumentation works with both CommonJS and native ESM, without an import hook or copied tracing helpers. ```typescript +import { register, ProjectType } from "@traceai/fi-core"; import { CohereInstrumentation } from "@traceai/cohere"; -import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; -import { registerInstrumentations } from "@opentelemetry/instrumentation"; - -// Set up the tracer provider -const provider = new NodeTracerProvider(); -provider.register(); - -// Register the Cohere instrumentation -registerInstrumentations({ - instrumentations: [new CohereInstrumentation()], -}); - -// Now use the Cohere SDK as normal -import { CohereClient } from "cohere-ai"; - -const cohere = new CohereClient({ - token: process.env.COHERE_API_KEY, -}); +import * as Cohere from "cohere-ai"; -// Chat example -const chatResponse = await cohere.chat({ - message: "Hello, how are you?", - model: "command-r-08-2024", -}); - -// Embed example -const embedResponse = await cohere.embed({ - texts: ["hello", "goodbye"], - model: "embed-english-v3.0", - inputType: "search_document", -}); - -// Rerank example -const rerankResponse = await cohere.rerank({ - query: "What is the capital of France?", - documents: [ - "Paris is the capital of France.", - "Berlin is the capital of Germany.", - "London is the capital of the United Kingdom.", - ], - model: "rerank-english-v3.0", +const provider = register({ + projectName: "cohere-app", + projectType: ProjectType.OBSERVE, }); -``` - -### Manual Instrumentation - -If automatic instrumentation doesn't work (e.g., with bundlers), you can manually instrument: -```typescript -import { CohereInstrumentation } from "@traceai/cohere"; -import * as CohereModule from "cohere-ai"; - -const instrumentation = new CohereInstrumentation(); -instrumentation.manuallyInstrument(CohereModule); -``` - -### Configuration Options - -```typescript const instrumentation = new CohereInstrumentation({ - // OpenTelemetry instrumentation config - instrumentationConfig: { - enabled: true, - }, - // TraceAI-specific config traceConfig: { - hideInputs: false, // Set to true to hide input content in traces - hideOutputs: false, // Set to true to hide output content in traces + hideInputs: true, + hideOutputs: true, + hideEmbeddingVectors: true, }, }); -``` - -### With TraceAI Core - -```typescript -import { register, ProjectType } from "@traceai/fi-core"; -import { CohereInstrumentation } from "@traceai/cohere"; -import { registerInstrumentations } from "@opentelemetry/instrumentation"; - -// 1. Register TraceAI Core TracerProvider -const tracerProvider = register({ - projectName: "my-cohere-app", - projectType: ProjectType.OBSERVE, - sessionName: "cohere-session-" + Date.now(), -}); - -// 2. Register instrumentation BEFORE importing Cohere SDK -registerInstrumentations({ - tracerProvider: tracerProvider, - instrumentations: [new CohereInstrumentation()], -}); +instrumentation.setTracerProvider(provider); +instrumentation.manuallyInstrument(Cohere); -// 3. NOW import and use Cohere -const { CohereClient } = await import("cohere-ai"); -const client = new CohereClient({ token: process.env.COHERE_API_KEY }); - -// 4. Use normally - all calls are traced -const response = await client.chat({ - message: "Hello!", - model: "command-r-08-2024", +const client = new Cohere.CohereClientV2({ + token: process.env.COHERE_API_KEY, }); -// 5. Shutdown when done -await tracerProvider.shutdown(); -``` - -## Real-World Use Cases - -### 1. Conversational AI with Chat History - -```typescript -const { CohereClient } = await import("cohere-ai"); -const client = new CohereClient({ token: process.env.COHERE_API_KEY }); - -async function chat(message: string, history: Array<{ role: string; message: string }>) { +try { const response = await client.chat({ - message, model: "command-r-08-2024", - chatHistory: history, - preamble: "You are a helpful customer support agent for an e-commerce platform.", + messages: [{ role: "user", content: "Explain semantic search in one sentence." }], }); - - // Add to history - history.push({ role: "USER", message }); - history.push({ role: "CHATBOT", message: response.text }); - - return response.text; + console.log(response.message); +} finally { + instrumentation.disable(); + await provider.shutdown(); } - -// Usage -const history: Array<{ role: string; message: string }> = []; -const answer1 = await chat("I have a problem with my order", history); -const answer2 = await chat("Order ID is 12345", history); ``` -### 2. Semantic Search with Embeddings +Open the **cohere-app** project in Future AGI to inspect the `Cohere Chat` span, model and token usage. This example hides request/response content, embedding vectors and error details. To inspect content in traces, explicitly turn off the corresponding hiding options only when it is safe for your application. -```typescript -async function semanticSearch(query: string, documents: string[]) { - // Get query embedding (use 'search_query' for queries) - const queryEmbedding = await client.embed({ - texts: [query], - model: "embed-english-v3.0", - inputType: "search_query", - }); +If you already have an OpenTelemetry tracer provider, pass it to `setTracerProvider` instead of creating a second provider. For CommonJS auto-instrumentation, register `CohereInstrumentation` with OpenTelemetry before requiring `cohere-ai`; do not also install another Cohere instrumentor. - // Get document embeddings (use 'search_document' for documents) - const docEmbeddings = await client.embed({ - texts: documents, - model: "embed-english-v3.0", - inputType: "search_document", - }); +## Stream a response - // Calculate cosine similarities - const queryVector = queryEmbedding.embeddings[0] as number[]; - const similarities = (docEmbeddings.embeddings as number[][]).map((docVector, i) => ({ - document: documents[i], - score: cosineSimilarity(queryVector, docVector), - })); - - return similarities.sort((a, b) => b.score - a.score); -} -``` - -### 3. Document Reranking for RAG +With the same initialized client, replace the chat call with: ```typescript -async function rerankDocuments(query: string, documents: string[]) { - const response = await client.rerank({ - query, - documents, - model: "rerank-english-v3.0", - topN: 5, - returnDocuments: true, - }); - - return response.results.map((result) => ({ - document: result.document?.text, - score: result.relevanceScore, - originalIndex: result.index, - })); -} - -// Usage in RAG pipeline -const retrievedDocs = await vectorSearch(userQuery, 20); // Get top 20 from vector DB -const rerankedDocs = await rerankDocuments(userQuery, retrievedDocs); -const topDocs = rerankedDocs.slice(0, 3); // Use top 3 for context -``` - -### 4. Tool-Augmented Chatbot - -```typescript -async function chatWithTools(message: string) { - const response = await client.chat({ - message, - model: "command-r-08-2024", - tools: [ - { - name: "search_database", - description: "Search the product database", - parameterDefinitions: { - query: { type: "str", description: "Search query", required: true }, - category: { type: "str", description: "Product category", required: false }, - }, - }, - { - name: "get_order_status", - description: "Get status of a customer order", - parameterDefinitions: { - orderId: { type: "str", description: "Order ID", required: true }, - }, - }, - ], - }); +const stream = await client.chatStream({ + model: "command-r-08-2024", + messages: [{ role: "user", content: "Explain semantic search in one sentence." }], +}); - if (response.toolCalls && response.toolCalls.length > 0) { - // Execute tool calls and continue conversation - for (const toolCall of response.toolCalls) { - console.log(`Tool: ${toolCall.name}, Args:`, toolCall.parameters); - } +for await (const event of stream) { + if (event.type === "content-start" || event.type === "content-delta") { + process.stdout.write(event.delta?.message?.content?.text ?? ""); } - - return response; } ``` -### 5. Multilingual Embeddings - -```typescript -async function multilingualSearch(query: string, documents: string[]) { - // Use multilingual model for cross-language search - const queryEmbedding = await client.embed({ - texts: [query], - model: "embed-multilingual-v3.0", - inputType: "search_query", - }); +The span stays open while you consume the stream. It finishes at the terminal event, iterator completion, failure, early return or controller abort. Breaking or returning from the iterator, aborting its controller, or receiving `USER_CANCEL` sets `cohere.stream.cancelled` and leaves the span status unset. A stream that ends without a terminal event is an error, not a cancellation. - const docEmbeddings = await client.embed({ - texts: documents, - model: "embed-multilingual-v3.0", - inputType: "search_document", - }); +For both chat and streaming chat, `ERROR`, `ERROR_TOXIC`, `ERROR_LIMIT` and `TIMEOUT` finish reasons mark the span as an error. A nonstreaming `USER_CANCEL` response sets `cohere.request.cancelled` with an unset status. Partial generated output is retained when content capture is enabled; token counts are available only if Cohere sends usage. These telemetry statuses do not change SDK results or turn resolved responses into exceptions. - // Works across languages! - // Query in English can match documents in French, German, Spanish, etc. - return calculateSimilarities(queryEmbedding, docEmbeddings); -} -``` +Consume the stream or explicitly close/abort it rather than abandoning it. SDK stream controllers, events and `withRawResponse()` remain available. -### 6. Streaming Chat Responses +## Embed and rerank ```typescript -async function streamChat(message: string) { - const stream = await client.chatStream({ - message, - model: "command-r-08-2024", - }); - - let fullText = ""; - for await (const event of stream) { - if (event.eventType === "text-generation") { - process.stdout.write(event.text); - fullText += event.text; - } else if (event.eventType === "stream-end") { - console.log("\n--- Stream complete ---"); - console.log("Token usage:", event.response?.meta?.tokens); - } - } +await client.embed({ + model: "embed-english-v3.0", + texts: ["Paris is the capital of France."], + inputType: "search_document", + embeddingTypes: ["float"], +}); - return fullText; -} +await client.rerank({ + model: "rerank-v3.5", + query: "What is the capital of France?", + documents: ["Paris is the capital of France.", "Berlin is the capital of Germany."], + topN: 1, +}); ``` -## Traced Operations - -### Chat API - -The instrumentation captures: - -- Model name -- Input messages and chat history -- Preamble and connectors -- Temperature and other generation parameters -- Response text and metadata -- Token usage (prompt and completion tokens) - -### Embed API - -The instrumentation captures: +Embeddings use the `EMBEDDING` span kind; reranking uses `RERANKER`, including result indices and relevance scores. Chat uses `LLM`. The instrumentor reads the SDK's camelCase properties, including `chatHistory`, `toolCalls`, `inputTokens` and `relevanceScore`, and retains support for v1 payloads. -- Model name -- Input texts -- Input type -- Truncation settings -- Output embeddings (unless hidden) -- Token usage +## Existing v1 applications -### Rerank API +Keep using `CohereClient`, `CohereInstrumentation` and `manuallyInstrument` with the same setup. Calls through `CohereClient.v2` and `CohereClientV2` are also traced when available in your SDK. Disable/enable applies to subsequent calls, including already-bound v2 methods; in-flight calls finish their spans normally. -The instrumentation captures: +This integration covers `chat`, `chatStream`, `embed` and `rerank`, not every endpoint on the Cohere SDK. It does not infer a default model when no model was supplied. -- Model name -- Query text -- Documents to rerank -- Top N parameter -- Reranked results with relevance scores +## Troubleshooting -## Semantic Conventions +- **No v2 spans:** initialize instrumentation before constructing the client. V2 binds its methods in the constructor; recreate clients constructed before setup. +- **Duplicate spans:** use one Cohere instrumentor per application and remove overlapping instrumentation. +- **A stream span never finishes:** consume the stream, break/return from its iterator, or abort its controller before shutting down the provider. +- **Missing content or vectors:** check your `traceConfig` and `FI_HIDE_*` environment settings. Input/output hiding also applies to Cohere-specific fields and raw payloads; vector hiding excludes the raw embedding response. Connector access tokens are always omitted from telemetry, regardless of content settings; requests to Cohere retain them. -This instrumentation follows OpenTelemetry semantic conventions and FutureAGI conventions for LLM observability: +## Development -- Span names follow the pattern: `cohere.{operation}` (e.g., `cohere.chat`, `cohere.embed`, `cohere.rerank`) -- Uses `LLM` span kind for chat operations -- Uses `EMBEDDING` span kind for embed operations -- Uses `RERANKER` span kind for rerank operations - -## Supported Models - -### Chat Models - -- command -- command-r -- command-r-08-2024 -- command-light -- command-nightly - -### Embedding Models - -- embed-english-v3.0 -- embed-multilingual-v3.0 -- embed-english-light-v3.0 -- embed-multilingual-light-v3.0 - -### Rerank Models - -- rerank-english-v3.0 -- rerank-multilingual-v3.0 -- rerank-english-v2.0 -- rerank-multilingual-v2.0 - -## Requirements - -- Node.js >= 18 -- cohere-ai SDK - -## Running Examples - -The `examples/` directory contains real-world examples: - -```bash -cd examples -cp .env.example .env -# Edit .env with your API keys - -pnpm install -pnpm run example -``` - -## E2E Testing - -Run E2E tests with real API keys: +From the TypeScript workspace, build shared packages and the instrumentor before running its offline suite: ```bash -COHERE_API_KEY=your_key pnpm test -- --testPathPattern=e2e -``` - -## Integration with TraceAI Platform - -```typescript -import { CohereInstrumentation } from "@traceai/cohere"; -import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; -import { SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base"; -import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto"; -import { Resource } from "@opentelemetry/resources"; - -const provider = new NodeTracerProvider({ - resource: new Resource({ - "service.name": "my-cohere-app", - "deployment.environment": "production", - }), -}); - -const exporter = new OTLPTraceExporter({ - url: "https://api.traceai.com/v1/traces", - headers: { - Authorization: `Bearer ${process.env.TRACEAI_API_KEY}`, - }, -}); - -provider.addSpanProcessor(new SimpleSpanProcessor(exporter)); -provider.register(); - -const instrumentation = new CohereInstrumentation(); -instrumentation.manuallyInstrument(require("cohere-ai")); +pnpm --filter '@traceai/cohere...' --workspace-concurrency=1 run build +pnpm --filter '@traceai/cohere' exec jest --runInBand ``` ## License Apache-2.0 - -## Related Packages - -- [@traceai/openai](../traceai_openai) - OpenAI instrumentation -- [@traceai/anthropic](../traceai_anthropic) - Anthropic instrumentation -- [@traceai/groq](../traceai_groq) - Groq instrumentation -- [@traceai/mistral](../traceai_mistral) - Mistral instrumentation -- [@traceai/google-genai](../traceai_google_genai) - Google GenAI instrumentation -- [@traceai/fi-core](../fi-core) - Core tracing utilities diff --git a/typescript/packages/traceai_cohere/package.json b/typescript/packages/traceai_cohere/package.json index 2bb212042..d7b290983 100644 --- a/typescript/packages/traceai_cohere/package.json +++ b/typescript/packages/traceai_cohere/package.json @@ -15,7 +15,7 @@ }, "scripts": { "prebuild": "rimraf dist", - "build": "tsc --build tsconfig.json tsconfig.esm.json tsconfig.esnext.json && tsc-alias -p tsconfig.esm.json && pnpm run postbuild", + "build": "tsc --build tsconfig.json tsconfig.esm.json tsconfig.esnext.json && tsc-alias -p tsconfig.esm.json --resolve-full-paths --resolve-full-extension .js && pnpm run postbuild", "postbuild": "echo '{\"type\": \"module\"}' > ./dist/esm/package.json", "test": "jest", "test:watch": "jest --watch", @@ -38,25 +38,28 @@ "homepage": "https://docs.futureagi.com/docs/tracing", "dependencies": { "@opentelemetry/api": "^1.9.0", - "@opentelemetry/instrumentation": "^0.202.0", "@opentelemetry/core": "^2.0.1", + "@opentelemetry/instrumentation": "^0.202.0", "@traceai/fi-core": "workspace:^", "@traceai/fi-semantic-conventions": "workspace:^" }, "peerDependencies": { - "cohere-ai": ">=7.0.0" + "cohere-ai": ">=7.0.0 <9" }, "devDependencies": { + "@opentelemetry/context-async-hooks": "2.2.0", + "@opentelemetry/exporter-trace-otlp-http": "0.202.0", "@opentelemetry/exporter-trace-otlp-proto": "^0.50.0", "@opentelemetry/resources": "^2.0.1", "@opentelemetry/sdk-trace-base": "^2.0.1", "@opentelemetry/sdk-trace-node": "^2.0.1", "@opentelemetry/semantic-conventions": "^1.34.0", "@types/jest": "^29.5.12", + "cohere-ai": "8.3.0", + "cohere-ai-v7": "npm:cohere-ai@7.0.0", "jest": "^29.7.0", - "ts-jest": "^29.1.4", - "cohere-ai": "^7.0.0", "rimraf": "^6.0.1", + "ts-jest": "^29.1.4", "tsc-alias": "^1.8.16" }, "keywords": [ diff --git a/typescript/packages/traceai_cohere/src/__tests__/attributes.test.ts b/typescript/packages/traceai_cohere/src/__tests__/attributes.test.ts new file mode 100644 index 000000000..0c83fe55e --- /dev/null +++ b/typescript/packages/traceai_cohere/src/__tests__/attributes.test.ts @@ -0,0 +1,16 @@ +import { generateTraceConfig } from "@traceai/fi-core"; +import { SemanticConventions as S } from "@traceai/fi-semantic-conventions"; +import { inputAttributes } from "../attributes"; + +it.each(["userAccessToken", "user_access_token"])("omits %s from telemetry without changing any connector", (field) => { + const authenticated = Object.freeze({ id: "private-connector", [field]: "synthetic-token" }); + const publicConnector = Object.freeze({ id: "web-search", options: { site: "example.test" } }); + const request = Object.freeze({ message: "question", connectors: Object.freeze([authenticated, publicConnector]) }); + const attrs = inputAttributes("chat", request, generateTraceConfig()); + expect(JSON.stringify(attrs)).not.toContain("synthetic-token"); + expect(JSON.parse(attrs[S.RAW_INPUT] as string).connectors).toEqual([ + { id: "private-connector" }, publicConnector, + ]); + expect(request.connectors[0]).toBe(authenticated); + expect(authenticated[field]).toBe("synthetic-token"); +}); diff --git a/typescript/packages/traceai_cohere/src/__tests__/auto.test.ts b/typescript/packages/traceai_cohere/src/__tests__/auto.test.ts new file mode 100644 index 000000000..5ae44a7fd --- /dev/null +++ b/typescript/packages/traceai_cohere/src/__tests__/auto.test.ts @@ -0,0 +1,61 @@ +import { execFile } from "node:child_process"; +import { promisify } from "node:util"; +import { resolve } from "node:path"; + +const run = promisify(execFile); +const cwd = resolve(__dirname, "../.."); + +it("automatically instruments a fresh CommonJS SDK import", async () => { + // Isolate the require hook from Jest's module loader; no SDK methods are mocked. + const source = ` + const { InMemorySpanExporter, SimpleSpanProcessor, BasicTracerProvider } = require('@opentelemetry/sdk-trace-base'); + const { CohereInstrumentation } = require('./dist/src/index.js'); + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); + const instrumentation = new CohereInstrumentation(); + instrumentation.setTracerProvider(provider); + const { CohereClientV2 } = require('cohere-ai'); + const client = new CohereClientV2({ token: 'offline-placeholder', fetcher: async () => ({ + ok: true, rawResponse: { status: 200 }, body: { + id: 'fixture', finish_reason: 'COMPLETE', message: { role: 'assistant', content: [{ type: 'text', text: 'hello' }] } + } + }) }); + (async () => { + await client.chat({ model: 'command-test', messages: [{ role: 'user', content: 'question' }] }); + const spans = exporter.getFinishedSpans(); + if (spans.length !== 1) throw new Error('Expected one automatically instrumented span, got ' + spans.length); + instrumentation.disable(); + await client.chat({ model: 'command-test', messages: [] }); + if (exporter.getFinishedSpans().length !== 1) throw new Error('Disabled bound method still traced'); + await provider.shutdown(); + console.log('automatic-v2-ok'); + })().catch((error) => { console.error(error); process.exitCode = 1; }); + `; + const { stdout } = await run(process.execPath, ["-e", source], { cwd }); + expect(stdout).toContain("automatic-v2-ok"); +}); + +it("loads the native ESM package entry and manually instruments a real v2 client", async () => { + const source = ` + import { CohereInstrumentation } from './dist/esm/index.js'; + import * as Cohere from 'cohere-ai'; + import { InMemorySpanExporter, SimpleSpanProcessor, BasicTracerProvider } from '@opentelemetry/sdk-trace-base'; + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); + const instrumentation = new CohereInstrumentation(); + instrumentation.setTracerProvider(provider); + instrumentation.manuallyInstrument(Cohere); + const client = new Cohere.CohereClientV2({ token: 'offline-placeholder', fetcher: async () => ({ + ok: true, rawResponse: { status: 200 }, body: { + id: 'fixture', finish_reason: 'COMPLETE', message: { role: 'assistant', content: [{ type: 'text', text: 'hello' }] } + } + }) }); + await client.chat({ model: 'command-test', messages: [{ role: 'user', content: 'question' }] }); + if (exporter.getFinishedSpans().length !== 1) throw new Error('Missing ESM v2 span'); + instrumentation.disable(); + await provider.shutdown(); + console.log('esm-v2-ok'); + `; + const { stdout } = await run(process.execPath, ["--input-type=module", "-e", source], { cwd }); + expect(stdout).toContain("esm-v2-ok"); +}); diff --git a/typescript/packages/traceai_cohere/src/__tests__/lifecycle.test.ts b/typescript/packages/traceai_cohere/src/__tests__/lifecycle.test.ts new file mode 100644 index 000000000..adf1d99e4 --- /dev/null +++ b/typescript/packages/traceai_cohere/src/__tests__/lifecycle.test.ts @@ -0,0 +1,149 @@ +import { context, trace, SpanStatusCode } from "@opentelemetry/api"; +import { BasicTracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base"; +import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; +import { CohereInstrumentation, isPatched } from "../instrumentation"; + +const exporter = new InMemorySpanExporter(); +const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); +let instrumentation: CohereInstrumentation; +beforeAll(() => context.setGlobalContextManager(new AsyncLocalStorageContextManager().enable())); +beforeEach(() => { + exporter.reset(); + instrumentation = new CohereInstrumentation(); + instrumentation.setTracerProvider(provider); +}); +afterEach(() => instrumentation.disable()); +afterAll(async () => { context.disable(); await provider.shutdown(); }); + +function instrument(CohereClient: T): T { + instrumentation.manuallyInstrument({ CohereClient }); + return CohereClient; +} + +it("does not flag an unsupported module as patched", () => { + instrumentation.manuallyInstrument({}); + expect(isPatched()).toBe(false); +}); + +it("preserves synchronous errors and ends exactly one span", () => { + const error = new Error("synchronous fixture failure"); + class Client { chat() { throw error; } } + instrument(Client); + expect(() => new Client().chat()).toThrow(error); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it("does not double trace nested Cohere methods", async () => { + class Client { + async chat() { return this.embed(); } + async embed() { return { embeddings: [[1]] }; } + } + instrument(Client); + await new Client().chat(); + expect(exporter.getFinishedSpans()).toHaveLength(1); +}); + +it("restores inherited methods without leaving shadow properties", () => { + class Base { async chat() { return { text: "hello" }; } } + class Client extends Base {} + const original = Base.prototype.chat; + instrument(Client); + expect(Object.hasOwn(Client.prototype, "chat")).toBe(true); + instrumentation.disable(); + expect(Object.hasOwn(Client.prototype, "chat")).toBe(false); + expect(Client.prototype.chat).toBe(original); +}); + +it("keeps SDK iterator work in the span context, but not consumer work", async () => { + const active: (string | undefined)[] = []; + let controller: AbortController; + class Client { + async chatStream() { + active.push(trace.getActiveSpan()?.spanContext().spanId); + controller = new AbortController(); + return { + controller, + async *[Symbol.asyncIterator]() { + active.push(trace.getActiveSpan()?.spanContext().spanId); + yield { eventType: "text-generation", text: "hello" }; + await Promise.resolve(); + active.push(trace.getActiveSpan()?.spanContext().spanId); + }, + }; + } + } + instrument(Client); + for await (const _event of await new Client().chatStream()) { + expect(trace.getActiveSpan()).toBeUndefined(); + } + const [span] = exporter.getFinishedSpans(); + expect(active).toEqual(Array(3).fill(span.spanContext().spanId)); + controller!.abort(); + expect(exporter.getFinishedSpans()).toHaveLength(1); +}); + +it("forwards iterator.throw and preserves error identity", async () => { + const failure = new Error("iterator failure"); + const cleanup = jest.fn(); + class Client { + async chatStream() { + return { + async *[Symbol.asyncIterator]() { + try { yield { eventType: "text-generation", text: "partial" }; } + finally { cleanup(); } + }, + }; + } + } + instrument(Client); + const stream = await new Client().chatStream(); + const iterator = stream[Symbol.asyncIterator](); + await iterator.next(); + await expect(iterator.throw(failure)).rejects.toBe(failure); + expect(cleanup).toHaveBeenCalledTimes(1); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it("marks terminal error events without changing their payload", async () => { + const event = { eventType: "stream-end", finishReason: "ERROR", response: { text: "partial" } }; + class Client { async chatStream() { return { async *[Symbol.asyncIterator]() { yield event; } }; } } + instrument(Client); + const values = []; + for await (const item of await new Client().chatStream()) values.push(item); + expect(values[0]).toBe(event); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it("keeps v1 usable if v2 discovery fails", async () => { + class Client { + get v2() { throw new Error("unsupported getter"); } + async chat() { return { text: "hello" }; } + } + instrument(Client); + await new Client().chat(); + expect(exporter.getFinishedSpans()).toHaveLength(1); +}); + +it("ends a span when iterator creation itself fails", async () => { + const failure = new Error("iterator creation failed"); + class Client { + async chatStream() { + return { [Symbol.asyncIterator](): AsyncIterator { throw failure; } }; + } + } + instrument(Client); + const stream = await new Client().chatStream(); + expect(() => stream[Symbol.asyncIterator]()).toThrow(failure); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it("marks an undefined rejection as an error without changing it", async () => { + class Client { chat() { return Promise.reject(undefined); } } + instrument(Client); + await expect(new Client().chat()).rejects.toBeUndefined(); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); diff --git a/typescript/packages/traceai_cohere/src/__tests__/sdk.test.ts b/typescript/packages/traceai_cohere/src/__tests__/sdk.test.ts new file mode 100644 index 000000000..15db663e4 --- /dev/null +++ b/typescript/packages/traceai_cohere/src/__tests__/sdk.test.ts @@ -0,0 +1,708 @@ +import { createServer, Server } from "node:http"; +import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node"; +import { InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base"; +import { SemanticConventions as S } from "@traceai/fi-semantic-conventions"; +import { CohereInstrumentation, isPatched } from "../instrumentation"; +import { context, trace, SpanStatusCode } from "@opentelemetry/api"; +import { suppressTracing } from "@opentelemetry/core"; +import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http"; + +// Real SDK serializers and transports, with an offline localhost API double. +const sdk8 = require("cohere-ai"); +const sdk7 = require("cohere-ai-v7"); +const exporter = new InMemorySpanExporter(); +const provider = new NodeTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); +let server: Server; +let url: string; +let response: unknown; +let contentType = "application/json"; +let status = 200; +let received: { path: string; body: any; headers: unknown }[] = []; +let wire: any[] = []; +let instrumentation: CohereInstrumentation; +const v1Response = { text: "hello", generation_id: "generation", meta: { tokens: { input_tokens: 5, output_tokens: 2 } } }; +const v2Response = { id: "generation", finish_reason: "COMPLETE", message: { role: "assistant", content: [{ type: "text", text: "hello" }] }, usage: { tokens: { input_tokens: 5, output_tokens: 2 } } }; + +beforeAll(async () => { + provider.register(); + server = createServer((request, reply) => { + const chunks: Buffer[] = []; + request.on("data", (chunk) => chunks.push(chunk)); + request.on("end", () => { + const body = JSON.parse(Buffer.concat(chunks).toString() || "{}"); + if (request.url === "/v1/traces") { + wire.push(body); + reply.setHeader("Content-Type", "application/json"); + reply.end("{}"); + return; + } + received.push({ path: request.url!, body, headers: request.headers }); + reply.statusCode = status; + reply.setHeader("Content-Type", contentType); + reply.end(typeof response === "string" ? response : JSON.stringify(response)); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + url = `http://127.0.0.1:${(server.address() as { port: number }).port}`; +}); +beforeEach(() => { + exporter.reset(); + wire = []; + received = []; + status = 200; + contentType = "application/json"; + instrumentation = new CohereInstrumentation(); + instrumentation.setTracerProvider(provider); +}); +afterEach(() => instrumentation.disable()); +afterAll(async () => { + await new Promise((resolve) => server.close(() => resolve())); + await provider.shutdown(); +}); + +it("traces CohereClientV2's generated v2 client", async () => { + instrumentation.manuallyInstrument(sdk8); + response = v2Response; + const client = new sdk8.CohereClientV2({ token: "offline-placeholder", environment: url }); + const result = await client.chat({ model: "command-test", messages: [{ role: "user", content: "question" }] }); + expect(result.message.content[0].text).toBe("hello"); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); +}); + +it.each([["7.0.0", sdk7], ["8.3.0", sdk8]])("reads the real %s SDK camelCase chat history", async (_version, sdk) => { + instrumentation.manuallyInstrument(sdk); + response = v1Response; + const client = new sdk.CohereClient({ token: "offline-placeholder", environment: url }); + await client.chat({ model: "command-test", message: "question", chatHistory: [{ role: "USER", message: "earlier" }] }); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].attributes[S.LLM_INPUT_MESSAGES]).toContain("earlier"); +}); + +function clientFor(sdk: any, v2 = false): any { + instrumentation.manuallyInstrument(sdk); + return new (v2 ? sdk.CohereClientV2 : sdk.CohereClient)({ token: "offline-placeholder", environment: url }); +} +const request = { model: "command-test", message: "question" }; +const requestV2 = { model: "command-test", messages: [{ role: "user", content: "question" }] }; +function streamResponse(v2: boolean) { + const events = v2 ? [ + { type: "content-delta", index: 0, delta: { message: { content: { text: "hel" } } } }, + { type: "content-delta", index: 0, delta: { message: { content: { text: "lo" } } } }, + { type: "tool-call-start", index: 0, delta: { message: { tool_calls: { id: "call-1", type: "function", function: { name: "lookup", arguments: "{\"key\":" } } } } }, + { type: "tool-call-delta", index: 0, delta: { message: { tool_calls: { function: { arguments: "\"value\"}" } } } } }, + { type: "message-end", delta: { finish_reason: "COMPLETE", usage: { tokens: { input_tokens: 5, output_tokens: 2 } } } }, + ] : [ + { event_type: "text-generation", text: "hel" }, + { event_type: "text-generation", text: "lo" }, + { event_type: "stream-end", finish_reason: "COMPLETE", response: v1Response }, + ]; + contentType = v2 ? "text/event-stream" : "application/x-ndjson"; + response = events.map((event) => v2 ? `data: ${JSON.stringify(event)}\n\n` : `${JSON.stringify(event)}\n`).join(""); + return events; +} + +it.each([false, true])("traces real v2=%s SDK streaming through terminal usage and preserves helpers", async (v2) => { + const client = clientFor(sdk8, v2); + const expected = streamResponse(v2); + const pending = client.chatStream(v2 ? requestV2 : request); + expect(typeof pending.withRawResponse).toBe("function"); + const { data: stream, rawResponse } = await pending.withRawResponse(); + expect(rawResponse.status).toBe(200); + expect(await pending).toBe(stream); + expect(stream.controller).toBeInstanceOf(AbortController); + expect(exporter.getFinishedSpans()).toHaveLength(0); + const events = []; + for await (const event of stream) events.push(event); + expect(events).toHaveLength(expected.length); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe("hello"); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_COMPLETION]).toBe(2); + expect(spans[0].status.code).toBe(SpanStatusCode.OK); + if (v2) { + const messages = JSON.parse(spans[0].attributes[S.LLM_OUTPUT_MESSAGES] as string); + expect(messages[0].tool_calls[0]).toEqual({ id: "call-1", type: "function", function: { name: "lookup", arguments: '{"key":"value"}' } }); + } +}); + +it.each([false, true])("closes a real v2=%s SDK stream when the consumer breaks early", async (v2) => { + const client = clientFor(sdk8, v2); + streamResponse(v2); + const stream = await client.chatStream(v2 ? requestV2 : request); + for await (const _event of stream) break; + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].attributes["cohere.stream.cancelled"]).toBe(true); + expect(exporter.getFinishedSpans()[0].attributes[S.OUTPUT_VALUE]).toBe("hel"); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.UNSET); +}); + +it("finalizes on abort before iteration, without consuming events", async () => { + const client = clientFor(sdk8, true); + streamResponse(true); + const stream = await client.chatStream(requestV2); + stream.controller.abort(); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].attributes["cohere.stream.cancelled"]).toBe(true); +}); + +it("finalizes on return before first next", async () => { + const client = clientFor(sdk8, true); + streamResponse(true); + const stream = await client.chatStream(requestV2); + await stream[Symbol.asyncIterator]().return(); + expect(exporter.getFinishedSpans()).toHaveLength(1); +}); + +it("records SDK stream parser errors and preserves the rejection", async () => { + const client = clientFor(sdk8, true); + contentType = "text/event-stream"; + response = "data: not-json\n\n"; + const stream = await client.chatStream(requestV2); + await expect(stream[Symbol.asyncIterator]().next()).rejects.toBeInstanceOf(Error); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it.each([["7.0.0", sdk7], ["8.3.0", sdk8]])("preserves real %s SDK embeddings and reranker types", async (_version, sdk) => { + const client = clientFor(sdk); + response = { id: "embedding", response_type: "embeddings_floats", texts: ["document"], embeddings: [[0.1, 0.2]], meta: { billed_units: { input_tokens: 3 } } }; + await client.embed({ texts: ["document"], model: "embed-test", inputType: "search_document" }); + response = { id: "rerank", results: [{ index: 0, relevance_score: 0.9, document: { text: "document" } }] }; + const result = await client.rerank({ model: "rerank-test", query: "query", documents: [{ text: "document" }], topN: 1 }); + expect(result.results[0].relevanceScore).toBe(0.9); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(2); + expect(spans[0].attributes[S.FI_SPAN_KIND]).toBe("EMBEDDING"); + expect(spans[0].attributes[`${S.EMBEDDING_EMBEDDINGS}.0.${S.EMBEDDING_VECTOR}`]).toEqual([0.1, 0.2]); + expect(spans[1].attributes[S.FI_SPAN_KIND]).toBe("RERANKER"); + expect(spans[1].attributes["cohere.rerank.results.0.relevance_score"]).toBe(0.9); + expect(spans[1].attributes["cohere.rerank.top_n"]).toBe(1); + expect(received[1].body.top_n).toBe(1); +}); + +it("instruments v2 embeddings, rerank and client.v2 without duplicate spans", async () => { + const client = clientFor(sdk8, true); + response = { id: "embedding", texts: ["document"], embeddings: { float: [[0.1, 0.2]] }, meta: { billed_units: { input_tokens: 3 } } }; + await client.embed({ texts: ["document"], model: "embed-test", inputType: "search_document", embeddingTypes: ["float"] }); + response = { id: "rerank", results: [{ index: 0, relevance_score: 0.9 }] }; + await client.rerank({ model: "rerank-test", query: "query", documents: ["document"] }); + response = v2Response; + const v1 = new sdk8.CohereClient({ token: "offline-placeholder", environment: url }); + const pending = v1.v2.chat(requestV2, { headers: { "X-Test-Forwarded": "yes" } }); + const raw = await pending.withRawResponse(); + expect(raw.data).toBe(await pending); + expect(raw.rawResponse.status).toBe(200); + expect(received[2].headers).toMatchObject({ "x-test-forwarded": "yes" }); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(3); + expect(spans.map((span) => span.attributes[S.FI_SPAN_KIND])).toEqual(["EMBEDDING", "RERANKER", "LLM"]); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBe(3); +}); + +it("does not invent a model name when the caller leaves it to Cohere", async () => { + const client = clientFor(sdk7); + response = v1Response; + await client.chat({ message: "question" }); + expect(exporter.getFinishedSpans()[0].attributes[S.LLM_MODEL_NAME]).toBeUndefined(); +}); + +it("supports multiple SDK copies, idempotent manual patching and enable/disable", async () => { + const client7 = clientFor(sdk7); + const client8 = clientFor(sdk8, true); + instrumentation.manuallyInstrument(sdk8); + expect(isPatched()).toBe(true); + response = v1Response; + await client7.chat(request); + response = v2Response; + await client8.chat(requestV2); + expect(exporter.getFinishedSpans()).toHaveLength(2); + instrumentation.disable(); + expect(isPatched()).toBe(false); + await client8.chat(requestV2); // already-bound method must respect disable + instrumentation.enable(); + await client8.chat(requestV2); + const newClient8 = clientFor(sdk8, true); + await newClient8.chat(requestV2); + expect(exporter.getFinishedSpans()).toHaveLength(4); +}); + +it("does not let a second instrumentor unwrap another owner's methods", async () => { + const client = clientFor(sdk8, true); + const other = new CohereInstrumentation(); + other.setTracerProvider(provider); + other.manuallyInstrument(sdk8); + other.disable(); + response = v2Response; + await client.chat(requestV2); + expect(exporter.getFinishedSpans()).toHaveLength(1); +}); + +it("honors suppression and propagates the parent across concurrent SDK calls", async () => { + const client = clientFor(sdk8, true); + response = v2Response; + await context.with(suppressTracing(context.active()), () => client.chat(requestV2)); + expect(exporter.getFinishedSpans()).toHaveLength(0); + const parent = provider.getTracer("parent").startSpan("parent"); + await context.with(trace.setSpan(context.active(), parent), () => Promise.all([client.chat(requestV2), client.chat(requestV2)])); + parent.end(); + const children = exporter.getFinishedSpans().filter((span) => span.name !== "parent"); + expect(children).toHaveLength(2); + for (const child of children) expect(child.parentSpanContext?.spanId).toBe(parent.spanContext().spanId); + expect(trace.getActiveSpan()).toBeUndefined(); +}); + +it.each(["normal", "raw", "stream"])("records real SDK HTTP errors using the %s path once", async (mode) => { + const client = clientFor(sdk8, true); + status = 400; + response = { message: "fixture-error" }; + const pending = mode === "stream" ? client.chatStream(requestV2) : client.chat(requestV2); + await expect(mode === "raw" ? pending.withRawResponse() : pending).rejects.toBeInstanceOf(sdk8.CohereError); + expect(exporter.getFinishedSpans()).toHaveLength(1); + expect(exporter.getFinishedSpans()[0].status.code).toBe(SpanStatusCode.ERROR); +}); + +it.each(["content", "hidden", "vectors"])("exports correct shape and %s privacy over actual OTLP/HTTP", async (mode) => { + instrumentation.disable(); + const wireProvider = new NodeTracerProvider({ spanProcessors: [new SimpleSpanProcessor(new OTLPTraceExporter({ url: `${url}/v1/traces` }))] }); + instrumentation = new CohereInstrumentation({ traceConfig: mode === "hidden" ? { hideInputs: true, hideOutputs: true } : { hideEmbeddingVectors: mode === "vectors" } }); + instrumentation.setTracerProvider(wireProvider); + try { + const client = clientFor(sdk8, true); + response = { id: "embedding", texts: ["secret-input"], embeddings: { float: [[0.123456789]] } }; + await client.embed({ texts: ["secret-input"], model: "embed-test", inputType: "search_document", embeddingTypes: ["float"] }); + response = { id: "rerank", results: [{ index: 0, relevance_score: 0.9, document: { text: "secret-output" } }] }; + await client.rerank({ model: "rerank-test", query: "secret-input", documents: ["secret-input"], topN: 1 }); + streamResponse(true); + response = String(response).replaceAll("hel", "secret-output"); + const stream = await client.chatStream({ ...requestV2, messages: [{ role: "user", content: "secret-input" }], tools: [{ type: "function", function: { name: "tool-secret-input", description: "secret-input", parameters: {} } }] }); + for await (const _event of stream) { /* consume */ } + status = 400; + response = { message: "secret-input secret-output" }; + await expect(client.chat(requestV2)).rejects.toBeInstanceOf(Error); + await wireProvider.forceFlush(); + const encoded = JSON.stringify(wire); + const spans = wire.flatMap((batch) => batch.resourceSpans.flatMap((resource: any) => resource.scopeSpans.flatMap((scope: any) => scope.spans))); + expect(spans).toHaveLength(4); + const attrs = (span: any) => Object.fromEntries(span.attributes.map((item: any) => [item.key, Object.values(item.value)[0]])); + expect(attrs(spans[0])[S.FI_SPAN_KIND]).toBe("EMBEDDING"); + expect(attrs(spans[1])[S.FI_SPAN_KIND]).toBe("RERANKER"); + expect(attrs(spans[2])[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); + expect(spans[3].status.code).toBe(SpanStatusCode.ERROR); + if (mode === "hidden") { + expect(encoded).not.toContain("secret-input"); + expect(encoded).not.toContain("secret-output"); + expect(encoded).not.toContain("0.123456789"); + } else { + expect(encoded).toContain("secret-input"); + expect(encoded).toContain("secret-output"); + } + if (mode === "vectors") expect(encoded).not.toContain("0.123456789"); + } finally { + await wireProvider.shutdown(); + } +}); + + +it.each([[false, false], [false, true], [true, false], [true, true]])( + "marks premature EOF as incomplete (v2=%s, empty=%s) without changing SDK iteration", + async (v2, empty) => { + const client = clientFor(sdk8, v2); + const events = streamResponse(v2); + response = empty ? "" : (v2 ? `data: ${JSON.stringify(events[0])}\n\n` : `${JSON.stringify(events[0])}\n`); + const delivered = []; + for await (const event of await client.chatStream(v2 ? requestV2 : request)) delivered.push(event); + expect(delivered).toHaveLength(empty ? 0 : 1); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0].status.code).toBe(SpanStatusCode.ERROR); + expect(spans[0].status.message).toBe("Cohere stream ended without a terminal event"); + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe(empty ? "" : "hel"); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBeUndefined(); + }, +); + +const privacyClients = [ + ["7.0.0 v1", sdk7, false], ["8.3.0 v1", sdk8, false], ["8.3.0 v2", sdk8, true], +] as const; +describe.each(privacyClients)("%s response privacy", (_label, sdk, v2) => { + it.each(["hideInputs", "hideInputMessages", "hideInputText", "hideInputImages"])( + "removes response-side input echoes with %s but retains generated output on the OTLP wire", + async (flag) => { + instrumentation.disable(); + const exporter = new InMemorySpanExporter(); + const wireProvider = new NodeTracerProvider({ spanProcessors: [ + new SimpleSpanProcessor(exporter), + new SimpleSpanProcessor(new OTLPTraceExporter({ url: `${url}/v1/traces` })), + ] }); + instrumentation = new CohereInstrumentation({ traceConfig: { [flag]: true } }); + instrumentation.setTracerProvider(wireProvider); + try { + const client = clientFor(sdk, v2); + response = v2 ? { ...v2Response, message: { role: "assistant", content: [{ type: "text", text: "generated-answer" }], + citations: [{ start: 0, end: 6, text: "generated-answer", sources: [{ type: "document", id: "doc", document: { text: "private-input-marker" } }] }] } } + : { ...v1Response, text: "generated-answer", chat_history: [ + { role: "USER", message: "private-input-marker" }, + { role: "CHATBOT", message: "private-input-marker-previous-answer" }, + { role: "CHATBOT", message: "generated-answer" }, + ] }; + const result = await client.chat(v2 ? { ...requestV2, messages: [{ role: "user", content: "private-input-marker" }] } : { ...request, message: "private-input-marker" }); + expect(v2 ? result.message.content[0].text : result.text).toBe("generated-answer"); + // Sanitizing telemetry must not alter the SDK's response object. + expect(JSON.stringify(result)).toContain("private-input-marker"); + response = { id: "embedding", response_type: "embeddings_floats", texts: ["private-input-marker"], embeddings: v2 ? { float: [[0.123456789]] } : [[0.123456789]] }; + await client.embed({ texts: ["private-input-marker"], model: "embed-test", inputType: "search_document", ...(v2 ? { embeddingTypes: ["float"] } : {}) }); + response = { id: "rerank", results: [{ index: 0, relevance_score: 0.9, document: { text: "private-input-marker" } }] }; + await client.rerank({ model: "rerank-test", query: "private-input-marker", documents: v2 ? ["private-input-marker"] : [{ text: "private-input-marker" }] }); + await wireProvider.forceFlush(); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(3); + expect(JSON.stringify(spans.map(span => span.attributes))).not.toContain("private-input-marker"); + expect(JSON.stringify(wire)).not.toContain("private-input-marker"); + expect(JSON.stringify(wire)).toContain("generated-answer"); + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe("generated-answer"); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); + expect(spans[1].attributes[`${S.EMBEDDING_EMBEDDINGS}.0.${S.EMBEDDING_VECTOR}`]).toEqual([0.123456789]); + expect(spans[2].attributes["cohere.rerank.results.0.relevance_score"]).toBe(0.9); + } finally { + await wireProvider.shutdown(); + } + }, + ); + + it.each(["hideOutputs", "hideOutputMessages", "hideOutputText"])("keeps inputs visible when only %s is set", async (flag) => { + instrumentation.disable(); + instrumentation = new CohereInstrumentation({ traceConfig: { [flag]: true } }); + instrumentation.setTracerProvider(provider); + const client = clientFor(sdk, v2); + response = v2 ? v2Response : v1Response; + await client.chat(v2 ? requestV2 : request); + const attrs = exporter.getFinishedSpans()[0].attributes; + expect(attrs[S.INPUT_VALUE]).toContain("question"); + expect(attrs[S.RAW_OUTPUT]).toBeUndefined(); + expect(attrs[S.OUTPUT_VALUE]).toBeUndefined(); + expect(attrs[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); + }); +}); + +it("hides v1 terminal stream response echoes without losing generated text", async () => { + instrumentation.disable(); + instrumentation = new CohereInstrumentation({ traceConfig: { hideInputs: true } }); + instrumentation.setTracerProvider(provider); + const client = clientFor(sdk8); + contentType = "application/x-ndjson"; + response = JSON.stringify({ event_type: "stream-end", finish_reason: "COMPLETE", response: { + ...v1Response, chat_history: [{ role: "USER", message: "private-input-marker" }], + } }) + "\n"; + for await (const _event of await client.chatStream({ ...request, message: "private-input-marker" })) { /* consume */ } + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(JSON.stringify(spans[0].attributes)).not.toContain("private-input-marker"); + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe("hello"); + expect(spans[0].status.code).toBe(SpanStatusCode.OK); +}); + +// Connector authentication is not user content and must never enter telemetry. +describe.each(["capture", "hideOutputs", "hideInputs"])("8.3.0 connector credentials (%s)", (mode) => { + it.each([false, true])("keeps credentials on the vendor request only (stream=%s)", async (streaming) => { + instrumentation.disable(); + const captured = new InMemorySpanExporter(); + const wireProvider = new NodeTracerProvider({ spanProcessors: [ + new SimpleSpanProcessor(captured), + new SimpleSpanProcessor(new OTLPTraceExporter({ url: `${url}/v1/traces` })), + ] }); + instrumentation = new CohereInstrumentation({ traceConfig: mode === "capture" ? {} : { [mode]: true } }); + instrumentation.setTracerProvider(wireProvider); + try { + const client = clientFor(sdk8); + const connector = Object.freeze({ id: "fixture-connector", userAccessToken: "synthetic-connector-credential", continueOnFailure: true, options: { site: "example.test" } }); + const input = Object.freeze({ ...request, connectors: Object.freeze([connector]) }); + if (streaming) { + streamResponse(false); + for await (const _event of await client.chatStream(input)) { /* consume */ } + } else { + response = v1Response; + await client.chat(input); + } + await wireProvider.forceFlush(); + expect(received).toHaveLength(1); + expect(received[0].body.connectors[0].user_access_token).toBe(connector.userAccessToken); + expect(input.connectors[0]).toBe(connector); + expect(input.connectors[0].userAccessToken).toBe("synthetic-connector-credential"); + const spans = captured.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(JSON.stringify(spans[0].attributes)).not.toContain(connector.userAccessToken); + expect(wire).not.toHaveLength(0); + expect(JSON.stringify(wire)).not.toContain(connector.userAccessToken); + if (mode === "hideInputs") { + expect(spans[0].attributes[S.RAW_INPUT]).toBeUndefined(); + } else { + expect(JSON.parse(spans[0].attributes[S.RAW_INPUT] as string).connectors).toEqual([ + { id: connector.id, continueOnFailure: true, options: { site: "example.test" } }, + ]); + } + } finally { + await wireProvider.shutdown(); + } + }); +}); + +const finishCases = [ + ...[false, true].flatMap(v2 => [false, true].flatMap(streaming => + (v2 ? ["COMPLETE", "STOP_SEQUENCE", "MAX_TOKENS", "TOOL_CALL", "ERROR", "TIMEOUT"] + : ["COMPLETE", "STOP_SEQUENCE", "MAX_TOKENS", "ERROR", "ERROR_TOXIC", "ERROR_LIMIT", "TIMEOUT", "USER_CANCEL"]) + .map(reason => ({ label: `8.3.0 v${v2 ? 2 : 1}`, sdk: sdk8, v2, streaming, reason })))), +]; +describe.each([false, true])("finish reasons with hideOutputs=%s", (hideOutputs) => { + it.each(finishCases)("classifies $label stream=$streaming $reason without changing SDK results", async ({ sdk, v2, streaming, reason }) => { + instrumentation.disable(); + instrumentation = new CohereInstrumentation({ traceConfig: { hideOutputs } }); + instrumentation.setTracerProvider(provider); + const client = clientFor(sdk, v2); + const failed = ["ERROR", "ERROR_TOXIC", "ERROR_LIMIT", "TIMEOUT"].includes(reason); + const cancelled = reason === "USER_CANCEL"; + if (streaming) { + const events = streamResponse(v2); + if (v2) (events[events.length - 1] as any).delta.finish_reason = reason; + else (events[events.length - 1] as any).finish_reason = reason; + response = events.map(event => v2 ? `data: ${JSON.stringify(event)}\n\n` : `${JSON.stringify(event)}\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream(v2 ? requestV2 : request)) delivered.push(event); + expect(delivered).toHaveLength(events.length); + const last = delivered[delivered.length - 1]; + expect(v2 ? last.delta.finishReason : last.finishReason).toBe(reason); + } else { + response = { ...(v2 ? v2Response : v1Response), finish_reason: reason }; + const pending = client.chat(v2 ? requestV2 : request); + const result = await pending; + expect(result.finishReason).toBe(reason); + expect(v2 ? result.message.content[0].text : result.text).toBe("hello"); + expect((await pending.withRawResponse()).data).toBe(result); + } + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + expect(spans[0].status.code).toBe(failed ? SpanStatusCode.ERROR : cancelled ? SpanStatusCode.UNSET : SpanStatusCode.OK); + expect(spans[0].attributes[streaming ? "cohere.stream.cancelled" : "cohere.request.cancelled"]).toBe(cancelled ? true : undefined); + expect(spans[0].events.filter(event => event.name === "exception")).toHaveLength(failed ? 1 : 0); + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe(hideOutputs ? undefined : "hello"); + expect(spans[0].attributes[S.LLM_TOKEN_COUNT_PROMPT]).toBe(5); + }); +}); + +it.each(["complete", "start-only", "cancelled", "eof"])("captures v2 content-start text for %s streams", async (mode) => { + const client = clientFor(sdk8, true); + const events: any[] = [{ type: "content-start", index: 0, delta: { message: { content: { type: "text", text: "initial-" } } } }]; + if (mode === "complete") events.push({ type: "content-delta", index: 0, delta: { message: { content: { text: "final" } } } }); + if (mode !== "eof") events.push({ type: "message-end", delta: { finish_reason: "COMPLETE" } }); + contentType = "text/event-stream"; + response = events.map(event => `data: ${JSON.stringify(event)}\n\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream(requestV2)) { + delivered.push(event); + if (mode === "cancelled") break; + } + expect(delivered[0].delta.message.content.text).toBe("initial-"); + expect(delivered).toHaveLength(mode === "cancelled" ? 1 : events.length); + const spans = exporter.getFinishedSpans(); + expect(spans).toHaveLength(1); + const text = mode === "complete" ? "initial-final" : "initial-"; + expect(spans[0].attributes[S.OUTPUT_VALUE]).toBe(text); + expect(JSON.parse(spans[0].attributes[S.LLM_OUTPUT_MESSAGES] as string)[0].content).toEqual([{ type: "text", text }]); + expect(JSON.parse(spans[0].attributes[S.RAW_OUTPUT] as string).text).toBe(text); + expect(spans[0].status.code).toBe(mode === "eof" ? SpanStatusCode.ERROR : mode === "cancelled" ? SpanStatusCode.UNSET : SpanStatusCode.OK); +}); + + +// Generated reasoning and partial tool arguments are output, not input echoes. +const outputModes = ["capture", "hideInputs", "hideOutputs", "hideOutputMessages", "hideOutputText"]; +describe.each(outputModes)("generated content with %s", (mode) => { + const hidden = mode.startsWith("hideOutput"); + async function captureOutput(action: (spans: InMemorySpanExporter) => Promise) { + instrumentation.disable(); + const captured = new InMemorySpanExporter(); + const wireProvider = new NodeTracerProvider({ spanProcessors: [ + new SimpleSpanProcessor(captured), + new SimpleSpanProcessor(new OTLPTraceExporter({ url: `${url}/v1/traces` })), + ] }); + instrumentation = new CohereInstrumentation({ traceConfig: mode === "capture" ? {} : { [mode]: true } }); + instrumentation.setTracerProvider(wireProvider); + try { + await action(captured); + await wireProvider.forceFlush(); + expect(wire).not.toHaveLength(0); + expect(JSON.stringify(wire).includes("synthetic-generated-")).toBe(!hidden); + } finally { + await wireProvider.shutdown(); + } + } + + it.each(["chat", "complete", "cancelled", "eof", "start-only"])("preserves v2 reasoning blocks (%s)", async (ending) => { + await captureOutput(async (captured) => { + const client = clientFor(sdk8, true); + const thinking = "synthetic-generated-reasoning-"; + const blocks = [ + { type: "thinking", thinking: thinking + (ending === "start-only" ? "" : "continued") }, + { type: "text", text: "visible-answer" }, + { type: "thinking", thinking: "second-thought" }, + ]; + if (ending === "chat") { + response = { ...v2Response, message: { role: "assistant", content: blocks } }; + const result = await client.chat({ ...requestV2, thinking: { type: "enabled" } }); + expect(result.message.content).toEqual(blocks); + } else { + const events: any[] = [ + { type: "content-start", index: 0, delta: { message: { content: { type: "thinking", thinking } } } }, + ]; + if (ending !== "start-only") events.push( + { type: "content-delta", index: 0, delta: { message: { content: { thinking: "continued" } } } }, + { type: "content-end", index: 0 }, + { type: "content-start", index: 1, delta: { message: { content: blocks[1] } } }, + // A delta without a start must also retain its generated content. + { type: "content-delta", index: 2, delta: { message: { content: { thinking: "second-thought" } } } }, + ); + if (ending === "complete") events.push({ type: "message-end", delta: { finish_reason: "COMPLETE" } }); + contentType = "text/event-stream"; + response = events.map(event => `data: ${JSON.stringify(event)}\n\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream({ ...requestV2, thinking: { type: "enabled" } })) { + delivered.push(event); + if (ending === "cancelled" && delivered.length === 2) break; + } + expect(delivered[0].delta.message.content.thinking).toBe(thinking); + expect(delivered).toHaveLength(ending === "cancelled" ? 2 : events.length); + } + const spans = captured.getFinishedSpans(); + expect(spans).toHaveLength(1); + const attrs = spans[0].attributes; + expect(spans[0].status.code).toBe(ending === "cancelled" ? SpanStatusCode.UNSET : ["eof", "start-only"].includes(ending) ? SpanStatusCode.ERROR : SpanStatusCode.OK); + if (hidden) { + expect(attrs[S.LLM_OUTPUT_MESSAGES]).toBeUndefined(); + expect(attrs[S.RAW_OUTPUT]).toBeUndefined(); + expect(JSON.stringify(attrs)).not.toContain(thinking); + } else { + const expectedBlocks = ["cancelled", "start-only"].includes(ending) ? blocks.slice(0, 1) : blocks; + expect(JSON.parse(attrs[S.LLM_OUTPUT_MESSAGES] as string)[0].content).toEqual(expectedBlocks); + expect(JSON.parse(attrs[S.RAW_OUTPUT] as string).message.content).toEqual(expectedBlocks); + expect(attrs[S.OUTPUT_VALUE]).toBe(expectedBlocks.length === 1 ? "" : "visible-answer"); + } + }); + }); + + it.each(["chat", "complete", "cancelled", "eof"])("retains v2 generated tool plans (%s)", async (ending) => { + await captureOutput(async (captured) => { + const client = clientFor(sdk8, true); + const plan = "synthetic-generated-tool-plan-continued"; + if (ending === "chat") { + response = { ...v2Response, message: { role: "assistant", tool_plan: plan, tool_calls: [{ id: "call-1", type: "function", function: { name: "lookup", arguments: "{}" } }] } }; + const result = await client.chat(requestV2); + expect(result.message.toolPlan).toBe(plan); + } else { + const events: any[] = [ + { type: "tool-plan-delta", delta: { message: { tool_plan: "synthetic-generated-tool-plan-" } } }, + { type: "tool-plan-delta", delta: { message: { tool_plan: "continued" } } }, + { type: "tool-call-start", index: 0, delta: { message: { tool_calls: { id: "call-1", type: "function", function: { name: "lookup", arguments: "{}" } } } } }, + ]; + if (ending === "complete") events.push({ type: "message-end", delta: { finish_reason: "TOOL_CALL" } }); + contentType = "text/event-stream"; + response = events.map(event => `data: ${JSON.stringify(event)}\n\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream(requestV2)) { + delivered.push(event); + if (ending === "cancelled" && delivered.length === 3) break; + } + expect(delivered[0].delta.message.toolPlan).toBe("synthetic-generated-tool-plan-"); + } + const spans = captured.getFinishedSpans(); + expect(spans).toHaveLength(1); + const attrs = spans[0].attributes; + expect(spans[0].status.code).toBe(ending === "cancelled" ? SpanStatusCode.UNSET : ending === "eof" ? SpanStatusCode.ERROR : SpanStatusCode.OK); + if (hidden) { + expect(attrs[S.LLM_OUTPUT_MESSAGES]).toBeUndefined(); + expect(attrs[S.RAW_OUTPUT]).toBeUndefined(); + expect(JSON.stringify(attrs)).not.toContain("synthetic-generated-"); + } else { + expect(JSON.parse(attrs[S.RAW_OUTPUT] as string).message.toolPlan).toBe(plan); + expect(JSON.parse(attrs[S.LLM_OUTPUT_MESSAGES] as string)[0].tool_plan).toBe(plan); + expect(attrs[S.OUTPUT_VALUE]).not.toContain(plan); + } + }); + }); + + it.each(["cancelled", "eof", "terminal", "generation"])("retains v1 generated tool plans (%s)", async (ending) => { + await captureOutput(async (captured) => { + const client = clientFor(sdk8); + const partial = "synthetic-generated-tool-plan-continued"; + const events: any[] = [ + { event_type: "tool-calls-chunk", tool_call_delta: { index: 0, text: "synthetic-generated-tool-plan-", name: "lookup", parameters: "{}" } }, + { event_type: "tool-calls-chunk", tool_call_delta: { index: 0, text: "continued" } }, + ]; + if (ending === "generation") events.push({ event_type: "tool-calls-generation", text: partial, tool_calls: [{ name: "lookup", parameters: {} }] }); + if (ending === "terminal") events.push({ event_type: "stream-end", finish_reason: "COMPLETE", response: { text: partial, tool_calls: [{ name: "lookup", parameters: {} }] } }); + contentType = "application/x-ndjson"; + response = events.map(event => `${JSON.stringify(event)}\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream(request)) { + delivered.push(event); + if (ending === "cancelled" && delivered.length === 2) break; + } + expect(delivered[0].toolCallDelta.text).toBe("synthetic-generated-tool-plan-"); + const spans = captured.getFinishedSpans(); + expect(spans).toHaveLength(1); + const attrs = spans[0].attributes; + expect(spans[0].status.code).toBe(ending === "terminal" ? SpanStatusCode.OK : ending === "cancelled" ? SpanStatusCode.UNSET : SpanStatusCode.ERROR); + if (hidden) { + expect(attrs[S.LLM_OUTPUT_MESSAGES]).toBeUndefined(); + expect(attrs[S.RAW_OUTPUT]).toBeUndefined(); + expect(JSON.stringify(attrs)).not.toContain("synthetic-generated-"); + } else { + expect(JSON.parse(attrs[S.RAW_OUTPUT] as string).message.toolPlan).toBe(partial); + expect(JSON.parse(attrs[S.LLM_OUTPUT_MESSAGES] as string)[0].tool_plan).toBe(partial); + } + }); + }); + + it.each(["cancelled", "eof", "terminal", "generation"])("preserves v1 tool chunks (%s)", async (ending) => { + await captureOutput(async (captured) => { + const client = clientFor(sdk8); + // Interleave calls and leave one argument incomplete: telemetry must not + // require JSON parsing, mix call indices, or quote an already-raw string. + const events: any[] = [ + { event_type: "tool-calls-chunk", tool_call_delta: { index: 1, name: "second", parameters: '{"synthetic-generated-second":' } }, + { event_type: "tool-calls-chunk", tool_call_delta: { index: 0, name: "lookup", parameters: '{"synthetic-generated-key":' } }, + { event_type: "tool-calls-chunk", tool_call_delta: { index: 0, parameters: '"value"}' } }, + ]; + const finalCalls = [{ name: "authoritative", parameters: { "synthetic-generated-final": true } }]; + if (ending === "terminal") events.push({ event_type: "stream-end", finish_reason: "COMPLETE", response: { text: "", tool_calls: finalCalls } }); + if (ending === "generation") events.push({ event_type: "tool-calls-generation", tool_calls: finalCalls }); + contentType = "application/x-ndjson"; + response = events.map(event => `${JSON.stringify(event)}\n`).join(""); + const delivered = []; + for await (const event of await client.chatStream(request)) { + delivered.push(event); + if (ending === "cancelled" && delivered.length === 3) break; + } + expect(delivered[1].toolCallDelta.name).toBe("lookup"); + expect(delivered[2].toolCallDelta.parameters).toBe('"value"}'); + expect(delivered).toHaveLength(events.length); + const spans = captured.getFinishedSpans(); + expect(spans).toHaveLength(1); + const attrs = spans[0].attributes; + expect(spans[0].status.code).toBe(ending === "terminal" ? SpanStatusCode.OK : ending === "cancelled" ? SpanStatusCode.UNSET : SpanStatusCode.ERROR); + if (hidden) { + expect(attrs[S.LLM_OUTPUT_MESSAGES]).toBeUndefined(); + expect(attrs[S.RAW_OUTPUT]).toBeUndefined(); + expect(JSON.stringify(attrs)).not.toContain("synthetic-generated-"); + } else { + const calls = ["terminal", "generation"].includes(ending) ? finalCalls : [ + { name: "lookup", parameters: '{"synthetic-generated-key":"value"}' }, + { name: "second", parameters: '{"synthetic-generated-second":' }, + ]; + expect(JSON.parse(attrs[S.RAW_OUTPUT] as string).toolCalls).toEqual(calls); + expect(JSON.parse(attrs[S.LLM_OUTPUT_MESSAGES] as string)[0].tool_calls).toEqual(calls.map(call => ({ + function: { name: call.name, arguments: typeof call.parameters === "string" ? call.parameters : JSON.stringify(call.parameters) }, + }))); + } + }); + }); +}); diff --git a/typescript/packages/traceai_cohere/src/attributes.ts b/typescript/packages/traceai_cohere/src/attributes.ts new file mode 100644 index 000000000..cb69e5a17 --- /dev/null +++ b/typescript/packages/traceai_cohere/src/attributes.ts @@ -0,0 +1,177 @@ +import { Attributes } from "@opentelemetry/api"; +import { safelyJSONStringify, TraceConfig } from "@traceai/fi-core"; +import { FISpanKind, LLMProvider, MimeType, SemanticConventions as S } from "@traceai/fi-semantic-conventions"; + +export type Operation = "chat" | "chatStream" | "embed" | "rerank"; +export type Data = Record; +export const object = (value: unknown): Data => value !== null && typeof value === "object" ? value as Data : {}; +export const list = (value: unknown): unknown[] => Array.isArray(value) ? value : []; +const json = (value: unknown) => safelyJSONStringify(value) ?? ""; +const field = (data: Data, camel: string, snake: string) => data[camel] ?? data[snake]; +export const hideInput = (config: TraceConfig) => config.hideInputs || config.hideInputMessages || config.hideInputText || config.hideInputImages; +export const hideOutput = (config: TraceConfig) => config.hideOutputs || config.hideOutputMessages || config.hideOutputText; + +/** Connector authentication is never telemetry content, even with input capture + * enabled. Project only the telemetry copy; the SDK still needs these tokens. + */ +function withoutConnectorCredentials(request: Data): Data { + if (!Array.isArray(request.connectors)) return request; + return { + ...request, + connectors: request.connectors.map((connector) => { + const { userAccessToken: _camel, user_access_token: _snake, ...content } = object(connector); + return content; + }), + }; +} + +export function inputAttributes(operation: Operation, request: Data, config: TraceConfig): Attributes { + const embedding = operation === "embed"; + const rerank = operation === "rerank"; + const attrs: Attributes = { + [S.FI_SPAN_KIND]: embedding ? FISpanKind.EMBEDDING : rerank ? FISpanKind.RERANKER : FISpanKind.LLM, + [S.LLM_PROVIDER]: LLMProvider.COHERE, + [S.GEN_AI_OPERATION_NAME]: embedding ? "embeddings" : rerank ? "rerank" : "chat", + }; + if (typeof request.model === "string") { + attrs[embedding ? S.EMBEDDING_MODEL_NAME : S.LLM_MODEL_NAME] = request.model; + if (rerank) attrs[S.RERANKER_MODEL_NAME] = request.model; + } + if (embedding) { + if (!hideInput(config)) { + attrs[S.INPUT_VALUE] = json(request.texts ?? request.inputs ?? request.images); + attrs[S.INPUT_MIME_TYPE] = MimeType.JSON; + list(request.texts).forEach((text, index) => { + if (typeof text === "string") attrs[`${S.EMBEDDING_EMBEDDINGS}.${index}.${S.EMBEDDING_TEXT}`] = text; + }); + } + } else if (rerank) { + attrs["cohere.rerank.documents_count"] = list(request.documents).length; + const topN = field(request, "topN", "top_n"); + if (typeof topN === "number") { + attrs[S.RERANKER_TOP_K] = topN; + attrs["cohere.rerank.top_n"] = topN; + } + if (!hideInput(config)) { + attrs[S.INPUT_VALUE] = typeof request.query === "string" ? request.query : ""; + attrs[S.INPUT_MIME_TYPE] = MimeType.TEXT; + attrs[S.RERANKER_QUERY] = attrs[S.INPUT_VALUE]; + attrs["cohere.rerank.query"] = attrs[S.INPUT_VALUE]; + attrs[S.RERANKER_INPUT_DOCUMENTS] = json(request.documents); + } + } else { + // Only generation settings belong here. Chat history, tools, documents and + // arbitrary SDK fields must not bypass the content privacy controls. + const params: Data = {}; + for (const key of ["temperature", "maxTokens", "max_tokens", "p", "k", "seed", "frequencyPenalty", "presencePenalty"]) { + if (typeof request[key] === "number") params[key] = request[key]; + } + attrs[S.LLM_INVOCATION_PARAMETERS] = json(params); + if (!hideInput(config)) { + const messages = Array.isArray(request.messages) ? request.messages : [ + ...(request.preamble ? [{ role: "system", content: request.preamble }] : []), + ...list(field(request, "chatHistory", "chat_history")).map((entry) => { + const item = object(entry); + return { role: item.role === "CHATBOT" ? "assistant" : String(item.role).toLowerCase(), content: item.message }; + }), + { role: "user", content: request.message }, + ]; + attrs[S.LLM_INPUT_MESSAGES] = json(messages); + attrs[S.INPUT_VALUE] = typeof request.message === "string" ? request.message : json(messages); + attrs[S.INPUT_MIME_TYPE] = typeof request.message === "string" ? MimeType.TEXT : MimeType.JSON; + if (request.tools) attrs[S.LLM_TOOLS] = json(request.tools); + } + } + if (!hideInput(config)) attrs[S.RAW_INPUT] = json(withoutConnectorCredentials(request)); + return attrs; +} + +/** Keep generated fields only when inputs are hidden. Responses can echo full + * chat history, embedding texts, reranked documents and citation sources. + * Project rather than mutate the SDK response or filter history by role: even + * earlier assistant messages are input to the current generation. + */ +function withoutInputEchoes(operation: Operation, response: Data): Data { + if (operation === "embed") return { embeddings: response.embeddings }; + if (operation === "rerank") return { + results: list(response.results).map((item) => { + const result = object(item); + return { index: result.index, relevanceScore: field(result, "relevanceScore", "relevance_score") }; + }), + }; + const message = object(response.message); + return { + text: response.text, + toolCalls: field(response, "toolCalls", "tool_calls"), + message: { + role: message.role, + content: Array.isArray(message.content) ? message.content.map((part) => { + const content = object(part); + return { type: content.type, text: content.text, thinking: content.thinking }; + }) : message.content, + toolCalls: field(message, "toolCalls", "tool_calls"), + toolPlan: field(message, "toolPlan", "tool_plan"), + }, + }; +} + +export function outputAttributes(operation: Operation, response: Data, config: TraceConfig): Attributes { + const attrs: Attributes = usageAttributes(response); + if (hideInput(config)) response = withoutInputEchoes(operation, response); + if (operation === "embed") { + // raw.output would otherwise bypass hideEmbeddingVectors. + if (!hideOutput(config) && !config.hideEmbeddingVectors) { + const vectors = Array.isArray(response.embeddings) ? response.embeddings : object(response.embeddings).float; + list(vectors).forEach((vector, index) => { + if (Array.isArray(vector) && vector.every((n) => typeof n === "number")) { + attrs[`${S.EMBEDDING_EMBEDDINGS}.${index}.${S.EMBEDDING_VECTOR}`] = vector; + } + }); + attrs[S.RAW_OUTPUT] = json(response); + } + } else if (operation === "rerank") { + const results = list(response.results); + attrs["cohere.rerank.results_count"] = results.length; + results.forEach((item, index) => { + const result = object(item); + const score = field(result, "relevanceScore", "relevance_score"); + if (typeof result.index === "number") attrs[`cohere.rerank.results.${index}.index`] = result.index; + if (typeof score === "number") attrs[`cohere.rerank.results.${index}.relevance_score`] = score; + const text = object(result.document).text; + if (!hideOutput(config) && typeof text === "string") attrs[`cohere.rerank.results.${index}.document`] = text; + }); + if (!hideOutput(config)) { + attrs[S.RERANKER_OUTPUT_DOCUMENTS] = json(results); + attrs[S.RAW_OUTPUT] = json(response); + } + } else if (!hideOutput(config)) { + const message = object(response.message); + const text = typeof response.text === "string" ? response.text : list(message.content).map((part) => object(part).text ?? "").join(""); + const calls = field(message, "toolCalls", "tool_calls") ?? field(response, "toolCalls", "tool_calls"); + const normalized: Data = { role: "assistant", content: message.content ?? text }; + const toolPlan = field(message, "toolPlan", "tool_plan"); + if (typeof toolPlan === "string") normalized.tool_plan = toolPlan; + if (calls) normalized.tool_calls = list(calls).map((item) => { + const call = object(item); + return call.function ? call : { function: { name: call.name, arguments: typeof call.parameters === "string" ? call.parameters : json(call.parameters) } }; + }); + attrs[S.OUTPUT_VALUE] = text; + attrs[S.OUTPUT_MIME_TYPE] = MimeType.TEXT; + attrs[S.LLM_OUTPUT_MESSAGES] = json([normalized]); + attrs[S.RAW_OUTPUT] = json(response); + } + return attrs; +} + +function usageAttributes(response: Data): Attributes { + const usage = object(response.usage ?? response.meta); + // Actual tokens, when available, are not the same as discounted billed units. + const tokens = object(usage.tokens ?? field(usage, "billedUnits", "billed_units")); + const input = field(tokens, "inputTokens", "input_tokens"); + const output = field(tokens, "outputTokens", "output_tokens"); + const attrs: Attributes = {}; + if (typeof input === "number") attrs[S.LLM_TOKEN_COUNT_PROMPT] = input; + if (typeof output === "number") attrs[S.LLM_TOKEN_COUNT_COMPLETION] = output; + if (typeof input === "number" || typeof output === "number") attrs[S.LLM_TOKEN_COUNT_TOTAL] = (typeof input === "number" ? input : 0) + (typeof output === "number" ? output : 0); + return attrs; +} diff --git a/typescript/packages/traceai_cohere/src/instrumentation.ts b/typescript/packages/traceai_cohere/src/instrumentation.ts index bcfc1de73..573ed41cd 100644 --- a/typescript/packages/traceai_cohere/src/instrumentation.ts +++ b/typescript/packages/traceai_cohere/src/instrumentation.ts @@ -1,592 +1,197 @@ -import { - InstrumentationBase, - InstrumentationConfig, - InstrumentationNodeModuleDefinition, - safeExecuteInTheMiddle, -} from "@opentelemetry/instrumentation"; -import { - diag, - context, - trace, - SpanKind, - Attributes, - SpanStatusCode, - Span, -} from "@opentelemetry/api"; -import { VERSION } from "./version"; -import { - SemanticConventions, - FISpanKind, - MimeType, - LLMSystem, - LLMProvider, -} from "@traceai/fi-semantic-conventions"; +import { InstrumentationBase, InstrumentationConfig, InstrumentationNodeModuleDefinition } from "@opentelemetry/instrumentation"; +import { context, createContextKey, diag, trace, SpanKind, SpanStatusCode, Span } from "@opentelemetry/api"; import { isTracingSuppressed } from "@opentelemetry/core"; -import { - FITracer, - safelyJSONStringify, - TraceConfigOptions, -} from "@traceai/fi-core"; -import { isString } from "./typeUtils"; - -const MODULE_NAME = "cohere-ai"; - -/** - * Flag to check if the Cohere module has been patched - */ -let _isFIPatched = false; - -/** - * Function to check if instrumentation is enabled / disabled - */ -export function isPatched() { - return _isFIPatched; -} - -/** - * Resolves the execution context for the current span - */ -function getExecContext(span: Span) { - const activeContext = context.active(); - const suppressTracing = isTracingSuppressed(activeContext); - const execContext = suppressTracing - ? trace.setSpan(context.active(), span) - : activeContext; - if (suppressTracing) { - trace.deleteSpan(activeContext); - } - return execContext; -} - -// Type definitions for Cohere SDK -interface CohereChatRequest { - model?: string; - message: string; - preamble?: string; - chat_history?: Array<{ - role: string; - message: string; - }>; - temperature?: number; - max_tokens?: number; - p?: number; - k?: number; - tools?: Array<{ - name: string; - description?: string; - parameterDefinitions?: Record; - }>; - [key: string]: unknown; -} - -interface CohereChatResponse { - text: string; - generation_id?: string; - chat_history?: Array<{ - role: string; - message: string; - }>; - tool_calls?: Array<{ - name: string; - parameters: Record; - }>; - meta?: { - api_version?: { version: string }; - billed_units?: { - input_tokens?: number; - output_tokens?: number; - }; - tokens?: { - input_tokens?: number; - output_tokens?: number; - }; - }; -} - -interface CohereEmbedRequest { - texts: string[]; - model?: string; - input_type?: string; - embedding_types?: string[]; - truncate?: string; -} - -interface CohereEmbedResponse { - id: string; - embeddings: number[][] | { float?: number[][] }; - texts: string[]; - meta?: { - api_version?: { version: string }; - billed_units?: { - input_tokens?: number; - }; - }; -} +import { FITracer, generateTraceConfig, TraceConfigOptions } from "@traceai/fi-core"; +import { VERSION } from "./version"; +import { Data, Operation, object, inputAttributes, outputAttributes, hideOutput } from "./attributes"; +import { wrapStream } from "./stream"; -interface CohereRerankRequest { - query: string; - documents: string[] | Array<{ text: string }>; - model?: string; - top_n?: number; - return_documents?: boolean; -} +const OPERATIONS: Operation[] = ["chat", "chatStream", "embed", "rerank"]; +const IN_COHERE_CALL = createContextKey("traceai.cohere.call"); +type Method = (this: unknown, ...args: unknown[]) => unknown; +type Patch = { target: Data; method: Operation; wrapped: Method; hadOwn: boolean }; +const owners = new WeakMap>(); +let patchCount = 0; -interface CohereRerankResponse { - id: string; - results: Array<{ - index: number; - relevance_score: number; - document?: { text: string }; - }>; - meta?: { - api_version?: { version: string }; - billed_units?: { - search_units?: number; - }; - }; -} +/** True when at least one SDK method is instrumented, across module copies. */ +export function isPatched(): boolean { return patchCount > 0; } -/** - * @param instrumentationConfig The config for the instrumentation - * @param traceConfig The FI trace configuration for masking sensitive information - */ +/** Cohere v1/v2 instrumentation. Install before constructing SDK clients. */ export class CohereInstrumentation extends InstrumentationBase { - private fiTracer!: FITracer; private _traceConfig?: TraceConfigOptions; + // InstrumentationBase may call enable/patch during super(), so initialize lazily. + private patches?: Patch[]; + private manualModules?: Set; public readonly instrumentationName = "@traceai/cohere"; - constructor({ - instrumentationConfig, - traceConfig, - }: { + constructor({ instrumentationConfig, traceConfig }: { instrumentationConfig?: InstrumentationConfig; traceConfig?: TraceConfigOptions; } = {}) { - super( - "@traceai/cohere", - VERSION, - Object.assign({}, instrumentationConfig), - ); + super("@traceai/cohere", VERSION, { ...instrumentationConfig }); this._traceConfig = traceConfig; } + protected init() { + return new InstrumentationNodeModuleDefinition("cohere-ai", [">=7.0.0 <9"], this.patch.bind(this), this.unpatch.bind(this)); + } + public override enable(): void { super.enable(); - this.fiTracer = new FITracer({ tracer: this.tracer, traceConfig: this._traceConfig }); + for (const module of this.manualModules ?? []) this.patch(module); } - protected init() { - const module = new InstrumentationNodeModuleDefinition( - "cohere-ai", - [">=7.0.0"], - this.patch.bind(this), - this.unpatch.bind(this), - ); - return module; + public override disable(): void { + super.disable(); + this.restorePatches(); } - /** - * Manually instruments the Cohere module - */ - manuallyInstrument(module: Record) { - this.patch(module); + manuallyInstrument(module: Data): void { + (this.manualModules ??= new Set()).add(module); + if (this.isEnabled()) this.patch(module); } - /** - * Patches the Cohere module - */ - private patch( - module: Record & { fiPatched?: boolean }, - moduleVersion?: string, - ) { - if (module?.fiPatched || _isFIPatched) { - return module; - } - - const instrumentation: CohereInstrumentation = this; - - // Ensure fiTracer uses the most up-to-date tracer - if (!(instrumentation.fiTracer as unknown as Record)?.tracer || - (instrumentation.fiTracer as unknown as Record)?.tracer !== instrumentation.tracer) { - instrumentation.fiTracer = new FITracer({ - tracer: instrumentation.tracer, - traceConfig: instrumentation._traceConfig - }); - } - - // Patch CohereClient class - const CohereClient = module.CohereClient as { - prototype?: Record; - } | undefined; - - if (CohereClient?.prototype) { - // Patch chat method - if (typeof CohereClient.prototype.chat === 'function') { - this._wrap( - CohereClient.prototype, - "chat", - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (original: any): any => { - return function patchedChat( - this: unknown, - ...args: [CohereChatRequest] - ) { - const request = args[0]; - const modelName = request.model || 'command'; - const { message, chat_history, preamble, ...invocationParameters } = request; - - const span = instrumentation.fiTracer.startSpan( - `Cohere Chat`, - { - kind: SpanKind.INTERNAL, - attributes: { - [SemanticConventions.FI_SPAN_KIND]: FISpanKind.LLM, - [SemanticConventions.LLM_MODEL_NAME]: modelName, - [SemanticConventions.INPUT_VALUE]: request.message, - [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT, - [SemanticConventions.LLM_INVOCATION_PARAMETERS]: - JSON.stringify(invocationParameters), - [SemanticConventions.LLM_PROVIDER]: LLMProvider.COHERE, - [SemanticConventions.GEN_AI_OPERATION_NAME]: "chat", - ...getChatInputMessagesAttributes(request), - ...getChatToolsAttributes(request), - [SemanticConventions.RAW_INPUT]: safelyJSONStringify(request) ?? "", - }, - }, - ); - - const execContext = getExecContext(span); - const execPromise = safeExecuteInTheMiddle( - () => { - return context.with(trace.setSpan(execContext, span), () => { - return original.apply(this, args); - }); - }, - (error) => { - if (error) { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - } - }, - ); - - const wrappedPromise = (execPromise as Promise).then((result) => { - span.setAttributes({ - [SemanticConventions.OUTPUT_VALUE]: result.text, - [SemanticConventions.OUTPUT_MIME_TYPE]: MimeType.TEXT, - ...getChatOutputMessagesAttributes(result), - ...getChatUsageAttributes(result), - [SemanticConventions.RAW_OUTPUT]: safelyJSONStringify(result) ?? "", - }); - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return result; - }).catch((error: Error) => { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - throw error; - }); - - return context.bind(execContext, wrappedPromise); - }; - }, - ); - } - - // Patch embed method - if (typeof CohereClient.prototype.embed === 'function') { - this._wrap( - CohereClient.prototype, - "embed", - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (original: any): any => { - return function patchedEmbed( - this: unknown, - ...args: [CohereEmbedRequest] - ) { - const request = args[0]; - const modelName = request.model || 'embed-english-v3.0'; - - const span = instrumentation.fiTracer.startSpan(`Cohere Embed`, { - kind: SpanKind.INTERNAL, - attributes: { - [SemanticConventions.FI_SPAN_KIND]: FISpanKind.EMBEDDING, - [SemanticConventions.EMBEDDING_MODEL_NAME]: modelName, - [SemanticConventions.INPUT_VALUE]: JSON.stringify(request.texts), - [SemanticConventions.INPUT_MIME_TYPE]: MimeType.JSON, - [SemanticConventions.LLM_PROVIDER]: LLMProvider.COHERE, - [SemanticConventions.GEN_AI_OPERATION_NAME]: "embeddings", - ...getEmbedTextAttributes(request), - [SemanticConventions.RAW_INPUT]: safelyJSONStringify(request) ?? "", - }, - }); - - const execContext = getExecContext(span); - const execPromise = safeExecuteInTheMiddle( - () => { - return context.with(trace.setSpan(execContext, span), () => { - return original.apply(this, args); - }); - }, - (error) => { - if (error) { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - } - }, - ); - - const wrappedPromise = (execPromise as Promise).then((result) => { - span.setAttributes({ - ...getEmbedVectorAttributes(result), - [SemanticConventions.RAW_OUTPUT]: safelyJSONStringify(result) ?? "", - }); - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return result; - }).catch((error: Error) => { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - throw error; - }); - - return context.bind(execContext, wrappedPromise); - }; - }, - ); - } - - // Patch rerank method - if (typeof CohereClient.prototype.rerank === 'function') { - this._wrap( - CohereClient.prototype, - "rerank", - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (original: any): any => { - return function patchedRerank( - this: unknown, - ...args: [CohereRerankRequest] - ) { - const request = args[0]; - const modelName = request.model || 'rerank-english-v3.0'; - - const span = instrumentation.fiTracer.startSpan(`Cohere Rerank`, { - kind: SpanKind.INTERNAL, - attributes: { - [SemanticConventions.FI_SPAN_KIND]: FISpanKind.LLM, - [SemanticConventions.LLM_MODEL_NAME]: modelName, - [SemanticConventions.INPUT_VALUE]: request.query, - [SemanticConventions.INPUT_MIME_TYPE]: MimeType.TEXT, - [SemanticConventions.LLM_PROVIDER]: LLMProvider.COHERE, - [SemanticConventions.GEN_AI_OPERATION_NAME]: "rerank", - "cohere.rerank.query": request.query, - "cohere.rerank.documents_count": request.documents.length, - "cohere.rerank.top_n": request.top_n, - [SemanticConventions.RAW_INPUT]: safelyJSONStringify(request) ?? "", - }, - }); - - const execContext = getExecContext(span); - const execPromise = safeExecuteInTheMiddle( - () => { - return context.with(trace.setSpan(execContext, span), () => { - return original.apply(this, args); - }); - }, - (error) => { - if (error) { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - } - }, - ); - - const wrappedPromise = (execPromise as Promise).then((result) => { - span.setAttributes({ - "cohere.rerank.results_count": result.results.length, - ...getRerankResultsAttributes(result), - [SemanticConventions.RAW_OUTPUT]: safelyJSONStringify(result) ?? "", - }); - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return result; - }).catch((error: Error) => { - span.recordException(error); - span.setStatus({ code: SpanStatusCode.ERROR, message: error.message }); - span.end(); - throw error; - }); - - return context.bind(execContext, wrappedPromise); - }; - }, - ); + private patch(module: Data): Data { + const Client = module.CohereClient as { prototype?: Data; new(options: { token: string }): { v2?: object } } | undefined; + const targets = [Client?.prototype]; + // V2Client is not exported at the module root. The public client.v2 getter + // exposes the same generated class used by CohereClientV2's bound methods. + // Construction performs no I/O; an explicit inert token avoids env lookup. + if (typeof Client === "function" && Client.prototype && "v2" in Client.prototype) { + try { + const v2 = new Client({ token: "traceai-instrumentation-placeholder" }).v2; + if (v2) targets.push(Object.getPrototypeOf(v2) as Data); + } catch { + // Telemetry must not prevent module loading; v1 remains usable. + diag.warn("Cohere instrumentation could not discover the v2 client"); } } - - _isFIPatched = true; - try { - module.fiPatched = true; - } catch (e) { - diag.warn(`Failed to set fiPatched flag on module '${MODULE_NAME}'. Error: ${e}`); + for (const target of targets) { + if (!target) continue; + for (const method of OPERATIONS) { + if (typeof target[method] !== "function" || owners.get(target)?.has(method)) continue; + const hadOwn = Object.prototype.hasOwnProperty.call(target, method); + this._wrap(target, method, (original) => this.wrapMethod(method, original as Method)); + const wrapped = target[method] as Method; + (this.patches ??= []).push({ target, method, wrapped, hadOwn }); + if (!owners.has(target)) owners.set(target, new Map()); + owners.get(target)!.set(method, this); + patchCount++; + } } - return module; } - /** - * Un-patches the Cohere module - */ - private unpatch( - moduleExports: Record & { fiPatched?: boolean }, - moduleVersion?: string, - ) { - const CohereClient = moduleExports.CohereClient as { - prototype?: Record; - } | undefined; + private unpatch(): void { this.restorePatches(); } - if (CohereClient?.prototype) { - if (CohereClient.prototype.chat) { - this._unwrap(CohereClient.prototype, "chat"); - } - if (CohereClient.prototype.embed) { - this._unwrap(CohereClient.prototype, "embed"); - } - if (CohereClient.prototype.rerank) { - this._unwrap(CohereClient.prototype, "rerank"); + private restorePatches(): void { + for (const { target, method, wrapped, hadOwn } of this.patches ?? []) { + if (target[method] === wrapped) { + this._unwrap(target, method); + if (!hadOwn) delete target[method]; } + owners.get(target)?.delete(method); + patchCount--; } - - _isFIPatched = false; - try { - moduleExports.fiPatched = false; - } catch (e) { - diag.warn(`Failed to unset fiPatched flag on module '${MODULE_NAME}'. Error: ${e}`); - } + this.patches = []; } -} - -/** - * Converts the chat request to LLM input messages - */ -function getChatInputMessagesAttributes(request: CohereChatRequest): Attributes { - const messages: Record[] = []; - - if (request.preamble) { - messages.push({ role: 'system', content: request.preamble }); - } - - if (request.chat_history) { - request.chat_history.forEach((msg) => { - messages.push({ role: msg.role, content: msg.message }); - }); - } - - messages.push({ role: 'user', content: request.message }); - - return { [SemanticConventions.LLM_INPUT_MESSAGES]: safelyJSONStringify(messages) ?? "[]" }; -} - -/** - * Gets chat tools attributes - */ -function getChatToolsAttributes(request: CohereChatRequest): Attributes { - if (!request.tools) { - return {}; - } - return { [SemanticConventions.LLM_TOOLS]: safelyJSONStringify(request.tools) ?? "[]" }; -} -/** - * Gets chat output messages attributes - */ -function getChatOutputMessagesAttributes(response: CohereChatResponse): Attributes { - const msg: Record = { role: 'assistant', content: response.text }; - - if (response.tool_calls) { - msg.tool_calls = response.tool_calls.map((tc) => ({ - function: { name: tc.name, arguments: JSON.stringify(tc.parameters) }, - })); - } - - return { [SemanticConventions.LLM_OUTPUT_MESSAGES]: safelyJSONStringify([msg]) ?? "[]" }; -} - -/** - * Gets chat usage attributes - */ -function getChatUsageAttributes(response: CohereChatResponse): Attributes { - const meta = response.meta; - if (!meta) { - return {}; - } - - const billedUnits = meta.billed_units; - const tokens = meta.tokens; - - if (billedUnits) { - return { - [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: billedUnits.input_tokens, - [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: billedUnits.output_tokens, - [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: - (billedUnits.input_tokens || 0) + (billedUnits.output_tokens || 0), - }; - } - - if (tokens) { - return { - [SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: tokens.input_tokens, - [SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: tokens.output_tokens, - [SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: - (tokens.input_tokens || 0) + (tokens.output_tokens || 0), + private wrapMethod(operation: Operation, original: Method): Method { + const instrumentation = this; + return function (this: unknown, ...args: unknown[]) { + const parent = context.active(); + // Bound v2 methods can outlive unpatch. Respect disable even for those. + if (!instrumentation.isEnabled() || isTracingSuppressed(parent) || parent.getValue(IN_COHERE_CALL)) { + return original.apply(this, args); + } + const config = generateTraceConfig(instrumentation._traceConfig); + let span: Span; + try { + const tracer = new FITracer({ tracer: instrumentation.tracer, traceConfig: config }); + const name = operation === "chatStream" ? "Chat Stream" : operation[0].toUpperCase() + operation.slice(1); + span = tracer.startSpan(`Cohere ${name}`, { kind: SpanKind.INTERNAL, attributes: inputAttributes(operation, object(args[0]), config) }, parent); + } catch { + diag.warn("Cohere instrumentation could not start span"); + return original.apply(this, args); + } + const execContext = trace.setSpan(parent, span).setValue(IN_COHERE_CALL, true); + let ended = false; + const finish = (response: Data, error?: unknown, cancelled = false) => { + if (ended) return; + ended = true; + try { + // Resolved HTTP-200 chat responses can still signal failure or user + // cancellation. Use the same classification for both chat paths, + // independent of content capture and without rejecting SDK results. + if ((operation === "chat" || operation === "chatStream") && error === undefined && !cancelled) { + const reason = response.finishReason ?? response.finish_reason; + switch (reason) { + case "ERROR": + case "ERROR_TOXIC": + case "ERROR_LIMIT": + case "TIMEOUT": + error = new Error(`Cohere generation failed (${reason})`); + break; + case "USER_CANCEL": + cancelled = true; + break; + } + } + span.setAttributes(outputAttributes(operation, response, config)); + if (error !== undefined) { + span.recordException(error instanceof Error ? error : new Error(String(error))); + span.setStatus({ code: SpanStatusCode.ERROR, message: error instanceof Error ? error.message : String(error) }); + } else if (cancelled) { + span.setAttribute(operation === "chatStream" ? "cohere.stream.cancelled" : "cohere.request.cancelled", true); + } else { + span.setStatus({ code: SpanStatusCode.OK }); + } + } catch { + diag.warn("Cohere instrumentation could not record response"); + } finally { + span.end(); + } + }; + const streams = new WeakMap(); + const resolved = (result: unknown): unknown => { + if (operation === "chatStream" && result && typeof (result as AsyncIterable)[Symbol.asyncIterator] === "function") { + if (!streams.has(result as object)) streams.set(result as object, wrapStream(result as object, execContext, !hideOutput(config), finish)); + return streams.get(result as object); + } + finish(object(result)); + return result; + }; + const rejected = (error: unknown): never => { + finish({}, error ?? new Error("Cohere request failed")); + throw error; + }; + let pending: unknown; + try { + pending = context.with(execContext, () => original.apply(this, args)); + } catch (error) { return rejected(error); } + return observePromise(pending, resolved, rejected); }; } - - return {}; } -/** - * Gets embed text attributes - */ -function getEmbedTextAttributes(request: CohereEmbedRequest): Attributes { - return request.texts.reduce((acc, text, index) => { - acc[`${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.${SemanticConventions.EMBEDDING_TEXT}`] = text; - return acc; - }, {} as Attributes); -} - -/** - * Gets embed vector attributes - */ -function getEmbedVectorAttributes(response: CohereEmbedResponse): Attributes { - const embeddings = Array.isArray(response.embeddings) - ? response.embeddings - : (response.embeddings as { float?: number[][] })?.float || []; - - return embeddings.reduce((acc, embedding, index) => { - acc[`${SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.${SemanticConventions.EMBEDDING_VECTOR}`] = embedding; - return acc; - }, {} as Attributes); -} - -/** - * Gets rerank results attributes - */ -function getRerankResultsAttributes(response: CohereRerankResponse): Attributes { - const attributes: Attributes = {}; - - response.results.forEach((result, index) => { - attributes[`cohere.rerank.results.${index}.index`] = result.index; - attributes[`cohere.rerank.results.${index}.relevance_score`] = result.relevance_score; - if (result.document) { - attributes[`cohere.rerank.results.${index}.document`] = result.document.text; - } +/** Keep SDK promise helpers (notably HttpResponsePromise.withRawResponse). */ +function observePromise(pending: unknown, resolved: (value: unknown) => unknown, rejected: (error: unknown) => never): unknown { + const promise = pending as Promise & { withRawResponse?: () => Promise<{ data: unknown; rawResponse: unknown }> }; + const observed = promise.then(resolved, rejected); + if (typeof promise.withRawResponse !== "function") return observed; + // A caller may exclusively use withRawResponse. Do not leave an unobserved + // rejected branch merely because telemetry also observes the normal promise. + void observed.catch(() => {}); + return new Proxy(promise, { + get(target, key) { + if (key === "then") return observed.then.bind(observed); + if (key === "catch") return observed.catch.bind(observed); + if (key === "finally") return observed.finally.bind(observed); + if (key === "withRawResponse") return () => target.withRawResponse!().then( + (raw) => ({ ...raw, data: resolved(raw.data) }), rejected, + ); + const value = Reflect.get(target, key, target); + return typeof value === "function" ? value.bind(target) : value; + }, }); - - return attributes; } diff --git a/typescript/packages/traceai_cohere/src/stream.ts b/typescript/packages/traceai_cohere/src/stream.ts new file mode 100644 index 000000000..814349453 --- /dev/null +++ b/typescript/packages/traceai_cohere/src/stream.ts @@ -0,0 +1,156 @@ +import { context, Context } from "@opentelemetry/api"; +import { Data, object, list } from "./attributes"; + +type Finish = (response: Data, error?: unknown, cancelled?: boolean) => void; + +/** Aggregate content, not the event log. Consumers still receive original events. */ +class ChatAccumulator { + private text = ""; + private toolPlan = ""; + private content = new Map(); + private tools = new Map(); + private response: Data = {}; + constructor(private captureContent: boolean) {} + + add(value: unknown): boolean { + const event = object(value); + const type = event.type ?? event.eventType ?? event.event_type; + const delta = object(event.delta); + const message = object(delta.message); + if (this.captureContent) { + if (type === "text-generation" && typeof event.text === "string") this.text += event.text; + if (type === "content-start" || type === "content-delta") { + const content = object(message.content); + const index = typeof event.index === "number" ? event.index : 0; + const previous = this.content.get(index) ?? {}; + const block: Data = { ...previous }; + // Keep reasoning separate from visible text, preserving block order and + // partial content even when a stream ends before its terminal event. + for (const key of ["text", "thinking"]) { + if (typeof content[key] === "string") { + block.type = content.type ?? block.type ?? key; + block[key] = String(previous[key] ?? "") + content[key]; + } + } + if (block.type) this.content.set(index, block); + if (typeof content.text === "string") this.text += content.text; + } + if (type === "tool-plan-delta") { + const plan = message.toolPlan ?? message.tool_plan; + if (typeof plan === "string") this.toolPlan += plan; + } + if (type === "tool-calls-chunk") { + const call = object(event.toolCallDelta ?? event.tool_call_delta); + if (typeof call.text === "string") this.toolPlan += call.text; + const index = typeof call.index === "number" ? call.index : 0; + const previous = this.tools.get(index) ?? {}; + this.tools.set(index, { + ...previous, + ...(typeof call.name === "string" ? { name: call.name } : {}), + // These are JSON fragments, not objects: retain unfinished arguments. + parameters: String(previous.parameters ?? "") + String(call.parameters ?? ""), + }); + } + if (type === "tool-calls-generation") { + if (typeof event.text === "string") this.toolPlan = event.text; + this.tools.clear(); + list(event.toolCalls ?? event.tool_calls).forEach((call, index) => this.tools.set(index, object(call))); + } + if (type === "tool-call-start" || type === "tool-call-delta") { + const call = object(message.toolCalls ?? message.tool_calls); + const index = typeof event.index === "number" ? event.index : 0; + const previous = this.tools.get(index) ?? {}; + const fn = object(call.function); + const previousFn = object(previous.function); + this.tools.set(index, { + ...previous, ...call, + function: { ...previousFn, ...fn, arguments: String(previousFn.arguments ?? "") + String(fn.arguments ?? "") }, + }); + } + } + if (type === "stream-end") { + const response = object(event.response); + // V1's terminal text is authoritative after provisional tool chunks. + if (this.toolPlan && typeof response.text === "string" && response.text) this.toolPlan = response.text; + this.response = { ...response, finishReason: event.finishReason ?? event.finish_reason ?? response.finishReason ?? response.finish_reason }; + } + if (type === "message-end") this.response = { usage: delta.usage, finishReason: delta.finishReason ?? delta.finish_reason }; + return type === "stream-end" || type === "message-end"; + } + + result(): Data { + if (!this.captureContent) return { usage: this.response.usage, meta: this.response.meta, finishReason: this.response.finishReason }; + // V1's terminal response has authoritative text/tool calls; v2 only sends usage. + return { + text: this.text, + toolCalls: [...this.tools].sort(([a], [b]) => a - b).map(([, call]) => call), + ...(this.content.size || this.toolPlan ? { message: { + role: "assistant", + ...(this.content.size ? { content: [...this.content].sort(([a], [b]) => a - b).map(([, block]) => block) } : {}), + ...(this.toolPlan ? { toolPlan: this.toolPlan } : {}), + } } : {}), + ...this.response, + }; + } +} + +export function wrapStream(stream: T, execContext: Context, captureContent: boolean, finish: Finish): T { + const accumulator = new ChatAccumulator(captureContent); + const signal = (stream as { controller?: AbortController }).controller?.signal; + let ended = false; + const end = (error?: unknown, cancelled = false) => { + if (ended) return; + ended = true; + signal?.removeEventListener("abort", onAbort); + finish(accumulator.result(), error, cancelled); + }; + const onAbort = () => end(undefined, true); + signal?.addEventListener("abort", onAbort, { once: true }); + if (signal?.aborted) onAbort(); + + return new Proxy(stream, { + get(target, key) { + if (key === Symbol.asyncIterator) return () => { + let iterator: AsyncIterator; + try { + iterator = context.with(execContext, () => (target as unknown as AsyncIterable)[Symbol.asyncIterator]()); + } catch (error) { + end(error ?? new Error("Cohere stream failed")); + throw error; + } + const invoke = async (method: "next" | "return" | "throw", value?: unknown): Promise> => { + try { + const fn = iterator[method]; + if (!fn) { + if (method === "throw") throw value; + end(undefined, method === "return"); + return { done: true, value }; + } + const result = await context.with(execContext, () => fn.call(iterator, value)); + if (!ended && !result.done) { + const terminal = accumulator.add(result.value); + if (terminal) end(); + } + if (!ended && (result.done || method === "return")) { + // Terminal events finish the span above. A clean transport EOF + // without one is still an incomplete generation, not success. + end(method === "return" ? undefined : new Error("Cohere stream ended without a terminal event"), method === "return"); + } + return result; + } catch (error) { + end(error ?? new Error("Cohere stream failed")); + throw error; + } + }; + return { + next: (value?: unknown) => invoke("next", value), + return: (value?: unknown) => invoke("return", value), + throw: (error?: unknown) => invoke("throw", error), + [Symbol.asyncIterator]() { return this; }, + }; + }; + const value = Reflect.get(target, key, target); + return typeof value === "function" ? value.bind(target) : value; + }, + }); +} diff --git a/typescript/pnpm-lock.yaml b/typescript/pnpm-lock.yaml index 45ff919a2..98f5b3035 100644 --- a/typescript/pnpm-lock.yaml +++ b/typescript/pnpm-lock.yaml @@ -596,6 +596,12 @@ importers: specifier: workspace:^ version: link:../fi-semantic-conventions devDependencies: + '@opentelemetry/context-async-hooks': + specifier: 2.2.0 + version: 2.2.0(@opentelemetry/api@1.9.0) + '@opentelemetry/exporter-trace-otlp-http': + specifier: 0.202.0 + version: 0.202.0(@opentelemetry/api@1.9.0) '@opentelemetry/exporter-trace-otlp-proto': specifier: ^0.50.0 version: 0.50.0(@opentelemetry/api@1.9.0) @@ -615,8 +621,11 @@ importers: specifier: ^29.5.12 version: 29.5.14 cohere-ai: - specifier: ^7.0.0 - version: 7.20.0 + specifier: 8.3.0 + version: 8.3.0(@aws-crypto/sha256-js@5.2.0)(@aws-sdk/credential-providers@3.975.0)(@smithy/protocol-http@5.3.8)(@smithy/signature-v4@5.3.8) + cohere-ai-v7: + specifier: npm:cohere-ai@7.0.0 + version: cohere-ai@7.0.0 jest: specifier: ^29.7.0 version: 29.7.0(@types/node@22.15.33)(ts-node@10.9.2(@types/node@22.15.33)(typescript@5.8.3)) @@ -4459,6 +4468,12 @@ packages: peerDependencies: '@opentelemetry/api': '>=1.0.0 <1.10.0' + '@opentelemetry/context-async-hooks@2.2.0': + resolution: {integrity: sha512-qRkLWiUEZNAmYapZ7KGS5C4OmBLcP/H2foXeOEaowYCR0wi89fHejrfYfbuLVCMLp/dWZXKvQusdbUEZjERfwQ==} + engines: {node: ^18.19.0 || >=20.6.0} + peerDependencies: + '@opentelemetry/api': '>=1.0.0 <1.10.0' + '@opentelemetry/context-async-hooks@2.7.1': resolution: {integrity: sha512-OPFBYuXEn1E4ja3Y6eeA7O+ZnLBNcXTV5Cgsn1VaqBZ6hC5FnpZPLBNme1LJY8ZtF4aOujPKFoeWN4ik487KuQ==} engines: {node: ^18.19.0 || >=20.6.0} @@ -6108,6 +6123,9 @@ packages: '@types/unist@3.0.3': resolution: {integrity: sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q==} + '@types/url-join@4.0.1': + resolution: {integrity: sha512-wDXw9LEEUHyV+7UWy7U315nrJGJ7p1BzaCxDpEoLr789Dk1WDVMMlf3iBfbG2F8NdWnYyFbtTxUn2ZNbm1Q4LQ==} + '@types/uuid@10.0.0': resolution: {integrity: sha512-7gqG38EyHgyP1S+7+xomFtL+ZNHcKv6DwNaCZmJmo1vgMugyF3TCnXVg4t1uk89mLNwnLtnY3TpOpCOyp1/xHQ==} @@ -6193,6 +6211,9 @@ packages: peerDependencies: typescript: '*' + '@ungap/url-search-params@0.2.2': + resolution: {integrity: sha512-qQsguKXZVKdCixOHX9jqnX/K/1HekPDpGKyEcXHT+zR6EjGA7S4boSuelL4uuPv6YfhN0n8c4UxW+v/Z3gM2iw==} + '@vercel/oidc@3.1.0': resolution: {integrity: sha512-Fw28YZpRnA3cAHHDlkt7xQHiJ0fcL+NRcIqsocZQUSmbzeIKRpwttJjik5ZGanXP+vlA4SbTg+AbA3bP363l+w==} engines: {node: '>= 20'} @@ -6426,6 +6447,9 @@ packages: aws4fetch@1.0.20: resolution: {integrity: sha512-/djoAN709iY65ETD6LKCtyyEI04XIBP5xVvfmNxsEP0uJB5tyaGBztSryRr4HqMStr9R06PisQE7m9zDTXKu6g==} + axios@0.27.2: + resolution: {integrity: sha512-t+yRIyySRTp/wua5xEr+z1q60QmLq8ABsS5O9Me1AsE5dfKqgnCFzwiCZZ/cGNd1lq4/7akDWMxdhVlucjmnOQ==} + axios@1.13.5: resolution: {integrity: sha512-cz4ur7Vb0xS4/KUN0tPWe44eqxrIu31me+fbang3ijiNscE129POzipJJA6zniq2C/Z6sJCjMimjS8Lc/GAs8Q==} @@ -6740,10 +6764,31 @@ packages: resolution: {integrity: sha512-YIFQQ1n2NSgwoB3sCe7RpkZzsrPxTMek6jc7wC9fXOm1wwfWAKja9gLOMEjlXOUd3LKV3o6Jci7n9BoHs5Z8Sg==} engines: {node: '>=14.18.0'} + cohere-ai@7.0.0: + resolution: {integrity: sha512-SBJJ6ESpKSy4pP7YzVt2Jp48Gt2jAIFtXnE4qpbDrQcgxWzchJOs3cp8XnSwJKZp0c4Z8GSMOZxbln2L2ItlmA==} + cohere-ai@7.20.0: resolution: {integrity: sha512-h/3h3pcLXRUmkzp/W+/FWViEMcAFtSZ8YayCTFQXpib112uNSj3feApOtJg7v9lreWR1t7gznhE6N9KNCX5FOA==} engines: {node: '>=18.0.0'} + cohere-ai@8.3.0: + resolution: {integrity: sha512-3EIpXefkHcjYgbPDxwJDPa+XgHRSotD1YNi+HrROtuezLLQ12seKZA2GBJSHllbOEBhv+j/TY+f1BGUBn9w0SQ==} + engines: {node: '>=18.0.0'} + peerDependencies: + '@aws-crypto/sha256-js': ^5.2.0 + '@aws-sdk/credential-providers': ^3.583.0 + '@smithy/protocol-http': ^5.1.2 + '@smithy/signature-v4': ^5.1.2 + peerDependenciesMeta: + '@aws-crypto/sha256-js': + optional: true + '@aws-sdk/credential-providers': + optional: true + '@smithy/protocol-http': + optional: true + '@smithy/signature-v4': + optional: true + collect-v8-coverage@1.0.2: resolution: {integrity: sha512-lHl4d5/ONEbLlJvaJNtsF/Lz+WvB07u2ycqTYbdrq7UypDXailES4valYb2eWiJFxZlVmpGekfqoxQhzyFdT4Q==} @@ -8044,6 +8089,9 @@ packages: resolution: {integrity: sha512-34wB/Y7MW7bzjKRjUKTa46I2Z7eV62Rkhva+KkopW7Qvv/OSWBqvkSY7vusOPrNuZcUG3tApvdVgNB8POj3SPw==} engines: {node: '>=10'} + js-base64@3.7.2: + resolution: {integrity: sha512-NnRs6dsyqUXejqk/yv2aiXlAvOs56sLkX6nUdeaNezI5LFFLlsZjOThmwnrcwh5ZZRwZlCMnVAY3CvhIhoVEKQ==} + js-tiktoken@1.0.20: resolution: {integrity: sha512-Xlaqhhs8VfCd6Sh7a1cFkZHQbYTLCwVJJWiHVxBYzLPxW0XsoxBy1hitmjkdIjD3Aon5BXLHFwU5O8WUx6HH+A==} @@ -10296,6 +10344,9 @@ packages: uri-templates@0.2.0: resolution: {integrity: sha512-EWkjYEN0L6KOfEoOH6Wj4ghQqU7eBZMJqRHQnxQAq+dSEzRPClkWjf8557HkWQXF6BrAUoLSAyy9i3RVTliaNg==} + url-join@4.0.1: + resolution: {integrity: sha512-jk1+QP6ZJqyOiuEI9AEWQfju/nB2Pw466kbA0LEZljHwKeMgd9WrAEgEGxjPDD2+TNbbb37rTyhEfrCXfuKXnA==} + url-template@2.0.8: resolution: {integrity: sha512-XdVKMF4SJ0nP/O7XIPB0JwAEuT9lDIYnNsK8yGVe43y0AWoKeJNdv3ZNWh7ksJ6KqQFjOO6ox/VEitLnaVNufw==} @@ -14010,6 +14061,10 @@ snapshots: dependencies: '@opentelemetry/api': 1.9.0 + '@opentelemetry/context-async-hooks@2.2.0(@opentelemetry/api@1.9.0)': + dependencies: + '@opentelemetry/api': 1.9.0 + '@opentelemetry/context-async-hooks@2.7.1(@opentelemetry/api@1.9.1)': dependencies: '@opentelemetry/api': 1.9.1 @@ -16299,6 +16354,8 @@ snapshots: '@types/unist@3.0.3': {} + '@types/url-join@4.0.1': {} + '@types/uuid@10.0.0': {} '@types/webidl-conversions@7.0.3': {} @@ -16419,6 +16476,8 @@ snapshots: transitivePeerDependencies: - supports-color + '@ungap/url-search-params@0.2.2': {} + '@vercel/oidc@3.1.0': {} '@vitest/expect@3.2.4': @@ -16720,6 +16779,13 @@ snapshots: aws4fetch@1.0.20: {} + axios@0.27.2: + dependencies: + follow-redirects: 1.16.1 + form-data: 4.0.6 + transitivePeerDependencies: + - debug + axios@1.13.5: dependencies: follow-redirects: 1.15.11 @@ -17131,6 +17197,16 @@ snapshots: dependencies: rfdc: 1.4.1 + cohere-ai@7.0.0: + dependencies: + '@types/url-join': 4.0.1 + '@ungap/url-search-params': 0.2.2 + axios: 0.27.2 + js-base64: 3.7.2 + url-join: 4.0.1 + transitivePeerDependencies: + - debug + cohere-ai@7.20.0: dependencies: '@aws-crypto/sha256-js': 5.2.0 @@ -17146,6 +17222,18 @@ snapshots: transitivePeerDependencies: - aws-crt + cohere-ai@8.3.0(@aws-crypto/sha256-js@5.2.0)(@aws-sdk/credential-providers@3.975.0)(@smithy/protocol-http@5.3.8)(@smithy/signature-v4@5.3.8): + dependencies: + form-data: 4.0.6 + form-data-encoder: 4.1.0 + formdata-node: 6.0.3 + readable-stream: 4.7.0 + optionalDependencies: + '@aws-crypto/sha256-js': 5.2.0 + '@aws-sdk/credential-providers': 3.975.0 + '@smithy/protocol-http': 5.3.8 + '@smithy/signature-v4': 5.3.8 + collect-v8-coverage@1.0.2: {} color-convert@2.0.1: @@ -19007,6 +19095,8 @@ snapshots: joycon@3.1.1: {} + js-base64@3.7.2: {} + js-tiktoken@1.0.20: dependencies: base64-js: 1.5.1 @@ -21594,6 +21684,8 @@ snapshots: uri-templates@0.2.0: {} + url-join@4.0.1: {} + url-template@2.0.8: {} util-deprecate@1.0.2: {}