Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
from .tool_call_details import ToolCallDetails
from .tool_type import ToolType
from .trace_processor.span_processor import SpanProcessor
from .utils import extract_context_from_headers, get_traceparent

__all__ = [
# Main SDK functions
Expand Down Expand Up @@ -71,6 +72,9 @@
"ExecutionType",
"InferenceOperationType",
"ToolType",
# Utility functions
"extract_context_from_headers",
"get_traceparent",
# Constants
# all constants from constants.py are exported via *
]
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from datetime import datetime

from opentelemetry.context import Context
from opentelemetry.trace import SpanKind

from .agent_details import AgentDetails
Expand Down Expand Up @@ -33,7 +34,7 @@ def start(
agent_details: AgentDetails,
tenant_details: TenantDetails,
request: Request | None = None,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
span_kind: SpanKind | None = None,
Expand All @@ -45,8 +46,9 @@ def start(
agent_details: The details of the agent making the call
tenant_details: The details of the tenant
request: Optional request details for additional context
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object. Useful when
recording a tool call after execution has already completed.
end_time: Optional explicit end time as a datetime object. When provided,
Expand All @@ -63,7 +65,7 @@ def start(
agent_details,
tenant_details,
request,
parent_id,
parent_context,
start_time,
end_time,
span_kind,
Expand All @@ -75,7 +77,7 @@ def __init__(
agent_details: AgentDetails,
tenant_details: TenantDetails,
request: Request | None = None,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
span_kind: SpanKind | None = None,
Expand All @@ -87,8 +89,9 @@ def __init__(
agent_details: The details of the agent making the call
tenant_details: The details of the tenant
request: Optional request details for additional context
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object. Useful when
recording a tool call after execution has already completed.
end_time: Optional explicit end time as a datetime object. When provided,
Expand All @@ -103,7 +106,7 @@ def __init__(
activity_name=f"{EXECUTE_TOOL_OPERATION_NAME} {details.tool_name}",
agent_details=agent_details,
tenant_details=tenant_details,
parent_id=parent_id,
parent_context=parent_context,
start_time=start_time,
end_time=end_time,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
from datetime import datetime
from typing import List

from opentelemetry.context import Context

from .agent_details import AgentDetails
from .constants import (
CHANNEL_LINK_KEY,
Expand Down Expand Up @@ -36,7 +38,7 @@ def start(
agent_details: AgentDetails,
tenant_details: TenantDetails,
request: Request | None = None,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
) -> "InferenceScope":
Expand All @@ -47,16 +49,17 @@ def start(
agent_details: The details of the agent making the call
tenant_details: The details of the tenant
request: Optional request details for additional context
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.

Returns:
A new InferenceScope instance
"""
return InferenceScope(
details, agent_details, tenant_details, request, parent_id, start_time, end_time
details, agent_details, tenant_details, request, parent_context, start_time, end_time
)

def __init__(
Expand All @@ -65,7 +68,7 @@ def __init__(
agent_details: AgentDetails,
tenant_details: TenantDetails,
request: Request | None = None,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
):
Expand All @@ -76,8 +79,9 @@ def __init__(
agent_details: The details of the agent making the call
tenant_details: The details of the tenant
request: Optional request details for additional context
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.
"""
Expand All @@ -88,7 +92,7 @@ def __init__(
activity_name=f"{details.operationName.value} {details.model}",
agent_details=agent_details,
tenant_details=tenant_details,
parent_id=parent_id,
parent_context=parent_context,
start_time=start_time,
end_time=end_time,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import logging
from datetime import datetime

from opentelemetry.context import Context
from opentelemetry.trace import SpanKind

from .agent_details import AgentDetails
Expand Down Expand Up @@ -50,6 +51,7 @@ def start(
request: Request | None = None,
caller_agent_details: AgentDetails | None = None,
caller_details: CallerDetails | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
span_kind: SpanKind | None = None,
Expand All @@ -63,6 +65,9 @@ def start(
request: Optional request details for additional context
caller_agent_details: Optional details of the caller agent
caller_details: Optional details of the non-agentic caller
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.
span_kind: Optional span kind override. Defaults to ``SpanKind.CLIENT``.
Expand All @@ -77,6 +82,7 @@ def start(
request,
caller_agent_details,
caller_details,
parent_context,
start_time,
end_time,
span_kind,
Expand All @@ -89,6 +95,7 @@ def __init__(
request: Request | None = None,
caller_agent_details: AgentDetails | None = None,
caller_details: CallerDetails | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
span_kind: SpanKind | None = None,
Expand All @@ -101,6 +108,9 @@ def __init__(
request: Optional request details for additional context
caller_agent_details: Optional details of the caller agent
caller_details: Optional details of the non-agentic caller
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.
span_kind: Optional span kind override. Defaults to ``SpanKind.CLIENT``.
Expand All @@ -118,6 +128,7 @@ def __init__(
activity_name=activity_name,
agent_details=invoke_agent_details.details,
tenant_details=tenant_details,
parent_context=parent_context,
start_time=start_time,
end_time=end_time,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
from typing import TYPE_CHECKING, Any

from opentelemetry import context, trace
from opentelemetry.context import Context
from opentelemetry.propagate import inject
from opentelemetry.trace import (
Span,
SpanKind,
Expand Down Expand Up @@ -43,7 +45,7 @@
TELEMETRY_SDK_VERSION_KEY,
TENANT_ID_KEY,
)
from .utils import get_sdk_version, parse_parent_id_to_context
from .utils import get_sdk_version

if TYPE_CHECKING:
from .agent_details import AgentDetails
Expand Down Expand Up @@ -97,7 +99,7 @@ def __init__(
activity_name: str,
agent_details: "AgentDetails | None" = None,
tenant_details: "TenantDetails | None" = None,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
):
Expand All @@ -111,8 +113,9 @@ def __init__(
activity_name: The name of the activity for display purposes
agent_details: Optional agent details
tenant_details: Optional tenant details
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to extract a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
Useful when recording an operation after it has already completed.
end_time: Optional explicit end time as a datetime object.
Expand Down Expand Up @@ -146,9 +149,8 @@ def __init__(
activity_kind = SpanKind.CONSUMER

# Get context for parent relationship
# If parent_id is provided, parse it and use it as the parent context
# If parent_context is provided, use it directly
# Otherwise, use the current context
parent_context = parse_parent_id_to_context(parent_id)
span_context = parent_context if parent_context else context.get_current()

# Convert custom start time to OTel-compatible format (nanoseconds since epoch)
Expand Down Expand Up @@ -286,6 +288,49 @@ def _end(self) -> None:
else:
self._span.end()

def get_context(self) -> Context | None:
"""Get the OpenTelemetry context for this scope's span.

This method returns a Context object containing this scope's span,
which can be used to propagate trace context to child operations
or downstream services.

Returns:
A Context containing this scope's span, or None if telemetry
is disabled or no span exists.
"""
if self._span and self._is_telemetry_enabled():
return set_span_in_context(self._span)
return None

def inject_context_to_headers(self) -> dict[str, str]:
"""Inject this span's trace context into W3C HTTP headers.

Returns a dictionary of headers containing ``traceparent`` and
optionally ``tracestate`` that can be forwarded to downstream services
or stored for later context propagation.

Example usage:

.. code-block:: python

scope = OpenTelemetryScope(...)
headers = scope.inject_context_to_headers()
# Add headers to outgoing HTTP request
requests.get("https://downstream-service/api", headers=headers)

Returns:
A dictionary containing W3C trace context headers. Returns an
empty dictionary if telemetry is disabled or no span exists.
"""
headers: dict[str, str] = {}
if self._span and self._is_telemetry_enabled():
# Create a context with the current span
ctx = set_span_in_context(self._span)
# Use the global propagator to inject trace context into headers
inject(headers, context=ctx)
return headers

def __enter__(self):
"""Enter the context manager and make span active."""
if self._span and self._is_telemetry_enabled():
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

from datetime import datetime

from opentelemetry.context import Context

from ..agent_details import AgentDetails
from ..constants import GEN_AI_OUTPUT_MESSAGES_KEY
from ..models.response import Response
Expand All @@ -23,7 +25,7 @@ def start(
agent_details: AgentDetails,
tenant_details: TenantDetails,
response: Response,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
) -> "OutputScope":
Expand All @@ -33,22 +35,25 @@ def start(
agent_details: The details of the agent
tenant_details: The details of the tenant
response: The response details from the agent
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.

Returns:
A new OutputScope instance
"""
return OutputScope(agent_details, tenant_details, response, parent_id, start_time, end_time)
return OutputScope(
agent_details, tenant_details, response, parent_context, start_time, end_time
)

def __init__(
self,
agent_details: AgentDetails,
tenant_details: TenantDetails,
response: Response,
parent_id: str | None = None,
parent_context: Context | None = None,
start_time: datetime | None = None,
end_time: datetime | None = None,
):
Expand All @@ -58,8 +63,9 @@ def __init__(
agent_details: The details of the agent
tenant_details: The details of the tenant
response: The response details from the agent
parent_id: Optional parent Activity ID used to link this span to an upstream
operation
parent_context: Optional OpenTelemetry Context used to link this span to an
upstream operation. Use ``extract_context_from_headers()`` to convert a
Context from HTTP headers containing W3C traceparent.
start_time: Optional explicit start time as a datetime object.
end_time: Optional explicit end time as a datetime object.
"""
Expand All @@ -69,7 +75,7 @@ def __init__(
activity_name=(f"{OUTPUT_OPERATION_NAME} {agent_details.agent_id}"),
agent_details=agent_details,
tenant_details=tenant_details,
parent_id=parent_id,
parent_context=parent_context,
start_time=start_time,
end_time=end_time,
)
Expand Down
Loading
Loading