Skip to content

Commit def8293

Browse files
authored
feat(integrations): trace openai and litellm decisions API (#927)
OpenAI `client.decisions.create()` and LiteLLM `decisions()` / `adecisions()` previously emitted no spans. These calls now emit `decision` spans for sync and async requests, using TypeSafe's `{state, questions}` input and `{answers}` output structure. Closes #921
1 parent 4117e44 commit def8293

26 files changed

Lines changed: 1665 additions & 17 deletions

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,18 @@
11
"""Test auto_instrument for LiteLLM."""
22

3+
import asyncio
4+
35
import litellm
46
from braintrust.auto import auto_instrument
57
from braintrust.integrations.litellm import LiteLLMIntegration
8+
from braintrust.integrations.test_decisions import (
9+
DECISIONS_INPUT,
10+
DECISIONS_MODEL,
11+
DECISIONS_QUESTIONS,
12+
assert_decisions_span,
13+
)
614
from braintrust.integrations.test_utils import autoinstrument_test_context
15+
from braintrust.integrations.versioning import detect_module_version, version_satisfies
716

817

918
# 1. Verify not patched initially
@@ -34,4 +43,22 @@
3443
span = spans[0]
3544
assert span["metadata"]["provider"] == "openai"
3645

46+
if version_satisfies(detect_module_version(litellm, ("litellm",)), ">=1.105.0"):
47+
with autoinstrument_test_context("test_litellm_decisions_metrics[async]", integration="litellm") as memory_logger:
48+
response = asyncio.run(
49+
litellm.adecisions(
50+
model=f"openai/{DECISIONS_MODEL}",
51+
input=DECISIONS_INPUT,
52+
questions=DECISIONS_QUESTIONS,
53+
safety_identifier="do-not-log-this-identifier",
54+
)
55+
)
56+
assert_decisions_span(
57+
memory_logger.pop(),
58+
response,
59+
name="Decision",
60+
model=f"openai/{DECISIONS_MODEL}",
61+
instrumentation="litellm-auto",
62+
)
63+
3764
print("SUCCESS")

‎py/src/braintrust/integrations/auto_test_scripts/test_auto_openai.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,19 @@
11
"""Test auto_instrument for OpenAI."""
22

3+
import asyncio
34
import inspect
45

56
import openai
67
from braintrust import wrap_openai
78
from braintrust.auto import auto_instrument
9+
from braintrust.integrations.test_decisions import (
10+
DECISIONS_INPUT,
11+
DECISIONS_MODEL,
12+
DECISIONS_QUESTIONS,
13+
assert_decisions_span,
14+
)
815
from braintrust.integrations.test_utils import autoinstrument_test_context
16+
from braintrust.integrations.versioning import version_satisfies
917
from wrapt import FunctionWrapper
1018

1119

@@ -15,6 +23,10 @@ def _is_braintrust_wrapped(resource) -> bool:
1523

1624

1725
resources = [openai.resources.chat.completions.Completions]
26+
if version_satisfies(openai.__version__, ">=3.26.0"):
27+
from openai.resources.decisions import AsyncDecisions, Decisions
28+
29+
resources.extend((Decisions, AsyncDecisions))
1830
try:
1931
from openai.resources.beta.agents.sessions.sessions import Sessions
2032

@@ -66,4 +78,23 @@ def _is_braintrust_wrapped(resource) -> bool:
6678
assert spans[0]["metrics"]["prompt_tokens"] == response.usage.input_tokens
6779
assert spans[0]["metrics"]["completion_tokens"] == response.usage.output_tokens
6880

81+
if version_satisfies(openai.__version__, ">=3.26.0"):
82+
with autoinstrument_test_context("test_openai_decisions_metrics[async]", integration="openai") as memory_logger:
83+
client = wrap_openai(wrap_openai(openai.AsyncOpenAI()))
84+
response = asyncio.run(
85+
client.decisions.create(
86+
model=DECISIONS_MODEL,
87+
input=DECISIONS_INPUT,
88+
questions=DECISIONS_QUESTIONS,
89+
safety_identifier="do-not-log-this-identifier",
90+
)
91+
)
92+
assert_decisions_span(
93+
memory_logger.pop(),
94+
response,
95+
name="openai.decisions.create",
96+
model=DECISIONS_MODEL,
97+
instrumentation="openai-auto",
98+
)
99+
69100
print("SUCCESS")

‎py/src/braintrust/integrations/auto_test_scripts/test_auto_typesafe.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
spans = memory_logger.pop()
2323
assert len(spans) == 1
2424
span = spans[0]
25+
assert span["span_attributes"]["type"] == "decision"
2526
assert span["metadata"]["provider"] == "typesafe"
2627
assert span["metadata"]["model"].startswith("jev-")
2728
assert span["input"]["questions"][0]["id"] == "positive"

‎py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_decisions_image[async].yaml‎

Lines changed: 79 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_decisions_image[sync].yaml‎

Lines changed: 79 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎py/src/braintrust/integrations/litellm/cassettes/latest/test_litellm_decisions_metrics[async].yaml‎

Lines changed: 82 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)