fix(litellm): trace decision media and systemone calls

py/src/braintrust/integrations/auto_test_scripts/test_auto_litellm.py

@@ -61,4 +61,22 @@
             instrumentation="litellm-auto",
         )
 
+    with autoinstrument_test_context("test_litellm_systemone_metrics[async]", integration="litellm") as memory_logger:
+        response = asyncio.run(
+            litellm.asystemone(
+                f"openai/{DECISIONS_MODEL}",
+                DECISIONS_INPUT,
+                {"positive": {"type": "noul", "instructions": "Is the sentiment positive?"}},
+            )
+        )
+        spans = memory_logger.pop()
+        assert len(spans) == 1
+        span = spans[0]
+        assert span["span_attributes"]["type"] == "decision"
+        assert span["input"]["state"] == DECISIONS_INPUT
+        assert span["output"]["answers"] == [
+            {"id": "positive", "type": "noul", "noul": response.answers["positive"].noul}
+        ]
+        assert span["context"]["span_origin"]["instrumentation"]["name"] == "litellm-auto"
+
 print("SUCCESS")

py/src/braintrust/integrations/litellm/patchers.py

@@ -81,6 +81,18 @@ class LiteLLMAdecisionsPatcher(FunctionWrapperPatcher):
     wrapper = _adecisions_wrapper_async
 
 
+class LiteLLMSystemOnePatcher(FunctionWrapperPatcher):
+    name = "litellm.systemone"
+    target_path = "systemone"
+    wrapper = _decisions_wrapper
+
+
+class LiteLLMASystemOnePatcher(FunctionWrapperPatcher):
+    name = "litellm.asystemone"
+    target_path = "asystemone"
+    wrapper = _adecisions_wrapper_async
+
+
 class LiteLLMImageGenerationPatcher(FunctionWrapperPatcher):
     name = "litellm.image_generation"
     target_path = "image_generation"
@@ -166,6 +178,8 @@ class LiteLLMArerankPatcher(FunctionWrapperPatcher):
     LiteLLMAresponsesPatcher,
     LiteLLMDecisionsPatcher,
     LiteLLMAdecisionsPatcher,
+    LiteLLMSystemOnePatcher,
+    LiteLLMASystemOnePatcher,
     LiteLLMImageGenerationPatcher,
     LiteLLMAimageGenerationPatcher,
     LiteLLMEmbeddingPatcher,