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
@@ -0,0 +1 @@
Add tool definitions attribute to interactions API spans/events. Use invoke remote agent span instead of inference span when an agent is specified.
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,11 @@ GenAI operations in aggregate.
Experimental
------------

This package is still experimental. The instrumentation may not be complete or correct just yet.
The ``interactions`` and ``embed_content`` methods are newly instrumented and may contain issues.
Please treat the telemetry produced by these methods as experimental.

Please see "TODOS.md" for a list of known defects/TODOs that are blockers to package stability.
The ``interactions`` API currently does not support automatic function calling, so no ``execute_tool`` spans
are generated.


Installation
Expand Down Expand Up @@ -66,7 +68,7 @@ Make sure to configure OpenTelemetry tracing, logging, and metrics to capture al
Limitations
***********

When using the Google GenAI SDK with automatic function calling enabled,
When using the Google GenAI SDK with automatic function calling enabled for ``generate_content``,
the OpenTelemetry instrumentation creates an ``execute_tool`` span for each tool call the SDK executes,
these spans are nested under the ``generate_content`` span.

Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
Expand Up @@ -83,19 +83,23 @@ class Stream:
)
from opentelemetry.util.genai.handler import TelemetryHandler
from opentelemetry.util.genai.invocation import (
AgentInvocation,
InferenceInvocation,
)
from opentelemetry.util.genai.stream import (
AsyncStreamWrapper,
SyncStreamWrapper,
)
from opentelemetry.util.genai.types import (
FunctionToolDefinition,
GenericPart,
GenericToolDefinition,
InputMessage,
OutputMessage,
Text,
ToolCallRequest,
ToolCallResponse,
ToolDefinition,
Uri,
)

Expand Down Expand Up @@ -127,18 +131,20 @@ def _set_co_filename(wrapped: object) -> None:

def _apply_interaction_response_attributes(
response: Interaction,
invocation: InferenceInvocation,
invocation: InferenceInvocation | AgentInvocation,
telemetry_handler: TelemetryHandler,
) -> None:
invocation.response_model_name = response.model
if getattr(response, "id", None):
invocation.response_id = response.id
if isinstance(invocation, InferenceInvocation):
invocation.response_model_name = response.model
if getattr(response, "id", None):
invocation.response_id = response.id

usage = response.usage or Usage()

invocation.input_tokens = usage.total_input_tokens
invocation.output_tokens = usage.total_output_tokens
invocation.thinking_tokens = usage.total_thought_tokens
if isinstance(invocation, InferenceInvocation):
invocation.thinking_tokens = usage.total_thought_tokens
invocation.cache_read_input_tokens = usage.total_cached_tokens

if telemetry_handler.should_capture_content():
Expand Down Expand Up @@ -245,7 +251,7 @@ class InteractionsStreamWrapper(SyncStreamWrapper[InteractionSSEEvent]):
def __init__(
self,
stream: Iterable[InteractionSSEEvent],
invocation: InferenceInvocation,
invocation: InferenceInvocation | AgentInvocation,
telemetry_handler: TelemetryHandler,
) -> None:
super().__init__(stream)
Expand Down Expand Up @@ -277,7 +283,7 @@ class AsyncInteractionsStreamWrapper(AsyncStreamWrapper[InteractionSSEEvent]):
def __init__(
self,
stream: AsyncIterable[InteractionSSEEvent],
invocation: InferenceInvocation,
invocation: InferenceInvocation | AgentInvocation,
telemetry_handler: TelemetryHandler,
) -> None:
super().__init__(stream)
Expand Down Expand Up @@ -305,6 +311,102 @@ def _on_stream_error(self, error: Exception) -> None:
self._self_invocation.fail(error)


# See https://ai.google.dev/gemini-api/docs/function-calling
def _maybe_get_tool_definitions(
tools: Any | None,
) -> list[ToolDefinition] | None:
if not isinstance(tools, Sequence) or isinstance(tools, (str, bytes)):
return None
definitions: list[ToolDefinition] = []
for tool in tools:
# Tool must be a list of dictionaries
if not isinstance(tool, dict):
continue
# This is currently a required field and the SDK will raise an error if it isn't present or
# takes on a type field that isn't supported.
tool_type = tool.get("type")
if tool_type in (
"bash",
"code_execution",
"filesystem",
"file_search",
"google_maps",
"google_search",
"url_context",
"computer_use",
):
definitions.append(
GenericToolDefinition(
name=tool.get("name") or tool_type,
type=tool_type,
)
)

elif tool_type == "mcp_server":
# Name and uri are optional.
name = tool.get("name") or tool.get("url") or "mcp_server"
# GenericToolDefinition only has 2 fields (name and type). It'd be useful to
# add support for extra fields somehow, so we can put the URI in here.
definitions.append(
GenericToolDefinition(
name=name,
type="mcp_server",
)
)
elif tool_type == "function":
definitions.append(
FunctionToolDefinition(
name=tool.get("name") or "function",
description=tool.get("description"),
parameters=tool.get("parameters"),
)
)
return definitions if definitions else None


def _start_interactions_invocation(
telemetry_handler: TelemetryHandler,
instance: InteractionsResource | AsyncInteractionsResource,
kwargs: dict[str, Any],
) -> InferenceInvocation | AgentInvocation:
# Vertex AI does not support the interactions API yet, but eventually will.
# SDK will raise an exception if model or agent is not passed or if input data is not passed.
is_vertex, server_address = _get_client_info(instance)
provider = (
GenAIAttributes.GenAiSystemValues.VERTEX_AI.value
if is_vertex
else GenAIAttributes.GenAiSystemValues.GEMINI.value
)
if agent := kwargs.get("agent"):
invocation: InferenceInvocation | AgentInvocation = (
telemetry_handler.invoke_remote_agent(
provider=provider,
request_model=kwargs.get("model"),
server_address=server_address,
agent_name=agent,
)
)
else:
invocation = telemetry_handler.inference(
provider=provider,
request_model=kwargs.get("model"),
operation_name="interactions.create",
server_address=server_address,
)
invocation.tool_definitions = _maybe_get_tool_definitions(
kwargs.get("tools")
)

if telemetry_handler.should_capture_content():
invocation.input_messages = _interactions_input_to_messages(
kwargs.get("input")
)
if system_instruction := kwargs.get("system_instruction"):
invocation.system_instruction = [Text(content=system_instruction)]

return invocation


def _create_instrumented_interactions_create(
telemetry_handler: TelemetryHandler,
) -> Callable[
Expand All @@ -322,32 +424,15 @@ def instrumented_interactions_create(
args: tuple[Any, ...],
kwargs: dict[str, Any],
) -> Interaction | InteractionsStreamWrapper:
# Vertex AI does not support the interactions API yet, but eventually will.
# SDK will raise an exception if model or agent is not passed or if input data is not passed.
is_vertex, server_address = _get_client_info(instance)
invocation = telemetry_handler.inference(
provider=(
GenAIAttributes.GenAiSystemValues.VERTEX_AI.value
if is_vertex
else GenAIAttributes.GenAiSystemValues.GEMINI.value
),
request_model=kwargs.get("model") or kwargs.get("agent"),
operation_name="interactions.create",
server_address=server_address,
invocation = _start_interactions_invocation(
telemetry_handler, instance, kwargs
)

if telemetry_handler.should_capture_content():
invocation.input_messages = _interactions_input_to_messages(
kwargs.get("input")
)
if system_instruction := kwargs.get("system_instruction"):
invocation.system_instruction = [
Text(content=system_instruction)
]

if kwargs.get("stream", False):
return InteractionsStreamWrapper(
wrapped(*args, **kwargs), invocation, telemetry_handler
wrapped(*args, **kwargs),
invocation,
telemetry_handler,
)
try:
response = wrapped(*args, **kwargs)
Expand Down Expand Up @@ -380,35 +465,21 @@ async def instrumented_interactions_create(
args: tuple[Any, ...],
kwargs: dict[str, Any],
) -> Interaction | AsyncInteractionsStreamWrapper:
is_vertex, server_address = _get_client_info(instance)
invocation = telemetry_handler.inference(
provider=(
GenAIAttributes.GenAiSystemValues.VERTEX_AI.value
if is_vertex
else GenAIAttributes.GenAiSystemValues.GEMINI.value
),
request_model=kwargs.get("model") or kwargs.get("agent"),
operation_name="interactions.create",
server_address=server_address,
invocation = _start_interactions_invocation(
telemetry_handler, instance, kwargs
)

if telemetry_handler.should_capture_content():
invocation.input_messages = _interactions_input_to_messages(
kwargs.get("input")
)
if system_instruction := kwargs.get("system_instruction"):
invocation.system_instruction = [
Text(content=system_instruction)
]

if kwargs.get("stream", False):
return AsyncInteractionsStreamWrapper(
await wrapped(*args, **kwargs),
invocation,
telemetry_handler,
)
try:
response = cast(Interaction, await wrapped(*args, **kwargs))
response = cast(
Interaction,
await wrapped(*args, **kwargs),
)
_apply_interaction_response_attributes(
response, invocation, telemetry_handler
)
Expand Down
Loading
Loading