fix tests

Author: alexander-alderman-webb

sentry_sdk/integrations/litellm.py

@@ -18,11 +18,6 @@
 except ImportError:
     raise DidNotEnable("LiteLLM not installed")
 
-_CALL_TYPE_TO_OPERATION = {
-    "embedding": "embeddings",
-    "completion": "chat",
-}
-
 
 def _get_metadata_dict(kwargs):
     # type: (Dict[str, Any]) -> Dict[str, Any]
@@ -54,10 +49,10 @@ def _input_callback(kwargs):
         provider = "unknown"
 
     call_type = kwargs.get("call_type", None)
-    if call_type not in _CALL_TYPE_TO_OPERATION:
-        return
-
-    operation = _CALL_TYPE_TO_OPERATION[call_type]
+    if call_type == "embedding":
+        operation = "embeddings"
+    else:
+        operation = "chat"
 
     # Start a new span/transaction
     span = get_start_span_function()(

tests/integrations/litellm/test_litellm.py

@@ -208,14 +208,15 @@ def test_embeddings_create(sentry_init, capture_events):
     )
     events = capture_events()
 
+    messages = [{"role": "user", "content": "Some text to test embeddings"}]
     mock_response = MockEmbeddingResponse()
 
     with start_transaction(name="litellm test"):
-        # For embeddings, messages would be empty
         kwargs = {
             "model": "text-embedding-ada-002",
             "input": "Hello!",
-            "messages": [],  # Empty for embeddings
+            "messages": messages,
+            "call_type": "embeddings",
         }
 
         _input_callback(kwargs)