From d4474d70242074412c3135c831c7b9a20ed1075c Mon Sep 17 00:00:00 2001 From: Nikhil Pareek <8347993+nik13@users.noreply.github.com> Date: Sat, 31 Oct 2026 05:41:00 -0700 Subject: [PATCH 1/7] fix(cohere): trace v2 clients and async chat streams Resolve each SDK method owner once, preserve v1 instrumentation, and capture v2 messages, tool deltas, token usage and typed embeddings. Keep streaming span context local to SDK iteration and finish on completion, errors, cancellation or close. Use canonical message attributes for privacy. Add offline SDK and OTLP-wire regression coverage for Cohere 5.0.0, 5.13.12 and 7.3.0. Update package examples to v2 without changing the public instrumentor API. Co-Authored-By: Claude Code --- python/frameworks/cohere/README.md | 345 +++--------- .../cohere/tests/test_sdk_coverage.py | 515 ++++++++++++++++++ .../cohere/traceai_cohere/__init__.py | 155 ++---- .../_request_attributes_extractor.py | 48 +- .../_response_attributes_extractor.py | 125 +++-- .../cohere/traceai_cohere/_stream.py | 191 +++++++ .../cohere/traceai_cohere/_utils.py | 34 ++ .../cohere/traceai_cohere/_with_span.py | 2 +- .../cohere/traceai_cohere/_wrappers.py | 122 ++--- 9 files changed, 1028 insertions(+), 509 deletions(-) create mode 100644 python/frameworks/cohere/tests/test_sdk_coverage.py create mode 100644 python/frameworks/cohere/traceai_cohere/_stream.py diff --git a/python/frameworks/cohere/README.md b/python/frameworks/cohere/README.md index cb50b729e..54829e3cb 100644 --- a/python/frameworks/cohere/README.md +++ b/python/frameworks/cohere/README.md @@ -1,315 +1,150 @@ -# TraceAI Cohere Instrumentation +# Cohere tracing -OpenTelemetry instrumentation for [Cohere](https://cohere.com/) - embeddings, chat, and rerank APIs. +Trace Cohere chat, streaming chat, embeddings, and reranking with OpenTelemetry. -## Installation +## Install ```bash pip install traceai-cohere ``` -## Features +The instrumentor supports `cohere>=5.0.0`. It instruments `Client` and +`AsyncClient`, plus `ClientV2` and `AsyncClientV2` when those classes are available +in your installed Cohere SDK. Use one Cohere instrumentor per application; do not +also enable another library's Cohere instrumentation. -- Automatic tracing of Cohere API calls -- Support for chat, embed, and rerank endpoints -- Streaming response support -- Token usage tracking -- Rerank relevance scores -- Full OpenTelemetry semantic conventions compliance +## Configure tracing -## Usage +Pass your application's OpenTelemetry tracer provider to the instrumentor. This +example prints spans locally; use your configured OTLP provider to export traces +to [Future AGI](https://docs.futureagi.com/docs/tracing). -### Basic Setup +Set `CO_API_KEY` to your Cohere API key and `COHERE_MODEL` to a chat model available +to your account before running the example. Keep both outside your source code. ```python +import os + import cohere -from opentelemetry import trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import ConsoleSpanExporter, SimpleSpanProcessor - from traceai_cohere import CohereInstrumentor -# Set up tracing provider = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter())) -trace.set_tracer_provider(provider) - -# Instrument Cohere CohereInstrumentor().instrument(tracer_provider=provider) -# Use Cohere -client = cohere.Client(api_key="your-api-key") -response = client.chat(message="Hello!") -print(response.text) -``` - -### Chat - -```python -import cohere - -client = cohere.Client() - -# Simple chat +client = cohere.ClientV2(api_key=os.environ["CO_API_KEY"]) response = client.chat( - model="command-r-plus", - message="What is machine learning?" -) -print(response.text) - -# With chat history -response = client.chat( - model="command-r-plus", - message="What's my name?", - chat_history=[ - {"role": "USER", "message": "My name is Alice"}, - {"role": "CHATBOT", "message": "Hello Alice!"} - ] -) - -# With preamble (system prompt) -response = client.chat( - model="command-r-plus", - message="Write a poem", - preamble="You are a creative poet who writes in haiku format." + model=os.environ["COHERE_MODEL"], + messages=[{"role": "user", "content": "Explain semantic search in one sentence."}], ) +for content in response.message.content or []: + if content.type == "text": + print(content.text) +provider.force_flush() ``` -### Streaming Chat - -```python -import cohere - -client = cohere.Client() +Existing v1 calls such as `client.chat(message="Hello", model=...)` continue to +work with the same `CohereInstrumentor().instrument(...)` setup. -# Streaming -for event in client.chat_stream( - model="command-r-plus", - message="Tell me a story" -): - if event.event_type == "text-generation": - print(event.text, end="", flush=True) -``` +## Stream responses -### Embeddings +The stream span finishes when you exhaust or close the iterator. Close a stream +explicitly if you stop reading early. ```python -import cohere - -client = cohere.Client() - -# Generate embeddings -response = client.embed( - model="embed-english-v3.0", - texts=["Hello world", "Machine learning is great"], - input_type="search_document" +stream = client.chat_stream( + model=os.environ["COHERE_MODEL"], + messages=[{"role": "user", "content": "Give three tips for semantic search."}], ) -print(f"Generated {len(response.embeddings)} embeddings") -print(f"Dimensions: {len(response.embeddings[0])}") +try: + for event in stream: + if event.type == "content-delta" and event.delta.message.content.text: + print(event.delta.message.content.text, end="", flush=True) +finally: + stream.close() ``` -### Rerank (for RAG) +For async applications, use `AsyncClientV2` and iterate directly over +`chat_stream(...)`; do not await the stream itself. ```python -import cohere - -client = cohere.Client() - -# Rerank documents for a query -query = "What is the capital of France?" -documents = [ - "Paris is the capital of France.", - "London is the capital of England.", - "The Eiffel Tower is in Paris.", - "France is a country in Europe." -] - -response = client.rerank( - model="rerank-english-v3.0", - query=query, - documents=documents, - top_n=3 -) - -for result in response.results: - print(f"Index: {result.index}, Score: {result.relevance_score:.4f}") - print(f" {documents[result.index]}") +import asyncio + +async def main(): + client = cohere.AsyncClientV2(api_key=os.environ["CO_API_KEY"]) + stream = client.chat_stream( + model=os.environ["COHERE_MODEL"], + messages=[{"role": "user", "content": "Explain reranking in one sentence."}], + ) + try: + async for event in stream: + if event.type == "content-delta" and event.delta.message.content.text: + print(event.delta.message.content.text, end="", flush=True) + finally: + await stream.aclose() + +asyncio.run(main()) +provider.force_flush() ``` -### RAG with Cohere - -```python -import cohere - -client = cohere.Client() - -# Documents for context -documents = [ - {"title": "Paris", "text": "Paris is the capital of France."}, - {"title": "London", "text": "London is the capital of England."}, -] - -response = client.chat( - model="command-r-plus", - message="What is the capital of France?", - documents=documents -) - -print(response.text) -if response.citations: - print("\nCitations:") - for citation in response.citations: - print(f" - {citation}") -``` +## Embed and rerank -### Tool Use +Both operations use the same instrumentor setup, for sync and async clients. ```python -import cohere - -client = cohere.Client() - -tools = [ - { - "name": "get_weather", - "description": "Get the weather for a location", - "parameter_definitions": { - "location": { - "type": "str", - "description": "The city name", - "required": True - } - } - } -] - -response = client.chat( - model="command-r-plus", - message="What's the weather in Paris?", - tools=tools +embedding = client.embed( + model="embed-v4.0", + texts=["Paris is the capital of France."], + input_type="search_document", + embedding_types=["float"], ) +print(len(embedding.embeddings.float[0])) -if response.tool_calls: - for tool_call in response.tool_calls: - print(f"Tool: {tool_call.name}") - print(f"Parameters: {tool_call.parameters}") +ranking = client.rerank( + model="rerank-v3.5", + query="What is the capital of France?", + documents=["Paris is the capital of France.", "London is the capital of England."], + top_n=1, +) +print(ranking.results[0].index) ``` -## Configuration Options +## Control content capture -### TraceConfig +By default, chat spans include prompts, responses, and tool-call arguments. +Configure privacy before instrumenting if you do not want that content exported. +Replace the earlier `instrument(...)` call with this one: ```python from fi_instrumentation import TraceConfig -from traceai_cohere import CohereInstrumentor - -config = TraceConfig( - hide_inputs=False, - hide_outputs=False, -) CohereInstrumentor().instrument( tracer_provider=provider, - config=config + config=TraceConfig(hide_inputs=True, hide_outputs=True), ) ``` -## Captured Attributes - -### Chat Attributes - -| Attribute | Description | -|-----------|-------------| -| `gen_ai.span.kind` | "LLM" | -| `gen_ai.provider.name` | "cohere" | -| `gen_ai.request.model` | Model name (command-r-plus, etc.) | -| `gen_ai.usage.input_tokens` | Input token count | -| `gen_ai.usage.output_tokens` | Output token count | -| `cohere.finish_reason` | Response finish reason | -| `cohere.citations_count` | Number of citations | -| `cohere.tool_calls_count` | Number of tool calls | - -### Embed Attributes - -| Attribute | Description | -|-----------|-------------| -| `gen_ai.span.kind` | "EMBEDDING" | -| `embedding.model_name` | Embedding model name | -| `cohere.texts_count` | Number of texts embedded | -| `cohere.embedding_dimensions` | Vector dimensions | - -### Rerank Attributes - -| Attribute | Description | -|-----------|-------------| -| `gen_ai.span.kind` | "RERANKER" | -| `reranker.model_name` | Rerank model name | -| `reranker.query` | Search query | -| `reranker.top_k` | Top N results requested | -| `cohere.results_count` | Number of results | -| `cohere.rerank.scores` | Relevance scores | - -## Available Models - -| Category | Models | -|----------|--------| -| Chat | `command-r-plus`, `command-r`, `command` | -| Embed | `embed-english-v3.0`, `embed-multilingual-v3.0` | -| Rerank | `rerank-english-v3.0`, `rerank-multilingual-v3.0` | - -## Real-World Use Cases - -### Semantic Search - -```python -import cohere - -client = cohere.Client() - -# Index documents -documents = ["doc1", "doc2", "doc3"] -doc_embeddings = client.embed( - model="embed-english-v3.0", - texts=documents, - input_type="search_document" -).embeddings - -# Search query -query_embedding = client.embed( - model="embed-english-v3.0", - texts=["search query"], - input_type="search_query" -).embeddings[0] - -# Compute similarity (not shown: use vector DB) -``` - -### Two-Stage RAG +You can also set `FI_HIDE_INPUTS=true` and `FI_HIDE_OUTPUTS=true` in your environment. +Explicit `TraceConfig` settings take precedence. Hiding content retains model, +usage, counts, and rerank scores. -```python -import cohere +## Span attributes -client = cohere.Client() +| Operation | Span name | Captured metadata | +| --- | --- | --- | +| Chat | `cohere.chat` | Model, input/output messages, tool calls, token usage, finish reason | +| Streaming chat | `cohere.chat_stream` | Accumulated text and tool-call arguments, final usage and finish reason | +| Embedding | `cohere.embed` | Model, text count, embedding count and dimensions | +| Rerank | `cohere.rerank` | Model, query, result count, relevance scores | -# Stage 1: Semantic search (not shown) -initial_results = ["doc1", "doc2", "doc3", "doc4", "doc5"] +Token usage prefers actual token counts when the response provides them, falling +back to billed units otherwise. A stream closed early may not have final usage. +Tool calls use canonical `gen_ai.input.messages` and `gen_ai.output.messages` +attributes so content privacy settings also apply to their arguments. -# Stage 2: Rerank for precision -reranked = client.rerank( - model="rerank-english-v3.0", - query="user question", - documents=initial_results, - top_n=3 -) - -# Stage 3: Generate with top results -top_docs = [initial_results[r.index] for r in reranked.results] -response = client.chat( - model="command-r-plus", - message="user question", - documents=[{"text": d} for d in top_docs] -) -``` +To remove the patches, call `CohereInstrumentor().uninstrument()`. ## License diff --git a/python/frameworks/cohere/tests/test_sdk_coverage.py b/python/frameworks/cohere/tests/test_sdk_coverage.py new file mode 100644 index 000000000..d521ce060 --- /dev/null +++ b/python/frameworks/cohere/tests/test_sdk_coverage.py @@ -0,0 +1,515 @@ +"""Real SDK dispatch and OTLP serialization, with synthetic HTTP only.""" + +import asyncio +import json + +import cohere +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.instrumentation.utils import suppress_instrumentation +from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ( + ExportTraceServiceRequest, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from traceai_cohere import CohereInstrumentor + +INPUT = "PRIVATE_INPUT_4829" +OUTPUT = "PRIVATE_OUTPUT_7251" +TOOL_ARGUMENT = "PRIVATE_TOOL_6723" +USAGE = { + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"input_tokens": 5, "output_tokens": 2}, +} + + +def payload(version, operation): + if operation == "embed": + embeddings = [[0.1, 0.2]] if version == 1 else {"float": [[0.1, 0.2]]} + return { + "id": "offline", + "texts": [INPUT], + "embeddings": embeddings, + "response_type": "embeddings_floats" + if version == 1 + else "embeddings_by_type", + "meta": {"billed_units": {"input_tokens": 7}}, + } + if operation == "rerank": + return { + "id": "offline", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": {"billed_units": {"search_units": 1}}, + } + if version == 1: + return { + "text": OUTPUT, + "generation_id": "00000000-0000-0000-0000-000000000001", + "meta": USAGE, + "finish_reason": "COMPLETE", + "tool_calls": [{"name": "lookup", "parameters": {"query": TOOL_ARGUMENT}}], + } + return { + "id": "offline", + "finish_reason": "COMPLETE", + "usage": USAGE, + "message": { + "role": "assistant", + "content": [{"type": "text", "text": OUTPUT}], + "tool_calls": [ + { + "id": "call-1", + "type": "function", + "function": { + "name": "lookup", + "arguments": json.dumps({"query": TOOL_ARGUMENT}), + }, + } + ], + }, + } + + +def events(version): + if version == 1: + return [ + {"event_type": "text-generation", "text": OUTPUT}, + { + "event_type": "stream-end", + "finish_reason": "COMPLETE", + "response": payload(1, "chat"), + }, + ] + return [ + { + "type": "message-start", + "id": "offline", + "delta": {"message": {"role": "assistant"}}, + }, + { + "type": "content-start", + "index": 0, + "delta": {"message": {"content": {"type": "text", "text": ""}}}, + }, + { + "type": "content-delta", + "index": 0, + "delta": {"message": {"content": {"text": OUTPUT}}}, + }, + { + "type": "tool-call-start", + "index": 0, + "delta": { + "message": { + "tool_calls": { + "id": "call-1", + "type": "function", + "function": {"name": "lookup", "arguments": ""}, + } + } + }, + }, + { + "type": "tool-call-delta", + "index": 0, + "delta": { + "message": {"tool_calls": {"function": {"arguments": '{"query": '}}} + }, + }, + { + "type": "tool-call-delta", + "index": 0, + "delta": { + "message": { + "tool_calls": { + "function": {"arguments": json.dumps(TOOL_ARGUMENT) + "}"} + } + } + }, + }, + {"type": "message-end", "delta": {"finish_reason": "COMPLETE", "usage": USAGE}}, + ] + + +def request_args(version, operation): + if operation == "embed": + args = dict(model="embed-v4.0", texts=[INPUT], input_type="search_document") + if version == 2: + args["embedding_types"] = ["float"] + return args + if operation == "rerank": + return dict(model="rerank-v3.5", query=INPUT, documents=["offline"]) + if version == 1: + return dict(model="command-r", message=INPUT) + return dict(model="command-r", messages=[{"role": "user", "content": INPUT}]) + + +def client_class(version, asynchronous): + name = ("AsyncClient" if asynchronous else "Client") + ( + "V2" if version == 2 else "" + ) + if not hasattr(cohere, name): + pytest.skip("Cohere 5.0.0 predates v2 clients") + return getattr(cohere, name) + + +@pytest.fixture +def tracing(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + + def setup(**privacy): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**privacy)) + return provider, exporter + + yield setup + instrumentor.uninstrument() + provider.shutdown() + + +def wire_attributes(exporter): + spans = exporter.get_finished_spans() + assert len(spans) == 1 + wire = encode_spans(spans).SerializeToString() + decoded = ExportTraceServiceRequest.FromString(wire) + span = decoded.resource_spans[0].scope_spans[0].spans[0] + attrs = { + a.key: getattr(a.value, a.value.WhichOneof("value")) for a in span.attributes + } + return spans[0], wire, attrs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("operation", ["chat", "chat_stream", "embed", "rerank"]) +@pytest.mark.parametrize("hidden", [False, True]) +async def test_sdk_surfaces_and_wire(tracing, version, asynchronous, operation, hidden): + cls = client_class(version, asynchronous) + _, exporter = tracing(hide_inputs=hidden, hide_outputs=hidden) + requests = [] + + def respond(request): + requests.append(request) + if operation == "chat_stream": + chunks = events(version) + text = "".join( + ("data: " if version == 2 else "") + + json.dumps(e) + + ("\n\n" if version == 2 else "\n") + for e in chunks + ) + return httpx.Response( + 200, text=text, headers={"content-type": "text/event-stream"} + ) + return httpx.Response(200, json=payload(version, operation)) + + args = request_args(version, operation) + if asynchronous: + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="offline-placeholder", httpx_client=http) + if operation == "chat_stream": + response = [e async for e in client.chat_stream(**args)] + else: + response = await getattr(client, operation)(**args) + else: + with httpx.Client(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="offline-placeholder", httpx_client=http) + response = getattr(client, operation)(**args) + if operation == "chat_stream": + response = list(response) + assert response is not None + assert len(requests) == 1 + span, wire, attrs = wire_attributes(exporter) + assert span.name == "cohere." + operation + assert span.status.status_code == trace.StatusCode.OK + assert attrs["gen_ai.provider.name"] == "cohere" + if hidden: + for private in (INPUT, OUTPUT, TOOL_ARGUMENT): + assert private.encode() not in wire + elif operation.startswith("chat"): + assert attrs["gen_ai.input.messages.0.message.content"] == INPUT + assert attrs["gen_ai.output.messages.0.message.content"] == OUTPUT + assert ( + TOOL_ARGUMENT + in attrs[ + "gen_ai.output.messages.0.message.tool_calls.0.tool_call.function.arguments" + ] + ) + if operation.startswith("chat"): + assert attrs["gen_ai.request.model"] == "command-r" + assert attrs["gen_ai.usage.input_tokens"] == 7 + assert attrs["gen_ai.usage.output_tokens"] == 3 + assert attrs["gen_ai.usage.total_tokens"] == 10 + elif operation == "embed": + assert attrs["cohere.embeddings_count"] == 1 + assert attrs["cohere.embedding_dimensions"] == 2 + else: + assert attrs["cohere.results_count"] == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_suppression_and_uninstrument(tracing, version, asynchronous): + cls = client_class(version, asynchronous) + originals = { + method: getattr(cls, method) + for method in ("chat", "chat_stream", "embed", "rerank") + } + _, exporter = tracing() + transport = httpx.MockTransport( + lambda r: httpx.Response(200, json=payload(version, "chat")) + ) + if asynchronous: + async with httpx.AsyncClient(transport=transport) as http: + client = cls(api_key="offline", httpx_client=http) + with suppress_instrumentation(): + await client.chat(**request_args(version, "chat")) + else: + with httpx.Client(transport=transport) as http: + client = cls(api_key="offline", httpx_client=http) + with suppress_instrumentation(): + client.chat(**request_args(version, "chat")) + assert not exporter.get_finished_spans() + CohereInstrumentor().uninstrument() + for method, original in originals.items(): + assert getattr(cls, method) is original + # A second cycle must neither retain wrappers nor double-instrument inherited methods. + tracing() + CohereInstrumentor().uninstrument() + for method, original in originals.items(): + assert getattr(cls, method) is original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + "termination", ["exhaust", "close", "unstarted_close", "error", "cancel"] +) +async def test_stream_lifecycle_context_and_passthrough( + tracing, asynchronous, termination +): + from traceai_cohere._wrappers import _AsyncChatStreamWrapper, _ChatStreamWrapper + + provider, exporter = tracing(hide_outputs=True) + seen = [] + closed = [] + event = {"event_type": "text-generation", "text": OUTPUT} + error = ( + asyncio.CancelledError(OUTPUT) + if termination == "cancel" + else RuntimeError(OUTPUT) + ) + + def source(): + try: + seen.append(trace.get_current_span().get_span_context().span_id) + yield event + if termination in ("error", "cancel"): + raise error + finally: + closed.append(True) + + async def async_source(): + try: + seen.append(trace.get_current_span().get_span_context().span_id) + yield event + if termination in ("error", "cancel"): + raise error + finally: + closed.append(True) + + tracer = CohereInstrumentor()._tracer + wrapper = (_AsyncChatStreamWrapper if asynchronous else _ChatStreamWrapper)( + tracer=tracer + ) + with provider.get_tracer("test").start_as_current_span("parent") as parent: + parent_id = parent.get_span_context().span_id + stream = wrapper(async_source if asynchronous else source, None, (), {}) + assert trace.get_current_span().get_span_context().span_id == parent_id + assert not exporter.get_finished_spans() + if termination != "unstarted_close": + got = await stream.__anext__() if asynchronous else next(stream) + assert got is event + assert trace.get_current_span().get_span_context().span_id == parent_id + assert not exporter.get_finished_spans() + if termination == "exhaust": + with pytest.raises(StopAsyncIteration if asynchronous else StopIteration): + if asynchronous: + await stream.__anext__() + else: + next(stream) + elif termination in ("error", "cancel"): + with pytest.raises(type(error)) as caught: + if asynchronous: + await stream.__anext__() + else: + next(stream) + assert caught.value is error + if asynchronous: + await stream.aclose() + await stream.aclose() + else: + stream.close() + stream.close() + assert trace.get_current_span().get_span_context().span_id == parent_id + span, wire, _ = wire_attributes(exporter) + assert span.parent.span_id == parent_id + if seen: + assert seen == [span.context.span_id] + assert closed == [True] + assert OUTPUT.encode() not in wire + expected = ( + trace.StatusCode.ERROR + if termination in ("error", "cancel") + else ( + trace.StatusCode.OK + if termination == "exhaust" + else trace.StatusCode.UNSET + ) + ) + assert span.status.status_code == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_sdk_errors_are_preserved_and_private( + tracing, version, asynchronous, streaming +): + cls = client_class(version, asynchronous) + _, exporter = tracing(hide_inputs=True, hide_outputs=True) + error = httpx.ConnectError(OUTPUT) + requests = [] + + def respond(request): + requests.append(request) + raise error + + args = request_args(version, "chat") + # Cohere 5.0.0 configures retries per request; no real network is used. + args["request_options"] = {"max_retries": 0} + with pytest.raises(httpx.ConnectError) as caught: + if asynchronous: + async with httpx.AsyncClient( + transport=httpx.MockTransport(respond) + ) as http: + client = cls(api_key="offline", httpx_client=http) + if streaming: + _ = [e async for e in client.chat_stream(**args)] + else: + await client.chat(**args) + else: + with httpx.Client(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="offline", httpx_client=http) + if streaming: + list(client.chat_stream(**args)) + else: + client.chat(**args) + assert caught.value is error + assert len(requests) == 1 + span, wire, _ = wire_attributes(exporter) + assert span.status.status_code == trace.StatusCode.ERROR + assert OUTPUT.encode() not in wire + + +def test_v2_typed_messages_and_null_usage(tracing): + from traceai_cohere._request_attributes_extractor import ( + _ChatRequestAttributesExtractor, + ) + from traceai_cohere._response_attributes_extractor import ( + _ChatResponseAttributesExtractor, + ) + + if not hasattr(cohere, "UserChatMessageV2"): + pytest.skip("Typed v2 messages are absent at the SDK floor") + _, exporter = tracing() + request = { + "messages": [ + cohere.UserChatMessageV2(role="user", content=INPUT), + { + "role": "assistant", + "tool_calls": [ + { + "id": "call-1", + "function": {"name": "lookup", "arguments": TOOL_ARGUMENT}, + } + ], + }, + {"role": "tool", "tool_call_id": "call-1", "content": "tool result"}, + ], + "documents": None, + "tools": None, + } + extractor = _ChatRequestAttributesExtractor() + with CohereInstrumentor()._tracer.start_as_current_span("typed-message") as span: + span.set_attributes(dict(extractor.get_extra_attributes_from_request(request))) + _, wire, attrs = wire_attributes(exporter) + assert attrs["gen_ai.input.messages.0.message.content"] == INPUT + assert ( + attrs["gen_ai.input.messages.1.message.tool_calls.0.tool_call.id"] == "call-1" + ) + assert attrs["gen_ai.input.messages.2.message.tool_call_id"] == "call-1" + assert TOOL_ARGUMENT.encode() in wire + response = _ChatResponseAttributesExtractor() + assert dict(response.get_attributes({"usage": None, "meta": None})) == {} + zero = dict( + response.get_attributes( + {"usage": {"tokens": {"input_tokens": 0, "output_tokens": 0}}} + ) + ) + assert zero["gen_ai.usage.total_tokens"] == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_abandoned_stream_finishes_span(tracing, asynchronous): + import gc + from traceai_cohere._wrappers import _AsyncChatStreamWrapper, _ChatStreamWrapper + + _, exporter = tracing() + + def source(): + yield {"event_type": "text-generation", "text": OUTPUT} + + async def async_source(): + yield {"event_type": "text-generation", "text": OUTPUT} + + wrapper = (_AsyncChatStreamWrapper if asynchronous else _ChatStreamWrapper)( + tracer=CohereInstrumentor()._tracer + ) + stream = wrapper(async_source if asynchronous else source, None, (), {}) + if asynchronous: + await stream.__anext__() + else: + next(stream) + assert not exporter.get_finished_spans() + del stream + gc.collect() + assert len(exporter.get_finished_spans()) == 1 + + +def test_unknown_event_extraction_does_not_break_stream(tracing): + from traceai_cohere._wrappers import _ChatStreamWrapper + + _, exporter = tracing() + + class UnknownEvent: + def model_dump(self): + raise ValueError("unrecognized event") + + event = UnknownEvent() + wrapper = _ChatStreamWrapper(tracer=CohereInstrumentor()._tracer) + stream = wrapper(lambda: iter([event]), None, (), {}) + assert list(stream) == [event] + span, _, _ = wire_attributes(exporter) + assert span.status.status_code == trace.StatusCode.OK diff --git a/python/frameworks/cohere/traceai_cohere/__init__.py b/python/frameworks/cohere/traceai_cohere/__init__.py index 630cb9c9f..0d4c5155d 100644 --- a/python/frameworks/cohere/traceai_cohere/__init__.py +++ b/python/frameworks/cohere/traceai_cohere/__init__.py @@ -17,6 +17,7 @@ from wrapt import wrap_function_wrapper from traceai_cohere._wrappers import ( + _AsyncChatStreamWrapper, _AsyncChatWrapper, _AsyncEmbedWrapper, _AsyncRerankWrapper, @@ -27,136 +28,66 @@ ) from traceai_cohere.version import __version__ -logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) _instruments = ("cohere >= 5.0.0",) class CohereInstrumentor(BaseInstrumentor): - """OpenTelemetry instrumentor for Cohere.""" - - __slots__ = ( - "_original_chat", - "_original_async_chat", - "_original_chat_stream", - "_original_async_chat_stream", - "_original_embed", - "_original_async_embed", - "_original_rerank", - "_original_async_rerank", - "_tracer", - ) + """Instrument v1/v2 Cohere chat, streaming, embeddings and reranking.""" def instrumentation_dependencies(self) -> Collection[str]: return _instruments def _instrument(self, **kwargs: Any) -> None: - if not (tracer_provider := kwargs.get("tracer_provider")): - tracer_provider = trace_api.get_tracer_provider() - if not (config := kwargs.get("config")): - config = TraceConfig() - else: - assert isinstance(config, TraceConfig) - + tracer_provider = ( + kwargs.get("tracer_provider") or trace_api.get_tracer_provider() + ) + config = kwargs.get("config") or TraceConfig() + assert isinstance(config, TraceConfig) self._tracer = FITracer( - trace_api.get_tracer(__name__, __version__, tracer_provider), - config=config, + trace_api.get_tracer(__name__, __version__, tracer_provider), config=config ) - - # Import cohere client + self._original_methods = [] cohere_module = import_module("cohere") - - # Wrap Client methods - # Chat - if hasattr(cohere_module.Client, "chat"): - self._original_chat = cohere_module.Client.chat - wrap_function_wrapper( - module="cohere", - name="Client.chat", - wrapper=_ChatWrapper(tracer=self._tracer), - ) - - # Chat stream - if hasattr(cohere_module.Client, "chat_stream"): - self._original_chat_stream = cohere_module.Client.chat_stream - wrap_function_wrapper( - module="cohere", - name="Client.chat_stream", - wrapper=_ChatStreamWrapper(tracer=self._tracer), - ) - - # Embed - if hasattr(cohere_module.Client, "embed"): - self._original_embed = cohere_module.Client.embed - wrap_function_wrapper( - module="cohere", - name="Client.embed", - wrapper=_EmbedWrapper(tracer=self._tracer), - ) - - # Rerank - if hasattr(cohere_module.Client, "rerank"): - self._original_rerank = cohere_module.Client.rerank - wrap_function_wrapper( - module="cohere", - name="Client.rerank", - wrapper=_RerankWrapper(tracer=self._tracer), - ) - - # Wrap AsyncClient methods - if hasattr(cohere_module, "AsyncClient"): - if hasattr(cohere_module.AsyncClient, "chat"): - self._original_async_chat = cohere_module.AsyncClient.chat - wrap_function_wrapper( - module="cohere", - name="AsyncClient.chat", - wrapper=_AsyncChatWrapper(tracer=self._tracer), - ) - - if hasattr(cohere_module.AsyncClient, "embed"): - self._original_async_embed = cohere_module.AsyncClient.embed - wrap_function_wrapper( - module="cohere", - name="AsyncClient.embed", - wrapper=_AsyncEmbedWrapper(tracer=self._tracer), - ) - - if hasattr(cohere_module.AsyncClient, "rerank"): - self._original_async_rerank = cohere_module.AsyncClient.rerank - wrap_function_wrapper( - module="cohere", - name="AsyncClient.rerank", - wrapper=_AsyncRerankWrapper(tracer=self._tracer), + sync = (_ChatWrapper, _ChatStreamWrapper, _EmbedWrapper, _RerankWrapper) + asynchronous = ( + _AsyncChatWrapper, + _AsyncChatStreamWrapper, + _AsyncEmbedWrapper, + _AsyncRerankWrapper, + ) + patched = set() + for client_name, wrappers in ( + ("Client", sync), + ("AsyncClient", asynchronous), + ("ClientV2", sync), + ("AsyncClientV2", asynchronous), + ): + client = getattr(cohere_module, client_name, None) + if client is None: # v2 does not exist at the supported SDK floor. + continue + for method, wrapper in zip( + ("chat", "chat_stream", "embed", "rerank"), wrappers + ): + # ClientV2 uses distinct v2 base methods; older SDKs inherit some + # v1 methods. Patch each implementation only once, at its owner. + owner = next( + (cls for cls in client.__mro__ if method in vars(cls)), None ) + if owner is None or (owner, method) in patched: + continue + patched.add((owner, method)) + self._original_methods.append((owner, method, vars(owner)[method])) + wrap_function_wrapper(owner, method, wrapper(tracer=self._tracer)) - # Wrap Protect if available if Protect is not None: - self._original_protect = Protect.protect + self._original_methods.append((Protect, "protect", Protect.protect)) wrap_function_wrapper( - module="fi.evals", - name="Protect.protect", - wrapper=GuardrailProtectWrapper(tracer=self._tracer), + Protect, "protect", GuardrailProtectWrapper(tracer=self._tracer) ) - else: - self._original_protect = None def _uninstrument(self, **kwargs: Any) -> None: - cohere_module = import_module("cohere") - - if hasattr(self, "_original_chat"): - cohere_module.Client.chat = self._original_chat - if hasattr(self, "_original_chat_stream"): - cohere_module.Client.chat_stream = self._original_chat_stream - if hasattr(self, "_original_embed"): - cohere_module.Client.embed = self._original_embed - if hasattr(self, "_original_rerank"): - cohere_module.Client.rerank = self._original_rerank - - if hasattr(cohere_module, "AsyncClient"): - if hasattr(self, "_original_async_chat"): - cohere_module.AsyncClient.chat = self._original_async_chat - if hasattr(self, "_original_async_embed"): - cohere_module.AsyncClient.embed = self._original_async_embed - if hasattr(self, "_original_async_rerank"): - cohere_module.AsyncClient.rerank = self._original_async_rerank + for owner, method, original in reversed(self._original_methods): + setattr(owner, method, original) + self._original_methods.clear() diff --git a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py index fae95f291..c98faabb4 100644 --- a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py @@ -1,16 +1,17 @@ import logging -from typing import Any, Dict, Iterator, List, Mapping, Optional, Tuple +from typing import Any, Dict, Iterator, Mapping, Optional, Tuple from fi_instrumentation import TraceConfig, safe_json_dumps from fi_instrumentation.fi_types import ( FiMimeTypeValues, FiSpanKindValues, - MessageAttributes, RerankerAttributes, SpanAttributes, ) from opentelemetry.util.types import AttributeValue +from traceai_cohere._utils import _message_attributes + logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -24,7 +25,6 @@ def get_attributes_from_request( """Extract primary attributes from request.""" yield SpanAttributes.GEN_AI_SPAN_KIND, FiSpanKindValues.LLM.value yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" - yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" model = request_parameters.get("model", "command-r-plus") yield SpanAttributes.GEN_AI_REQUEST_MODEL, model @@ -35,7 +35,10 @@ def get_attributes_from_request( # Invocation parameters invocation_params = self._extract_invocation_parameters(request_parameters) if invocation_params: - yield SpanAttributes.GEN_AI_REQUEST_PARAMETERS, safe_json_dumps(invocation_params) + yield ( + SpanAttributes.GEN_AI_REQUEST_PARAMETERS, + safe_json_dumps(invocation_params), + ) def get_extra_attributes_from_request( self, request_parameters: Mapping[str, Any] @@ -45,25 +48,31 @@ def get_extra_attributes_from_request( # indices. Canonical input-message keys also honor FI_HIDE_INPUTS. messages = [] if "preamble" in request_parameters: - messages.append({"role": "system", "message": request_parameters["preamble"]}) - messages.extend(request_parameters.get("chat_history", [])) + messages.append( + {"role": "system", "message": request_parameters["preamble"]} + ) + messages.extend((request_parameters.get("chat_history") or [])) if "message" in request_parameters: message = request_parameters.get("message", "") messages.append({"role": "user", "message": message}) yield SpanAttributes.INPUT_VALUE, message + if "messages" in request_parameters: + messages = request_parameters.get("messages") or [] + yield SpanAttributes.INPUT_VALUE, safe_json_dumps(messages) for i, msg in enumerate(messages): - yield f"{SpanAttributes.GEN_AI_INPUT_MESSAGES}.{i}.{MessageAttributes.MESSAGE_ROLE}", msg.get("role", "") - yield f"{SpanAttributes.GEN_AI_INPUT_MESSAGES}.{i}.{MessageAttributes.MESSAGE_CONTENT}", msg.get("message", "") + yield from _message_attributes( + msg, f"{SpanAttributes.GEN_AI_INPUT_MESSAGES}.{i}" + ) # Documents for RAG if "documents" in request_parameters: - docs = request_parameters.get("documents", []) + docs = request_parameters.get("documents") or [] yield "cohere.documents_count", len(docs) # Tools if "tools" in request_parameters: - tools = request_parameters.get("tools", []) + tools = request_parameters.get("tools") or [] yield "cohere.tools_count", len(tools) def _extract_invocation_parameters( @@ -72,8 +81,14 @@ def _extract_invocation_parameters( """Extract model invocation parameters.""" params = {} param_keys = [ - "temperature", "max_tokens", "k", "p", "stop_sequences", - "frequency_penalty", "presence_penalty", "seed" + "temperature", + "max_tokens", + "k", + "p", + "stop_sequences", + "frequency_penalty", + "presence_penalty", + "seed", ] for key in param_keys: if key in request_parameters: @@ -90,7 +105,6 @@ def get_attributes_from_request( """Extract primary attributes from request.""" yield SpanAttributes.GEN_AI_SPAN_KIND, FiSpanKindValues.EMBEDDING.value yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" - yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" model = request_parameters.get("model", "embed-english-v3.0") yield SpanAttributes.EMBEDDING_MODEL_NAME, model @@ -104,7 +118,10 @@ def get_attributes_from_request( # Embedding types if "embedding_types" in request_parameters: - yield "cohere.embedding_types", safe_json_dumps(request_parameters.get("embedding_types")) + yield ( + "cohere.embedding_types", + safe_json_dumps(request_parameters.get("embedding_types")), + ) def get_extra_attributes_from_request( self, request_parameters: Mapping[str, Any] @@ -130,7 +147,6 @@ def get_attributes_from_request( """Extract primary attributes from request.""" yield SpanAttributes.GEN_AI_SPAN_KIND, FiSpanKindValues.RERANKER.value yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" - yield SpanAttributes.GEN_AI_PROVIDER_NAME, "cohere" model = request_parameters.get("model", "rerank-english-v3.0") yield RerankerAttributes.RERANKER_MODEL_NAME, model @@ -153,5 +169,5 @@ def get_extra_attributes_from_request( yield RerankerAttributes.RERANKER_QUERY, masked_query yield SpanAttributes.INPUT_VALUE, query - documents = request_parameters.get("documents", []) + documents = request_parameters.get("documents") or [] yield "cohere.documents_count", len(documents) diff --git a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py index 212832b66..740366e74 100644 --- a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py @@ -4,11 +4,12 @@ from fi_instrumentation import safe_json_dumps from fi_instrumentation.fi_types import ( EmbeddingAttributes, - MessageAttributes, SpanAttributes, ) from opentelemetry.util.types import AttributeValue +from traceai_cohere._utils import _message_attributes + logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -25,25 +26,23 @@ def get_attributes( if response is None: return - # Token counts from meta - meta = response.get("meta", {}) - if "billed_units" in meta: - billed = meta["billed_units"] - if "input_tokens" in billed: - yield SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS, billed["input_tokens"] - if "output_tokens" in billed: - yield SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, billed["output_tokens"] - - total = billed.get("input_tokens", 0) + billed.get("output_tokens", 0) - if total: - yield SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS, total - - if "tokens" in meta: - tokens = meta["tokens"] - if "input_tokens" in tokens: - yield SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS, tokens["input_tokens"] - if "output_tokens" in tokens: - yield SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, tokens["output_tokens"] + # v1 reports meta; v2 reports usage. Prefer actual tokens over billed + # units, and compute the total from the same source (including zeroes). + usage = response.get("usage") or response.get("meta") or {} + tokens = usage.get("tokens") or usage.get("billed_units") or {} + for field, attribute in ( + ("input_tokens", SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS), + ("output_tokens", SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS), + ): + if tokens.get(field) is not None: + yield attribute, tokens[field] + if any( + tokens.get(field) is not None for field in ("input_tokens", "output_tokens") + ): + yield ( + SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS, + (tokens.get("input_tokens") or 0) + (tokens.get("output_tokens") or 0), + ) def get_extra_attributes( self, @@ -55,24 +54,30 @@ def get_extra_attributes( if response is None: return - # Text response - if "text" in response: - text = response.get("text", "") - yield f"{SpanAttributes.GEN_AI_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}", "assistant" - yield f"{SpanAttributes.GEN_AI_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}", text - yield SpanAttributes.OUTPUT_VALUE, text - - # Tool calls - if "tool_calls" in response and response["tool_calls"]: - tool_calls = response["tool_calls"] - yield "cohere.tool_calls_count", len(tool_calls) - for i, tc in enumerate(tool_calls[:5]): # Limit - yield f"cohere.tool_calls.{i}.name", tc.get("name", "") - yield f"cohere.tool_calls.{i}.parameters", safe_json_dumps(tc.get("parameters", {})) - - # Citations - if "citations" in response and response["citations"]: - yield "cohere.citations_count", len(response["citations"]) + message = response.get("message") + if isinstance(message, dict): + message = dict(message) + content = message.get("content") + if isinstance(content, list): + message["content"] = "".join( + block.get("text") or "" + for block in content + if isinstance(block, dict) + ) + else: + message = { + "role": "assistant", + "content": response.get("text"), + "tool_calls": response.get("tool_calls"), + } + yield from _message_attributes( + message, f"{SpanAttributes.GEN_AI_OUTPUT_MESSAGES}.0" + ) + if message.get("tool_calls"): + yield "cohere.tool_calls_count", len(message["tool_calls"]) + citations = message.get("citations") or response.get("citations") + if citations: + yield "cohere.citations_count", len(citations) # Finish reason if "finish_reason" in response: @@ -95,9 +100,9 @@ def get_attributes( return # Token count from meta - meta = response.get("meta", {}) + meta = response.get("meta") or {} if "billed_units" in meta: - billed = meta["billed_units"] + billed = meta["billed_units"] or {} if "input_tokens" in billed: yield SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS, billed["input_tokens"] @@ -111,17 +116,21 @@ def get_extra_attributes( if response is None: return - # Embeddings - embeddings = response.get("embeddings", []) - if embeddings: - yield "cohere.embeddings_count", len(embeddings) - # First embedding info - if len(embeddings) > 0 and isinstance(embeddings[0], list): - yield "cohere.embedding_dimensions", len(embeddings[0]) - yield f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.0.{EmbeddingAttributes.EMBEDDING_VECTOR}", safe_json_dumps(embeddings[0][:10]) - - yield SpanAttributes.OUTPUT_VALUE, f"[{len(embeddings)} embeddings]" - yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps({"embeddings_count": len(embeddings)}) + # Typed embeddings (v1 optional, v2 required) group vectors by encoding. + embeddings = response.get("embeddings") or [] + if isinstance(embeddings, dict): + embeddings = next( + (vectors for vectors in embeddings.values() if vectors), [] + ) + count = len(embeddings) + yield "cohere.embeddings_count", count + if count and isinstance(embeddings[0], list): + yield "cohere.embedding_dimensions", len(embeddings[0]) + yield ( + f"{SpanAttributes.EMBEDDING_EMBEDDINGS}.0.{EmbeddingAttributes.EMBEDDING_VECTOR}", + safe_json_dumps(embeddings[0][:10]), + ) + yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps({"embeddings_count": count}) class _RerankResponseAttributesExtractor: @@ -137,9 +146,9 @@ def get_attributes( return # Token count from meta - meta = response.get("meta", {}) + meta = response.get("meta") or {} if "billed_units" in meta: - billed = meta["billed_units"] + billed = meta["billed_units"] or {} if "search_units" in billed: yield "cohere.search_units", billed["search_units"] @@ -169,5 +178,13 @@ def get_extra_attributes( if scores: yield "cohere.rerank.scores", safe_json_dumps(scores) - yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps([{"index": r.get("index"), "score": r.get("relevance_score")} for r in results[:5]]) + yield ( + SpanAttributes.OUTPUT_VALUE, + safe_json_dumps( + [ + {"index": r.get("index"), "score": r.get("relevance_score")} + for r in results[:5] + ] + ), + ) yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps(response) diff --git a/python/frameworks/cohere/traceai_cohere/_stream.py b/python/frameworks/cohere/traceai_cohere/_stream.py new file mode 100644 index 000000000..659598cae --- /dev/null +++ b/python/frameworks/cohere/traceai_cohere/_stream.py @@ -0,0 +1,191 @@ +"""Pass-through stream lifecycles and v1/v2 event accumulation.""" + +import logging +from typing import Any, Dict, Mapping +from weakref import finalize + +from opentelemetry import trace + +from traceai_cohere._response_attributes_extractor import ( + _ChatResponseAttributesExtractor, +) +from traceai_cohere._utils import _finish_tracing, _to_dict +from traceai_cohere._with_span import _WithSpan + +logger = logging.getLogger(__name__) + + +class _Accumulator: + def __init__(self) -> None: + self.response: Dict[str, Any] = {} + self.text = [] + self.tools: Dict[int, Dict[str, Any]] = {} + self.v2 = False + + def add(self, event: Any) -> None: + data = _to_dict(event) + kind = data.get("type") or data.get("event_type") + if kind == "text-generation": + self.text.append(data.get("text") or "") + elif kind == "stream-end": + self.response.update(data.get("response") or {}) + if data.get("finish_reason"): + self.response["finish_reason"] = data["finish_reason"] + elif kind in ( + "message-start", + "content-start", + "content-delta", + "tool-call-start", + "tool-call-delta", + "message-end", + ): + self.v2 = True + delta = data.get("delta") or {} + message = delta.get("message") or {} + if kind == "message-start": + if data.get("id"): + self.response["id"] = data["id"] + elif kind in ("content-start", "content-delta"): + self.text.append((message.get("content") or {}).get("text") or "") + elif kind in ("tool-call-start", "tool-call-delta"): + incoming = message.get("tool_calls") or {} + tool = self.tools.setdefault( + data.get("index") or 0, {"function": {"arguments": ""}} + ) + for key in ("id", "type"): + if incoming.get(key) is not None: + tool[key] = incoming[key] + function = incoming.get("function") or {} + if function.get("name") is not None: + tool["function"]["name"] = function["name"] + tool["function"]["arguments"] += function.get("arguments") or "" + elif kind == "message-end": + self.response.update(delta) + + def result(self) -> Dict[str, Any]: + result = dict(self.response) + if self.v2: + result["message"] = { + "role": "assistant", + "content": [{"type": "text", "text": "".join(self.text)}], + "tool_calls": [self.tools[i] for i in sorted(self.tools)], + } + elif "text" not in result: + result["text"] = "".join(self.text) + return result + + +class _StreamBase: + def __init__( + self, response: Any, span: _WithSpan, request: Mapping[str, Any] + ) -> None: + self._response = response + self._span = span + self._request = request + self._accumulator = _Accumulator() + self._finished = False + # Match generator cleanup when callers abandon a stream. Explicit + # close/aclose is still required for deterministic transport cleanup. + finalize(self, span.finish_tracing) + + def _use_span(self): + # Scope the span to SDK work, never to the caller between yielded events. + return trace.use_span( + self._span._span, + end_on_exit=False, + record_exception=False, + set_status_on_exception=False, + ) + + def _add(self, event: Any) -> Any: + try: + self._accumulator.add(event) + except Exception: + logger.debug("Could not extract Cohere stream event", exc_info=True) + return event + + def _finish(self, status: trace.StatusCode = trace.StatusCode.UNSET) -> None: + if self._finished: + return + self._finished = True + try: + result = self._accumulator.result() + extractor = _ChatResponseAttributesExtractor() + _finish_tracing( + status=trace.Status(status), + with_span=self._span, + attributes=extractor.get_attributes(result, is_streaming=True), + extra_attributes=extractor.get_extra_attributes( + result, self._request, is_streaming=True + ), + ) + except Exception: + logger.debug("Could not finalize Cohere stream", exc_info=True) + finally: + self._span.finish_tracing(status=trace.Status(status)) + + def _error(self, error: BaseException) -> None: + self._span.record_exception(error) + self._finish(trace.StatusCode.ERROR) + + +class _Stream(_StreamBase): + def __iter__(self): + return self + + def __next__(self): + if self._finished: + raise StopIteration + try: + with self._use_span(): + event = next(self._response) + except StopIteration: + self._finish(trace.StatusCode.OK) + raise + except BaseException as error: + self._error(error) + raise + return self._add(event) + + def close(self) -> None: + try: + with self._use_span(): + close = getattr(self._response, "close", None) + if close: + close() + except BaseException as error: + self._error(error) + raise + finally: + self._finish() + + +class _AsyncStream(_StreamBase): + def __aiter__(self): + return self + + async def __anext__(self): + if self._finished: + raise StopAsyncIteration + try: + with self._use_span(): + event = await self._response.__anext__() + except StopAsyncIteration: + self._finish(trace.StatusCode.OK) + raise + except BaseException as error: + self._error(error) + raise + return self._add(event) + + async def aclose(self) -> None: + try: + with self._use_span(): + close = getattr(self._response, "aclose", None) + if close: + await close() + except BaseException as error: + self._error(error) + raise + finally: + self._finish() diff --git a/python/frameworks/cohere/traceai_cohere/_utils.py b/python/frameworks/cohere/traceai_cohere/_utils.py index f19f364c5..c0e3c3765 100644 --- a/python/frameworks/cohere/traceai_cohere/_utils.py +++ b/python/frameworks/cohere/traceai_cohere/_utils.py @@ -2,6 +2,7 @@ from typing import Any, Dict, Iterator, Mapping, Optional, Tuple from fi_instrumentation import safe_json_dumps +from fi_instrumentation.fi_types import MessageAttributes, ToolCallAttributes from opentelemetry import trace as trace_api from opentelemetry.util.types import AttributeValue @@ -62,3 +63,36 @@ def _finish_tracing( if value is not None: with_span._span.set_attribute(key, value) with_span.finish_tracing(status=status) + + +def _message_attributes( + message: Any, prefix: str +) -> Iterator[Tuple[str, AttributeValue]]: + """Use canonical keys so the core masks message content and tool arguments.""" + message = _to_dict(message) + yield f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", message.get("role") or "" + content = message.get("content", message.get("message")) + if content is not None: + if not isinstance(content, str): + content = safe_json_dumps(content) + yield f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}", content + if message.get("tool_call_id"): + yield ( + f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALL_ID}", + message["tool_call_id"], + ) + for index, item in enumerate(message.get("tool_calls") or []): + tool = _to_dict(item) + function = tool.get("function") or {} + key = f"{prefix}.{MessageAttributes.MESSAGE_TOOL_CALLS}.{index}" + if tool.get("id"): + yield f"{key}.{ToolCallAttributes.TOOL_CALL_ID}", tool["id"] + yield ( + f"{key}.{ToolCallAttributes.TOOL_CALL_FUNCTION_NAME}", + function.get("name", tool.get("name", "")), + ) + arguments = function.get("arguments", tool.get("parameters", {})) + yield ( + f"{key}.{ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON}", + (arguments if isinstance(arguments, str) else safe_json_dumps(arguments)), + ) diff --git a/python/frameworks/cohere/traceai_cohere/_with_span.py b/python/frameworks/cohere/traceai_cohere/_with_span.py index dd0d014dd..d470d9de6 100644 --- a/python/frameworks/cohere/traceai_cohere/_with_span.py +++ b/python/frameworks/cohere/traceai_cohere/_with_span.py @@ -18,7 +18,7 @@ def __init__( self._context_attributes = context_attributes or {} self._extra_attributes = extra_attributes or {} - def record_exception(self, exception: Exception) -> None: + def record_exception(self, exception: BaseException) -> None: """Record an exception on the span.""" if self._span.is_recording(): self._span.record_exception(exception) diff --git a/python/frameworks/cohere/traceai_cohere/_wrappers.py b/python/frameworks/cohere/traceai_cohere/_wrappers.py index 31e20a82c..46ca176be 100644 --- a/python/frameworks/cohere/traceai_cohere/_wrappers.py +++ b/python/frameworks/cohere/traceai_cohere/_wrappers.py @@ -1,11 +1,10 @@ import logging from abc import ABC from contextlib import contextmanager -from typing import Any, Callable, Dict, Iterable, Iterator, List, Mapping, Tuple +from typing import Any, Callable, Iterable, Iterator, Mapping, Tuple import opentelemetry.context as context_api -from fi_instrumentation import TraceConfig, get_attributes_from_context, safe_json_dumps -from fi_instrumentation.fi_types import SpanAttributes +from fi_instrumentation import TraceConfig, get_attributes_from_context from opentelemetry import trace as trace_api from opentelemetry.trace import INVALID_SPAN from opentelemetry.util.types import AttributeValue @@ -20,6 +19,7 @@ _EmbedResponseAttributesExtractor, _RerankResponseAttributesExtractor, ) +from traceai_cohere._stream import _AsyncStream, _Stream from traceai_cohere._utils import _finish_tracing, _to_dict from traceai_cohere._with_span import _WithSpan @@ -88,9 +88,13 @@ def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = wrapped(*args, **kwargs) @@ -142,9 +146,13 @@ async def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = await wrapped(*args, **kwargs) @@ -196,9 +204,13 @@ def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = wrapped(*args, **kwargs) @@ -213,63 +225,15 @@ def __call__( return self._handle_streaming_response(span, response, request_parameters) - def _handle_streaming_response( - self, span: _WithSpan, response: Any, request_parameters: Dict[str, Any] - ): - """Handle streaming response.""" - content_parts: List[str] = [] - raw_chunks: List[Dict[str, Any]] = [] - final_response: Dict[str, Any] = {} + def _handle_streaming_response(self, span, response, request_parameters): + return _Stream(response, span, request_parameters) - def streaming_wrapper(): - nonlocal final_response - try: - for event in response: - event_dict = _to_dict(event) - raw_chunks.append(event_dict) - - event_type = event_dict.get("event_type", "") - - if event_type == "text-generation": - text = event_dict.get("text", "") - if text: - content_parts.append(text) - if event_type == "stream-end": - final_response = event_dict.get("response", {}) +class _AsyncChatStreamWrapper(_ChatStreamWrapper): + """Async SDK streams return an async iterator directly, not an awaitable.""" - yield event - - except Exception as exception: - span.record_exception(exception) - status = trace_api.Status( - status_code=trace_api.StatusCode.ERROR, - description=f"{type(exception).__name__}: {exception}", - ) - span.finish_tracing(status=status) - raise - else: - try: - final_response["text"] = "".join(content_parts) - span._span.set_attribute(SpanAttributes.OUTPUT_VALUE, safe_json_dumps(raw_chunks[-5:])) # Last 5 chunks - span._span.set_attribute(SpanAttributes.OUTPUT_VALUE, "".join(content_parts)) - - _finish_tracing( - status=trace_api.Status(status_code=trace_api.StatusCode.OK), - with_span=span, - attributes=self._response_extractor.get_attributes(final_response, is_streaming=True), - extra_attributes=self._response_extractor.get_extra_attributes( - final_response, request_parameters, is_streaming=True - ), - ) - except Exception: - logger.exception("Failed to finalize streaming response") - span.finish_tracing() - finally: - if span._span.is_recording(): - span.finish_tracing() - - return streaming_wrapper() + def _handle_streaming_response(self, span, response, request_parameters): + return _AsyncStream(response, span, request_parameters) class _EmbedWrapper(_WithTracer): @@ -295,9 +259,13 @@ def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = wrapped(*args, **kwargs) @@ -349,9 +317,13 @@ async def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = await wrapped(*args, **kwargs) @@ -403,9 +375,13 @@ def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = wrapped(*args, **kwargs) @@ -457,9 +433,13 @@ async def __call__( with self._start_as_current_span( span_name=span_name, - attributes=self._request_extractor.get_attributes_from_request(request_parameters), + attributes=self._request_extractor.get_attributes_from_request( + request_parameters + ), context_attributes=get_attributes_from_context(), - extra_attributes=self._request_extractor.get_extra_attributes_from_request(request_parameters), + extra_attributes=self._request_extractor.get_extra_attributes_from_request( + request_parameters + ), ) as span: try: response = await wrapped(*args, **kwargs) From b254c97f69028d72c36abf1e16676e85eb8d44fc Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 06:39:58 -0800 Subject: [PATCH 2/7] fix(cohere): mask images before serializing chat inputs Apply the configured core image mask to telemetry-only message copies so hide_input_images and base64 size limits cover both input.value and message.content. Preserve original requests and explicit config priority. Add failing-first real-SDK OTLP regressions for sync/async chat and streams, typed/dict messages, environment options, and direct floor-SDK extraction. Co-Authored-By: Claude Code --- .../cohere/tests/test_sdk_coverage.py | 123 ++++++++++++++++++ .../_request_attributes_extractor.py | 40 +++++- .../cohere/traceai_cohere/_wrappers.py | 6 +- 3 files changed, 164 insertions(+), 5 deletions(-) diff --git a/python/frameworks/cohere/tests/test_sdk_coverage.py b/python/frameworks/cohere/tests/test_sdk_coverage.py index d521ce060..2d75e10a5 100644 --- a/python/frameworks/cohere/tests/test_sdk_coverage.py +++ b/python/frameworks/cohere/tests/test_sdk_coverage.py @@ -498,6 +498,129 @@ async def async_source(): assert len(exporter.get_finished_spans()) == 1 +IMAGE_URL = "https://example.invalid/PRIVATE_IMAGE_346.png" +IMAGE_DATA = "data:image/png;base64,U1lOVEhFVElDX1BSSVZBVEVfSU1BR0VfMzQ2" + + +def multimodal_message(): + return { + "role": "user", + "content": [ + {"type": "text", "text": INPUT}, + {"type": "image_url", "image_url": {"url": IMAGE_URL}}, + {"type": "image_url", "image_url": {"url": IMAGE_DATA}}, + ], + } + + +def assert_image_privacy(exporter, privacy): + _, wire, attrs = wire_attributes(exporter) + # Check both copies and serialized OTLP, not just canonical attribute keys. + for value in ( + attrs["input.value"], + attrs["gen_ai.input.messages.0.message.content"], + wire.decode("utf-8", errors="replace"), + ): + assert (IMAGE_URL in value) == (privacy != "hidden") + assert (IMAGE_DATA in value) == (privacy == "visible") + assert INPUT in value + assert attrs["gen_ai.input.messages.0.message.role"] == "user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("typed", [False, True]) +@pytest.mark.parametrize("configuration", ["explicit", "environment"]) +@pytest.mark.parametrize("privacy", ["visible", "hidden", "bounded"]) +async def test_multimodal_image_privacy_on_wire( + tracing, monkeypatch, asynchronous, streaming, typed, configuration, privacy +): + import copy + + cls = client_class(2, asynchronous) + if typed and not hasattr(cohere, "ImageUrlContent"): + pytest.skip("This SDK predates typed multimodal chat messages") + hidden = privacy == "hidden" + limit = 16 if privacy == "bounded" else 32000 + monkeypatch.setenv( + "FI_HIDE_INPUT_IMAGES", + str(hidden if configuration == "environment" else not hidden).lower(), + ) + monkeypatch.setenv( + "FI_BASE64_IMAGE_MAX_LENGTH", + str( + limit if configuration == "environment" else (32000 if limit == 16 else 16) + ), + ) + _, exporter = tracing( + **( + {"hide_input_images": hidden, "base64_image_max_length": limit} + if configuration == "explicit" + else {} + ) + ) + message = multimodal_message() + if typed: + message = cohere.UserChatMessageV2(**message) + original = copy.deepcopy(message) + requests = [] + + def respond(request): + requests.append(json.loads(request.content)) + if streaming: + return httpx.Response( + 200, + text="".join("data: " + json.dumps(e) + "\n\n" for e in events(2)), + headers={"content-type": "text/event-stream"}, + ) + return httpx.Response(200, json=payload(2, "chat")) + + args = {"model": "command-a-vision-07-2025", "messages": [message]} + if asynchronous: + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="offline-placeholder", httpx_client=http) + if streaming: + response = [e async for e in client.chat_stream(**args)] + else: + response = await client.chat(**args) + else: + with httpx.Client(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="offline-placeholder", httpx_client=http) + response = ( + list(client.chat_stream(**args)) if streaming else client.chat(**args) + ) + assert response is not None + assert message == original + assert len(requests) == 1 + content = requests[0]["messages"][0]["content"] + assert content[0]["text"] == INPUT + assert content[1]["image_url"]["url"] == IMAGE_URL + assert content[2]["image_url"]["url"] == IMAGE_DATA + assert_image_privacy(exporter, privacy) + + +@pytest.mark.parametrize("privacy", ["visible", "hidden", "bounded"]) +def test_multimodal_extractor_privacy_at_sdk_floor(tracing, privacy): + from traceai_cohere._request_attributes_extractor import ( + _ChatRequestAttributesExtractor, + ) + + options = { + "hide_input_images": privacy == "hidden", + "base64_image_max_length": 16 if privacy == "bounded" else 32000, + } + _, exporter = tracing(**options) + extractor = _ChatRequestAttributesExtractor(config=TraceConfig(**options)) + message = multimodal_message() + with CohereInstrumentor()._tracer.start_as_current_span("image-privacy") as span: + span.set_attributes( + dict(extractor.get_extra_attributes_from_request({"messages": [message]})) + ) + assert message == multimodal_message() + assert_image_privacy(exporter, privacy) + + def test_unknown_event_extraction_does_not_break_stream(tracing): from traceai_cohere._wrappers import _ChatStreamWrapper diff --git a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py index c98faabb4..b5966bd70 100644 --- a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py @@ -5,12 +5,15 @@ from fi_instrumentation.fi_types import ( FiMimeTypeValues, FiSpanKindValues, + ImageAttributes, + MessageAttributes, + MessageContentAttributes, RerankerAttributes, SpanAttributes, ) from opentelemetry.util.types import AttributeValue -from traceai_cohere._utils import _message_attributes +from traceai_cohere._utils import _message_attributes, _to_dict logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -19,6 +22,36 @@ class _ChatRequestAttributesExtractor: """Extract span attributes from Cohere chat request.""" + def __init__(self, config: Optional[TraceConfig] = None) -> None: + self._config = config if config is not None else TraceConfig() + + def _mask_message_images(self, message: Any, index: int) -> Dict[str, Any]: + # Both input.value and message.content serialize entire content blocks. + # Apply the core's image mask before either copy loses its semantic keys. + # Copy only telemetry data: never redact the request sent to Cohere. + message = dict(_to_dict(message)) + content = message.get("content") + if not isinstance(content, (list, tuple)): + return message + masked_content = [] + for block_index, item in enumerate(content): + block = _to_dict(item) + if block.get("type") == "image_url": + image = _to_dict(block.get("image_url")) + key = ( + f"{SpanAttributes.GEN_AI_INPUT_MESSAGES}.{index}." + f"{MessageAttributes.MESSAGE_CONTENTS}.{block_index}." + f"{MessageContentAttributes.MESSAGE_CONTENT_IMAGE}." + f"{ImageAttributes.IMAGE_URL}" + ) + url = self._config.mask(key, image.get("url")) + if url is None: + continue + block = {**block, "image_url": {**image, "url": url}} + masked_content.append(block) + message["content"] = masked_content + return message + def get_attributes_from_request( self, request_parameters: Mapping[str, Any] ) -> Iterator[Tuple[str, AttributeValue]]: @@ -58,7 +91,10 @@ def get_extra_attributes_from_request( yield SpanAttributes.INPUT_VALUE, message if "messages" in request_parameters: - messages = request_parameters.get("messages") or [] + messages = [ + self._mask_message_images(message, i) + for i, message in enumerate(request_parameters.get("messages") or []) + ] yield SpanAttributes.INPUT_VALUE, safe_json_dumps(messages) for i, msg in enumerate(messages): yield from _message_attributes( diff --git a/python/frameworks/cohere/traceai_cohere/_wrappers.py b/python/frameworks/cohere/traceai_cohere/_wrappers.py index 46ca176be..56b7617c8 100644 --- a/python/frameworks/cohere/traceai_cohere/_wrappers.py +++ b/python/frameworks/cohere/traceai_cohere/_wrappers.py @@ -69,7 +69,7 @@ class _ChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) self._response_extractor = _ChatResponseAttributesExtractor() def __call__( @@ -128,7 +128,7 @@ class _AsyncChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) self._response_extractor = _ChatResponseAttributesExtractor() async def __call__( @@ -186,7 +186,7 @@ class _ChatStreamWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) self._response_extractor = _ChatResponseAttributesExtractor() def __call__( From 58b4577f8de400099213b813890ed781ec8f14a0 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 07:25:20 -0800 Subject: [PATCH 3/7] fix(cohere): enforce content privacy and finalize failed calls Honor selective content flags on raw telemetry copies, mask native stream error details with the effective config, finish cancelled calls, and classify terminal provider failures. Require tested core 1.2.2 and exercise SDK/core floors and latest distributions in compatibility CI. Co-Authored-By: Claude Code --- .github/workflows/cohere-python-tests.yml | 71 ++++ python/frameworks/cohere/pyproject.toml | 2 +- .../cohere/tests/test_review_regressions.py | 379 ++++++++++++++++++ .../_request_attributes_extractor.py | 8 +- .../_response_attributes_extractor.py | 27 +- .../cohere/traceai_cohere/_stream.py | 23 +- .../cohere/traceai_cohere/_wrappers.py | 24 +- 7 files changed, 514 insertions(+), 20 deletions(-) create mode 100644 .github/workflows/cohere-python-tests.yml create mode 100644 python/frameworks/cohere/tests/test_review_regressions.py diff --git a/.github/workflows/cohere-python-tests.yml b/.github/workflows/cohere-python-tests.yml new file mode 100644 index 000000000..c5a2213cb --- /dev/null +++ b/.github/workflows/cohere-python-tests.yml @@ -0,0 +1,71 @@ +name: Cohere Python contracts + +on: + pull_request: + branches: [main, dev] + paths: + - 'python/frameworks/cohere/**' + - '.github/workflows/cohere-python-tests.yml' + push: + branches: [dev] + paths: + - 'python/frameworks/cohere/**' + - '.github/workflows/cohere-python-tests.yml' + +permissions: + contents: read + +jobs: + contract: + name: Cohere (${{ matrix.label }}) + runs-on: ubuntu-24.04 + timeout-minutes: 15 + strategy: + fail-fast: false + matrix: + include: + - label: SDK and core floors + python: '3.10' + cohere: 'cohere==5.0.0' + core: 'fi-instrumentation-otel==1.2.2' + httpx: 'httpx==0.27.2' + - label: early v2 + python: '3.12' + cohere: 'cohere==5.13.12' + core: 'fi-instrumentation-otel==1.2.2' + httpx: 'httpx' + - label: latest SDK and core + python: '3.13' + cohere: 'cohere' + core: 'fi-instrumentation-otel' + httpx: 'httpx' + env: + FI_BASE_URL: http://127.0.0.1:9 + FI_GRPC_URL: http://127.0.0.1:9 + COHERE_REQUIREMENT: ${{ matrix.cohere }} + CORE_REQUIREMENT: ${{ matrix.core }} + HTTPX_REQUIREMENT: ${{ matrix.httpx }} + steps: + - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ matrix.python }} + - name: Install distribution and compatibility dependencies + shell: bash + run: | + python -m venv .cohere-test-venv + .cohere-test-venv/bin/pip install ./python/frameworks/cohere \ + "$COHERE_REQUIREMENT" "$CORE_REQUIREMENT" "$HTTPX_REQUIREMENT" \ + pytest pytest-asyncio opentelemetry-exporter-otlp-proto-http + .cohere-test-venv/bin/pip check + - name: Test with only loopback networking + shell: bash + run: | + sudo --preserve-env=FI_BASE_URL,FI_GRPC_URL unshare --net bash -c ' + set -e + ip link set lo up + "$1" -I -m pytest python/frameworks/cohere/tests \ + --import-mode=importlib -q -p no:cacheprovider + ' _ "$PWD/.cohere-test-venv/bin/python" diff --git a/python/frameworks/cohere/pyproject.toml b/python/frameworks/cohere/pyproject.toml index 5e857a798..2bd25b77a 100644 --- a/python/frameworks/cohere/pyproject.toml +++ b/python/frameworks/cohere/pyproject.toml @@ -27,7 +27,7 @@ classifiers = [ [tool.poetry.dependencies] python = ">=3.10,<3.14" -fi-instrumentation-otel = ">=0.1.11" +fi-instrumentation-otel = ">=1.2.2" cohere = ">=5.0.0" [tool.poetry.urls] diff --git a/python/frameworks/cohere/tests/test_review_regressions.py b/python/frameworks/cohere/tests/test_review_regressions.py new file mode 100644 index 000000000..e15aad4ad --- /dev/null +++ b/python/frameworks/cohere/tests/test_review_regressions.py @@ -0,0 +1,379 @@ +"""Privacy and lifecycle contracts using real SDK classes and synthetic HTTP.""" + +import asyncio +import json +from importlib.metadata import requires + +import cohere +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from packaging.requirements import Requirement + +from traceai_cohere import CohereInstrumentor + +INPUT = "SYNTHETIC_PRIVATE_INPUT_346" +OUTPUT = "SYNTHETIC_PRIVATE_OUTPUT_346" +ERROR = "SYNTHETIC_PRIVATE_ERROR_346" + + +class Started(SpanProcessor): + def __init__(self): + self.spans = [] + + def on_start(self, span, parent_context=None): + self.spans.append(span) + + +@pytest.fixture +def tracing(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + started = Started() + provider.add_span_processor(started) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + + def setup(**options): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**options)) + return provider, exporter, started + + yield setup + instrumentor.uninstrument() + provider.shutdown() + + +def client_class(version, asynchronous=False): + name = ("AsyncClient" if asynchronous else "Client") + ( + "V2" if version == 2 else "" + ) + if not hasattr(cohere, name): + pytest.skip("v2 not available at SDK floor") + return getattr(cohere, name) + + +def request_args(version, operation="chat"): + if operation == "embed": + return dict( + model="embed-v4.0", + texts=[INPUT], + input_type="search_document", + embedding_types=["float"], + ) + if operation == "rerank": + return dict(model="rerank-v3.5", query=INPUT, documents=["synthetic"]) + return dict( + model="command-r", + **( + {"message": INPUT} + if version == 1 + else {"messages": [{"role": "user", "content": INPUT}]} + ), + ) + + +def payload(version): + if version == 1: + return { + "text": OUTPUT, + "generation_id": "00000000-0000-0000-0000-000000000001", + "finish_reason": "COMPLETE", + # v1 returns an echoed copy of the input, not just output text. + "chat_history": [ + {"role": "USER", "message": INPUT}, + {"role": "CHATBOT", "message": OUTPUT}, + ], + } + return { + "id": "synthetic", + "finish_reason": "COMPLETE", + "message": { + "role": "assistant", + "content": [{"type": "text", "text": OUTPUT}], + }, + } + + +def event_response(version, finish="COMPLETE", error=None): + if version == 1: + response = payload(version) + response["finish_reason"] = finish + chunks = [ + {"event_type": "text-generation", "text": OUTPUT}, + { + "event_type": "stream-end", + "finish_reason": finish, + "response": response, + }, + ] + else: + delta = {"finish_reason": finish} + if error is not None: + delta["error"] = error + chunks = [ + { + "type": "message-start", + "id": "synthetic", + "delta": {"message": {"role": "assistant"}}, + }, + { + "type": "content-delta", + "index": 0, + "delta": {"message": {"content": {"text": OUTPUT}}}, + }, + {"type": "message-end", "delta": delta}, + ] + return httpx.Response( + 200, + text="".join( + ("data: " if version == 2 else "") + + json.dumps(event) + + ("\n\n" if version == 2 else "\n") + for event in chunks + ), + headers={"content-type": "text/event-stream"}, + ) + + +async def call_chat(version, asynchronous, streaming, transport): + cls = client_class(version, asynchronous) + if asynchronous: + async with httpx.AsyncClient(transport=transport) as http: + client = cls(api_key="synthetic-only", httpx_client=http) + if streaming: + return [ + event async for event in client.chat_stream(**request_args(version)) + ] + return await client.chat(**request_args(version)) + with httpx.Client(transport=transport) as http: + client = cls(api_key="synthetic-only", httpx_client=http) + if streaming: + return list(client.chat_stream(**request_args(version))) + return client.chat(**request_args(version)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize("configuration", ["explicit", "environment"]) +@pytest.mark.parametrize( + "flag", + [ + "hide_input_messages", + "hide_input_text", + "hide_output_messages", + "hide_output_text", + ], +) +async def test_selective_content_flags_do_not_leak_raw_copies( + tracing, monkeypatch, version, asynchronous, streaming, hidden, configuration, flag +): + client_class(version, asynchronous) + # Exercise both explicit overrides, not just a true environment setting. + monkeypatch.setenv( + "FI_" + flag.upper(), + str(not hidden if configuration == "explicit" else hidden).lower(), + ) + _, exporter, _ = tracing(**({flag: hidden} if configuration == "explicit" else {})) + requests = [] + + def respond(request): + requests.append(json.loads(request.content)) + return ( + event_response(version) + if streaming + else httpx.Response(200, json=payload(version)) + ) + + response = await call_chat( + version, asynchronous, streaming, httpx.MockTransport(respond) + ) + assert len(requests) == 1 + assert INPUT in json.dumps(requests[0]) + final = response[-1] if streaming else response + if streaming: + # Caller-visible events must retain the same native text and history. + assert OUTPUT in json.dumps([event.dict() for event in response]) + if version == 1: + assert final.response.chat_history[0].message == INPUT + elif version == 1: + assert final.text == OUTPUT + assert final.chat_history[0].message == INPUT + else: + assert final.message.content[0].text == OUTPUT + spans = exporter.get_finished_spans() + assert len(spans) == 1 + wire = encode_spans(spans).SerializeToString() + sentinel = INPUT if "input" in flag else OUTPUT + other = OUTPUT if "input" in flag else INPUT + assert other.encode() in wire # Preserve capture for the other direction. + assert (sentinel.encode() not in wire) == hidden + assert spans[0].attributes["cohere.finish_reason"] == "COMPLETE" + assert spans[0].status.status_code == trace.StatusCode.OK + + +@pytest.mark.asyncio +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize("operation", ["chat", "embed", "rerank"]) +@pytest.mark.parametrize("hidden", [False, True]) +async def test_cancelled_unary_request_ends_span(tracing, version, operation, hidden): + cls = client_class(version, True) + provider, exporter, started = tracing(hide_error_details=hidden) + entered = asyncio.Event() + cancellations = [] + + async def respond(request): + entered.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError as error: + cancellations.append(error) + raise + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http: + client = cls(api_key="synthetic-only", httpx_client=http) + + async def call(): + try: + return await getattr(client, operation)( + **request_args(version, operation) + ) + except asyncio.CancelledError as error: + # Check inside the task: Python 3.10 can drop the message at + # the awaiting-task boundary, independently of instrumentation. + cancellations.append(error) + raise + + task = asyncio.create_task(call()) + await asyncio.wait_for(entered.wait(), timeout=5) + task.cancel(ERROR) + with pytest.raises(asyncio.CancelledError): + await task + assert len(cancellations) == 2 + assert cancellations[0] is cancellations[1] + assert cancellations[1].args == (ERROR,) + assert len(started.spans) == 1 + provider.force_flush() + ended = exporter.get_finished_spans() + assert len(ended) == 1 + assert not started.spans[0].is_recording() + assert ended[0].status.status_code == trace.StatusCode.ERROR + assert "CancelledError" in ended[0].attributes["error.type"] + assert (ERROR.encode() not in encode_spans(ended).SerializeToString()) == hidden + assert not trace.get_current_span().is_recording() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("close_at_terminal", [False, True]) +@pytest.mark.parametrize( + "version,finish,error,failed", + [ + (1, "ERROR", None, True), + (1, "ERROR_TOXIC", None, True), + (2, "ERROR", None, True), + (2, "TIMEOUT", None, True), + (2, None, ERROR, True), + (1, "COMPLETE", None, False), + (1, "MAX_TOKENS", None, False), + (2, "COMPLETE", None, False), + (2, "MAX_TOKENS", None, False), + (2, "TOOL_CALL", None, False), + ], +) +async def test_native_stream_terminal_failure_status( + tracing, version, asynchronous, close_at_terminal, finish, error, failed +): + cls = client_class(version, asynchronous) + _, exporter, _ = tracing() + transport = httpx.MockTransport(lambda r: event_response(version, finish, error)) + response = [] + if asynchronous: + async with httpx.AsyncClient(transport=transport) as http: + client = cls(api_key="synthetic-only", httpx_client=http) + stream = client.chat_stream(**request_args(version)) + async for event in stream: + response.append(event) + if close_at_terminal and ( + getattr(event, "type", None) == "message-end" + or getattr(event, "event_type", None) == "stream-end" + ): + break + await stream.aclose() + await stream.aclose() + else: + with httpx.Client(transport=transport) as http: + client = cls(api_key="synthetic-only", httpx_client=http) + stream = client.chat_stream(**request_args(version)) + for event in stream: + response.append(event) + if close_at_terminal and ( + getattr(event, "type", None) == "message-end" + or getattr(event, "event_type", None) == "stream-end" + ): + break + stream.close() + stream.close() + assert len(response) == (2 if version == 1 else 3) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + assert spans[0].attributes.get("cohere.finish_reason") == finish + expected = ( + trace.StatusCode.ERROR + if failed + else (trace.StatusCode.UNSET if close_at_terminal else trace.StatusCode.OK) + ) + assert spans[0].status.status_code == expected + if failed: + assert spans[0].attributes["error.type"] == (finish or "_OTHER") + else: + assert "error.type" not in spans[0].attributes + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("configuration", ["explicit", "environment", "derived"]) +@pytest.mark.parametrize("hidden", [False, True]) +async def test_native_stream_error_details_are_private( + tracing, monkeypatch, asynchronous, configuration, hidden +): + client_class(2, asynchronous) + if configuration == "derived": + monkeypatch.delenv("FI_HIDE_ERROR_DETAILS", raising=False) + options = {"hide_input_text": hidden} + else: + monkeypatch.setenv( + "FI_HIDE_ERROR_DETAILS", + str(not hidden if configuration == "explicit" else hidden).lower(), + ) + options = {"hide_error_details": hidden} if configuration == "explicit" else {} + _, exporter, _ = tracing(**options) + response = await call_chat( + 2, + asynchronous, + True, + httpx.MockTransport(lambda r: event_response(2, "ERROR", ERROR)), + ) + assert response[-1].delta.error == ERROR + spans = exporter.get_finished_spans() + assert len(spans) == 1 + wire = encode_spans(spans).SerializeToString() + assert (ERROR.encode() not in wire) == hidden + assert OUTPUT.encode() in wire + assert spans[0].status.status_code == trace.StatusCode.ERROR + assert spans[0].attributes["error.type"] == "ERROR" + + +def test_distribution_requires_compatible_core(): + """The installed distribution must exclude the reproduced broken core floor.""" + dependencies = [Requirement(value) for value in requires("traceAI-cohere")] + core = next(dep for dep in dependencies if dep.name == "fi-instrumentation-otel") + assert not core.specifier.contains("0.1.11") + assert core.specifier.contains("1.2.2") diff --git a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py index b5966bd70..76aa82086 100644 --- a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py @@ -88,14 +88,18 @@ def get_extra_attributes_from_request( if "message" in request_parameters: message = request_parameters.get("message", "") messages.append({"role": "user", "message": message}) - yield SpanAttributes.INPUT_VALUE, message + if not (self._config.hide_input_messages or self._config.hide_input_text): + yield SpanAttributes.INPUT_VALUE, message if "messages" in request_parameters: messages = [ self._mask_message_images(message, i) for i, message in enumerate(request_parameters.get("messages") or []) ] - yield SpanAttributes.INPUT_VALUE, safe_json_dumps(messages) + # Raw JSON is a second copy of message content. Omit it under + # selective content privacy; canonical fields retain safe metadata. + if not (self._config.hide_input_messages or self._config.hide_input_text): + yield SpanAttributes.INPUT_VALUE, safe_json_dumps(messages) for i, msg in enumerate(messages): yield from _message_attributes( msg, f"{SpanAttributes.GEN_AI_INPUT_MESSAGES}.{i}" diff --git a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py index 740366e74..bf29659cc 100644 --- a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py @@ -1,7 +1,7 @@ import logging from typing import Any, Dict, Iterator, Mapping, Optional, Tuple -from fi_instrumentation import safe_json_dumps +from fi_instrumentation import TraceConfig, safe_json_dumps from fi_instrumentation.fi_types import ( EmbeddingAttributes, SpanAttributes, @@ -17,6 +17,9 @@ class _ChatResponseAttributesExtractor: """Extract span attributes from Cohere chat response.""" + def __init__(self, config: Optional[TraceConfig] = None) -> None: + self._config = config if config is not None else TraceConfig() + def get_attributes( self, response: Optional[Dict[str, Any]], @@ -83,7 +86,27 @@ def get_extra_attributes( if "finish_reason" in response: yield "cohere.finish_reason", response.get("finish_reason", "") - # Raw response + # Full responses duplicate message content (including history and tool + # payloads). Do not export that copy when selective content is hidden. + if self._config.hide_output_messages or self._config.hide_output_text: + return + # v1 echoes the request in chat_history; input privacy also applies to + # that response-side copy. Keep the actual reply and usage available. + if ( + self._config.hide_inputs + or self._config.hide_input_messages + or self._config.hide_input_text + ): + response = { + key: value for key, value in response.items() if key != "chat_history" + } + # Native message-end errors are data, not raised exceptions. Mask their + # telemetry copy before JSON serialization, never the caller's event. + if "error" in response: + response = dict(response) + error = self._config.mask("error.message", response.pop("error")) + if error is not None: + response["error"] = error yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps(response) diff --git a/python/frameworks/cohere/traceai_cohere/_stream.py b/python/frameworks/cohere/traceai_cohere/_stream.py index 659598cae..346ea6635 100644 --- a/python/frameworks/cohere/traceai_cohere/_stream.py +++ b/python/frameworks/cohere/traceai_cohere/_stream.py @@ -1,9 +1,10 @@ """Pass-through stream lifecycles and v1/v2 event accumulation.""" import logging -from typing import Any, Dict, Mapping +from typing import Any, Dict, Mapping, Optional from weakref import finalize +from fi_instrumentation import TraceConfig from opentelemetry import trace from traceai_cohere._response_attributes_extractor import ( @@ -77,12 +78,17 @@ def result(self) -> Dict[str, Any]: class _StreamBase: def __init__( - self, response: Any, span: _WithSpan, request: Mapping[str, Any] + self, + response: Any, + span: _WithSpan, + request: Mapping[str, Any], + config: Optional[TraceConfig] = None, ) -> None: self._response = response self._span = span self._request = request self._accumulator = _Accumulator() + self._response_extractor = _ChatResponseAttributesExtractor(config=config) self._finished = False # Match generator cleanup when callers abandon a stream. Explicit # close/aclose is still required for deterministic transport cleanup. @@ -110,7 +116,18 @@ def _finish(self, status: trace.StatusCode = trace.StatusCode.UNSET) -> None: self._finished = True try: result = self._accumulator.result() - extractor = _ChatResponseAttributesExtractor() + # Cohere yields native failure events without raising. Exhaustion + # (or close after the terminal event) must not mark those as OK. + finish_reason = result.get("finish_reason") + failure_reason = finish_reason in ("ERROR", "ERROR_TOXIC", "TIMEOUT") + if failure_reason or result.get("error"): + status = trace.StatusCode.ERROR + self._span._span.set_attribute( + "error.type", finish_reason if failure_reason else "_OTHER" + ) + if result.get("error"): + self._span._span.set_attribute("error.message", result["error"]) + extractor = self._response_extractor _finish_tracing( status=trace.Status(status), with_span=self._span, diff --git a/python/frameworks/cohere/traceai_cohere/_wrappers.py b/python/frameworks/cohere/traceai_cohere/_wrappers.py index 56b7617c8..79a2dcf77 100644 --- a/python/frameworks/cohere/traceai_cohere/_wrappers.py +++ b/python/frameworks/cohere/traceai_cohere/_wrappers.py @@ -70,7 +70,7 @@ class _ChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) - self._response_extractor = _ChatResponseAttributesExtractor() + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) def __call__( self, @@ -98,7 +98,7 @@ def __call__( ) as span: try: response = wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -129,7 +129,7 @@ class _AsyncChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) - self._response_extractor = _ChatResponseAttributesExtractor() + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) async def __call__( self, @@ -156,7 +156,7 @@ async def __call__( ) as span: try: response = await wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -187,7 +187,7 @@ class _ChatStreamWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) - self._response_extractor = _ChatResponseAttributesExtractor() + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) def __call__( self, @@ -214,7 +214,7 @@ def __call__( ) as span: try: response = wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -226,14 +226,14 @@ def __call__( return self._handle_streaming_response(span, response, request_parameters) def _handle_streaming_response(self, span, response, request_parameters): - return _Stream(response, span, request_parameters) + return _Stream(response, span, request_parameters, config=self._config) class _AsyncChatStreamWrapper(_ChatStreamWrapper): """Async SDK streams return an async iterator directly, not an awaitable.""" def _handle_streaming_response(self, span, response, request_parameters): - return _AsyncStream(response, span, request_parameters) + return _AsyncStream(response, span, request_parameters, config=self._config) class _EmbedWrapper(_WithTracer): @@ -269,7 +269,7 @@ def __call__( ) as span: try: response = wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -327,7 +327,7 @@ async def __call__( ) as span: try: response = await wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -385,7 +385,7 @@ def __call__( ) as span: try: response = wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -443,7 +443,7 @@ async def __call__( ) as span: try: response = await wrapped(*args, **kwargs) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, From 5e6faa6b3d32560ce95922235a435fed047e8ef3 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 08:07:02 -0800 Subject: [PATCH 4/7] fix(cohere): protect echoed inputs and classify native chat failures Strip document and tool-result citation payloads only from telemetry copies under effective input privacy. Share native response failure classification across unary and streaming chat, including ERROR_LIMIT. Add SDK-to-OTLP regression coverage for configuration precedence, immutable caller data, and terminal status. Co-Authored-By: Claude Code --- .../cohere/tests/test_response_regressions.py | 323 ++++++++++++++++++ .../_response_attributes_extractor.py | 49 ++- .../cohere/traceai_cohere/_stream.py | 12 +- .../cohere/traceai_cohere/_utils.py | 17 + .../cohere/traceai_cohere/_wrappers.py | 6 +- 5 files changed, 389 insertions(+), 18 deletions(-) create mode 100644 python/frameworks/cohere/tests/test_response_regressions.py diff --git a/python/frameworks/cohere/tests/test_response_regressions.py b/python/frameworks/cohere/tests/test_response_regressions.py new file mode 100644 index 000000000..9ede57f3f --- /dev/null +++ b/python/frameworks/cohere/tests/test_response_regressions.py @@ -0,0 +1,323 @@ +"""Response privacy and native failure contracts using real SDKs and OTLP.""" + +from copy import deepcopy +import json + +import cohere +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceai_cohere import CohereInstrumentor + +SECRET = "SYNTHETIC_PRIVATE_SOURCE_346_ROUND2" +OUTPUT = "A safe short answer" + + +@pytest.fixture +def tracing(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + + def setup(**privacy): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**privacy)) + return exporter + + yield setup + instrumentor.uninstrument() + provider.shutdown() + + +def client_cls(version, asynchronous): + name = ("AsyncClient" if asynchronous else "Client") + ( + "V2" if version == 2 else "" + ) + if not hasattr(cohere, name): + pytest.skip("v2 not present at SDK floor") + return getattr(cohere, name) + + +def args(version): + return dict( + model="synthetic-model", + **( + {"message": "Summarize the document"} + if version == 1 + else {"messages": [{"role": "user", "content": "Summarize the document"}]} + ), + ) + + +def response_body(version, finish="COMPLETE"): + if version == 1: + return {"text": OUTPUT, "finish_reason": finish} + return { + "id": "synthetic-response", + "finish_reason": finish, + "message": {"role": "assistant", "content": [{"type": "text", "text": OUTPUT}]}, + } + + +def response(version, body, streaming): + if not streaming: + return httpx.Response(200, json=body) + if version == 1: + events = [ + {"event_type": "text-generation", "text": OUTPUT}, + { + "event_type": "stream-end", + "finish_reason": body["finish_reason"], + "response": body, + }, + ] + else: + events = [ + { + "type": "message-start", + "id": "synthetic-response", + "delta": {"message": {"role": "assistant"}}, + }, + { + "type": "content-delta", + "index": 0, + "delta": {"message": {"content": {"text": OUTPUT}}}, + }, + {"type": "message-end", "delta": {"finish_reason": body["finish_reason"]}}, + ] + text = "".join( + ("data: " if version == 2 else "") + + json.dumps(e) + + ("\n\n" if version == 2 else "\n") + for e in events + ) + return httpx.Response(200, text=text, headers={"content-type": "text/event-stream"}) + + +async def invoke(version, asynchronous, streaming, request, handler): + cls = client_cls(version, asynchronous) + transport = httpx.MockTransport(handler) + if asynchronous: + async with httpx.AsyncClient(transport=transport) as http: + client = cls(api_key="synthetic-no-real-credential", httpx_client=http) + return ( + [e async for e in client.chat_stream(**request)] + if streaming + else await client.chat(**request) + ) + with httpx.Client(transport=transport) as http: + client = cls(api_key="synthetic-no-real-credential", httpx_client=http) + return ( + list(client.chat_stream(**request)) if streaming else client.chat(**request) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("version,streaming", [(1, False), (1, True), (2, False)]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize( + "flag", ["hide_inputs", "hide_input_messages", "hide_input_text"] +) +@pytest.mark.parametrize("configuration", ["explicit", "environment"]) +async def test_input_document_echo_is_hidden( + tracing, monkeypatch, version, asynchronous, streaming, hidden, flag, configuration +): + client_cls(version, asynchronous) + monkeypatch.setenv( + "FI_" + flag.upper(), + str(hidden if configuration == "environment" else not hidden).lower(), + ) + exporter = tracing(**({flag: hidden} if configuration == "explicit" else {})) + request = args(version) + document = {"id": "doc-1", "text": SECRET} + request["documents"] = [ + document if version == 1 else cohere.Document(id="doc-1", data={"text": SECRET}) + ] + body = response_body(version) + if version == 1: + body["documents"] = [document] + else: + body["message"]["citations"] = [ + { + "start": 0, + "end": len(OUTPUT), + "text": OUTPUT, + "sources": [{"type": "document", "id": "doc-1", "document": document}], + } + ] + seen = [] + + def handle(r): + seen.append(json.loads(r.content)) + return response(version, body, streaming) + + result = await invoke(version, asynchronous, streaming, request, handle) + assert len(seen) == 1 and SECRET in json.dumps(seen[0]) + assert SECRET in str(result) # SDK-visible data must remain unchanged. + spans = exporter.get_finished_spans() + assert len(spans) == 1 + wire = encode_spans(spans).SerializeToString() + assert OUTPUT.encode() in wire + leaking = [k for k, v in spans[0].attributes.items() if SECRET in str(v)] + assert (SECRET.encode() not in wire) == hidden, ( + f"hide_inputs={hidden}; echoed input leaked via {leaking}" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize( + "flag", ["hide_inputs", "hide_input_messages", "hide_input_text"] +) +@pytest.mark.parametrize("configuration", ["explicit", "environment"]) +async def test_tool_result_citation_echo_is_hidden( + tracing, monkeypatch, asynchronous, hidden, flag, configuration +): + client_cls(2, asynchronous) + monkeypatch.setenv( + "FI_" + flag.upper(), + str(hidden if configuration == "environment" else not hidden).lower(), + ) + exporter = tracing(**({flag: hidden} if configuration == "explicit" else {})) + request = args(2) + request["messages"] = [ + {"role": "user", "content": "Look up the record"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "tool-1", + "type": "function", + "function": {"name": "lookup", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "tool-1", + "content": json.dumps({"text": SECRET}), + }, + ] + body = response_body(2) + body["message"]["citations"] = [ + { + "start": 0, + "end": len(OUTPUT), + "text": OUTPUT, + "sources": [ + {"type": "tool", "id": "tool-1", "tool_output": {"text": SECRET}} + ], + } + ] + seen = [] + + def handle(r): + seen.append(json.loads(r.content)) + return response(2, body, False) + + result = await invoke(2, asynchronous, False, request, handle) + assert len(seen) == 1 and SECRET in json.dumps(seen[0]) + assert SECRET in str(result) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + wire = encode_spans(spans).SerializeToString() + leaking = [k for k, v in spans[0].attributes.items() if SECRET in str(v)] + assert (SECRET.encode() not in wire) == hidden, ( + f"hide_inputs={hidden}; input tool result leaked via {leaking}" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize( + "version,finish,failed", + [ + (1, "ERROR", True), + (1, "ERROR_TOXIC", True), + (1, "ERROR_LIMIT", True), + (2, "ERROR", True), + (2, "TIMEOUT", True), + (1, "COMPLETE", False), + (2, "COMPLETE", False), + (1, "MAX_TOKENS", False), + (2, "MAX_TOKENS", False), + (2, "STOP_SEQUENCE", False), + (2, "TOOL_CALL", False), + ], +) +async def test_native_failure_status( + tracing, version, asynchronous, streaming, hidden, finish, failed +): + client_cls(version, asynchronous) + exporter = tracing(hide_error_details=hidden) + body = response_body(version, finish) + result = await invoke( + version, + asynchronous, + streaming, + args(version), + lambda r: response(version, body, streaming), + ) + terminal = result[-1] if streaming else result + if streaming and version == 2: + terminal = terminal.delta + assert terminal.finish_reason == finish + spans = exporter.get_finished_spans() + assert len(spans) == 1 + expected = trace.StatusCode.ERROR if failed else trace.StatusCode.OK + assert spans[0].status.status_code == expected, ( + f"{finish} => {spans[0].status.status_code}; error.type={spans[0].attributes.get('error.type')}" + ) + assert spans[0].attributes.get("error.type") == (finish if failed else None) + assert spans[0].attributes["cohere.finish_reason"] == finish + assert not spans[0].events # Native failures do not raise SDK exceptions. + assert not trace.get_current_span().is_recording() + + +@pytest.mark.parametrize("version", [1, 2]) +@pytest.mark.parametrize( + "flag", ["hide_inputs", "hide_input_messages", "hide_input_text"] +) +def test_input_echo_sanitization_never_mutates_response(version, flag): + from traceai_cohere._response_attributes_extractor import ( + _ChatResponseAttributesExtractor, + ) + + body = response_body(version) + body["documents"] = [{"text": SECRET}] + body["chat_history"] = [{"role": "USER", "message": SECRET}] + container = body if version == 1 else body["message"] + container["citations"] = [ + { + "start": 0, + "end": len(OUTPUT), + "text": OUTPUT, + "sources": [ + {"type": "document", "id": "doc-1", "document": {"text": SECRET}}, + {"type": "tool", "id": "tool-1", "tool_output": {"text": SECRET}}, + ], + } + ] + original = deepcopy(body) + extractor = _ChatResponseAttributesExtractor(TraceConfig(**{flag: True})) + attributes = dict(extractor.get_extra_attributes(body, {})) + assert body == original + assert SECRET not in json.dumps(attributes) + output = json.loads(attributes["output.value"]) + assert "documents" not in output and "chat_history" not in output + citation = (output if version == 1 else output["message"])["citations"][0] + assert citation["text"] == OUTPUT + assert citation["sources"] == [ + {"type": "document", "id": "doc-1"}, + {"type": "tool", "id": "tool-1"}, + ] + assert attributes["cohere.citations_count"] == 1 diff --git a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py index bf29659cc..7ed782689 100644 --- a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py @@ -14,6 +14,46 @@ logger.addHandler(logging.NullHandler()) +def _without_citation_inputs(citations: Any) -> Any: + """Remove source payloads, retaining generated citation text and references.""" + if not isinstance(citations, list): + return citations + result = [] + for citation in citations: + if isinstance(citation, dict) and isinstance(citation.get("sources"), list): + citation = dict(citation) + citation["sources"] = [ + { + key: value + for key, value in source.items() + if key not in ("document", "tool_output") + } + if isinstance(source, dict) + else source + for source in citation["sources"] + ] + result.append(citation) + return result + + +def _without_input_echoes(response: Dict[str, Any]) -> Dict[str, Any]: + """Copy only telemetry paths containing known v1/v2 input echoes.""" + result = { + key: value + for key, value in response.items() + if key not in ("chat_history", "documents") + } + if "citations" in result: + result["citations"] = _without_citation_inputs(result["citations"]) + message = result.get("message") + if isinstance(message, dict) and "citations" in message: + result["message"] = { + **message, + "citations": _without_citation_inputs(message["citations"]), + } + return result + + class _ChatResponseAttributesExtractor: """Extract span attributes from Cohere chat response.""" @@ -90,16 +130,15 @@ def get_extra_attributes( # payloads). Do not export that copy when selective content is hidden. if self._config.hide_output_messages or self._config.hide_output_text: return - # v1 echoes the request in chat_history; input privacy also applies to - # that response-side copy. Keep the actual reply and usage available. + # v1 echoes history/documents; v2 citation sources echo documents and + # tool results. Input privacy applies to those structured copies, not + # to the generated answer. Never mutate the caller's response. if ( self._config.hide_inputs or self._config.hide_input_messages or self._config.hide_input_text ): - response = { - key: value for key, value in response.items() if key != "chat_history" - } + response = _without_input_echoes(response) # Native message-end errors are data, not raised exceptions. Mask their # telemetry copy before JSON serialization, never the caller's event. if "error" in response: diff --git a/python/frameworks/cohere/traceai_cohere/_stream.py b/python/frameworks/cohere/traceai_cohere/_stream.py index 346ea6635..fa580d564 100644 --- a/python/frameworks/cohere/traceai_cohere/_stream.py +++ b/python/frameworks/cohere/traceai_cohere/_stream.py @@ -10,7 +10,7 @@ from traceai_cohere._response_attributes_extractor import ( _ChatResponseAttributesExtractor, ) -from traceai_cohere._utils import _finish_tracing, _to_dict +from traceai_cohere._utils import _chat_response_status, _finish_tracing, _to_dict from traceai_cohere._with_span import _WithSpan logger = logging.getLogger(__name__) @@ -118,15 +118,7 @@ def _finish(self, status: trace.StatusCode = trace.StatusCode.UNSET) -> None: result = self._accumulator.result() # Cohere yields native failure events without raising. Exhaustion # (or close after the terminal event) must not mark those as OK. - finish_reason = result.get("finish_reason") - failure_reason = finish_reason in ("ERROR", "ERROR_TOXIC", "TIMEOUT") - if failure_reason or result.get("error"): - status = trace.StatusCode.ERROR - self._span._span.set_attribute( - "error.type", finish_reason if failure_reason else "_OTHER" - ) - if result.get("error"): - self._span._span.set_attribute("error.message", result["error"]) + status = _chat_response_status(result, self._span, default=status) extractor = self._response_extractor _finish_tracing( status=trace.Status(status), diff --git a/python/frameworks/cohere/traceai_cohere/_utils.py b/python/frameworks/cohere/traceai_cohere/_utils.py index c0e3c3765..b207ae3b7 100644 --- a/python/frameworks/cohere/traceai_cohere/_utils.py +++ b/python/frameworks/cohere/traceai_cohere/_utils.py @@ -47,6 +47,23 @@ def _flatten( yield full_key, value +def _chat_response_status( + response: Mapping[str, Any], + with_span: _WithSpan, + default: trace_api.StatusCode = trace_api.StatusCode.OK, +) -> trace_api.StatusCode: + """Classify native generation failures shared by unary and stream calls.""" + reason = response.get("finish_reason") + failed = reason in ("ERROR", "ERROR_TOXIC", "ERROR_LIMIT", "TIMEOUT") + if not (failed or response.get("error")): + return default + with_span._span.set_attribute("error.type", reason if failed else "_OTHER") + if response.get("error"): + # FITracer applies the effective error-detail privacy policy. + with_span._span.set_attribute("error.message", response["error"]) + return trace_api.StatusCode.ERROR + + def _finish_tracing( status: trace_api.Status, with_span: _WithSpan, diff --git a/python/frameworks/cohere/traceai_cohere/_wrappers.py b/python/frameworks/cohere/traceai_cohere/_wrappers.py index 79a2dcf77..4dd35cdf7 100644 --- a/python/frameworks/cohere/traceai_cohere/_wrappers.py +++ b/python/frameworks/cohere/traceai_cohere/_wrappers.py @@ -20,7 +20,7 @@ _RerankResponseAttributesExtractor, ) from traceai_cohere._stream import _AsyncStream, _Stream -from traceai_cohere._utils import _finish_tracing, _to_dict +from traceai_cohere._utils import _chat_response_status, _finish_tracing, _to_dict from traceai_cohere._with_span import _WithSpan logger = logging.getLogger(__name__) @@ -110,7 +110,7 @@ def __call__( try: response_dict = _to_dict(response) _finish_tracing( - status=trace_api.Status(status_code=trace_api.StatusCode.OK), + status=trace_api.Status(_chat_response_status(response_dict, span)), with_span=span, attributes=self._response_extractor.get_attributes(response_dict), extra_attributes=self._response_extractor.get_extra_attributes( @@ -168,7 +168,7 @@ async def __call__( try: response_dict = _to_dict(response) _finish_tracing( - status=trace_api.Status(status_code=trace_api.StatusCode.OK), + status=trace_api.Status(_chat_response_status(response_dict, span)), with_span=span, attributes=self._response_extractor.get_attributes(response_dict), extra_attributes=self._response_extractor.get_extra_attributes( From bf0e5644c475cebc21b10a3eaf50658552084e8e Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:05:41 -0800 Subject: [PATCH 5/7] fix(cohere): redact rerank document echoes under input privacy Co-Authored-By: Claude Code --- .../cohere/tests/test_rerank_input_echo.py | 106 ++++++++++++++++++ .../_response_attributes_extractor.py | 22 ++-- .../cohere/traceai_cohere/_wrappers.py | 8 +- 3 files changed, 125 insertions(+), 11 deletions(-) create mode 100644 python/frameworks/cohere/tests/test_rerank_input_echo.py diff --git a/python/frameworks/cohere/tests/test_rerank_input_echo.py b/python/frameworks/cohere/tests/test_rerank_input_echo.py new file mode 100644 index 000000000..a5de30848 --- /dev/null +++ b/python/frameworks/cohere/tests/test_rerank_input_echo.py @@ -0,0 +1,106 @@ +"""Independent round-3 SDK and OTLP probes. Synthetic data and transport only.""" + +import copy +import inspect +import json + +import cohere +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceai_cohere import CohereInstrumentor + +SECRET = "SYNTHETIC_PRIVATE_DOCUMENT_ROUND3_346" +API_KEY = "synthetic-only-not-a-real-api-key-346" + + +@pytest.fixture +def tracing(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + + def setup(**privacy): + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(**privacy)) + return provider, exporter + + yield setup + instrumentor.uninstrument() + provider.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("configuration", ["explicit", "environment"]) +@pytest.mark.parametrize("hidden", [False, True]) +@pytest.mark.parametrize("version", [1, 2]) +async def test_rerank_return_documents_respects_input_privacy( + tracing, monkeypatch, asynchronous, configuration, hidden, version +): + name = ("AsyncClient" if asynchronous else "Client") + ( + "V2" if version == 2 else "" + ) + if not hasattr(cohere, name): + pytest.skip("v2 not available at SDK floor") + cls = getattr(cohere, name) + if "return_documents" not in inspect.signature(cls.rerank).parameters: + pytest.skip("this SDK removed v2 return_documents") + # This is an SDK-declared input echo, not generated output text. + monkeypatch.setenv( + "FI_HIDE_INPUTS", + str(hidden if configuration == "environment" else not hidden).lower(), + ) + provider, exporter = tracing( + **({"hide_inputs": hidden} if configuration == "explicit" else {}) + ) + document = {"text": SECRET, "title": "synthetic document"} + original = copy.deepcopy(document) + body = { + "id": "synthetic", + "results": [{"index": 0, "relevance_score": 0.9, "document": document}], + "meta": {"billed_units": {"search_units": 1}}, + } + seen = [] + + def handle(request): + seen.append(json.loads(request.content)) + assert request.headers["authorization"] == "Bearer " + API_KEY + return httpx.Response(200, json=body) + + request = dict( + model="rerank-v3.5", + query="Find the relevant document", + documents=[document if version == 1 else SECRET], + return_documents=True, + ) + with provider.get_tracer("review").start_as_current_span("parent") as parent: + if asynchronous: + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as http: + client = cls(api_key=API_KEY, httpx_client=http) + result = await client.rerank(**request) + else: + with httpx.Client(transport=httpx.MockTransport(handle)) as http: + client = cls(api_key=API_KEY, httpx_client=http) + result = client.rerank(**request) + assert trace.get_current_span() is parent + assert len(seen) == 1 and seen[0]["documents"] == request["documents"] + assert result.results[0].document.text == SECRET + assert document == original + spans = exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.parent.span_id == parent.get_span_context().span_id + assert span.attributes["cohere.results_count"] == 1 + assert span.attributes["cohere.rerank.0.relevance_score"] == 0.9 + wire = encode_spans(spans).SerializeToString() + assert API_KEY.encode() not in wire + leaked = [key for key, value in span.attributes.items() if SECRET in str(value)] + assert (SECRET.encode() not in wire) == hidden, ( + f"hide_inputs={hidden}; copied input leaked via {leaked}" + ) diff --git a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py index 7ed782689..b197c2b08 100644 --- a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py @@ -198,6 +198,9 @@ def get_extra_attributes( class _RerankResponseAttributesExtractor: """Extract span attributes from Cohere rerank response.""" + def __init__(self, config: Optional[TraceConfig] = None) -> None: + self._config = config if config is not None else TraceConfig() + def get_attributes( self, response: Optional[Dict[str, Any]], @@ -240,13 +243,14 @@ def get_extra_attributes( if scores: yield "cohere.rerank.scores", safe_json_dumps(scores) - yield ( - SpanAttributes.OUTPUT_VALUE, - safe_json_dumps( - [ - {"index": r.get("index"), "score": r.get("relevance_score")} - for r in results[:5] - ] - ), - ) + # return_documents echoes the submitted input. Preserve the ranking + # metadata in telemetry without changing the SDK response for callers. + if self._config.hide_inputs: + response = dict(response) + response["results"] = [ + {key: value for key, value in result.items() if key != "document"} + if isinstance(result, dict) + else result + for result in results + ] yield SpanAttributes.OUTPUT_VALUE, safe_json_dumps(response) diff --git a/python/frameworks/cohere/traceai_cohere/_wrappers.py b/python/frameworks/cohere/traceai_cohere/_wrappers.py index 4dd35cdf7..b66144def 100644 --- a/python/frameworks/cohere/traceai_cohere/_wrappers.py +++ b/python/frameworks/cohere/traceai_cohere/_wrappers.py @@ -358,7 +358,9 @@ class _RerankWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._request_extractor = _RerankRequestAttributesExtractor(config=self._config) - self._response_extractor = _RerankResponseAttributesExtractor() + self._response_extractor = _RerankResponseAttributesExtractor( + config=self._config + ) def __call__( self, @@ -416,7 +418,9 @@ class _AsyncRerankWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._request_extractor = _RerankRequestAttributesExtractor(config=self._config) - self._response_extractor = _RerankResponseAttributesExtractor() + self._response_extractor = _RerankResponseAttributesExtractor( + config=self._config + ) async def __call__( self, From b70238f2f494659dd76edcea35877cc922dc49a6 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:29:18 -0800 Subject: [PATCH 6/7] fix(cohere): count multimodal embedding inputs safely Co-Authored-By: Claude Code --- .../cohere/tests/test_multimodal_embed.py | 62 +++++++++++++++++++ .../_request_attributes_extractor.py | 39 ++++++++++-- 2 files changed, 97 insertions(+), 4 deletions(-) create mode 100644 python/frameworks/cohere/tests/test_multimodal_embed.py diff --git a/python/frameworks/cohere/tests/test_multimodal_embed.py b/python/frameworks/cohere/tests/test_multimodal_embed.py new file mode 100644 index 000000000..3e0fb17c6 --- /dev/null +++ b/python/frameworks/cohere/tests/test_multimodal_embed.py @@ -0,0 +1,62 @@ +"""Offline real-SDK coverage for embedding modality counts and private images.""" +import json + +import cohere +import httpx +import pytest +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +from traceai_cohere import CohereInstrumentor + + +IMAGE = "data:image/png;base64,c3ludGhldGljLWltYWdl" + + +@pytest.mark.parametrize("mode", ["images", "mixed"]) +def test_embed_multimodal_sdk_counts_without_image_data(mode): + if not hasattr(cohere, "ClientV2"): + pytest.skip("Cohere 5.0.0 predates v2 images") + if mode == "mixed": + import inspect + if "inputs" not in inspect.signature(cohere.ClientV2.embed).parameters: + pytest.skip("This SDK predates mixed embedding inputs") + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + instrumentor.instrument(tracer_provider=provider) + sent = [] + + def respond(request): + sent.append(json.loads(request.content)) + return httpx.Response(200, json={ + "id": "synthetic", "response_type": "embeddings_by_type", + "embeddings": {"float": [[0.1, 0.2]]}, + }) + + try: + kwargs = {"images": [IMAGE]} if mode == "images" else { + "inputs": [{"content": [ + {"type": "text", "text": "synthetic caption"}, + {"type": "image_url", "image_url": {"url": IMAGE}}, + ]}], + } + with httpx.Client(transport=httpx.MockTransport(respond)) as http: + client = cohere.ClientV2(api_key="synthetic-only", httpx_client=http) + response = client.embed(model="embed-v4.0", input_type="search_document", embedding_types=["float"], **kwargs) + assert response.embeddings.float == [[0.1, 0.2]] + assert len(sent) == 1 and ("images" if mode == "images" else "inputs") in sent[0] + attrs = exporter.get_finished_spans()[0].attributes + assert attrs["cohere.images_count"] == 1 + assert attrs["cohere.inputs_count"] == 1 + if mode == "images": + assert "cohere.texts_count" not in attrs + assert attrs.get("input.value") != "[]" + else: + assert attrs["cohere.texts_count"] == 1 + assert IMAGE not in json.dumps(dict(attrs)) + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py index 76aa82086..f60b458dc 100644 --- a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py @@ -167,11 +167,42 @@ def get_extra_attributes_from_request( self, request_parameters: Mapping[str, Any] ) -> Iterator[Tuple[str, AttributeValue]]: """Extract extra attributes.""" - texts = request_parameters.get("texts", []) - if isinstance(texts, list): + texts = request_parameters.get("texts") + images = request_parameters.get("images") + inputs = request_parameters.get("inputs") + + # Cohere v2 accepts image-only requests and, in newer SDKs, inputs + # containing both text and image components. Never serialize input + # objects here: they may contain private image data URIs. + plain_text_count = len(texts) if isinstance(texts, (list, tuple)) else 0 + plain_image_count = len(images) if isinstance(images, (list, tuple)) else 0 + mixed_count = len(inputs) if isinstance(inputs, (list, tuple)) else 0 + text_count = plain_text_count + image_count = plain_image_count + for item in inputs if isinstance(inputs, (list, tuple)) else (): + content = _to_dict(item).get("content") + for block in content if isinstance(content, (list, tuple)) else (): + modality = _to_dict(block).get("type") + if modality == "text": + text_count += 1 + elif modality == "image_url": + image_count += 1 + + if texts is not None or text_count: + yield "cohere.texts_count", text_count + if images is not None or image_count: + yield "cohere.images_count", image_count + if images is not None or inputs is not None: + # Mixed items each produce one embedding even with two components. + input_count = plain_text_count + plain_image_count + mixed_count + yield "cohere.inputs_count", input_count + yield SpanAttributes.INPUT_VALUE, safe_json_dumps( + {"texts_count": text_count, "images_count": image_count, + "inputs_count": input_count} + ) + elif isinstance(texts, (list, tuple)): yield SpanAttributes.INPUT_VALUE, safe_json_dumps(texts[:5]) # Limit - yield "cohere.texts_count", len(texts) - else: + elif texts is not None: yield SpanAttributes.INPUT_VALUE, str(texts) From af20be025f0985b7986df9fbc71bc89854ca9fb9 Mon Sep 17 00:00:00 2001 From: Nikhil Pareek Date: Mon, 2 Nov 2026 13:56:13 -0800 Subject: [PATCH 7/7] fix(cohere): retain streamed v2 citations with input redaction Co-Authored-By: Claude Code --- .../cohere/tests/test_stream_citations.py | 111 ++++++++++++++++++ .../cohere/traceai_cohere/_stream.py | 8 ++ 2 files changed, 119 insertions(+) create mode 100644 python/frameworks/cohere/tests/test_stream_citations.py diff --git a/python/frameworks/cohere/tests/test_stream_citations.py b/python/frameworks/cohere/tests/test_stream_citations.py new file mode 100644 index 000000000..c2a98f333 --- /dev/null +++ b/python/frameworks/cohere/tests/test_stream_citations.py @@ -0,0 +1,111 @@ +"""Citation-start parity and privacy, exercised through the real v2 SDK.""" + +import json + +import cohere +import httpx +import pytest +from fi_instrumentation import TraceConfig +from opentelemetry.exporter.otlp.proto.common.trace_encoder import encode_spans +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from traceai_cohere import CohereInstrumentor + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source_kind", ["document", "tool"]) +@pytest.mark.parametrize("citation_count", [1, 2]) +@pytest.mark.parametrize( + "hide_flag", [None, "hide_inputs", "hide_input_messages", "hide_input_text"] +) +async def test_v2_stream_citations( + asynchronous, source_kind, citation_count, hide_flag +): + client_name = "AsyncClientV2" if asynchronous else "ClientV2" + if not hasattr(cohere, client_name): + pytest.skip("v2 unavailable at the declared SDK floor") + secret = "SYNTHETIC_PRIVATE_CITATION_SOURCE_346" + source = {"type": source_kind, "id": "source-1"} + source["document" if source_kind == "document" else "tool_output"] = { + "text": secret + } + citation = {"start": 0, "end": 5, "text": "hello", "sources": [source]} + events = [ + { + "type": "message-start", + "id": "synthetic", + "delta": {"message": {"role": "assistant"}}, + }, + { + "type": "content-delta", + "index": 0, + "delta": {"message": {"content": {"text": "hello"}}}, + }, + ] + for index in range(citation_count): + events.extend( + [ + { + "type": "citation-start", + "index": index, + "delta": {"message": {"citations": citation}}, + }, + {"type": "citation-end", "index": index}, + ] + ) + events.append({"type": "message-end", "delta": {"finish_reason": "COMPLETE"}}) + original_events = json.loads(json.dumps(events)) + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = CohereInstrumentor() + instrumentor.instrument( + tracer_provider=provider, + config=TraceConfig(**({hide_flag: True} if hide_flag else {})), + ) + + def handle(request): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join("data: " + json.dumps(event) + "\n\n" for event in events), + ) + + try: + request = { + "model": "synthetic-model", + "messages": [{"role": "user", "content": "hello"}], + } + if asynchronous: + async with httpx.AsyncClient(transport=httpx.MockTransport(handle)) as http: + client = cohere.AsyncClientV2( + api_key="synthetic-only", httpx_client=http + ) + returned = [event async for event in client.chat_stream(**request)] + else: + with httpx.Client(transport=httpx.MockTransport(handle)) as http: + client = cohere.ClientV2(api_key="synthetic-only", httpx_client=http) + returned = list(client.chat_stream(**request)) + assert events == original_events + returned_citations = [ + event for event in returned if event.type == "citation-start" + ] + assert len(returned_citations) == citation_count + assert all(secret in str(event) for event in returned_citations) + spans = exporter.get_finished_spans() + assert len(spans) == 1 + attributes = spans[0].attributes + assert attributes["cohere.citations_count"] == citation_count + output = json.loads(attributes["output.value"]) + assert len(output["message"]["citations"]) == citation_count + assert all( + item["sources"][0]["id"] == "source-1" + for item in output["message"]["citations"] + ) + wire = encode_spans(spans).SerializeToString() + assert (secret.encode() not in wire) == bool(hide_flag) + finally: + instrumentor.uninstrument() + provider.shutdown() diff --git a/python/frameworks/cohere/traceai_cohere/_stream.py b/python/frameworks/cohere/traceai_cohere/_stream.py index fa580d564..e9b87a20a 100644 --- a/python/frameworks/cohere/traceai_cohere/_stream.py +++ b/python/frameworks/cohere/traceai_cohere/_stream.py @@ -21,6 +21,7 @@ def __init__(self) -> None: self.response: Dict[str, Any] = {} self.text = [] self.tools: Dict[int, Dict[str, Any]] = {} + self.citations = [] self.v2 = False def add(self, event: Any) -> None: @@ -38,6 +39,7 @@ def add(self, event: Any) -> None: "content-delta", "tool-call-start", "tool-call-delta", + "citation-start", "message-end", ): self.v2 = True @@ -48,6 +50,10 @@ def add(self, event: Any) -> None: self.response["id"] = data["id"] elif kind in ("content-start", "content-delta"): self.text.append((message.get("content") or {}).get("text") or "") + elif kind == "citation-start": + citation = message.get("citations") + if isinstance(citation, dict): + self.citations.append(citation) elif kind in ("tool-call-start", "tool-call-delta"): incoming = message.get("tool_calls") or {} tool = self.tools.setdefault( @@ -71,6 +77,8 @@ def result(self) -> Dict[str, Any]: "content": [{"type": "text", "text": "".join(self.text)}], "tool_calls": [self.tools[i] for i in sorted(self.tools)], } + if self.citations: + result["message"]["citations"] = self.citations elif "text" not in result: result["text"] = "".join(self.text) return result