sentry-sdk 2.34.1__py2.py3-none-any.whl → 2.35.0__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of sentry-sdk might be problematic. Click here for more details.

@@ -3,55 +3,59 @@ from collections import OrderedDict
3
3
  from functools import wraps
4
4
 
5
5
  import sentry_sdk
6
- from sentry_sdk.ai.monitoring import set_ai_pipeline_name, record_token_usage
7
- from sentry_sdk.consts import OP, SPANDATA
6
+ from sentry_sdk.ai.monitoring import set_ai_pipeline_name
8
7
  from sentry_sdk.ai.utils import set_data_normalized
8
+ from sentry_sdk.consts import OP, SPANDATA
9
+ from sentry_sdk.integrations import DidNotEnable, Integration
9
10
  from sentry_sdk.scope import should_send_default_pii
10
11
  from sentry_sdk.tracing import Span
11
- from sentry_sdk.integrations import DidNotEnable, Integration
12
+ from sentry_sdk.tracing_utils import _get_value
12
13
  from sentry_sdk.utils import logger, capture_internal_exceptions
13
14
 
14
15
  from typing import TYPE_CHECKING
15
16
 
16
17
  if TYPE_CHECKING:
17
- from typing import Any, List, Callable, Dict, Union, Optional
18
+ from typing import (
19
+ Any,
20
+ AsyncIterator,
21
+ Callable,
22
+ Dict,
23
+ Iterator,
24
+ List,
25
+ Optional,
26
+ Union,
27
+ )
18
28
  from uuid import UUID
19
29
 
30
+
20
31
  try:
21
- from langchain_core.messages import BaseMessage
22
- from langchain_core.outputs import LLMResult
32
+ from langchain.agents import AgentExecutor
33
+ from langchain_core.agents import AgentFinish
23
34
  from langchain_core.callbacks import (
24
- manager,
25
35
  BaseCallbackHandler,
26
36
  BaseCallbackManager,
27
37
  Callbacks,
38
+ manager,
28
39
  )
29
- from langchain_core.agents import AgentAction, AgentFinish
40
+ from langchain_core.messages import BaseMessage
41
+ from langchain_core.outputs import LLMResult
42
+
30
43
  except ImportError:
31
44
  raise DidNotEnable("langchain not installed")
32
45
 
33
46
 
34
47
  DATA_FIELDS = {
35
- "temperature": SPANDATA.AI_TEMPERATURE,
36
- "top_p": SPANDATA.AI_TOP_P,
37
- "top_k": SPANDATA.AI_TOP_K,
38
- "function_call": SPANDATA.AI_FUNCTION_CALL,
39
- "tool_calls": SPANDATA.AI_TOOL_CALLS,
40
- "tools": SPANDATA.AI_TOOLS,
41
- "response_format": SPANDATA.AI_RESPONSE_FORMAT,
42
- "logit_bias": SPANDATA.AI_LOGIT_BIAS,
43
- "tags": SPANDATA.AI_TAGS,
48
+ "frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
49
+ "function_call": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
50
+ "max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
51
+ "presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
52
+ "temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
53
+ "tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
54
+ "tools": SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
55
+ "top_k": SPANDATA.GEN_AI_REQUEST_TOP_K,
56
+ "top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
44
57
  }
45
58
 
46
- # To avoid double collecting tokens, we do *not* measure
47
- # token counts for models for which we have an explicit integration
48
- NO_COLLECT_TOKEN_MODELS = [
49
- "openai-chat",
50
- "anthropic-chat",
51
- "cohere-chat",
52
- "huggingface_endpoint",
53
- ]
54
-
55
59
 
56
60
  class LangchainIntegration(Integration):
57
61
  identifier = "langchain"
@@ -60,25 +64,23 @@ class LangchainIntegration(Integration):
60
64
  # The most number of spans (e.g., LLM calls) that can be processed at the same time.
61
65
  max_spans = 1024
62
66
 
63
- def __init__(
64
- self, include_prompts=True, max_spans=1024, tiktoken_encoding_name=None
65
- ):
66
- # type: (LangchainIntegration, bool, int, Optional[str]) -> None
67
+ def __init__(self, include_prompts=True, max_spans=1024):
68
+ # type: (LangchainIntegration, bool, int) -> None
67
69
  self.include_prompts = include_prompts
68
70
  self.max_spans = max_spans
69
- self.tiktoken_encoding_name = tiktoken_encoding_name
70
71
 
71
72
  @staticmethod
72
73
  def setup_once():
73
74
  # type: () -> None
74
75
  manager._configure = _wrap_configure(manager._configure)
75
76
 
77
+ if AgentExecutor is not None:
78
+ AgentExecutor.invoke = _wrap_agent_executor_invoke(AgentExecutor.invoke)
79
+ AgentExecutor.stream = _wrap_agent_executor_stream(AgentExecutor.stream)
80
+
76
81
 
77
82
  class WatchedSpan:
78
83
  span = None # type: Span
79
- num_completion_tokens = 0 # type: int
80
- num_prompt_tokens = 0 # type: int
81
- no_collect_tokens = False # type: bool
82
84
  children = [] # type: List[WatchedSpan]
83
85
  is_pipeline = False # type: bool
84
86
 
@@ -88,26 +90,14 @@ class WatchedSpan:
88
90
 
89
91
 
90
92
  class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
91
- """Base callback handler that can be used to handle callbacks from langchain."""
93
+ """Callback handler that creates Sentry spans."""
92
94
 
93
- def __init__(self, max_span_map_size, include_prompts, tiktoken_encoding_name=None):
94
- # type: (int, bool, Optional[str]) -> None
95
+ def __init__(self, max_span_map_size, include_prompts):
96
+ # type: (int, bool) -> None
95
97
  self.span_map = OrderedDict() # type: OrderedDict[UUID, WatchedSpan]
96
98
  self.max_span_map_size = max_span_map_size
97
99
  self.include_prompts = include_prompts
98
100
 
99
- self.tiktoken_encoding = None
100
- if tiktoken_encoding_name is not None:
101
- import tiktoken # type: ignore
102
-
103
- self.tiktoken_encoding = tiktoken.get_encoding(tiktoken_encoding_name)
104
-
105
- def count_tokens(self, s):
106
- # type: (str) -> int
107
- if self.tiktoken_encoding is not None:
108
- return len(self.tiktoken_encoding.encode_ordinary(s))
109
- return 0
110
-
111
101
  def gc_span_map(self):
112
102
  # type: () -> None
113
103
 
@@ -117,39 +107,37 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
117
107
 
118
108
  def _handle_error(self, run_id, error):
119
109
  # type: (UUID, Any) -> None
120
- if not run_id or run_id not in self.span_map:
121
- return
110
+ with capture_internal_exceptions():
111
+ if not run_id or run_id not in self.span_map:
112
+ return
122
113
 
123
- span_data = self.span_map[run_id]
124
- if not span_data:
125
- return
126
- sentry_sdk.capture_exception(error, span_data.span.scope)
127
- span_data.span.__exit__(None, None, None)
128
- del self.span_map[run_id]
114
+ span_data = self.span_map[run_id]
115
+ span = span_data.span
116
+ span.set_status("unknown")
117
+
118
+ sentry_sdk.capture_exception(error, span.scope)
119
+
120
+ span.__exit__(None, None, None)
121
+ del self.span_map[run_id]
129
122
 
130
123
  def _normalize_langchain_message(self, message):
131
124
  # type: (BaseMessage) -> Any
132
- parsed = {"content": message.content, "role": message.type}
125
+ parsed = {"role": message.type, "content": message.content}
133
126
  parsed.update(message.additional_kwargs)
134
127
  return parsed
135
128
 
136
129
  def _create_span(self, run_id, parent_id, **kwargs):
137
130
  # type: (SentryLangchainCallback, UUID, Optional[Any], Any) -> WatchedSpan
138
-
139
131
  watched_span = None # type: Optional[WatchedSpan]
140
132
  if parent_id:
141
133
  parent_span = self.span_map.get(parent_id) # type: Optional[WatchedSpan]
142
134
  if parent_span:
143
135
  watched_span = WatchedSpan(parent_span.span.start_child(**kwargs))
144
136
  parent_span.children.append(watched_span)
137
+
145
138
  if watched_span is None:
146
139
  watched_span = WatchedSpan(sentry_sdk.start_span(**kwargs))
147
140
 
148
- if kwargs.get("op", "").startswith("ai.pipeline."):
149
- if kwargs.get("name"):
150
- set_ai_pipeline_name(kwargs.get("name"))
151
- watched_span.is_pipeline = True
152
-
153
141
  watched_span.span.__enter__()
154
142
  self.span_map[run_id] = watched_span
155
143
  self.gc_span_map()
@@ -157,7 +145,6 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
157
145
 
158
146
  def _exit_span(self, span_data, run_id):
159
147
  # type: (SentryLangchainCallback, WatchedSpan, UUID) -> None
160
-
161
148
  if span_data.is_pipeline:
162
149
  set_ai_pipeline_name(None)
163
150
 
@@ -180,21 +167,44 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
180
167
  with capture_internal_exceptions():
181
168
  if not run_id:
182
169
  return
170
+
183
171
  all_params = kwargs.get("invocation_params", {})
184
172
  all_params.update(serialized.get("kwargs", {}))
173
+
174
+ model = (
175
+ all_params.get("model")
176
+ or all_params.get("model_name")
177
+ or all_params.get("model_id")
178
+ or ""
179
+ )
180
+
185
181
  watched_span = self._create_span(
186
182
  run_id,
187
- kwargs.get("parent_run_id"),
188
- op=OP.LANGCHAIN_RUN,
183
+ parent_run_id,
184
+ op=OP.GEN_AI_PIPELINE,
189
185
  name=kwargs.get("name") or "Langchain LLM call",
190
186
  origin=LangchainIntegration.origin,
191
187
  )
192
188
  span = watched_span.span
189
+
190
+ if model:
191
+ span.set_data(
192
+ SPANDATA.GEN_AI_REQUEST_MODEL,
193
+ model,
194
+ )
195
+
196
+ ai_type = all_params.get("_type", "")
197
+ if "anthropic" in ai_type:
198
+ span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic")
199
+ elif "openai" in ai_type:
200
+ span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
201
+
202
+ for key, attribute in DATA_FIELDS.items():
203
+ if key in all_params and all_params[key] is not None:
204
+ set_data_normalized(span, attribute, all_params[key], unpack=False)
205
+
193
206
  if should_send_default_pii() and self.include_prompts:
194
- set_data_normalized(span, SPANDATA.AI_INPUT_MESSAGES, prompts)
195
- for k, v in DATA_FIELDS.items():
196
- if k in all_params:
197
- set_data_normalized(span, v, all_params[k])
207
+ set_data_normalized(span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompts)
198
208
 
199
209
  def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
200
210
  # type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
@@ -202,170 +212,150 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
202
212
  with capture_internal_exceptions():
203
213
  if not run_id:
204
214
  return
215
+
205
216
  all_params = kwargs.get("invocation_params", {})
206
217
  all_params.update(serialized.get("kwargs", {}))
218
+
219
+ model = (
220
+ all_params.get("model")
221
+ or all_params.get("model_name")
222
+ or all_params.get("model_id")
223
+ or ""
224
+ )
225
+
207
226
  watched_span = self._create_span(
208
227
  run_id,
209
228
  kwargs.get("parent_run_id"),
210
- op=OP.LANGCHAIN_CHAT_COMPLETIONS_CREATE,
211
- name=kwargs.get("name") or "Langchain Chat Model",
229
+ op=OP.GEN_AI_CHAT,
230
+ name=f"chat {model}".strip(),
212
231
  origin=LangchainIntegration.origin,
213
232
  )
214
233
  span = watched_span.span
215
- model = all_params.get(
216
- "model", all_params.get("model_name", all_params.get("model_id"))
217
- )
218
- watched_span.no_collect_tokens = any(
219
- x in all_params.get("_type", "") for x in NO_COLLECT_TOKEN_MODELS
220
- )
221
234
 
222
- if not model and "anthropic" in all_params.get("_type"):
223
- model = "claude-2"
235
+ span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
224
236
  if model:
225
- span.set_data(SPANDATA.AI_MODEL_ID, model)
237
+ span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
238
+
239
+ ai_type = all_params.get("_type", "")
240
+ if "anthropic" in ai_type:
241
+ span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic")
242
+ elif "openai" in ai_type:
243
+ span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
244
+
245
+ for key, attribute in DATA_FIELDS.items():
246
+ if key in all_params and all_params[key] is not None:
247
+ set_data_normalized(span, attribute, all_params[key], unpack=False)
248
+
226
249
  if should_send_default_pii() and self.include_prompts:
227
250
  set_data_normalized(
228
251
  span,
229
- SPANDATA.AI_INPUT_MESSAGES,
252
+ SPANDATA.GEN_AI_REQUEST_MESSAGES,
230
253
  [
231
254
  [self._normalize_langchain_message(x) for x in list_]
232
255
  for list_ in messages
233
256
  ],
234
257
  )
235
- for k, v in DATA_FIELDS.items():
236
- if k in all_params:
237
- set_data_normalized(span, v, all_params[k])
238
- if not watched_span.no_collect_tokens:
239
- for list_ in messages:
240
- for message in list_:
241
- self.span_map[run_id].num_prompt_tokens += self.count_tokens(
242
- message.content
243
- ) + self.count_tokens(message.type)
244
-
245
- def on_llm_new_token(self, token, *, run_id, **kwargs):
246
- # type: (SentryLangchainCallback, str, UUID, Any) -> Any
247
- """Run on new LLM token. Only available when streaming is enabled."""
258
+
259
+ def on_chat_model_end(self, response, *, run_id, **kwargs):
260
+ # type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
261
+ """Run when Chat Model ends running."""
248
262
  with capture_internal_exceptions():
249
263
  if not run_id or run_id not in self.span_map:
250
264
  return
265
+
251
266
  span_data = self.span_map[run_id]
252
- if not span_data or span_data.no_collect_tokens:
253
- return
254
- span_data.num_completion_tokens += self.count_tokens(token)
267
+ span = span_data.span
268
+
269
+ if should_send_default_pii() and self.include_prompts:
270
+ set_data_normalized(
271
+ span,
272
+ SPANDATA.GEN_AI_RESPONSE_TEXT,
273
+ [[x.text for x in list_] for list_ in response.generations],
274
+ )
275
+
276
+ _record_token_usage(span, response)
277
+ self._exit_span(span_data, run_id)
255
278
 
256
279
  def on_llm_end(self, response, *, run_id, **kwargs):
257
280
  # type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
258
281
  """Run when LLM ends running."""
259
282
  with capture_internal_exceptions():
260
- if not run_id:
283
+ if not run_id or run_id not in self.span_map:
261
284
  return
262
285
 
263
- token_usage = (
264
- response.llm_output.get("token_usage") if response.llm_output else None
265
- )
266
-
267
286
  span_data = self.span_map[run_id]
268
- if not span_data:
269
- return
287
+ span = span_data.span
288
+
289
+ try:
290
+ generation = response.generations[0][0]
291
+ except IndexError:
292
+ generation = None
293
+
294
+ if generation is not None:
295
+ try:
296
+ response_model = generation.generation_info.get("model_name")
297
+ if response_model is not None:
298
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
299
+ except AttributeError:
300
+ pass
301
+
302
+ try:
303
+ finish_reason = generation.generation_info.get("finish_reason")
304
+ if finish_reason is not None:
305
+ span.set_data(
306
+ SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reason
307
+ )
308
+ except AttributeError:
309
+ pass
310
+
311
+ try:
312
+ tool_calls = getattr(generation.message, "tool_calls", None)
313
+ if tool_calls is not None and tool_calls != []:
314
+ set_data_normalized(
315
+ span,
316
+ SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
317
+ tool_calls,
318
+ unpack=False,
319
+ )
320
+ except AttributeError:
321
+ pass
270
322
 
271
323
  if should_send_default_pii() and self.include_prompts:
272
324
  set_data_normalized(
273
- span_data.span,
274
- SPANDATA.AI_RESPONSES,
325
+ span,
326
+ SPANDATA.GEN_AI_RESPONSE_TEXT,
275
327
  [[x.text for x in list_] for list_ in response.generations],
276
328
  )
277
329
 
278
- if not span_data.no_collect_tokens:
279
- if token_usage:
280
- record_token_usage(
281
- span_data.span,
282
- input_tokens=token_usage.get("prompt_tokens"),
283
- output_tokens=token_usage.get("completion_tokens"),
284
- total_tokens=token_usage.get("total_tokens"),
285
- )
286
- else:
287
- record_token_usage(
288
- span_data.span,
289
- input_tokens=span_data.num_prompt_tokens,
290
- output_tokens=span_data.num_completion_tokens,
291
- )
292
-
330
+ _record_token_usage(span, response)
293
331
  self._exit_span(span_data, run_id)
294
332
 
295
333
  def on_llm_error(self, error, *, run_id, **kwargs):
296
334
  # type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
297
335
  """Run when LLM errors."""
298
- with capture_internal_exceptions():
299
- self._handle_error(run_id, error)
300
-
301
- def on_chain_start(self, serialized, inputs, *, run_id, **kwargs):
302
- # type: (SentryLangchainCallback, Dict[str, Any], Dict[str, Any], UUID, Any) -> Any
303
- """Run when chain starts running."""
304
- with capture_internal_exceptions():
305
- if not run_id:
306
- return
307
- watched_span = self._create_span(
308
- run_id,
309
- kwargs.get("parent_run_id"),
310
- op=(
311
- OP.LANGCHAIN_RUN
312
- if kwargs.get("parent_run_id") is not None
313
- else OP.LANGCHAIN_PIPELINE
314
- ),
315
- name=kwargs.get("name") or "Chain execution",
316
- origin=LangchainIntegration.origin,
317
- )
318
- metadata = kwargs.get("metadata")
319
- if metadata:
320
- set_data_normalized(watched_span.span, SPANDATA.AI_METADATA, metadata)
321
-
322
- def on_chain_end(self, outputs, *, run_id, **kwargs):
323
- # type: (SentryLangchainCallback, Dict[str, Any], UUID, Any) -> Any
324
- """Run when chain ends running."""
325
- with capture_internal_exceptions():
326
- if not run_id or run_id not in self.span_map:
327
- return
328
-
329
- span_data = self.span_map[run_id]
330
- if not span_data:
331
- return
332
- self._exit_span(span_data, run_id)
336
+ self._handle_error(run_id, error)
333
337
 
334
- def on_chain_error(self, error, *, run_id, **kwargs):
338
+ def on_chat_model_error(self, error, *, run_id, **kwargs):
335
339
  # type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
336
- """Run when chain errors."""
340
+ """Run when Chat Model errors."""
337
341
  self._handle_error(run_id, error)
338
342
 
339
- def on_agent_action(self, action, *, run_id, **kwargs):
340
- # type: (SentryLangchainCallback, AgentAction, UUID, Any) -> Any
341
- with capture_internal_exceptions():
342
- if not run_id:
343
- return
344
- watched_span = self._create_span(
345
- run_id,
346
- kwargs.get("parent_run_id"),
347
- op=OP.LANGCHAIN_AGENT,
348
- name=action.tool or "AI tool usage",
349
- origin=LangchainIntegration.origin,
350
- )
351
- if action.tool_input and should_send_default_pii() and self.include_prompts:
352
- set_data_normalized(
353
- watched_span.span, SPANDATA.AI_INPUT_MESSAGES, action.tool_input
354
- )
355
-
356
343
  def on_agent_finish(self, finish, *, run_id, **kwargs):
357
344
  # type: (SentryLangchainCallback, AgentFinish, UUID, Any) -> Any
358
345
  with capture_internal_exceptions():
359
- if not run_id:
346
+ if not run_id or run_id not in self.span_map:
360
347
  return
361
348
 
362
349
  span_data = self.span_map[run_id]
363
- if not span_data:
364
- return
350
+ span = span_data.span
351
+
365
352
  if should_send_default_pii() and self.include_prompts:
366
353
  set_data_normalized(
367
- span_data.span, SPANDATA.AI_RESPONSES, finish.return_values.items()
354
+ span,
355
+ SPANDATA.GEN_AI_RESPONSE_TEXT,
356
+ finish.return_values.items(),
368
357
  )
358
+
369
359
  self._exit_span(span_data, run_id)
370
360
 
371
361
  def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
@@ -374,23 +364,31 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
374
364
  with capture_internal_exceptions():
375
365
  if not run_id:
376
366
  return
367
+
368
+ tool_name = serialized.get("name") or kwargs.get("name") or ""
369
+
377
370
  watched_span = self._create_span(
378
371
  run_id,
379
372
  kwargs.get("parent_run_id"),
380
- op=OP.LANGCHAIN_TOOL,
381
- name=serialized.get("name") or kwargs.get("name") or "AI tool usage",
373
+ op=OP.GEN_AI_EXECUTE_TOOL,
374
+ name=f"execute_tool {tool_name}".strip(),
382
375
  origin=LangchainIntegration.origin,
383
376
  )
377
+ span = watched_span.span
378
+
379
+ span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
380
+ span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
381
+
382
+ tool_description = serialized.get("description")
383
+ if tool_description is not None:
384
+ span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description)
385
+
384
386
  if should_send_default_pii() and self.include_prompts:
385
387
  set_data_normalized(
386
- watched_span.span,
387
- SPANDATA.AI_INPUT_MESSAGES,
388
+ span,
389
+ SPANDATA.GEN_AI_TOOL_INPUT,
388
390
  kwargs.get("inputs", [input_str]),
389
391
  )
390
- if kwargs.get("metadata"):
391
- set_data_normalized(
392
- watched_span.span, SPANDATA.AI_METADATA, kwargs.get("metadata")
393
- )
394
392
 
395
393
  def on_tool_end(self, output, *, run_id, **kwargs):
396
394
  # type: (SentryLangchainCallback, str, UUID, Any) -> Any
@@ -400,10 +398,11 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
400
398
  return
401
399
 
402
400
  span_data = self.span_map[run_id]
403
- if not span_data:
404
- return
401
+ span = span_data.span
402
+
405
403
  if should_send_default_pii() and self.include_prompts:
406
- set_data_normalized(span_data.span, SPANDATA.AI_RESPONSES, output)
404
+ set_data_normalized(span, SPANDATA.GEN_AI_TOOL_OUTPUT, output)
405
+
407
406
  self._exit_span(span_data, run_id)
408
407
 
409
408
  def on_tool_error(self, error, *args, run_id, **kwargs):
@@ -412,6 +411,126 @@ class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
412
411
  self._handle_error(run_id, error)
413
412
 
414
413
 
414
+ def _extract_tokens(token_usage):
415
+ # type: (Any) -> tuple[Optional[int], Optional[int], Optional[int]]
416
+ if not token_usage:
417
+ return None, None, None
418
+
419
+ input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value(
420
+ token_usage, "input_tokens"
421
+ )
422
+ output_tokens = _get_value(token_usage, "completion_tokens") or _get_value(
423
+ token_usage, "output_tokens"
424
+ )
425
+ total_tokens = _get_value(token_usage, "total_tokens")
426
+
427
+ return input_tokens, output_tokens, total_tokens
428
+
429
+
430
+ def _extract_tokens_from_generations(generations):
431
+ # type: (Any) -> tuple[Optional[int], Optional[int], Optional[int]]
432
+ """Extract token usage from response.generations structure."""
433
+ if not generations:
434
+ return None, None, None
435
+
436
+ total_input = 0
437
+ total_output = 0
438
+ total_total = 0
439
+
440
+ for gen_list in generations:
441
+ for gen in gen_list:
442
+ token_usage = _get_token_usage(gen)
443
+ input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
444
+ total_input += input_tokens if input_tokens is not None else 0
445
+ total_output += output_tokens if output_tokens is not None else 0
446
+ total_total += total_tokens if total_tokens is not None else 0
447
+
448
+ return (
449
+ total_input if total_input > 0 else None,
450
+ total_output if total_output > 0 else None,
451
+ total_total if total_total > 0 else None,
452
+ )
453
+
454
+
455
+ def _get_token_usage(obj):
456
+ # type: (Any) -> Optional[Dict[str, Any]]
457
+ """
458
+ Check multiple paths to extract token usage from different objects.
459
+ """
460
+ possible_names = ("usage", "token_usage", "usage_metadata")
461
+
462
+ message = _get_value(obj, "message")
463
+ if message is not None:
464
+ for name in possible_names:
465
+ usage = _get_value(message, name)
466
+ if usage is not None:
467
+ return usage
468
+
469
+ llm_output = _get_value(obj, "llm_output")
470
+ if llm_output is not None:
471
+ for name in possible_names:
472
+ usage = _get_value(llm_output, name)
473
+ if usage is not None:
474
+ return usage
475
+
476
+ # check for usage in the object itself
477
+ for name in possible_names:
478
+ usage = _get_value(obj, name)
479
+ if usage is not None:
480
+ return usage
481
+
482
+ # no usage found anywhere
483
+ return None
484
+
485
+
486
+ def _record_token_usage(span, response):
487
+ # type: (Span, Any) -> None
488
+ token_usage = _get_token_usage(response)
489
+ if token_usage:
490
+ input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
491
+ else:
492
+ input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations(
493
+ response.generations
494
+ )
495
+
496
+ if input_tokens is not None:
497
+ span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
498
+
499
+ if output_tokens is not None:
500
+ span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
501
+
502
+ if total_tokens is not None:
503
+ span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
504
+
505
+
506
+ def _get_request_data(obj, args, kwargs):
507
+ # type: (Any, Any, Any) -> tuple[Optional[str], Optional[List[Any]]]
508
+ """
509
+ Get the agent name and available tools for the agent.
510
+ """
511
+ agent = getattr(obj, "agent", None)
512
+ runnable = getattr(agent, "runnable", None)
513
+ runnable_config = getattr(runnable, "config", {})
514
+ tools = (
515
+ getattr(obj, "tools", None)
516
+ or getattr(agent, "tools", None)
517
+ or runnable_config.get("tools")
518
+ or runnable_config.get("available_tools")
519
+ )
520
+ tools = tools if tools and len(tools) > 0 else None
521
+
522
+ try:
523
+ agent_name = None
524
+ if len(args) > 1:
525
+ agent_name = args[1].get("run_name")
526
+ if agent_name is None:
527
+ agent_name = runnable_config.get("run_name")
528
+ except Exception:
529
+ pass
530
+
531
+ return (agent_name, tools)
532
+
533
+
415
534
  def _wrap_configure(f):
416
535
  # type: (Callable[..., Any]) -> Callable[..., Any]
417
536
 
@@ -473,7 +592,6 @@ def _wrap_configure(f):
473
592
  sentry_handler = SentryLangchainCallback(
474
593
  integration.max_spans,
475
594
  integration.include_prompts,
476
- integration.tiktoken_encoding_name,
477
595
  )
478
596
  if isinstance(local_callbacks, BaseCallbackManager):
479
597
  local_callbacks = local_callbacks.copy()
@@ -495,3 +613,158 @@ def _wrap_configure(f):
495
613
  )
496
614
 
497
615
  return new_configure
616
+
617
+
618
+ def _wrap_agent_executor_invoke(f):
619
+ # type: (Callable[..., Any]) -> Callable[..., Any]
620
+
621
+ @wraps(f)
622
+ def new_invoke(self, *args, **kwargs):
623
+ # type: (Any, Any, Any) -> Any
624
+ integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
625
+ if integration is None:
626
+ return f(self, *args, **kwargs)
627
+
628
+ agent_name, tools = _get_request_data(self, args, kwargs)
629
+
630
+ with sentry_sdk.start_span(
631
+ op=OP.GEN_AI_INVOKE_AGENT,
632
+ name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
633
+ origin=LangchainIntegration.origin,
634
+ ) as span:
635
+ if agent_name:
636
+ span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
637
+
638
+ span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
639
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
640
+
641
+ if tools:
642
+ set_data_normalized(
643
+ span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools, unpack=False
644
+ )
645
+
646
+ # Run the agent
647
+ result = f(self, *args, **kwargs)
648
+
649
+ input = result.get("input")
650
+ if (
651
+ input is not None
652
+ and should_send_default_pii()
653
+ and integration.include_prompts
654
+ ):
655
+ set_data_normalized(
656
+ span,
657
+ SPANDATA.GEN_AI_REQUEST_MESSAGES,
658
+ [
659
+ input,
660
+ ],
661
+ )
662
+
663
+ output = result.get("output")
664
+ if (
665
+ output is not None
666
+ and should_send_default_pii()
667
+ and integration.include_prompts
668
+ ):
669
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
670
+
671
+ return result
672
+
673
+ return new_invoke
674
+
675
+
676
+ def _wrap_agent_executor_stream(f):
677
+ # type: (Callable[..., Any]) -> Callable[..., Any]
678
+
679
+ @wraps(f)
680
+ def new_stream(self, *args, **kwargs):
681
+ # type: (Any, Any, Any) -> Any
682
+ integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
683
+ if integration is None:
684
+ return f(self, *args, **kwargs)
685
+
686
+ agent_name, tools = _get_request_data(self, args, kwargs)
687
+
688
+ span = sentry_sdk.start_span(
689
+ op=OP.GEN_AI_INVOKE_AGENT,
690
+ name=f"invoke_agent {agent_name}".strip(),
691
+ origin=LangchainIntegration.origin,
692
+ )
693
+ span.__enter__()
694
+
695
+ if agent_name:
696
+ span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
697
+
698
+ span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
699
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
700
+
701
+ if tools:
702
+ set_data_normalized(
703
+ span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools, unpack=False
704
+ )
705
+
706
+ input = args[0].get("input") if len(args) >= 1 else None
707
+ if (
708
+ input is not None
709
+ and should_send_default_pii()
710
+ and integration.include_prompts
711
+ ):
712
+ set_data_normalized(
713
+ span,
714
+ SPANDATA.GEN_AI_REQUEST_MESSAGES,
715
+ [
716
+ input,
717
+ ],
718
+ )
719
+
720
+ # Run the agent
721
+ result = f(self, *args, **kwargs)
722
+
723
+ old_iterator = result
724
+
725
+ def new_iterator():
726
+ # type: () -> Iterator[Any]
727
+ for event in old_iterator:
728
+ yield event
729
+
730
+ try:
731
+ output = event.get("output")
732
+ except Exception:
733
+ output = None
734
+
735
+ if (
736
+ output is not None
737
+ and should_send_default_pii()
738
+ and integration.include_prompts
739
+ ):
740
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
741
+
742
+ span.__exit__(None, None, None)
743
+
744
+ async def new_iterator_async():
745
+ # type: () -> AsyncIterator[Any]
746
+ async for event in old_iterator:
747
+ yield event
748
+
749
+ try:
750
+ output = event.get("output")
751
+ except Exception:
752
+ output = None
753
+
754
+ if (
755
+ output is not None
756
+ and should_send_default_pii()
757
+ and integration.include_prompts
758
+ ):
759
+ span.set_data(SPANDATA.GEN_AI_RESPONSE_TEXT, output)
760
+
761
+ span.__exit__(None, None, None)
762
+
763
+ if str(type(result)) == "<class 'async_generator'>":
764
+ result = new_iterator_async()
765
+ else:
766
+ result = new_iterator()
767
+
768
+ return result
769
+
770
+ return new_stream