forked from getsentry/sentry-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathlanggraph.py
More file actions
515 lines (417 loc) · 18.4 KB
/
Copy pathlanggraph.py
File metadata and controls
515 lines (417 loc) · 18.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
from functools import wraps
from typing import Any, Callable, List, Optional
import sentry_sdk
from sentry_sdk.ai.utils import (
get_start_span_function,
normalize_message_roles,
set_data_normalized,
truncate_and_annotate_messages,
)
from sentry_sdk.consts import OP, SPANDATA
from sentry_sdk.integrations import DidNotEnable, Integration
# This is fine because langgraph depends on langchain-base, and LangchainIntegration only imports from langchain-base.
from sentry_sdk.integrations.langchain import LangchainIntegration
from sentry_sdk.scope import should_send_default_pii
from sentry_sdk.traces import StreamedSpan
from sentry_sdk.tracing_utils import (
has_span_streaming_enabled,
should_truncate_gen_ai_input,
)
from sentry_sdk.utils import safe_serialize
try:
from langgraph.errors import GraphBubbleUp
from langgraph.graph import StateGraph
from langgraph.pregel import Pregel
except ImportError:
raise DidNotEnable("langgraph not installed")
class LanggraphIntegration(Integration):
identifier = "langgraph"
origin = f"auto.ai.{identifier}"
def __init__(self: "LanggraphIntegration", include_prompts: bool = True) -> None:
self.include_prompts = include_prompts
@staticmethod
def setup_once() -> None:
LangchainIntegration._ignored_exceptions.add(GraphBubbleUp)
# LangGraph lets users create agents using a StateGraph or the Functional API.
# StateGraphs are then compiled to a CompiledStateGraph. Both CompiledStateGraph and
# the functional API execute on a Pregel instance. Pregel is the runtime for the graph
# and the invocation happens on Pregel, so patching the invoke methods takes care of both.
# The streaming methods are not patched, because due to some internal reasons, LangGraph
# will automatically patch the streaming methods to run through invoke, and by doing this
# we prevent duplicate spans for invocations.
StateGraph.compile = _wrap_state_graph_compile(StateGraph.compile)
if hasattr(Pregel, "invoke"):
Pregel.invoke = _wrap_pregel_invoke(Pregel.invoke)
if hasattr(Pregel, "ainvoke"):
Pregel.ainvoke = _wrap_pregel_ainvoke(Pregel.ainvoke)
def _get_graph_name(graph_obj: "Any") -> "Optional[str]":
for attr in ["name", "graph_name", "__name__", "_name"]:
if hasattr(graph_obj, attr):
name = getattr(graph_obj, attr)
if name and isinstance(name, str):
return name
return None
def _normalize_langgraph_message(message: "Any") -> "Any":
if not hasattr(message, "content"):
return None
parsed = {"role": getattr(message, "type", None), "content": message.content}
for attr in [
"name",
"tool_calls",
"function_call",
"tool_call_id",
"response_metadata",
]:
if hasattr(message, attr):
value = getattr(message, attr)
if value is not None:
parsed[attr] = value
return parsed
def _parse_langgraph_messages(state: "Any") -> "Optional[List[Any]]":
if not state:
return None
messages = None
if isinstance(state, dict):
messages = state.get("messages")
elif hasattr(state, "messages"):
messages = state.messages
elif hasattr(state, "get") and callable(state.get):
try:
messages = state.get("messages")
except Exception:
pass
if not messages or not isinstance(messages, (list, tuple)):
return None
normalized_messages = []
for message in messages:
try:
normalized = _normalize_langgraph_message(message)
if normalized:
normalized_messages.append(normalized)
except Exception:
continue
return normalized_messages if normalized_messages else None
def _wrap_state_graph_compile(f: "Callable[..., Any]") -> "Callable[..., Any]":
@wraps(f)
def new_compile(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(LanggraphIntegration)
if integration is None or has_span_streaming_enabled(client.options):
return f(self, *args, **kwargs)
with sentry_sdk.start_span(
op=OP.GEN_AI_CREATE_AGENT,
origin=LanggraphIntegration.origin,
) as span:
compiled_graph = f(self, *args, **kwargs)
compiled_graph_name = getattr(compiled_graph, "name", None)
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "create_agent")
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, compiled_graph_name)
if compiled_graph_name:
span.description = f"create_agent {compiled_graph_name}"
else:
span.description = "create_agent"
if kwargs.get("model", None) is not None:
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, kwargs.get("model"))
tools = None
get_graph = getattr(compiled_graph, "get_graph", None)
if get_graph and callable(get_graph):
graph_obj = compiled_graph.get_graph()
nodes = getattr(graph_obj, "nodes", None)
if nodes and isinstance(nodes, dict):
tools_node = nodes.get("tools")
if tools_node:
data = getattr(tools_node, "data", None)
if data and hasattr(data, "tools_by_name"):
tools = list(data.tools_by_name.keys())
if tools is not None:
span.set_data(SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, tools)
return compiled_graph
return new_compile
def _wrap_pregel_invoke(f: "Callable[..., Any]") -> "Callable[..., Any]":
@wraps(f)
def new_invoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(LanggraphIntegration)
if integration is None:
return f(self, *args, **kwargs)
graph_name = _get_graph_name(self)
span_name = (
f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent"
)
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=span_name,
attributes={
"sentry.op": OP.GEN_AI_INVOKE_AGENT,
"sentry.origin": LanggraphIntegration.origin,
SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
},
) as span:
if graph_name:
span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name)
span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name)
# Store input messages to later compare with output
input_messages = None
if (
len(args) > 0
and should_send_default_pii()
and integration.include_prompts
):
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(
input_messages
)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(
normalized_input_messages, span, scope
)
if should_truncate_gen_ai_input(client.options)
else normalized_input_messages
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)
result = f(self, *args, **kwargs)
_set_response_attributes(span, input_messages, result, integration)
return result
else:
with get_start_span_function()(
op=OP.GEN_AI_INVOKE_AGENT,
name=span_name,
origin=LanggraphIntegration.origin,
) as span:
if graph_name:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name)
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name)
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
# Store input messages to later compare with output
input_messages = None
if (
len(args) > 0
and should_send_default_pii()
and integration.include_prompts
):
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(
input_messages
)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(
normalized_input_messages, span, scope
)
if should_truncate_gen_ai_input(client.options)
else normalized_input_messages
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)
result = f(self, *args, **kwargs)
_set_response_attributes(span, input_messages, result, integration)
return result
return new_invoke
def _wrap_pregel_ainvoke(f: "Callable[..., Any]") -> "Callable[..., Any]":
@wraps(f)
async def new_ainvoke(self: "Any", *args: "Any", **kwargs: "Any") -> "Any":
client = sentry_sdk.get_client()
integration = client.get_integration(LanggraphIntegration)
if integration is None:
return await f(self, *args, **kwargs)
graph_name = _get_graph_name(self)
span_name = (
f"invoke_agent {graph_name}".strip() if graph_name else "invoke_agent"
)
if has_span_streaming_enabled(client.options):
with sentry_sdk.traces.start_span(
name=span_name,
attributes={
"sentry.op": OP.GEN_AI_INVOKE_AGENT,
"sentry.origin": LanggraphIntegration.origin,
SPANDATA.GEN_AI_OPERATION_NAME: "invoke_agent",
},
) as span:
if graph_name:
span.set_attribute(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name)
span.set_attribute(SPANDATA.GEN_AI_AGENT_NAME, graph_name)
input_messages = None
if (
len(args) > 0
and should_send_default_pii()
and integration.include_prompts
):
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(
input_messages
)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(
normalized_input_messages, span, scope
)
if should_truncate_gen_ai_input(client.options)
else normalized_input_messages
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)
result = await f(self, *args, **kwargs)
_set_response_attributes(span, input_messages, result, integration)
return result
with get_start_span_function()(
op=OP.GEN_AI_INVOKE_AGENT,
name=span_name,
origin=LanggraphIntegration.origin,
) as span:
if graph_name:
span.set_data(SPANDATA.GEN_AI_PIPELINE_NAME, graph_name)
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, graph_name)
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
input_messages = None
if (
len(args) > 0
and should_send_default_pii()
and integration.include_prompts
):
input_messages = _parse_langgraph_messages(args[0])
if input_messages:
normalized_input_messages = normalize_message_roles(input_messages)
client = sentry_sdk.get_client()
scope = sentry_sdk.get_current_scope()
messages_data = (
truncate_and_annotate_messages(
normalized_input_messages, span, scope
)
if should_truncate_gen_ai_input(client.options)
else normalized_input_messages
)
if messages_data is not None:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_MESSAGES,
messages_data,
unpack=False,
)
result = await f(self, *args, **kwargs)
_set_response_attributes(span, input_messages, result, integration)
return result
return new_ainvoke
def _get_new_messages(
input_messages: "Optional[List[Any]]", output_messages: "Optional[List[Any]]"
) -> "Optional[List[Any]]":
"""Extract only the new messages added during this invocation."""
if not output_messages:
return None
if not input_messages:
return output_messages
# only return the new messages, aka the output messages that are not in the input messages
input_count = len(input_messages)
new_messages = (
output_messages[input_count:] if len(output_messages) > input_count else []
)
return new_messages if new_messages else None
def _extract_llm_response_text(messages: "Optional[List[Any]]") -> "Optional[str]":
if not messages:
return None
for message in reversed(messages):
if isinstance(message, dict):
role = message.get("role")
if role in ["assistant", "ai"]:
content = message.get("content")
if content and isinstance(content, str):
return content
return None
def _extract_tool_calls(messages: "Optional[List[Any]]") -> "Optional[List[Any]]":
if not messages:
return None
tool_calls = []
for message in messages:
if isinstance(message, dict):
msg_tool_calls = message.get("tool_calls")
if msg_tool_calls and isinstance(msg_tool_calls, list):
tool_calls.extend(msg_tool_calls)
return tool_calls if tool_calls else None
def _set_usage_data(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
input_tokens = 0
output_tokens = 0
total_tokens = 0
for message in messages:
response_metadata = message.get("response_metadata")
if response_metadata is None:
continue
token_usage = response_metadata.get("token_usage")
if not token_usage:
continue
input_tokens += int(token_usage.get("prompt_tokens", 0))
output_tokens += int(token_usage.get("completion_tokens", 0))
total_tokens += int(token_usage.get("total_tokens", 0))
set_on_span = (
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
)
if input_tokens > 0:
set_on_span(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
if output_tokens > 0:
set_on_span(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
if total_tokens > 0:
set_on_span(
SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS,
total_tokens,
)
def _set_response_model_name(span: "sentry_sdk.tracing.Span", messages: "Any") -> None:
if len(messages) == 0:
return
last_message = messages[-1]
response_metadata = last_message.get("response_metadata")
if response_metadata is None:
return
model_name = response_metadata.get("model_name")
if model_name is None:
return
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, model_name)
def _set_response_attributes(
span: "Any",
input_messages: "Optional[List[Any]]",
result: "Any",
integration: "LanggraphIntegration",
) -> None:
parsed_response_messages = _parse_langgraph_messages(result)
new_messages = _get_new_messages(input_messages, parsed_response_messages)
if new_messages is None:
return
_set_usage_data(span, new_messages)
_set_response_model_name(span, new_messages)
if not (should_send_default_pii() and integration.include_prompts):
return
llm_response_text = _extract_llm_response_text(new_messages)
if llm_response_text:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, llm_response_text)
elif new_messages:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, new_messages)
else:
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, result)
tool_calls = _extract_tool_calls(new_messages)
if tool_calls:
set_data_normalized(
span,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
safe_serialize(tool_calls),
unpack=False,
)