Skip to content

Commit d79c3a7

Browse files
sabrennersabrenner
committed
llmobs_set_tags accepts extra kwargs
1 parent 17570f8 commit d79c3a7

File tree

22 files changed

+180
-91
lines changed

22 files changed

+180
-91
lines changed

‎ddtrace/contrib/internal/crewai/patch.py‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -48,8 +48,9 @@ def traced_kickoff(crewai, pin, func, instance, args, kwargs):
4848
span.set_exc_info(*sys.exc_info())
4949
raise
5050
finally:
51-
kwargs["_dd.instance"] = instance
52-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="crew")
51+
integration.llmobs_set_tags(
52+
span, args=args, kwargs=kwargs, response=result, operation="crew", instance=instance
53+
)
5354
span.finish()
5455
return result
5556

@@ -75,8 +76,9 @@ def traced_task_execute(crewai, pin, func, instance, args, kwargs):
7576
finally:
7677
if getattr(instance, "_ddtrace_ctx", None):
7778
delattr(instance, "_ddtrace_ctx")
78-
kwargs["_dd.instance"] = instance
79-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="task")
79+
integration.llmobs_set_tags(
80+
span, args=args, kwargs=kwargs, response=result, operation="task", instance=instance
81+
)
8082
span.finish()
8183
return result
8284

@@ -111,8 +113,9 @@ def traced_agent_execute(crewai, pin, func, instance, args, kwargs):
111113
span.set_exc_info(*sys.exc_info())
112114
raise
113115
finally:
114-
kwargs["_dd.instance"] = instance
115-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="agent")
116+
integration.llmobs_set_tags(
117+
span, args=args, kwargs=kwargs, response=result, operation="agent", instance=instance
118+
)
116119
span.finish()
117120
return result
118121

@@ -130,8 +133,9 @@ def traced_tool_run(crewai, pin, func, instance, args, kwargs):
130133
span.set_exc_info(*sys.exc_info())
131134
raise
132135
finally:
133-
kwargs["_dd.instance"] = instance
134-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="tool")
136+
integration.llmobs_set_tags(
137+
span, args=args, kwargs=kwargs, response=result, operation="tool", instance=instance
138+
)
135139
span.finish()
136140
return result
137141

@@ -142,7 +146,9 @@ async def traced_flow_kickoff(crewai, pin, func, instance, args, kwargs):
142146
span_name = getattr(type(instance), "__name__", "CrewAI Flow")
143147
with integration.trace(pin, "CrewAI Flow", span_name=span_name, operation="flow", submit_to_llmobs=True) as span:
144148
result = await func(*args, **kwargs)
145-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="flow")
149+
integration.llmobs_set_tags(
150+
span, args=args, kwargs=kwargs, response=result, operation="flow", instance=instance
151+
)
146152
return result
147153

148154

@@ -165,9 +171,15 @@ async def traced_flow_method(crewai, pin, func, instance, args, kwargs):
165171
elif hasattr(flow_state, "model_dump"):
166172
initial_flow_state = flow_state.model_dump()
167173
result = await func(*args, **kwargs)
168-
kwargs["_dd.instance"] = instance
169-
kwargs["_dd.initial_flow_state"] = initial_flow_state
170-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="flow_method")
174+
integration.llmobs_set_tags(
175+
span,
176+
args=args,
177+
kwargs=kwargs,
178+
response=result,
179+
operation="flow_method",
180+
instance=instance,
181+
initial_flow_state=initial_flow_state,
182+
)
171183
return result
172184

173185

‎ddtrace/contrib/internal/google_adk/patch.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,10 +65,10 @@ async def _generator():
6565
span.set_exc_info(*sys.exc_info())
6666
raise
6767
finally:
68-
kwargs["instance"] = instance.agent
69-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=response_events, operation="agent")
68+
integration.llmobs_set_tags(
69+
span, args=args, kwargs=kwargs, response=response_events, operation="agent", instance=instance.agent
70+
)
7071
span.finish()
71-
del kwargs["instance"]
7272

7373
return _generator()
7474

‎ddtrace/contrib/internal/langchain/patch.py‎

Lines changed: 48 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -81,8 +81,14 @@ def traced_llm_generate(langchain_core, pin, func, instance, args, kwargs):
8181
span.set_exc_info(*sys.exc_info())
8282
raise
8383
finally:
84-
kwargs["_dd.identifying_params"] = instance._identifying_params
85-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=completions, operation="llm")
84+
integration.llmobs_set_tags(
85+
span,
86+
args=args,
87+
kwargs=kwargs,
88+
response=completions,
89+
operation="llm",
90+
identifying_params=instance._identifying_params,
91+
)
8692
span.finish()
8793
return completions
8894

@@ -115,8 +121,14 @@ async def traced_llm_agenerate(langchain_core, pin, func, instance, args, kwargs
115121
span.set_exc_info(*sys.exc_info())
116122
raise
117123
finally:
118-
kwargs["_dd.identifying_params"] = instance._identifying_params
119-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=completions, operation="llm")
124+
integration.llmobs_set_tags(
125+
span,
126+
args=args,
127+
kwargs=kwargs,
128+
response=completions,
129+
operation="llm",
130+
identifying_params=instance._identifying_params,
131+
)
120132
span.finish()
121133
return completions
122134

@@ -148,8 +160,14 @@ def traced_chat_model_generate(langchain_core, pin, func, instance, args, kwargs
148160
span.set_exc_info(*sys.exc_info())
149161
raise
150162
finally:
151-
kwargs["_dd.identifying_params"] = instance._identifying_params
152-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=chat_completions, operation="chat")
163+
integration.llmobs_set_tags(
164+
span,
165+
args=args,
166+
kwargs=kwargs,
167+
response=chat_completions,
168+
operation="chat",
169+
identifying_params=instance._identifying_params,
170+
)
153171
span.finish()
154172
return chat_completions
155173

@@ -181,8 +199,14 @@ async def traced_chat_model_agenerate(langchain_core, pin, func, instance, args,
181199
span.set_exc_info(*sys.exc_info())
182200
raise
183201
finally:
184-
kwargs["_dd.identifying_params"] = instance._identifying_params
185-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=chat_completions, operation="chat")
202+
integration.llmobs_set_tags(
203+
span,
204+
args=args,
205+
kwargs=kwargs,
206+
response=chat_completions,
207+
operation="chat",
208+
identifying_params=instance._identifying_params,
209+
)
186210
span.finish()
187211
return chat_completions
188212

@@ -320,14 +344,20 @@ def _on_span_started(span: Span):
320344
integration.record_instance(instance, span)
321345

322346
def _on_span_finished(span: Span, streamed_chunks):
323-
kwargs["_dd.identifying_params"] = instance._identifying_params
324347
if len(streamed_chunks):
325348
joined_chunks = streamed_chunks[0]
326349
for chunk in streamed_chunks[1:]:
327350
joined_chunks += chunk # base message types support __add__ for concatenation
328351
else:
329352
joined_chunks = []
330-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=joined_chunks, operation="chat")
353+
integration.llmobs_set_tags(
354+
span,
355+
args=args,
356+
kwargs=kwargs,
357+
response=joined_chunks,
358+
operation="chat",
359+
identifying_params=instance._identifying_params,
360+
)
331361

332362
return shared_stream(
333363
integration=integration,
@@ -364,8 +394,14 @@ def _on_span_start(span: Span):
364394

365395
def _on_span_finished(span: Span, streamed_chunks):
366396
content = "".join([str(chunk) for chunk in streamed_chunks])
367-
kwargs["_dd.identifying_params"] = instance._identifying_params
368-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=content, operation="llm")
397+
integration.llmobs_set_tags(
398+
span,
399+
args=args,
400+
kwargs=kwargs,
401+
response=content,
402+
operation="llm",
403+
identifying_params=instance._identifying_params,
404+
)
369405

370406
return shared_stream(
371407
integration=integration,

‎ddtrace/contrib/internal/litellm/patch.py‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
from ddtrace.contrib.trace_utils import with_traced_module
1313
from ddtrace.contrib.trace_utils import wrap
1414
from ddtrace.internal.utils import get_argument_value
15-
from ddtrace.llmobs._constants import LITELLM_ROUTER_INSTANCE_KEY
1615
from ddtrace.llmobs._integrations import LiteLLMIntegration
1716
from ddtrace.llmobs._integrations.base_stream_handler import make_traced_stream
1817

@@ -117,8 +116,9 @@ def traced_router_completion(litellm, pin, func, instance, args, kwargs):
117116
raise
118117
finally:
119118
if not stream:
120-
kwargs[LITELLM_ROUTER_INSTANCE_KEY] = instance
121-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=resp, operation=operation)
119+
integration.llmobs_set_tags(
120+
span, args=args, kwargs=kwargs, response=resp, operation=operation, instance=instance
121+
)
122122
span.finish()
123123

124124

@@ -148,8 +148,9 @@ async def traced_router_acompletion(litellm, pin, func, instance, args, kwargs):
148148
raise
149149
finally:
150150
if not stream:
151-
kwargs[LITELLM_ROUTER_INSTANCE_KEY] = instance
152-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=resp, operation=operation)
151+
integration.llmobs_set_tags(
152+
span, args=args, kwargs=kwargs, response=resp, operation=operation, instance=instance
153+
)
153154
span.finish()
154155

155156

‎ddtrace/contrib/internal/pydantic_ai/patch.py‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,7 @@ def traced_agent_run_stream(pydantic_ai, pin, func, instance, args, kwargs):
3939
span.name = getattr(instance, "name", None) or "Pydantic Agent"
4040

4141
result = func(*args, **kwargs)
42-
kwargs["instance"] = instance
43-
return TracedPydanticRunStream(result, span, integration, args, kwargs)
42+
return TracedPydanticRunStream(result, span, instance, integration, args, kwargs)
4443

4544

4645
@with_traced_module
@@ -56,7 +55,6 @@ def traced_agent_iter(pydantic_ai, pin, func, instance, args, kwargs):
5655
span.name = getattr(instance, "name", None) or "Pydantic Agent"
5756

5857
result = func(*args, **kwargs)
59-
kwargs["instance"] = instance
6058
return TracedPydanticAsyncContextManager(result, span, instance, integration, args, kwargs)
6159

6260

@@ -87,8 +85,7 @@ async def traced_tool_run(pydantic_ai, pin, func, instance, args, kwargs, tool_n
8785
span.set_exc_info(*sys.exc_info())
8886
raise
8987
finally:
90-
kwargs["instance"] = instance
91-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=resp)
88+
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=resp, instance=instance)
9289
span.finish()
9390

9491

‎ddtrace/contrib/internal/pydantic_ai/utils.py‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,11 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
2727
self._dd_span.set_exc_info(exc_type, exc_val, exc_tb)
2828
elif self._dd_integration.is_pc_sampled_llmobs(self._dd_span):
2929
self._dd_integration.llmobs_set_tags(
30-
self._dd_span, args=self._args, kwargs=self._kwargs, response=self._agent_run
30+
self._dd_span,
31+
args=self._args,
32+
kwargs=self._kwargs,
33+
response=self._agent_run,
34+
instance=self._dd_instance,
3135
)
3236
finally:
3337
if exc_type:
@@ -36,18 +40,19 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
3640

3741

3842
class TracedPydanticRunStream(wrapt.ObjectProxy):
39-
def __init__(self, wrapped, span, integration, args, kwargs):
43+
def __init__(self, wrapped, span, instance, integration, args, kwargs):
4044
super().__init__(wrapped)
4145
self._dd_span = span
4246
self._dd_integration = integration
47+
self._dd_instance = instance
4348
self._args = args
4449
self._kwargs = kwargs
4550
self._streamed_run_result = None
4651

4752
async def __aenter__(self):
4853
result = await self.__wrapped__.__aenter__()
4954
self._streamed_run_result = TracedPydanticStreamedRunResult(
50-
result, self._dd_span, self._dd_integration, self._args, self._kwargs
55+
result, self._dd_span, self._dd_instance, self._dd_integration, self._args, self._kwargs
5156
)
5257
return self._streamed_run_result
5358

@@ -63,10 +68,11 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
6368

6469

6570
class TracedPydanticStreamedRunResult(wrapt.ObjectProxy):
66-
def __init__(self, wrapped, span, integration, args, kwargs):
71+
def __init__(self, wrapped, span, instance, integration, args, kwargs):
6772
super().__init__(wrapped)
6873
self._dd_span = span
6974
self._dd_integration = integration
75+
self._dd_instance = instance
7076
self._args = args
7177
self._kwargs = kwargs
7278
# needed for extracting usage metrics from the streamed run result
@@ -103,7 +109,9 @@ def stream_structured(self, *args, **kwargs):
103109

104110
async def get_output(self):
105111
result = await self.__wrapped__.get_output()
106-
self._dd_integration.llmobs_set_tags(self._dd_span, args=self._args, kwargs=self._kwargs, response=result)
112+
self._dd_integration.llmobs_set_tags(
113+
self._dd_span, args=self._args, kwargs=self._kwargs, response=result, instance=self._dd_instance
114+
)
107115
self._dd_span.finish()
108116
return result
109117

‎ddtrace/contrib/internal/vertexai/patch.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,9 @@ def _traced_generate(vertexai, pin, func, instance, args, kwargs, model_instance
7979
finally:
8080
# streamed spans will be finished separately once the stream generator is exhausted
8181
if span.error or not stream:
82-
kwargs["instance"] = model_instance
83-
kwargs["history"] = history
84-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=generations)
82+
integration.llmobs_set_tags(
83+
span, args=args, kwargs=kwargs, response=generations, instance=model_instance, history=history
84+
)
8585
span.finish()
8686
return generations
8787

@@ -115,9 +115,9 @@ async def _traced_agenerate(vertexai, pin, func, instance, args, kwargs, model_i
115115
finally:
116116
# streamed spans will be finished separately once the stream generator is exhausted
117117
if span.error or not stream:
118-
kwargs["instance"] = model_instance
119-
kwargs["history"] = history
120-
integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=generations)
118+
integration.llmobs_set_tags(
119+
span, args=args, kwargs=kwargs, response=generations, instance=model_instance, history=history
120+
)
121121
span.finish()
122122
return generations
123123

‎ddtrace/llmobs/_constants.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -100,8 +100,6 @@
100100
# When there are no tool call args, we use this as a place-holder lookup key
101101
OAI_HANDOFF_TOOL_ARG = "{}"
102102

103-
LITELLM_ROUTER_INSTANCE_KEY = "_dd.router_instance"
104-
105103
PROXY_REQUEST = "llmobs.proxy_request"
106104

107105
EXPERIMENT_ID_KEY = "_ml_obs.experiment_id"

‎ddtrace/llmobs/_integrations/anthropic.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ def _llmobs_set_tags(
5858
kwargs: Dict[str, Any],
5959
response: Optional[Any] = None,
6060
operation: str = "",
61+
**extra_kwargs: Any,
6162
) -> None:
6263
"""Extract prompt/response tags from a completion and set them as temporary "_ml_obs.*" tags."""
6364
parameters = {}

‎ddtrace/llmobs/_integrations/base.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,12 +77,13 @@ def llmobs_set_tags(
7777
kwargs: Dict[str, Any],
7878
response: Optional[Any] = None,
7979
operation: str = "",
80+
**extra_kwargs: Any,
8081
) -> None:
8182
"""Extract input/output information from the request and response to be submitted to LLMObs."""
8283
if not self.llmobs_enabled or not self.is_pc_sampled_llmobs(span):
8384
return
8485
try:
85-
self._llmobs_set_tags(span, args, kwargs, response, operation)
86+
self._llmobs_set_tags(span, args, kwargs, response, operation, **extra_kwargs)
8687
except Exception:
8788
log.error("Error extracting LLMObs fields for span %s, likely due to malformed data", span, exc_info=True)
8889

@@ -94,6 +95,7 @@ def _llmobs_set_tags(
9495
kwargs: Dict[str, Any],
9596
response: Optional[Any] = None,
9697
operation: str = "",
98+
**extra_kwargs: Any,
9799
) -> None:
98100
raise NotImplementedError()
99101

0 commit comments

Comments
 (0)