feat(langchain): Record run_name in on_chat_model_start (#5924)

Set the `run_name` LangChain attribute as the `gen_ai.function_id` attribute in `on_chat_model_start`.

Author: alexander-alderman-webb

sentry_sdk/consts.py

@@ -549,6 +549,12 @@ class SPANDATA:
     Example: "Hello!"
     """
 
+    GEN_AI_FUNCTION_ID = "gen_ai.function_id"
+    """
+    Framework-specific tracing label for the execution of a function or other unit of execution in a generative AI system.
+    Example: "my-awesome-function"
+    """
+
     GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
     """
     The name of the operation being performed.

sentry_sdk/integrations/langchain.py

@@ -432,6 +432,13 @@ def on_chat_model_start(
                     SPANDATA.GEN_AI_AGENT_NAME, agent_metadata["lc_agent_name"]
                 )
 
+            run_name = kwargs.get("name")
+            if run_name:
+                span.set_data(
+                    SPANDATA.GEN_AI_FUNCTION_ID,
+                    run_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)

tests/conftest.py

@@ -1265,26 +1265,31 @@ def streaming_chat_completions_model_response():
 
 @pytest.fixture
 def nonstreaming_chat_completions_model_response():
-    return openai.types.chat.ChatCompletion(
-        id="chatcmpl-test",
-        choices=[
-            openai.types.chat.chat_completion.Choice(
-                index=0,
-                finish_reason="stop",
-                message=openai.types.chat.ChatCompletionMessage(
-                    role="assistant", content="Test response"
-                ),
-            )
-        ],
-        created=1234567890,
-        model="gpt-3.5-turbo",
-        object="chat.completion",
-        usage=openai.types.CompletionUsage(
-            prompt_tokens=10,
-            completion_tokens=20,
-            total_tokens=30,
-        ),
-    )
+    def inner(
+        response_id: str,
+        response_model: str,
+        message_content: str,
+        created: int,
+        usage: openai.types.CompletionUsage,
+    ):
+        return openai.types.chat.ChatCompletion(
+            id=response_id,
+            choices=[
+                openai.types.chat.chat_completion.Choice(
+                    index=0,
+                    finish_reason="stop",
+                    message=openai.types.chat.ChatCompletionMessage(
+                        role="assistant", content=message_content
+                    ),
+                )
+            ],
+            created=created,
+            model=response_model,
+            object="chat.completion",
+            usage=usage,
+        )
+
+    return inner
 
 
 @pytest.fixture

tests/integrations/langchain/test_langchain.py

@@ -67,6 +67,7 @@
 )
 
 LANGCHAIN_VERSION = package_version("langchain")
+LANGCHAIN_OPENAI_VERSION = package_version("langchain-openai")
 
 
 @tool
@@ -170,6 +171,68 @@ def test_langchain_text_completion(
     assert llm_span["data"]["gen_ai.usage.output_tokens"] == 15
 
 
+def test_langchain_chat_with_run_name(
+    sentry_init,
+    capture_events,
+    get_model_response,
+    nonstreaming_chat_completions_model_response,
+):
+    sentry_init(
+        integrations=[
+            LangchainIntegration(
+                include_prompts=True,
+            )
+        ],
+        traces_sample_rate=1.0,
+        send_default_pii=True,
+    )
+    events = capture_events()
+
+    request_headers = {}
+    # Changed in https://github.com/langchain-ai/langchain/pull/32655
+    if LANGCHAIN_OPENAI_VERSION >= (0, 3, 32):
+        request_headers["X-Stainless-Raw-Response"] = "True"
+
+    model_response = get_model_response(
+        nonstreaming_chat_completions_model_response(
+            response_id="chat-id",
+            response_model="response-model-id",
+            message_content="the model response",
+            created=10000000,
+            usage=CompletionUsage(
+                prompt_tokens=20,
+                completion_tokens=10,
+                total_tokens=30,
+            ),
+        ),
+        serialize_pydantic=True,
+        request_headers=request_headers,
+    )
+
+    llm = ChatOpenAI(
+        model_name="gpt-3.5-turbo",
+        temperature=0,
+        openai_api_key="badkey",
+    )
+
+    with patch.object(
+        llm.client._client._client,
+        "send",
+        return_value=model_response,
+    ) as _:
+        with start_transaction():
+            llm.invoke(
+                "How many letters in the word eudca",
+                config={"run_name": "my-snazzy-pipeline"},
+            )
+
+    tx = events[0]
+
+    chat_spans = list(x for x in tx["spans"] if x["op"] == "gen_ai.chat")
+    assert len(chat_spans) == 1
+    assert chat_spans[0]["data"][SPANDATA.GEN_AI_FUNCTION_ID] == "my-snazzy-pipeline"
+
+
 @pytest.mark.skipif(
     LANGCHAIN_VERSION < (1,),
     reason="LangChain 1.0+ required (ONE AGENT refactor)",