Skip to content

fix(langchain): Use gen_ai.tool.definitions attribute when data collection is enabled #7204

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ericapisani wants to merge 1 commit into from
+41 −27
Open

fix(langchain): Use gen_ai.tool.definitions attribute when data collection is enabled #7204

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 17 additions & 6 deletions sentry_sdk/integrations/langchain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ def _get_ai_system(all_params: "Dict[str, Any]") -> "Optional[str]":

DATA_FIELDS = {
"frequency_penalty": SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY,
# "function_call" is an OpenAI convention for the now-legacy Chat Completions API field
"function_call": SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
"max_tokens": SPANDATA.GEN_AI_REQUEST_MAX_TOKENS,
"presence_penalty": SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY,
Expand Down Expand Up @@ -417,6 +418,12 @@ def on_llm_start(

for key, attribute in DATA_FIELDS.items():
if key in all_params and all_params[key] is not None:
# This is correctly gated on "inputs" at the moment because the
# "on_llm_start" method is the start of a request.
#
# TODO: GEN_AI_RESPONSE_TOOL_CALLS will need to be
# transitioned to non-deprecated tool call attributes

if (
attribute == SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS
and has_data_collection_enabled(client.options)
Expand Down Expand Up @@ -516,6 +523,11 @@ def on_chat_model_start(
for key, attribute in DATA_FIELDS.items():
if key in all_params and all_params[key] is not None:
if (
# This is correctly gated on "inputs" at the moment because the
# "on_chat_model_start" method is the start of a request.
#
# TODO: GEN_AI_RESPONSE_TOOL_CALLS will need to be
# transitioned to non-deprecated tool call attributes
attribute == SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS
and has_data_collection_enabled(client.options)
and not client.options["data_collection"]["gen_ai"]["inputs"]
Expand Down Expand Up @@ -615,14 +627,11 @@ def on_llm_end(

client = sentry_sdk.get_client()

record_inputs = False
record_outputs = False
if has_data_collection_enabled(client.options):
record_inputs = client.options["data_collection"]["gen_ai"]["inputs"]
record_outputs = client.options["data_collection"]["gen_ai"]["outputs"]
elif should_send_default_pii() and self.include_prompts:
# TODO: Remove this branch once `send_default_pii` is deprecated
record_inputs = True
record_outputs = True

try:
Expand All @@ -647,7 +656,7 @@ def on_llm_end(
if response_model is not None:
set_on_span(SPANDATA.GEN_AI_RESPONSE_MODEL, response_model)

if record_inputs:
if record_outputs:
tool_calls = getattr(generation.message, "tool_calls", None)
if tool_calls is not None and tool_calls != []:
set_data_normalized(
Expand Down Expand Up @@ -1048,17 +1057,19 @@ def _set_tools_on_span(span: "Union[Span, StreamedSpan]", tools: "Any") -> None:
return

client = sentry_sdk.get_client()
attribute_name = SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS
if has_data_collection_enabled(client.options):
if not client.options["data_collection"]["gen_ai"]["inputs"]:
return

else:
attribute_name = SPANDATA.GEN_AI_TOOL_DEFINITIONS
# Before data collection was introduced this was set unconditionally, so it
# stays that way when data collection is not configured.
simplified_tools = _simplify_langchain_tools(tools)
if simplified_tools:
set_data_normalized(
span,
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
attribute_name,
simplified_tools,
unpack=False,
)
Expand Down
45 changes: 24 additions & 21 deletions tests/integrations/langchain/test_langchain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6497,10 +6497,10 @@ def test_langchain_text_completion_data_collection(
False,
False,
[
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
[],
[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS],
id="gen-ai-inputs-and-outputs-enabled-tools-collected",
),
pytest.param(
Expand All @@ -6509,6 +6509,7 @@ def test_langchain_text_completion_data_collection(
True,
[],
[
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
Expand All @@ -6518,33 +6519,34 @@ def test_langchain_text_completion_data_collection(
{"gen_ai": {"inputs": True, "outputs": False}},
False,
False,
[SPANDATA.GEN_AI_TOOL_DEFINITIONS],
[
# REQUEST_AVAILABLE_TOOLS is the legacy value set when data collection is not enabled
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
[],
id="gen-ai-inputs-enabled-outputs-disabled-tools-collected",
id="gen-ai-inputs-enabled-outputs-disabled-only-tool-definitions-collected",
),
pytest.param(
{"gen_ai": {"inputs": False, "outputs": True}},
True,
True,
[],
[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS],
[
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
id="gen-ai-outputs-enabled-inputs-disabled-tools-not-collected",
id="gen-ai-outputs-enabled-inputs-disabled-only-response-tool-calls-collected",
),
pytest.param(
{"gen_ai": {}},
False,
False,
[
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
[],
[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS],
id="gen-ai-inputs-and-outputs-omitted-default-to-enabled",
),
pytest.param(
Expand All @@ -6555,15 +6557,18 @@ def test_langchain_text_completion_data_collection(
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
[],
[SPANDATA.GEN_AI_TOOL_DEFINITIONS],
id="no-gen-ai-config-legacy-pii-and-include-prompts-enabled",
),
pytest.param(
None,
False,
False,
[SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS],
[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS],
[
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS,
],
id="no-gen-ai-config-available-tools-collected-regardless-of-pii",
),
],
Expand Down Expand Up @@ -6664,16 +6669,14 @@ def test_langchain_data_collection_tools(
for key in expected_absent:
assert key not in chat_spans[0], f"{key} should not have been collected"

# Available tools are gated the same way on every span they are set on
available_tools_collected = (
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in expected_present
)
assert (
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in chat_spans[1]
) is available_tools_collected
assert (
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS in invoke_agent_span
) is available_tools_collected
# Tool definitions are gated the same way on every span they are set on
for key in (
SPANDATA.GEN_AI_TOOL_DEFINITIONS,
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
):
collected = key in expected_present
assert (key in chat_spans[1]) is collected
assert (key in invoke_agent_span) is collected


@pytest.mark.parametrize("span_streaming", [True, False])
Expand Down
Loading