format

Author: shellmayr

sentry_sdk/integrations/cohere/__init__.py

@@ -46,7 +46,9 @@ def _normalize_embedding_input(texts):
     "system": "cohere",
     "operation": "embeddings",
     "params": {"model": SPANDATA.GEN_AI_REQUEST_MODEL},
-    "extract_messages": lambda kw: _normalize_embedding_input(kw["texts"]) if "texts" in kw else None,
+    "extract_messages": lambda kw: (
+        _normalize_embedding_input(kw["texts"]) if "texts" in kw else None
+    ),
     "message_target": SPANDATA.GEN_AI_EMBEDDINGS_INPUT,
     "truncation_fn": None,
 }

sentry_sdk/integrations/cohere/v1.py

@@ -100,7 +100,6 @@ def _wrap_chat(f, streaming):
     if not _has_chat_types:
         return f
 
-
     @wraps(f)
     def new_chat(*args, **kwargs):
         # type: (*Any, **Any) -> Any
@@ -129,10 +128,15 @@ def new_chat(*args, **kwargs):
                 reraise(*exc_info)
 
             with capture_internal_exceptions():
-                set_input_span_data(span, kwargs, integration, {
-                    **COHERE_V1_CHAT_CONFIG,
-                    "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming},
-                })
+                set_input_span_data(
+                    span,
+                    kwargs,
+                    integration,
+                    {
+                        **COHERE_V1_CHAT_CONFIG,
+                        "extra_static": {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming},
+                    },
+                )
 
                 if streaming:
                     old_iterator = res
@@ -194,4 +198,5 @@ def collect_chat_response_fields(span, res, include_pii):
                     input_tokens=res.meta.tokens.input_tokens,
                     output_tokens=res.meta.tokens.output_tokens,
                 )
+
     return new_chat

sentry_sdk/integrations/cohere/v2.py

@@ -189,10 +189,15 @@ def new_chat(*args, **kwargs):
                 extra = {SPANDATA.GEN_AI_RESPONSE_STREAMING: streaming}
                 if model:
                     extra[SPANDATA.GEN_AI_RESPONSE_MODEL] = model
-                set_input_span_data(span, kwargs, integration, {
-                    **COHERE_V2_CHAT_CONFIG,
-                    "extra_static": extra,
-                })
+                set_input_span_data(
+                    span,
+                    kwargs,
+                    integration,
+                    {
+                        **COHERE_V2_CHAT_CONFIG,
+                        "extra_static": extra,
+                    },
+                )
 
                 if streaming:
                     old_iterator = res