|
| 1 | +"""Test auto_instrument for Cohere.""" |
| 2 | + |
| 3 | +import os |
| 4 | + |
| 5 | + |
| 6 | +os.environ.setdefault("CO_API_KEY", "co-test-dummy-api-key-for-vcr-tests") |
| 7 | +os.environ.setdefault("COHERE_API_KEY", os.environ["CO_API_KEY"]) |
| 8 | + |
| 9 | +import cohere |
| 10 | +from braintrust.auto import auto_instrument |
| 11 | +from braintrust.integrations.test_utils import autoinstrument_test_context |
| 12 | + |
| 13 | + |
| 14 | +results = auto_instrument() |
| 15 | +assert results.get("cohere") is True |
| 16 | + |
| 17 | +results2 = auto_instrument() |
| 18 | +assert results2.get("cohere") is True |
| 19 | + |
| 20 | + |
| 21 | +with autoinstrument_test_context("test_auto_cohere", integration="cohere") as memory_logger: |
| 22 | + use_v2 = hasattr(cohere, "ClientV2") and hasattr(cohere.ClientV2, "chat") |
| 23 | + if use_v2: |
| 24 | + client = cohere.ClientV2(api_key=os.environ["CO_API_KEY"]) |
| 25 | + response = client.chat( |
| 26 | + model="command-a-03-2025", |
| 27 | + messages=[{"role": "user", "content": "Say hi in one word."}], |
| 28 | + max_tokens=10, |
| 29 | + ) |
| 30 | + assert response.message.role == "assistant" |
| 31 | + else: |
| 32 | + client = cohere.Client(api_key=os.environ["CO_API_KEY"]) |
| 33 | + response = client.chat( |
| 34 | + model="command-a-03-2025", |
| 35 | + message="Say hi in one word.", |
| 36 | + max_tokens=10, |
| 37 | + ) |
| 38 | + assert isinstance(response.text, str) |
| 39 | + |
| 40 | + spans = memory_logger.pop() |
| 41 | + assert len(spans) == 1, f"Expected 1 span, got {len(spans)}" |
| 42 | + span = spans[0] |
| 43 | + assert span["metadata"]["provider"] == "cohere" |
| 44 | + assert span["metadata"]["model"] == "command-a-03-2025" |
| 45 | + assert span["span_attributes"]["name"] == "cohere.chat" |
| 46 | + |
| 47 | +print("SUCCESS") |
0 commit comments