feat(ai): Langchain integration (#2911)

Integration for Langchain. 

---------

Co-authored-by: Anton Pirker <[email protected]>

Author: colin-sentry

Diff for: .github/workflows/test-integrations-data-processing.yml

@@ -58,6 +58,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh "py${{ matrix.python-version }}-huey-latest" --cov=tests --cov=sentry_sdk --cov-report= --cov-branch
+      - name: Test langchain latest
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh "py${{ matrix.python-version }}-langchain-latest" --cov=tests --cov=sentry_sdk --cov-report= --cov-branch
       - name: Test openai latest
         run: |
           set -x # print commands that are executed
@@ -114,6 +118,10 @@ jobs:
         run: |
           set -x # print commands that are executed
           ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-huey" --cov=tests --cov=sentry_sdk --cov-report= --cov-branch
+      - name: Test langchain pinned
+        run: |
+          set -x # print commands that are executed
+          ./scripts/runtox.sh --exclude-latest "py${{ matrix.python-version }}-langchain" --cov=tests --cov=sentry_sdk --cov-report= --cov-branch
       - name: Test openai pinned
         run: |
           set -x # print commands that are executed

Diff for: mypy.ini

@@ -48,6 +48,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-asgiref.*]
 ignore_missing_imports = True
+[mypy-langchain_core.*]
+ignore_missing_imports = True
 [mypy-executing.*]
 ignore_missing_imports = True
 [mypy-asttokens.*]

Diff for: scripts/split-tox-gh-actions/split-tox-gh-actions.py

@@ -70,6 +70,7 @@
         "beam",
         "celery",
         "huey",
+        "langchain",
         "openai",
         "rq",
     ],

Diff for: sentry_sdk/ai/__init__.py

@@ -0,0 +1,77 @@
+from functools import wraps
+
+import sentry_sdk.utils
+from sentry_sdk import start_span
+from sentry_sdk.tracing import Span
+from sentry_sdk.utils import ContextVar
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Optional, Callable, Any
+
+_ai_pipeline_name = ContextVar("ai_pipeline_name", default=None)
+
+
+def set_ai_pipeline_name(name):
+    # type: (Optional[str]) -> None
+    _ai_pipeline_name.set(name)
+
+
+def get_ai_pipeline_name():
+    # type: () -> Optional[str]
+    return _ai_pipeline_name.get()
+
+
+def ai_track(description, **span_kwargs):
+    # type: (str, Any) -> Callable[..., Any]
+    def decorator(f):
+        # type: (Callable[..., Any]) -> Callable[..., Any]
+        @wraps(f)
+        def wrapped(*args, **kwargs):
+            # type: (Any, Any) -> Any
+            curr_pipeline = _ai_pipeline_name.get()
+            op = span_kwargs.get("op", "ai.run" if curr_pipeline else "ai.pipeline")
+            with start_span(description=description, op=op, **span_kwargs) as span:
+                if curr_pipeline:
+                    span.set_data("ai.pipeline.name", curr_pipeline)
+                    return f(*args, **kwargs)
+                else:
+                    _ai_pipeline_name.set(description)
+                    try:
+                        res = f(*args, **kwargs)
+                    except Exception as e:
+                        event, hint = sentry_sdk.utils.event_from_exception(
+                            e,
+                            client_options=sentry_sdk.get_client().options,
+                            mechanism={"type": "ai_monitoring", "handled": False},
+                        )
+                        sentry_sdk.capture_event(event, hint=hint)
+                        raise e from None
+                    finally:
+                        _ai_pipeline_name.set(None)
+                    return res
+
+        return wrapped
+
+    return decorator
+
+
+def record_token_usage(
+    span, prompt_tokens=None, completion_tokens=None, total_tokens=None
+):
+    # type: (Span, Optional[int], Optional[int], Optional[int]) -> None
+    ai_pipeline_name = get_ai_pipeline_name()
+    if ai_pipeline_name:
+        span.set_data("ai.pipeline.name", ai_pipeline_name)
+    if prompt_tokens is not None:
+        span.set_measurement("ai_prompt_tokens_used", value=prompt_tokens)
+    if completion_tokens is not None:
+        span.set_measurement("ai_completion_tokens_used", value=completion_tokens)
+    if (
+        total_tokens is None
+        and prompt_tokens is not None
+        and completion_tokens is not None
+    ):
+        total_tokens = prompt_tokens + completion_tokens
+    if total_tokens is not None:
+        span.set_measurement("ai_total_tokens_used", total_tokens)

Diff for: sentry_sdk/ai/monitoring.py

@@ -0,0 +1,32 @@
+from sentry_sdk._types import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+
+from sentry_sdk.tracing import Span
+from sentry_sdk.utils import logger
+
+
+def _normalize_data(data):
+    # type: (Any) -> Any
+
+    # convert pydantic data (e.g. OpenAI v1+) to json compatible format
+    if hasattr(data, "model_dump"):
+        try:
+            return data.model_dump()
+        except Exception as e:
+            logger.warning("Could not convert pydantic data to JSON: %s", e)
+            return data
+    if isinstance(data, list):
+        if len(data) == 1:
+            return _normalize_data(data[0])  # remove empty dimensions
+        return list(_normalize_data(x) for x in data)
+    if isinstance(data, dict):
+        return {k: _normalize_data(v) for (k, v) in data.items()}
+    return data
+
+
+def set_data_normalized(span, key, value):
+    # type: (Span, str, Any) -> None
+    normalized = _normalize_data(value)
+    span.set_data(key, normalized)

Diff for: sentry_sdk/ai/utils.py

@@ -91,6 +91,85 @@ class SPANDATA:
     See: https://develop.sentry.dev/sdk/performance/span-data-conventions/
     """
 
+    AI_INPUT_MESSAGES = "ai.input_messages"
+    """
+    The input messages to an LLM call.
+    Example: [{"role": "user", "message": "hello"}]
+    """
+
+    AI_MODEL_ID = "ai.model_id"
+    """
+    The unique descriptor of the model being execugted
+    Example: gpt-4
+    """
+
+    AI_METADATA = "ai.metadata"
+    """
+    Extra metadata passed to an AI pipeline step.
+    Example: {"executed_function": "add_integers"}
+    """
+
+    AI_TAGS = "ai.tags"
+    """
+    Tags that describe an AI pipeline step.
+    Example: {"executed_function": "add_integers"}
+    """
+
+    AI_STREAMING = "ai.streaming"
+    """
+    Whether or not the AI model call's repsonse was streamed back asynchronously
+    Example: true
+    """
+
+    AI_TEMPERATURE = "ai.temperature"
+    """
+    For an AI model call, the temperature parameter. Temperature essentially means how random the output will be.
+    Example: 0.5
+    """
+
+    AI_TOP_P = "ai.top_p"
+    """
+    For an AI model call, the top_p parameter. Top_p essentially controls how random the output will be.
+    Example: 0.5
+    """
+
+    AI_TOP_K = "ai.top_k"
+    """
+    For an AI model call, the top_k parameter. Top_k essentially controls how random the output will be.
+    Example: 35
+    """
+
+    AI_FUNCTION_CALL = "ai.function_call"
+    """
+    For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls
+    """
+
+    AI_TOOL_CALLS = "ai.tool_calls"
+    """
+    For an AI model call, the function that was called. This is deprecated for OpenAI, and replaced by tool_calls
+    """
+
+    AI_TOOLS = "ai.tools"
+    """
+    For an AI model call, the functions that are available
+    """
+
+    AI_RESPONSE_FORMAT = "ai.response_format"
+    """
+    For an AI model call, the format of the response
+    """
+
+    AI_LOGIT_BIAS = "ai.response_format"
+    """
+    For an AI model call, the logit bias
+    """
+
+    AI_RESPONSES = "ai.responses"
+    """
+    The responses to an AI model call. Always as a list.
+    Example: ["hello", "world"]
+    """
+
     DB_NAME = "db.name"
     """
     The name of the database being accessed. For commands that switch the database, this should be set to the target database (even if the command fails).
@@ -245,6 +324,11 @@ class OP:
     MIDDLEWARE_STARLITE_SEND = "middleware.starlite.send"
     OPENAI_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.openai"
     OPENAI_EMBEDDINGS_CREATE = "ai.embeddings.create.openai"
+    LANGCHAIN_PIPELINE = "ai.pipeline.langchain"
+    LANGCHAIN_RUN = "ai.run.langchain"
+    LANGCHAIN_TOOL = "ai.tool.langchain"
+    LANGCHAIN_AGENT = "ai.agent.langchain"
+    LANGCHAIN_CHAT_COMPLETIONS_CREATE = "ai.chat_completions.create.langchain"
     QUEUE_SUBMIT_ARQ = "queue.submit.arq"
     QUEUE_TASK_ARQ = "queue.task.arq"
     QUEUE_SUBMIT_CELERY = "queue.submit.celery"

Diff for: sentry_sdk/consts.py

@@ -85,6 +85,7 @@ def iter_default_integrations(with_auto_enabling_integrations):
     "sentry_sdk.integrations.graphene.GrapheneIntegration",
     "sentry_sdk.integrations.httpx.HttpxIntegration",
     "sentry_sdk.integrations.huey.HueyIntegration",
+    "sentry_sdk.integrations.langchain.LangchainIntegration",
     "sentry_sdk.integrations.loguru.LoguruIntegration",
     "sentry_sdk.integrations.openai.OpenAIIntegration",
     "sentry_sdk.integrations.pymongo.PyMongoIntegration",