feat(cohere): upgrade integration from ai to gen_ai#5597
feat(cohere): upgrade integration from ai to gen_ai#5597
Conversation
Semver Impact of This PR🟡 Minor (new features) 📋 Changelog PreviewThis is how your changes will appear in the changelog. New Features ✨
Bug Fixes 🐛
Documentation 📚
Internal Changes 🔧Openai Agents
Other
🤖 This preview updates automatically when you update the PR. |
Codecov Results 📊✅ 5 passed | Total: 5 | Pass Rate: 100% | Execution Time: 1.12s All tests are passing successfully. ❌ Patch coverage is 9.09%. Project has 15703 uncovered lines. Files with missing lines (2)
Generated by Codecov Action |
![sentry[bot]](https://avatars.githubusercontent.com/in/12637?s=60&v=4){kind=small}
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| from sentry_sdk.ai.monitoring import record_token_usage | ||
| from sentry_sdk.consts import SPANDATA | ||
| from sentry_sdk.ai.utils import set_data_normalized | ||
| from sentry_sdk.consts import OP, SPANDATA |
There was a problem hiding this comment.
Unused consts import after direct OP import added
Low Severity
The from sentry_sdk import consts import is now dead code. The old code accessed consts.OP.COHERE_CHAT_COMPLETIONS_CREATE and consts.OP.COHERE_EMBEDDINGS_CREATE, but those were replaced with OP.GEN_AI_CHAT and OP.GEN_AI_EMBEDDINGS using the new direct import from sentry_sdk.consts import OP, SPANDATA on line 6. There are no remaining consts. references in the file.
| if integration is None: | ||
| return f(*args, **kwargs) | ||
|
|
||
| model = kwargs.get("model", "") |
There was a problem hiding this comment.
the documentation currently does not specify the AI Client Span's description when there is no request model.
I'll bring this up in the sync, update the devdocs, and ping you once there's a concrete spec.
https://develop.sentry.dev/sdk/telemetry/traces/modules/ai-agents/
sentry_sdk/integrations/cohere.py
Outdated
| messages = [] | ||
| for x in kwargs.get("chat_history", []): | ||
| role = getattr(x, "role", "").lower() | ||
| if role == "chatbot": | ||
| role = "assistant" | ||
| messages.append( | ||
| { | ||
| "role": role, | ||
| "content": getattr(x, "message", ""), | ||
| } | ||
| ) | ||
| messages.append({"role": "user", "content": message}) | ||
| messages = normalize_message_roles(messages) |
There was a problem hiding this comment.
Either you should explicitly convert message roles in cohere inputs or rely on solely on normalize_message_roles().
Also, we are developing against a library with specific schemas and types.
Heuristics like normalize_message_roles should be avoided as they effectively add dead code.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
Autofix Details
Bugbot Autofix prepared fixes for both issues found in the latest run.
- ✅ Fixed: System/operation attributes missing on error spans
- Moved GEN_AI_SYSTEM and GEN_AI_OPERATION_NAME set_data calls before the try block so they are always present on spans, even when exceptions occur.
- ✅ Fixed: Cohere-specific role added to global mapping utility
- Removed 'chatbot' from the global GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING and handled the chatbot->assistant role conversion locally in the Cohere integration.
Or push these changes by commenting:
@cursor push 5089f11b7f
Preview (5089f11b7f)
diff --git a/sentry_sdk/ai/utils.py b/sentry_sdk/ai/utils.py
--- a/sentry_sdk/ai/utils.py
+++ b/sentry_sdk/ai/utils.py
@@ -30,7 +30,7 @@
GEN_AI_MESSAGE_ROLE_REVERSE_MAPPING = {
GEN_AI_ALLOWED_MESSAGE_ROLES.SYSTEM: ["system"],
GEN_AI_ALLOWED_MESSAGE_ROLES.USER: ["user", "human"],
- GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai", "chatbot"],
+ GEN_AI_ALLOWED_MESSAGE_ROLES.ASSISTANT: ["assistant", "ai"],
GEN_AI_ALLOWED_MESSAGE_ROLES.TOOL: ["tool", "tool_call"],
}
diff --git a/sentry_sdk/integrations/cohere.py b/sentry_sdk/integrations/cohere.py
--- a/sentry_sdk/integrations/cohere.py
+++ b/sentry_sdk/integrations/cohere.py
@@ -141,6 +141,8 @@
origin=CohereIntegration.origin,
)
span.__enter__()
+ span.set_data(SPANDATA.GEN_AI_SYSTEM, "cohere")
+ span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
try:
res = f(*args, **kwargs)
except Exception as e:
@@ -151,15 +153,16 @@
reraise(*exc_info)
with capture_internal_exceptions():
- span.set_data(SPANDATA.GEN_AI_SYSTEM, "cohere")
- span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
if should_send_default_pii() and integration.include_prompts:
messages = []
for x in kwargs.get("chat_history", []):
+ role = getattr(x, "role", "")
+ if role == "chatbot":
+ role = "assistant"
messages.append(
{
- "role": getattr(x, "role", ""),
+ "role": role,
"content": getattr(x, "message", ""),
}
)|
|
||
| with capture_internal_exceptions(): | ||
| span.set_data(SPANDATA.GEN_AI_SYSTEM, "cohere") | ||
| span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat") |
There was a problem hiding this comment.
System/operation attributes missing on error spans
Medium Severity
GEN_AI_SYSTEM and GEN_AI_OPERATION_NAME are set after the try/except block inside capture_internal_exceptions(). When f(*args, **kwargs) throws, the except block exits the span and reraises, so these attributes are never set on error spans. Other integrations like Google GenAI set these attributes before the try block, ensuring they appear even on failed spans. This means Cohere error spans won't be identifiable by system or operation name.
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
| "is_search_required": SPANDATA.AI_SEARCH_REQUIRED, | ||
| "finish_reason": SPANDATA.AI_FINISH_REASON, | ||
| "generation_id": SPANDATA.GEN_AI_RESPONSE_ID, | ||
| "finish_reason": SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, |
There was a problem hiding this comment.
Single finish_reason stored for plural finish_reasons attribute
Medium Severity
Cohere's finish_reason response attribute is a single string (e.g. "COMPLETE"), but it's mapped to SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS (plural), which semantically expects a list. The Google GenAI integration consistently stores this attribute as a list of strings via extract_finish_reasons(). The value needs to be wrapped in a list (e.g., [getattr(res, attr)]) to match the plural attribute's expected schema and maintain cross-integration consistency.
Additional Locations (1)
| origin=CohereIntegration.origin, | ||
| ) as span: | ||
| set_data_normalized(span, SPANDATA.GEN_AI_SYSTEM, "cohere") | ||
| set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") |
There was a problem hiding this comment.
Embed uses set_data_normalized instead of span.set_data
Low Severity
The embed function uses set_data_normalized for GEN_AI_SYSTEM and GEN_AI_OPERATION_NAME, while the chat function at lines 162-163 uses span.set_data directly for the same attributes with identical simple string values. The PR reviewer explicitly suggested using span.set_data here, and the chat function was updated accordingly, but the embed function wasn't.
2 participants
Closes TET-2023