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
13 changes: 11 additions & 2 deletions src/google/adk/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@
from .sessions.base_session_service import BaseSessionService
from .sessions.base_session_service import GetSessionConfig
from .sessions.session import Session
from .telemetry import _instrumentation
from .telemetry.tracing import tracer
from .tools.base_toolset import BaseToolset
from .utils._debug_output import print_event
Expand All @@ -66,6 +67,10 @@

logger = logging.getLogger('google_adk.' + __name__)

# Silence unused warning.
# tracer is imported for backwards compatibility, to avoid breaking change in the API.
_ = tracer


def _find_active_task_isolation_scope(session) -> Optional[str]:
"""Walk session backwards; find the active paused task agent's scope.
Expand Down Expand Up @@ -454,7 +459,9 @@ async def _run_node_async(
Events flow through ic._event_queue via NodeRunner.
"""

with tracer.start_as_current_span('invocation'):
with _instrumentation.record_invocation(
entrypoint_node=node or self.agent, conversation_id=session_id
):
# 1. Setup
session = await self._get_or_create_session(
user_id=user_id, session_id=session_id
Expand Down Expand Up @@ -1040,7 +1047,9 @@ async def _run_with_trace(
new_message: Optional[types.Content] = None,
invocation_id: Optional[str] = None,
) -> AsyncGenerator[Event, None]:
with tracer.start_as_current_span('invocation'):
with _instrumentation.record_invocation(
entrypoint_node=self.agent, conversation_id=session_id
):
session = await self._get_or_create_session(
user_id=user_id,
session_id=session_id,
Expand Down
21 changes: 21 additions & 0 deletions src/google/adk/telemetry/_instrumentation.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import sys
import time
from typing import AsyncIterator
from typing import Iterator
from typing import TYPE_CHECKING

from opentelemetry import trace
Expand All @@ -35,6 +36,7 @@
from ..models.llm_request import LlmRequest
from ..models.llm_response import LlmResponse
from ..tools.base_tool import BaseTool
from ..workflow._base_node import BaseNode

logger = logging.getLogger("google_adk." + __name__)

Expand Down Expand Up @@ -69,6 +71,25 @@ def _get_elapsed_s(
return time.monotonic() - fallback_start


@contextlib.contextmanager
def record_invocation(
entrypoint_node: BaseNode | None,
conversation_id: str,
) -> Iterator[None]:
"""Top-level ``invocation`` span for a runner invocation.

Args:
entrypoint_node: The runner's root agent/node.
conversation_id: Session/conversation id.

Yields:
Nothing; the span is active for the duration of the block.
"""
del entrypoint_node, conversation_id # Unused until schema v2 lands.
with tracing.tracer.start_as_current_span("invocation"):
yield


@dataclasses.dataclass
class TelemetryContext:
"""Stores all telemetry related state."""
Expand Down
Loading