1132 lines
39 KiB
Python
1132 lines
39 KiB
Python
import contextvars
|
|
import itertools
|
|
import warnings
|
|
from collections import OrderedDict
|
|
from functools import wraps
|
|
import sys
|
|
|
|
import sentry_sdk
|
|
from sentry_sdk.ai.monitoring import set_ai_pipeline_name
|
|
from sentry_sdk.ai.utils import (
|
|
GEN_AI_ALLOWED_MESSAGE_ROLES,
|
|
normalize_message_roles,
|
|
set_data_normalized,
|
|
get_start_span_function,
|
|
truncate_and_annotate_messages,
|
|
)
|
|
from sentry_sdk.consts import OP, SPANDATA
|
|
from sentry_sdk.integrations import DidNotEnable, Integration
|
|
from sentry_sdk.scope import should_send_default_pii
|
|
from sentry_sdk.tracing_utils import _get_value, set_span_errored
|
|
from sentry_sdk.utils import logger, capture_internal_exceptions
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from typing import (
|
|
Any,
|
|
AsyncIterator,
|
|
Callable,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Union,
|
|
)
|
|
from uuid import UUID
|
|
from sentry_sdk.tracing import Span
|
|
|
|
|
|
try:
|
|
from langchain_core.agents import AgentFinish
|
|
from langchain_core.callbacks import (
|
|
BaseCallbackHandler,
|
|
BaseCallbackManager,
|
|
Callbacks,
|
|
manager,
|
|
)
|
|
from langchain_core.messages import BaseMessage
|
|
from langchain_core.outputs import LLMResult
|
|
|
|
except ImportError:
|
|
raise DidNotEnable("langchain not installed")
|
|
|
|
|
|
try:
|
|
# >=v1
|
|
from langchain_classic.agents import AgentExecutor # type: ignore[import-not-found]
|
|
except ImportError:
|
|
try:
|
|
# <v1
|
|
from langchain.agents import AgentExecutor
|
|
except ImportError:
|
|
AgentExecutor = None
|
|
|
|
|
|
# Conditional imports for embeddings providers
|
|
try:
|
|
from langchain_openai import OpenAIEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
OpenAIEmbeddings = None
|
|
|
|
try:
|
|
from langchain_openai import AzureOpenAIEmbeddings
|
|
except ImportError:
|
|
AzureOpenAIEmbeddings = None
|
|
|
|
try:
|
|
from langchain_google_vertexai import VertexAIEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
VertexAIEmbeddings = None
|
|
|
|
try:
|
|
from langchain_aws import BedrockEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
BedrockEmbeddings = None
|
|
|
|
try:
|
|
from langchain_cohere import CohereEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
CohereEmbeddings = None
|
|
|
|
try:
|
|
from langchain_mistralai import MistralAIEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
MistralAIEmbeddings = None
|
|
|
|
try:
|
|
from langchain_huggingface import HuggingFaceEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
HuggingFaceEmbeddings = None
|
|
|
|
try:
|
|
from langchain_ollama import OllamaEmbeddings # type: ignore[import-not-found]
|
|
except ImportError:
|
|
OllamaEmbeddings = None
|
|
|
|
|
|
DATA_FIELDS = {
|
|
"frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
|
|
"function_call": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
|
|
"max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
|
|
"presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
|
|
"temperature": SPANDATA.GEN_AI_REQUEST_TEMPERATURE,
|
|
"tool_calls": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
|
|
"top_k": SPANDATA.GEN_AI_REQUEST_TOP_K,
|
|
"top_p": SPANDATA.GEN_AI_REQUEST_TOP_P,
|
|
}
|
|
|
|
|
|
# Contextvar to track agent names in a stack for re-entrant agent support
|
|
_agent_stack = contextvars.ContextVar("langchain_agent_stack", default=None) # type: contextvars.ContextVar[Optional[List[Optional[str]]]]
|
|
|
|
|
|
def _push_agent(agent_name):
|
|
# type: (Optional[str]) -> None
|
|
"""Push an agent name onto the stack."""
|
|
stack = _agent_stack.get()
|
|
if stack is None:
|
|
stack = []
|
|
else:
|
|
# Copy the list to maintain contextvar isolation across async contexts
|
|
stack = stack.copy()
|
|
stack.append(agent_name)
|
|
_agent_stack.set(stack)
|
|
|
|
|
|
def _pop_agent():
|
|
# type: () -> Optional[str]
|
|
"""Pop an agent name from the stack and return it."""
|
|
stack = _agent_stack.get()
|
|
if stack:
|
|
# Copy the list to maintain contextvar isolation across async contexts
|
|
stack = stack.copy()
|
|
agent_name = stack.pop()
|
|
_agent_stack.set(stack)
|
|
return agent_name
|
|
return None
|
|
|
|
|
|
def _get_current_agent():
|
|
# type: () -> Optional[str]
|
|
"""Get the current agent name (top of stack) without removing it."""
|
|
stack = _agent_stack.get()
|
|
if stack:
|
|
return stack[-1]
|
|
return None
|
|
|
|
|
|
class LangchainIntegration(Integration):
|
|
identifier = "langchain"
|
|
origin = f"auto.ai.{identifier}"
|
|
|
|
def __init__(self, include_prompts=True, max_spans=None):
|
|
# type: (LangchainIntegration, bool, Optional[int]) -> None
|
|
self.include_prompts = include_prompts
|
|
self.max_spans = max_spans
|
|
|
|
if max_spans is not None:
|
|
warnings.warn(
|
|
"The `max_spans` parameter of `LangchainIntegration` is "
|
|
"deprecated and will be removed in version 3.0 of sentry-sdk.",
|
|
DeprecationWarning,
|
|
stacklevel=2,
|
|
)
|
|
|
|
@staticmethod
|
|
def setup_once():
|
|
# type: () -> None
|
|
manager._configure = _wrap_configure(manager._configure)
|
|
|
|
if AgentExecutor is not None:
|
|
AgentExecutor.invoke = _wrap_agent_executor_invoke(AgentExecutor.invoke)
|
|
AgentExecutor.stream = _wrap_agent_executor_stream(AgentExecutor.stream)
|
|
|
|
# Patch embeddings providers
|
|
_patch_embeddings_provider(OpenAIEmbeddings)
|
|
_patch_embeddings_provider(AzureOpenAIEmbeddings)
|
|
_patch_embeddings_provider(VertexAIEmbeddings)
|
|
_patch_embeddings_provider(BedrockEmbeddings)
|
|
_patch_embeddings_provider(CohereEmbeddings)
|
|
_patch_embeddings_provider(MistralAIEmbeddings)
|
|
_patch_embeddings_provider(HuggingFaceEmbeddings)
|
|
_patch_embeddings_provider(OllamaEmbeddings)
|
|
|
|
|
|
class WatchedSpan:
|
|
span = None # type: Span
|
|
children = [] # type: List[WatchedSpan]
|
|
is_pipeline = False # type: bool
|
|
|
|
def __init__(self, span):
|
|
# type: (Span) -> None
|
|
self.span = span
|
|
|
|
|
|
class SentryLangchainCallback(BaseCallbackHandler): # type: ignore[misc]
|
|
"""Callback handler that creates Sentry spans."""
|
|
|
|
def __init__(self, max_span_map_size, include_prompts):
|
|
# type: (Optional[int], bool) -> None
|
|
self.span_map = OrderedDict() # type: OrderedDict[UUID, WatchedSpan]
|
|
self.max_span_map_size = max_span_map_size
|
|
self.include_prompts = include_prompts
|
|
|
|
def gc_span_map(self):
|
|
# type: () -> None
|
|
|
|
if self.max_span_map_size is not None:
|
|
while len(self.span_map) > self.max_span_map_size:
|
|
run_id, watched_span = self.span_map.popitem(last=False)
|
|
self._exit_span(watched_span, run_id)
|
|
|
|
def _handle_error(self, run_id, error):
|
|
# type: (UUID, Any) -> None
|
|
with capture_internal_exceptions():
|
|
if not run_id or run_id not in self.span_map:
|
|
return
|
|
|
|
span_data = self.span_map[run_id]
|
|
span = span_data.span
|
|
set_span_errored(span)
|
|
|
|
sentry_sdk.capture_exception(error, span.scope)
|
|
|
|
span.__exit__(None, None, None)
|
|
del self.span_map[run_id]
|
|
|
|
def _normalize_langchain_message(self, message):
|
|
# type: (BaseMessage) -> Any
|
|
parsed = {"role": message.type, "content": message.content}
|
|
parsed.update(message.additional_kwargs)
|
|
return parsed
|
|
|
|
def _create_span(self, run_id, parent_id, **kwargs):
|
|
# type: (SentryLangchainCallback, UUID, Optional[Any], Any) -> WatchedSpan
|
|
watched_span = None # type: Optional[WatchedSpan]
|
|
if parent_id:
|
|
parent_span = self.span_map.get(parent_id) # type: Optional[WatchedSpan]
|
|
if parent_span:
|
|
watched_span = WatchedSpan(parent_span.span.start_child(**kwargs))
|
|
parent_span.children.append(watched_span)
|
|
|
|
if watched_span is None:
|
|
watched_span = WatchedSpan(sentry_sdk.start_span(**kwargs))
|
|
|
|
watched_span.span.__enter__()
|
|
self.span_map[run_id] = watched_span
|
|
self.gc_span_map()
|
|
return watched_span
|
|
|
|
def _exit_span(self, span_data, run_id):
|
|
# type: (SentryLangchainCallback, WatchedSpan, UUID) -> None
|
|
if span_data.is_pipeline:
|
|
set_ai_pipeline_name(None)
|
|
|
|
span_data.span.__exit__(None, None, None)
|
|
del self.span_map[run_id]
|
|
|
|
def on_llm_start(
|
|
self,
|
|
serialized,
|
|
prompts,
|
|
*,
|
|
run_id,
|
|
tags=None,
|
|
parent_run_id=None,
|
|
metadata=None,
|
|
**kwargs,
|
|
):
|
|
# type: (SentryLangchainCallback, Dict[str, Any], List[str], UUID, Optional[List[str]], Optional[UUID], Optional[Dict[str, Any]], Any) -> Any
|
|
"""Run when LLM starts running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id:
|
|
return
|
|
|
|
all_params = kwargs.get("invocation_params", {})
|
|
all_params.update(serialized.get("kwargs", {}))
|
|
|
|
model = (
|
|
all_params.get("model")
|
|
or all_params.get("model_name")
|
|
or all_params.get("model_id")
|
|
or ""
|
|
)
|
|
|
|
watched_span = self._create_span(
|
|
run_id,
|
|
parent_run_id,
|
|
op=OP.GEN_AI_PIPELINE,
|
|
name=kwargs.get("name") or "Langchain LLM call",
|
|
origin=LangchainIntegration.origin,
|
|
)
|
|
span = watched_span.span
|
|
|
|
if model:
|
|
span.set_data(
|
|
SPANDATA.GEN_AI_REQUEST_MODEL,
|
|
model,
|
|
)
|
|
|
|
ai_type = all_params.get("_type", "")
|
|
if "anthropic" in ai_type:
|
|
span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic")
|
|
elif "openai" in ai_type:
|
|
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
|
|
|
|
for key, attribute in DATA_FIELDS.items():
|
|
if key in all_params and all_params[key] is not None:
|
|
set_data_normalized(span, attribute, all_params[key], unpack=False)
|
|
|
|
_set_tools_on_span(span, all_params.get("tools"))
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
normalized_messages = [
|
|
{
|
|
"role": GEN_AI_ALLOWED_MESSAGE_ROLES.USER,
|
|
"content": {"type": "text", "text": prompt},
|
|
}
|
|
for prompt in prompts
|
|
]
|
|
scope = sentry_sdk.get_current_scope()
|
|
messages_data = truncate_and_annotate_messages(
|
|
normalized_messages, span, scope
|
|
)
|
|
if messages_data is not None:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_REQUEST_MESSAGES,
|
|
messages_data,
|
|
unpack=False,
|
|
)
|
|
|
|
def on_chat_model_start(self, serialized, messages, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, Dict[str, Any], List[List[BaseMessage]], UUID, Any) -> Any
|
|
"""Run when Chat Model starts running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id:
|
|
return
|
|
|
|
all_params = kwargs.get("invocation_params", {})
|
|
all_params.update(serialized.get("kwargs", {}))
|
|
|
|
model = (
|
|
all_params.get("model")
|
|
or all_params.get("model_name")
|
|
or all_params.get("model_id")
|
|
or ""
|
|
)
|
|
|
|
watched_span = self._create_span(
|
|
run_id,
|
|
kwargs.get("parent_run_id"),
|
|
op=OP.GEN_AI_CHAT,
|
|
name=f"chat {model}".strip(),
|
|
origin=LangchainIntegration.origin,
|
|
)
|
|
span = watched_span.span
|
|
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "chat")
|
|
if model:
|
|
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model)
|
|
|
|
ai_type = all_params.get("_type", "")
|
|
if "anthropic" in ai_type:
|
|
span.set_data(SPANDATA.GEN_AI_SYSTEM, "anthropic")
|
|
elif "openai" in ai_type:
|
|
span.set_data(SPANDATA.GEN_AI_SYSTEM, "openai")
|
|
|
|
agent_name = _get_current_agent()
|
|
if agent_name:
|
|
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
|
|
|
|
for key, attribute in DATA_FIELDS.items():
|
|
if key in all_params and all_params[key] is not None:
|
|
set_data_normalized(span, attribute, all_params[key], unpack=False)
|
|
|
|
_set_tools_on_span(span, all_params.get("tools"))
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
normalized_messages = []
|
|
for list_ in messages:
|
|
for message in list_:
|
|
normalized_messages.append(
|
|
self._normalize_langchain_message(message)
|
|
)
|
|
normalized_messages = normalize_message_roles(normalized_messages)
|
|
scope = sentry_sdk.get_current_scope()
|
|
messages_data = truncate_and_annotate_messages(
|
|
normalized_messages, span, scope
|
|
)
|
|
if messages_data is not None:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_REQUEST_MESSAGES,
|
|
messages_data,
|
|
unpack=False,
|
|
)
|
|
|
|
def on_chat_model_end(self, response, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
|
|
"""Run when Chat Model ends running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id or run_id not in self.span_map:
|
|
return
|
|
|
|
span_data = self.span_map[run_id]
|
|
span = span_data.span
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_RESPONSE_TEXT,
|
|
[[x.text for x in list_] for list_ in response.generations],
|
|
)
|
|
|
|
_record_token_usage(span, response)
|
|
self._exit_span(span_data, run_id)
|
|
|
|
def on_llm_end(self, response, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, LLMResult, UUID, Any) -> Any
|
|
"""Run when LLM ends running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id or run_id not in self.span_map:
|
|
return
|
|
|
|
span_data = self.span_map[run_id]
|
|
span = span_data.span
|
|
|
|
try:
|
|
generation = response.generations[0][0]
|
|
except IndexError:
|
|
generation = None
|
|
|
|
if generation is not None:
|
|
try:
|
|
response_model = generation.generation_info.get("model_name")
|
|
if response_model is not None:
|
|
span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)
|
|
except AttributeError:
|
|
pass
|
|
|
|
try:
|
|
finish_reason = generation.generation_info.get("finish_reason")
|
|
if finish_reason is not None:
|
|
span.set_data(
|
|
SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, finish_reason
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
try:
|
|
if should_send_default_pii() and self.include_prompts:
|
|
tool_calls = getattr(generation.message, "tool_calls", None)
|
|
if tool_calls is not None and tool_calls != []:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
|
|
tool_calls,
|
|
unpack=False,
|
|
)
|
|
except AttributeError:
|
|
pass
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_RESPONSE_TEXT,
|
|
[[x.text for x in list_] for list_ in response.generations],
|
|
)
|
|
|
|
_record_token_usage(span, response)
|
|
self._exit_span(span_data, run_id)
|
|
|
|
def on_llm_error(self, error, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
|
|
"""Run when LLM errors."""
|
|
self._handle_error(run_id, error)
|
|
|
|
def on_chat_model_error(self, error, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
|
|
"""Run when Chat Model errors."""
|
|
self._handle_error(run_id, error)
|
|
|
|
def on_agent_finish(self, finish, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, AgentFinish, UUID, Any) -> Any
|
|
with capture_internal_exceptions():
|
|
if not run_id or run_id not in self.span_map:
|
|
return
|
|
|
|
span_data = self.span_map[run_id]
|
|
span = span_data.span
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
set_data_normalized(
|
|
span, SPANDATA.GEN_AI_RESPONSE_TEXT, finish.return_values.items()
|
|
)
|
|
|
|
self._exit_span(span_data, run_id)
|
|
|
|
def on_tool_start(self, serialized, input_str, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, Dict[str, Any], str, UUID, Any) -> Any
|
|
"""Run when tool starts running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id:
|
|
return
|
|
|
|
tool_name = serialized.get("name") or kwargs.get("name") or ""
|
|
|
|
watched_span = self._create_span(
|
|
run_id,
|
|
kwargs.get("parent_run_id"),
|
|
op=OP.GEN_AI_EXECUTE_TOOL,
|
|
name=f"execute_tool {tool_name}".strip(),
|
|
origin=LangchainIntegration.origin,
|
|
)
|
|
span = watched_span.span
|
|
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
|
|
span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool_name)
|
|
|
|
tool_description = serialized.get("description")
|
|
if tool_description is not None:
|
|
span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool_description)
|
|
|
|
agent_name = _get_current_agent()
|
|
if agent_name:
|
|
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_TOOL_INPUT,
|
|
kwargs.get("inputs", [input_str]),
|
|
)
|
|
|
|
def on_tool_end(self, output, *, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, str, UUID, Any) -> Any
|
|
"""Run when tool ends running."""
|
|
with capture_internal_exceptions():
|
|
if not run_id or run_id not in self.span_map:
|
|
return
|
|
|
|
span_data = self.span_map[run_id]
|
|
span = span_data.span
|
|
|
|
if should_send_default_pii() and self.include_prompts:
|
|
set_data_normalized(span, SPANDATA.GEN_AI_TOOL_OUTPUT, output)
|
|
|
|
self._exit_span(span_data, run_id)
|
|
|
|
def on_tool_error(self, error, *args, run_id, **kwargs):
|
|
# type: (SentryLangchainCallback, Union[Exception, KeyboardInterrupt], UUID, Any) -> Any
|
|
"""Run when tool errors."""
|
|
self._handle_error(run_id, error)
|
|
|
|
|
|
def _extract_tokens(token_usage):
|
|
# type: (Any) -> tuple[Optional[int], Optional[int], Optional[int]]
|
|
if not token_usage:
|
|
return None, None, None
|
|
|
|
input_tokens = _get_value(token_usage, "prompt_tokens") or _get_value(
|
|
token_usage, "input_tokens"
|
|
)
|
|
output_tokens = _get_value(token_usage, "completion_tokens") or _get_value(
|
|
token_usage, "output_tokens"
|
|
)
|
|
total_tokens = _get_value(token_usage, "total_tokens")
|
|
|
|
return input_tokens, output_tokens, total_tokens
|
|
|
|
|
|
def _extract_tokens_from_generations(generations):
|
|
# type: (Any) -> tuple[Optional[int], Optional[int], Optional[int]]
|
|
"""Extract token usage from response.generations structure."""
|
|
if not generations:
|
|
return None, None, None
|
|
|
|
total_input = 0
|
|
total_output = 0
|
|
total_total = 0
|
|
|
|
for gen_list in generations:
|
|
for gen in gen_list:
|
|
token_usage = _get_token_usage(gen)
|
|
input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
|
|
total_input += input_tokens if input_tokens is not None else 0
|
|
total_output += output_tokens if output_tokens is not None else 0
|
|
total_total += total_tokens if total_tokens is not None else 0
|
|
|
|
return (
|
|
total_input if total_input > 0 else None,
|
|
total_output if total_output > 0 else None,
|
|
total_total if total_total > 0 else None,
|
|
)
|
|
|
|
|
|
def _get_token_usage(obj):
|
|
# type: (Any) -> Optional[Dict[str, Any]]
|
|
"""
|
|
Check multiple paths to extract token usage from different objects.
|
|
"""
|
|
possible_names = ("usage", "token_usage", "usage_metadata")
|
|
|
|
message = _get_value(obj, "message")
|
|
if message is not None:
|
|
for name in possible_names:
|
|
usage = _get_value(message, name)
|
|
if usage is not None:
|
|
return usage
|
|
|
|
llm_output = _get_value(obj, "llm_output")
|
|
if llm_output is not None:
|
|
for name in possible_names:
|
|
usage = _get_value(llm_output, name)
|
|
if usage is not None:
|
|
return usage
|
|
|
|
for name in possible_names:
|
|
usage = _get_value(obj, name)
|
|
if usage is not None:
|
|
return usage
|
|
|
|
return None
|
|
|
|
|
|
def _record_token_usage(span, response):
|
|
# type: (Span, Any) -> None
|
|
token_usage = _get_token_usage(response)
|
|
if token_usage:
|
|
input_tokens, output_tokens, total_tokens = _extract_tokens(token_usage)
|
|
else:
|
|
input_tokens, output_tokens, total_tokens = _extract_tokens_from_generations(
|
|
response.generations
|
|
)
|
|
|
|
if input_tokens is not None:
|
|
span.set_data(SPANDATA.GEN_AI_USAGE_INPUT_TOKENS, input_tokens)
|
|
|
|
if output_tokens is not None:
|
|
span.set_data(SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS, output_tokens)
|
|
|
|
if total_tokens is not None:
|
|
span.set_data(SPANDATA.GEN_AI_USAGE_TOTAL_TOKENS, total_tokens)
|
|
|
|
|
|
def _get_request_data(obj, args, kwargs):
|
|
# type: (Any, Any, Any) -> tuple[Optional[str], Optional[List[Any]]]
|
|
"""
|
|
Get the agent name and available tools for the agent.
|
|
"""
|
|
agent = getattr(obj, "agent", None)
|
|
runnable = getattr(agent, "runnable", None)
|
|
runnable_config = getattr(runnable, "config", {})
|
|
tools = (
|
|
getattr(obj, "tools", None)
|
|
or getattr(agent, "tools", None)
|
|
or runnable_config.get("tools")
|
|
or runnable_config.get("available_tools")
|
|
)
|
|
tools = tools if tools and len(tools) > 0 else None
|
|
|
|
try:
|
|
agent_name = None
|
|
if len(args) > 1:
|
|
agent_name = args[1].get("run_name")
|
|
if agent_name is None:
|
|
agent_name = runnable_config.get("run_name")
|
|
except Exception:
|
|
pass
|
|
|
|
return (agent_name, tools)
|
|
|
|
|
|
def _simplify_langchain_tools(tools):
|
|
# type: (Any) -> Optional[List[Any]]
|
|
"""Parse and simplify tools into a cleaner format."""
|
|
if not tools:
|
|
return None
|
|
|
|
if not isinstance(tools, (list, tuple)):
|
|
return None
|
|
|
|
simplified_tools = []
|
|
for tool in tools:
|
|
try:
|
|
if isinstance(tool, dict):
|
|
if "function" in tool and isinstance(tool["function"], dict):
|
|
func = tool["function"]
|
|
simplified_tool = {
|
|
"name": func.get("name"),
|
|
"description": func.get("description"),
|
|
}
|
|
if simplified_tool["name"]:
|
|
simplified_tools.append(simplified_tool)
|
|
elif "name" in tool:
|
|
simplified_tool = {
|
|
"name": tool.get("name"),
|
|
"description": tool.get("description"),
|
|
}
|
|
simplified_tools.append(simplified_tool)
|
|
else:
|
|
name = (
|
|
tool.get("name")
|
|
or tool.get("tool_name")
|
|
or tool.get("function_name")
|
|
)
|
|
if name:
|
|
simplified_tools.append(
|
|
{
|
|
"name": name,
|
|
"description": tool.get("description")
|
|
or tool.get("desc"),
|
|
}
|
|
)
|
|
elif hasattr(tool, "name"):
|
|
simplified_tool = {
|
|
"name": getattr(tool, "name", None),
|
|
"description": getattr(tool, "description", None)
|
|
or getattr(tool, "desc", None),
|
|
}
|
|
if simplified_tool["name"]:
|
|
simplified_tools.append(simplified_tool)
|
|
elif hasattr(tool, "__name__"):
|
|
simplified_tools.append(
|
|
{
|
|
"name": tool.__name__,
|
|
"description": getattr(tool, "__doc__", None),
|
|
}
|
|
)
|
|
else:
|
|
tool_str = str(tool)
|
|
if tool_str and tool_str != "":
|
|
simplified_tools.append({"name": tool_str, "description": None})
|
|
except Exception:
|
|
continue
|
|
|
|
return simplified_tools if simplified_tools else None
|
|
|
|
|
|
def _set_tools_on_span(span, tools):
|
|
# type: (Span, Any) -> None
|
|
"""Set available tools data on a span if tools are provided."""
|
|
if tools is not None:
|
|
simplified_tools = _simplify_langchain_tools(tools)
|
|
if simplified_tools:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
|
|
simplified_tools,
|
|
unpack=False,
|
|
)
|
|
|
|
|
|
def _wrap_configure(f):
|
|
# type: (Callable[..., Any]) -> Callable[..., Any]
|
|
|
|
@wraps(f)
|
|
def new_configure(
|
|
callback_manager_cls, # type: type
|
|
inheritable_callbacks=None, # type: Callbacks
|
|
local_callbacks=None, # type: Callbacks
|
|
*args, # type: Any
|
|
**kwargs, # type: Any
|
|
):
|
|
# type: (...) -> Any
|
|
|
|
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
|
|
if integration is None:
|
|
return f(
|
|
callback_manager_cls,
|
|
inheritable_callbacks,
|
|
local_callbacks,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
local_callbacks = local_callbacks or []
|
|
|
|
# Handle each possible type of local_callbacks. For each type, we
|
|
# extract the list of callbacks to check for SentryLangchainCallback,
|
|
# and define a function that would add the SentryLangchainCallback
|
|
# to the existing callbacks list.
|
|
if isinstance(local_callbacks, BaseCallbackManager):
|
|
callbacks_list = local_callbacks.handlers
|
|
elif isinstance(local_callbacks, BaseCallbackHandler):
|
|
callbacks_list = [local_callbacks]
|
|
elif isinstance(local_callbacks, list):
|
|
callbacks_list = local_callbacks
|
|
else:
|
|
logger.debug("Unknown callback type: %s", local_callbacks)
|
|
# Just proceed with original function call
|
|
return f(
|
|
callback_manager_cls,
|
|
inheritable_callbacks,
|
|
local_callbacks,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
# Handle each possible type of inheritable_callbacks.
|
|
if isinstance(inheritable_callbacks, BaseCallbackManager):
|
|
inheritable_callbacks_list = inheritable_callbacks.handlers
|
|
elif isinstance(inheritable_callbacks, list):
|
|
inheritable_callbacks_list = inheritable_callbacks
|
|
else:
|
|
inheritable_callbacks_list = []
|
|
|
|
if not any(
|
|
isinstance(cb, SentryLangchainCallback)
|
|
for cb in itertools.chain(callbacks_list, inheritable_callbacks_list)
|
|
):
|
|
sentry_handler = SentryLangchainCallback(
|
|
integration.max_spans,
|
|
integration.include_prompts,
|
|
)
|
|
if isinstance(local_callbacks, BaseCallbackManager):
|
|
local_callbacks = local_callbacks.copy()
|
|
local_callbacks.handlers = [
|
|
*local_callbacks.handlers,
|
|
sentry_handler,
|
|
]
|
|
elif isinstance(local_callbacks, BaseCallbackHandler):
|
|
local_callbacks = [local_callbacks, sentry_handler]
|
|
else:
|
|
local_callbacks = [*local_callbacks, sentry_handler]
|
|
|
|
return f(
|
|
callback_manager_cls,
|
|
inheritable_callbacks,
|
|
local_callbacks,
|
|
*args,
|
|
**kwargs,
|
|
)
|
|
|
|
return new_configure
|
|
|
|
|
|
def _wrap_agent_executor_invoke(f):
|
|
# type: (Callable[..., Any]) -> Callable[..., Any]
|
|
|
|
@wraps(f)
|
|
def new_invoke(self, *args, **kwargs):
|
|
# type: (Any, Any, Any) -> Any
|
|
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
|
|
if integration is None:
|
|
return f(self, *args, **kwargs)
|
|
|
|
agent_name, tools = _get_request_data(self, args, kwargs)
|
|
start_span_function = get_start_span_function()
|
|
|
|
with start_span_function(
|
|
op=OP.GEN_AI_INVOKE_AGENT,
|
|
name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
|
|
origin=LangchainIntegration.origin,
|
|
) as span:
|
|
_push_agent(agent_name)
|
|
try:
|
|
if agent_name:
|
|
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
|
|
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
|
|
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, False)
|
|
|
|
_set_tools_on_span(span, tools)
|
|
|
|
# Run the agent
|
|
result = f(self, *args, **kwargs)
|
|
|
|
input = result.get("input")
|
|
if (
|
|
input is not None
|
|
and should_send_default_pii()
|
|
and integration.include_prompts
|
|
):
|
|
normalized_messages = normalize_message_roles([input])
|
|
scope = sentry_sdk.get_current_scope()
|
|
messages_data = truncate_and_annotate_messages(
|
|
normalized_messages, span, scope
|
|
)
|
|
if messages_data is not None:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_REQUEST_MESSAGES,
|
|
messages_data,
|
|
unpack=False,
|
|
)
|
|
|
|
output = result.get("output")
|
|
if (
|
|
output is not None
|
|
and should_send_default_pii()
|
|
and integration.include_prompts
|
|
):
|
|
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
|
|
|
|
return result
|
|
finally:
|
|
# Ensure agent is popped even if an exception occurs
|
|
_pop_agent()
|
|
|
|
return new_invoke
|
|
|
|
|
|
def _wrap_agent_executor_stream(f):
|
|
# type: (Callable[..., Any]) -> Callable[..., Any]
|
|
|
|
@wraps(f)
|
|
def new_stream(self, *args, **kwargs):
|
|
# type: (Any, Any, Any) -> Any
|
|
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
|
|
if integration is None:
|
|
return f(self, *args, **kwargs)
|
|
|
|
agent_name, tools = _get_request_data(self, args, kwargs)
|
|
start_span_function = get_start_span_function()
|
|
|
|
span = start_span_function(
|
|
op=OP.GEN_AI_INVOKE_AGENT,
|
|
name=f"invoke_agent {agent_name}" if agent_name else "invoke_agent",
|
|
origin=LangchainIntegration.origin,
|
|
)
|
|
span.__enter__()
|
|
|
|
_push_agent(agent_name)
|
|
|
|
if agent_name:
|
|
span.set_data(SPANDATA.GEN_AI_AGENT_NAME, agent_name)
|
|
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
|
|
span.set_data(SPANDATA.GEN_AI_RESPONSE_STREAMING, True)
|
|
|
|
_set_tools_on_span(span, tools)
|
|
|
|
input = args[0].get("input") if len(args) >= 1 else None
|
|
if (
|
|
input is not None
|
|
and should_send_default_pii()
|
|
and integration.include_prompts
|
|
):
|
|
normalized_messages = normalize_message_roles([input])
|
|
scope = sentry_sdk.get_current_scope()
|
|
messages_data = truncate_and_annotate_messages(
|
|
normalized_messages, span, scope
|
|
)
|
|
if messages_data is not None:
|
|
set_data_normalized(
|
|
span,
|
|
SPANDATA.GEN_AI_REQUEST_MESSAGES,
|
|
messages_data,
|
|
unpack=False,
|
|
)
|
|
|
|
# Run the agent
|
|
result = f(self, *args, **kwargs)
|
|
|
|
old_iterator = result
|
|
|
|
def new_iterator():
|
|
# type: () -> Iterator[Any]
|
|
exc_info = (None, None, None) # type: tuple[Any, Any, Any]
|
|
try:
|
|
for event in old_iterator:
|
|
yield event
|
|
|
|
try:
|
|
output = event.get("output")
|
|
except Exception:
|
|
output = None
|
|
|
|
if (
|
|
output is not None
|
|
and should_send_default_pii()
|
|
and integration.include_prompts
|
|
):
|
|
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
|
|
except Exception:
|
|
exc_info = sys.exc_info()
|
|
set_span_errored(span)
|
|
raise
|
|
finally:
|
|
# Ensure cleanup happens even if iterator is abandoned or fails
|
|
_pop_agent()
|
|
span.__exit__(*exc_info)
|
|
|
|
async def new_iterator_async():
|
|
# type: () -> AsyncIterator[Any]
|
|
exc_info = (None, None, None) # type: tuple[Any, Any, Any]
|
|
try:
|
|
async for event in old_iterator:
|
|
yield event
|
|
|
|
try:
|
|
output = event.get("output")
|
|
except Exception:
|
|
output = None
|
|
|
|
if (
|
|
output is not None
|
|
and should_send_default_pii()
|
|
and integration.include_prompts
|
|
):
|
|
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output)
|
|
except Exception:
|
|
exc_info = sys.exc_info()
|
|
set_span_errored(span)
|
|
raise
|
|
finally:
|
|
# Ensure cleanup happens even if iterator is abandoned or fails
|
|
_pop_agent()
|
|
span.__exit__(*exc_info)
|
|
|
|
if str(type(result)) == "<class 'async_generator'>":
|
|
result = new_iterator_async()
|
|
else:
|
|
result = new_iterator()
|
|
|
|
return result
|
|
|
|
return new_stream
|
|
|
|
|
|
def _patch_embeddings_provider(provider_class):
|
|
# type: (Any) -> None
|
|
"""Patch an embeddings provider class with monitoring wrappers."""
|
|
if provider_class is None:
|
|
return
|
|
|
|
if hasattr(provider_class, "embed_documents"):
|
|
provider_class.embed_documents = _wrap_embedding_method(
|
|
provider_class.embed_documents
|
|
)
|
|
if hasattr(provider_class, "embed_query"):
|
|
provider_class.embed_query = _wrap_embedding_method(provider_class.embed_query)
|
|
if hasattr(provider_class, "aembed_documents"):
|
|
provider_class.aembed_documents = _wrap_async_embedding_method(
|
|
provider_class.aembed_documents
|
|
)
|
|
if hasattr(provider_class, "aembed_query"):
|
|
provider_class.aembed_query = _wrap_async_embedding_method(
|
|
provider_class.aembed_query
|
|
)
|
|
|
|
|
|
def _wrap_embedding_method(f):
|
|
# type: (Callable[..., Any]) -> Callable[..., Any]
|
|
"""Wrap sync embedding methods (embed_documents and embed_query)."""
|
|
|
|
@wraps(f)
|
|
def new_embedding_method(self, *args, **kwargs):
|
|
# type: (Any, Any, Any) -> Any
|
|
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
|
|
if integration is None:
|
|
return f(self, *args, **kwargs)
|
|
|
|
model_name = getattr(self, "model", None) or getattr(self, "model_name", None)
|
|
with sentry_sdk.start_span(
|
|
op=OP.GEN_AI_EMBEDDINGS,
|
|
name=f"embeddings {model_name}" if model_name else "embeddings",
|
|
origin=LangchainIntegration.origin,
|
|
) as span:
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
|
|
if model_name:
|
|
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
|
|
|
|
# Capture input if PII is allowed
|
|
if (
|
|
should_send_default_pii()
|
|
and integration.include_prompts
|
|
and len(args) > 0
|
|
):
|
|
input_data = args[0]
|
|
# Normalize to list format
|
|
texts = input_data if isinstance(input_data, list) else [input_data]
|
|
set_data_normalized(
|
|
span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, texts, unpack=False
|
|
)
|
|
|
|
result = f(self, *args, **kwargs)
|
|
return result
|
|
|
|
return new_embedding_method
|
|
|
|
|
|
def _wrap_async_embedding_method(f):
|
|
# type: (Callable[..., Any]) -> Callable[..., Any]
|
|
"""Wrap async embedding methods (aembed_documents and aembed_query)."""
|
|
|
|
@wraps(f)
|
|
async def new_async_embedding_method(self, *args, **kwargs):
|
|
# type: (Any, Any, Any) -> Any
|
|
integration = sentry_sdk.get_client().get_integration(LangchainIntegration)
|
|
if integration is None:
|
|
return await f(self, *args, **kwargs)
|
|
|
|
model_name = getattr(self, "model", None) or getattr(self, "model_name", None)
|
|
with sentry_sdk.start_span(
|
|
op=OP.GEN_AI_EMBEDDINGS,
|
|
name=f"embeddings {model_name}" if model_name else "embeddings",
|
|
origin=LangchainIntegration.origin,
|
|
) as span:
|
|
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "embeddings")
|
|
if model_name:
|
|
span.set_data(SPANDATA.GEN_AI_REQUEST_MODEL, model_name)
|
|
|
|
# Capture input if PII is allowed
|
|
if (
|
|
should_send_default_pii()
|
|
and integration.include_prompts
|
|
and len(args) > 0
|
|
):
|
|
input_data = args[0]
|
|
# Normalize to list format
|
|
texts = input_data if isinstance(input_data, list) else [input_data]
|
|
set_data_normalized(
|
|
span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, texts, unpack=False
|
|
)
|
|
|
|
result = await f(self, *args, **kwargs)
|
|
return result
|
|
|
|
return new_async_embedding_method
|