diff --git a/sentry_sdk/__init__.py b/sentry_sdk/__init__.py index a900f73c0b..50ab118d0d 100644 --- a/sentry_sdk/__init__.py +++ b/sentry_sdk/__init__.py @@ -6,6 +6,8 @@ from sentry_sdk.api import * # noqa +from sentry_sdk.tracing import new_trace, update_current_span + from sentry_sdk.consts import VERSION # noqa __all__ = [ # noqa @@ -43,6 +45,8 @@ "start_span", "start_transaction", "trace", + "new_trace", + "update_current_span", "monitor", "logger", "start_session", diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 4cdec031e1..3dba03f6b8 100644 --- a/sentry_sdk/tracing.py +++ b/sentry_sdk/tracing.py @@ -1,8 +1,12 @@ from __future__ import annotations from datetime import datetime +import functools +import inspect import json import warnings +import sentry_sdk + from opentelemetry import trace as otel_trace, context from opentelemetry.trace import ( format_trace_id, @@ -20,6 +24,7 @@ DEFAULT_SPAN_NAME, DEFAULT_SPAN_ORIGIN, BAGGAGE_HEADER_NAME, + OP, SENTRY_TRACE_HEADER_NAME, SPANSTATUS, SPANDATA, @@ -574,3 +579,134 @@ async def my_async_function(): return start_child_span_decorator(func) else: return start_child_span_decorator + + +def set_input_attributes(span, as_type, args, kwargs): + # depending on `as_type` set some attributes on the span derived from args and kwargs + pass + + +def set_output_attributes(span, as_type, result): + # depending on `as_type` set some attributes on the span derived from result + pass + + +DEFAULT_SPAN_OP = "function" + + +def new_trace(func=None, *, as_type=None, name=None): + def span_decorator(f, *a, **kw): + @functools.wraps(f) + async def async_wrapper(*args, **kwargs): + op = kw.get("op", DEFAULT_SPAN_OP) + span_name = kw.get("name", f.__name__) + attributes = kw.get("attributes", {}) + + with sentry_sdk.start_span( + op=op, + name=span_name, + ) as span: + for key, value in attributes.items(): + span.set_attribute(key, value) + + set_input_attributes(span, as_type, args, kwargs) + + # run wrapped function + result = await f(*args, **kwargs) + + set_output_attributes(span, as_type, result) + return result + + @functools.wraps(f) + def sync_wrapper(*args, **kwargs): + op = kw.get("op", DEFAULT_SPAN_OP) + span_name = kw.get("name", f.__name__) + attributes = kw.get("attributes", {}) + + with sentry_sdk.start_span( + op=op, + name=span_name, + ) as span: + for key, value in attributes.items(): + span.set_attribute(key, value) + + set_input_attributes(span, as_type, args, kwargs) + + # run wrapped function + result = f(*args, **kwargs) + + set_output_attributes(span, as_type, result) + return result + + if inspect.iscoroutinefunction(f): + wrapper = async_wrapper + else: + wrapper = sync_wrapper + + return wrapper + + def ai_chat_decorator(f): + attributes = { + "gen_ai.operation.name": "chat", + } + + return span_decorator( + f, + op=OP.GEN_AI_CHAT, + name="chat [MODEL]", + attributes=attributes, + ) + + def ai_agent_decorator(f): + span_name = name or f.__name__ + attributes = { + "gen_ai.agent.name": span_name, + "gen_ai.operation.name": "invoke_agent", + } + + return span_decorator( + f, + op=OP.GEN_AI_INVOKE_AGENT, + name=f"invoke_agent {span_name}", + attributes=attributes, + ) + + def ai_tool_decorator(f): + span_name = name or f.__name__ + attributes = { + "gen_ai.tool.name": span_name, + "gen_ai.operation.name": "execute_tool", + } + + return span_decorator( + f, + op=OP.GEN_AI_EXECUTE_TOOL, + name=f"execute_tool {span_name}", + attributes=attributes, + ) + + decorator_for_type = { + "ai_chat": ai_chat_decorator, + "ai_agent": ai_agent_decorator, + "ai_tool": ai_tool_decorator, + } + decorator = decorator_for_type.get(as_type, span_decorator) + + if func: + return decorator(func) + else: + return decorator + + +def update_current_span(op=None, name=None, attributes=None) -> None: + current_span = sentry_sdk.get_current_span() + + if op is not None: + current_span.op = op + + if name is not None: + current_span.name = name + + if attributes is not None: + for key, value in attributes.items(): + current_span.set_attribute(key, value)