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/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/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_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/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/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/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/tests/test_sdk_coverage.py b/python/frameworks/cohere/tests/test_sdk_coverage.py new file mode 100644 index 000000000..2d75e10a5 --- /dev/null +++ b/python/frameworks/cohere/tests/test_sdk_coverage.py @@ -0,0 +1,638 @@ +"""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 + + +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 + + _, 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/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/__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..f60b458dc 100644 --- a/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_request_attributes_extractor.py @@ -1,16 +1,20 @@ 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, + ImageAttributes, MessageAttributes, + MessageContentAttributes, RerankerAttributes, SpanAttributes, ) from opentelemetry.util.types import AttributeValue +from traceai_cohere._utils import _message_attributes, _to_dict + logger = logging.getLogger(__name__) logger.addHandler(logging.NullHandler()) @@ -18,13 +22,42 @@ 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]]: """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 +68,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 +81,38 @@ 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 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 []) + ] + # 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 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 +121,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 +145,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,17 +158,51 @@ 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] ) -> 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) @@ -130,7 +218,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 +240,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..b197c2b08 100644 --- a/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py +++ b/python/frameworks/cohere/traceai_cohere/_response_attributes_extractor.py @@ -1,21 +1,65 @@ 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, - MessageAttributes, SpanAttributes, ) from opentelemetry.util.types import AttributeValue +from traceai_cohere._utils import _message_attributes + logger = logging.getLogger(__name__) 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.""" + 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]], @@ -25,25 +69,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,30 +97,55 @@ 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: 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 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 = _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: + 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) @@ -95,9 +162,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,22 +178,29 @@ 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: """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]], @@ -137,9 +211,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 +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/_stream.py b/python/frameworks/cohere/traceai_cohere/_stream.py new file mode 100644 index 000000000..e9b87a20a --- /dev/null +++ b/python/frameworks/cohere/traceai_cohere/_stream.py @@ -0,0 +1,208 @@ +"""Pass-through stream lifecycles and v1/v2 event accumulation.""" + +import logging +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 ( + _ChatResponseAttributesExtractor, +) +from traceai_cohere._utils import _chat_response_status, _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.citations = [] + 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", + "citation-start", + "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 == "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( + 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)], + } + if self.citations: + result["message"]["citations"] = self.citations + 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], + 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. + 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() + # Cohere yields native failure events without raising. Exhaustion + # (or close after the terminal event) must not mark those as OK. + status = _chat_response_status(result, self._span, default=status) + extractor = self._response_extractor + _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..b207ae3b7 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 @@ -46,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, @@ -62,3 +80,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..b66144def 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,7 +19,8 @@ _EmbedResponseAttributesExtractor, _RerankResponseAttributesExtractor, ) -from traceai_cohere._utils import _finish_tracing, _to_dict +from traceai_cohere._stream import _AsyncStream, _Stream +from traceai_cohere._utils import _chat_response_status, _finish_tracing, _to_dict from traceai_cohere._with_span import _WithSpan logger = logging.getLogger(__name__) @@ -69,8 +69,8 @@ class _ChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() - self._response_extractor = _ChatResponseAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) def __call__( self, @@ -88,13 +88,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -106,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( @@ -124,8 +128,8 @@ class _AsyncChatWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() - self._response_extractor = _ChatResponseAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) async def __call__( self, @@ -142,13 +146,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -160,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( @@ -178,8 +186,8 @@ class _ChatStreamWrapper(_WithTracer): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) - self._request_extractor = _ChatRequestAttributesExtractor() - self._response_extractor = _ChatResponseAttributesExtractor() + self._request_extractor = _ChatRequestAttributesExtractor(config=self._config) + self._response_extractor = _ChatResponseAttributesExtractor(config=self._config) def __call__( self, @@ -196,13 +204,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -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 streaming_wrapper(): - nonlocal final_response - try: - for event in response: - event_dict = _to_dict(event) - raw_chunks.append(event_dict) + def _handle_streaming_response(self, span, response, request_parameters): + return _Stream(response, span, request_parameters, config=self._config) - event_type = event_dict.get("event_type", "") - if event_type == "text-generation": - text = event_dict.get("text", "") - if text: - content_parts.append(text) +class _AsyncChatStreamWrapper(_ChatStreamWrapper): + """Async SDK streams return an async iterator directly, not an awaitable.""" - if event_type == "stream-end": - final_response = event_dict.get("response", {}) - - 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, config=self._config) class _EmbedWrapper(_WithTracer): @@ -295,13 +259,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -349,13 +317,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -386,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, @@ -403,13 +377,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR, @@ -440,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, @@ -457,13 +437,17 @@ 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) - except Exception as exception: + except BaseException as exception: span.record_exception(exception) status = trace_api.Status( status_code=trace_api.StatusCode.ERROR,