2020import logging
2121import os
2222import time
23+ import typing
2324from typing import (
2425 Any ,
2526 AsyncIterator ,
4243 GenerateContentConfig ,
4344 GenerateContentConfigOrDict ,
4445 GenerateContentResponse ,
46+ Tool ,
47+ ToolListUnionDict ,
48+ ToolUnionDict ,
4549)
4650
4751from opentelemetry import context as context_api
8286from .otel_wrapper import OTelWrapper
8387from .tool_call_wrapper import wrapped as wrapped_tool
8488
89+ _is_mcp_imported = False
90+ if typing .TYPE_CHECKING :
91+ from mcp import ClientSession as McpClientSession
92+ from mcp import Tool as McpTool
93+
94+ is_mcp_imported = True
95+ else :
96+ try :
97+ from mcp import ClientSession as McpClientSession
98+ from mcp import Tool as McpTool
99+
100+ _is_mcp_imported = True
101+ except ImportError :
102+ McpClientSession = None
103+ McpTool = None
104+
85105_logger = logging .getLogger (__name__ )
86106
107+ GEN_AI_TOOL_DEFINITIONS = getattr (
108+ gen_ai_attributes , "GEN_AI_TOOL_DEFINITIONS" , "gen_ai.tool.definitions"
109+ )
87110
88111# Constant used to make the absence of content more understandable.
89112_CONTENT_ELIDED = "<elided>"
@@ -185,6 +208,68 @@ def _to_dict(value: object):
185208 return json .loads (json .dumps (value ))
186209
187210
211+ def _tool_to_tool_definition (tool : ToolUnionDict ) -> MessagePart :
212+ if hasattr (tool , "model_dump" ):
213+ return tool .model_dump (exclude_none = True )
214+
215+ return str (tool )
216+
217+
218+ def _callable_tool_to_tool_definition (tool : Any ) -> MessagePart :
219+ doc = getattr (tool , "__doc__" , "" ) or ""
220+ return {
221+ "name" : getattr (tool , "__name__" , type (tool ).__name__ ),
222+ "description" : doc .strip (),
223+ }
224+
225+
226+ def _mcp_tool_to_tool_definition (tool : McpTool ) -> MessagePart :
227+ if hasattr (tool , "model_dump" ):
228+ return tool .model_dump (exclude_none = True )
229+
230+ return {
231+ "name" : getattr (tool , "name" , type (tool ).__name__ ),
232+ "description" : getattr (tool , "description" , "" ) or "" ,
233+ "input_schema" : getattr (tool , "input_schema" , {}),
234+ }
235+
236+
237+ def _to_tool_definition_common (tool : ToolUnionDict ) -> MessagePart :
238+ if isinstance (tool , dict ):
239+ return tool
240+
241+ if isinstance (tool , Tool ):
242+ return _tool_to_tool_definition (tool )
243+
244+ if callable (tool ):
245+ return _callable_tool_to_tool_definition (tool )
246+
247+ if _is_mcp_imported and isinstance (tool , McpTool ):
248+ return _mcp_tool_to_tool_definition (tool )
249+
250+ try :
251+ return {"raw_definition" : json .loads (json .dumps (tool ))}
252+ except Exception : # pylint: disable=broad-exception-caught
253+ return {
254+ "error" : f"failed to serialize tool definition, tool type={ type (tool ).__name__ } "
255+ }
256+
257+
258+ def _to_tool_definition (tool : ToolUnionDict ) -> MessagePart :
259+ if _is_mcp_imported and isinstance (tool , McpClientSession ):
260+ return None
261+
262+ return _to_tool_definition_common (tool )
263+
264+
265+ async def _to_tool_definition_async (tool : ToolUnionDict ) -> MessagePart :
266+ if _is_mcp_imported and isinstance (tool , McpClientSession ):
267+ result = await tool .list_tools ()
268+ return [t .model_dump (exclude_none = True ) for t in result .tools ]
269+
270+ return _to_tool_definition_common (tool )
271+
272+
188273def _create_request_attributes (
189274 config : Optional [GenerateContentConfigOrDict ],
190275 allow_list : AllowList ,
@@ -285,10 +370,22 @@ def _config_to_system_instruction(
285370 return config .system_instruction
286371
287372
373+ def _config_to_tools (
374+ config : Union [GenerateContentConfigOrDict , None ],
375+ ) -> Union [ToolListUnionDict , None ]:
376+ if not config :
377+ return None
378+
379+ if isinstance (config , dict ):
380+ return GenerateContentConfig .model_validate (config ).tools
381+ return config .tools
382+
383+
288384def _create_completion_details_attributes (
289385 input_messages : list [InputMessage ],
290386 output_messages : list [OutputMessage ],
291387 system_instructions : list [MessagePart ],
388+ tool_definitions : list [MessagePart ],
292389 as_str : bool = False ,
293390) -> dict [str , AttributeValue ]:
294391 attributes : dict [str , AttributeValue ] = {
@@ -306,6 +403,9 @@ def _create_completion_details_attributes(
306403 dataclasses .asdict (sys_instr ) for sys_instr in system_instructions
307404 ]
308405
406+ if tool_definitions :
407+ attributes [GEN_AI_TOOL_DEFINITIONS ] = tool_definitions
408+
309409 return attributes
310410
311411
@@ -324,6 +424,7 @@ def __init__(
324424 model : str ,
325425 completion_hook : CompletionHook ,
326426 generate_content_config_key_allowlist : Optional [AllowList ] = None ,
427+ is_async : bool = False ,
327428 ):
328429 self ._start_time = time .time_ns ()
329430 self ._otel_wrapper = otel_wrapper
@@ -345,6 +446,7 @@ def __init__(
345446 self ._generate_content_config_key_allowlist = (
346447 generate_content_config_key_allowlist or AllowList ()
347448 )
449+ self ._is_async = is_async
348450
349451 def wrapped_config (
350452 self , config : Optional [GenerateContentConfigOrDict ]
@@ -461,6 +563,44 @@ def _maybe_update_error_type(self, response: GenerateContentResponse):
461563 block_reason = response .prompt_feedback .block_reason .name .upper ()
462564 self ._error_type = f"BLOCKED_{ block_reason } "
463565
566+ def _maybe_get_tool_definitions (self , config ):
567+ if (
568+ self .sem_conv_opt_in_mode
569+ != _StabilityMode .GEN_AI_LATEST_EXPERIMENTAL
570+ ):
571+ return None
572+
573+ tool_definitions = []
574+ if tools := _config_to_tools (config ):
575+ for tool in tools :
576+ definition = _to_tool_definition (tool )
577+ if definition is None :
578+ continue
579+ if isinstance (definition , list ):
580+ tool_definitions .extend (definition )
581+ else :
582+ tool_definitions .append (definition )
583+ return tool_definitions
584+
585+ async def _maybe_get_tool_definitions_async (self , config ):
586+ if (
587+ self .sem_conv_opt_in_mode
588+ != _StabilityMode .GEN_AI_LATEST_EXPERIMENTAL
589+ ):
590+ return None
591+
592+ tool_definitions = []
593+ if tools := _config_to_tools (config ):
594+ for tool in tools :
595+ definition = await _to_tool_definition_async (tool )
596+ if definition is None :
597+ continue
598+ if isinstance (definition , list ):
599+ tool_definitions .extend (definition )
600+ else :
601+ tool_definitions .append (definition )
602+ return tool_definitions
603+
464604 def _maybe_log_completion_details (
465605 self ,
466606 extra_attributes : dict [str , AttributeValue ],
@@ -469,6 +609,7 @@ def _maybe_log_completion_details(
469609 request : Union [ContentListUnion , ContentListUnionDict ],
470610 candidates : list [Candidate ],
471611 config : Optional [GenerateContentConfigOrDict ] = None ,
612+ tool_definitions : list [MessagePart ] = None ,
472613 ):
473614 if (
474615 self .sem_conv_opt_in_mode
@@ -503,6 +644,7 @@ def _maybe_log_completion_details(
503644 input_messages ,
504645 output_messages ,
505646 system_instructions ,
647+ tool_definitions ,
506648 )
507649 if self ._content_recording_enabled in [
508650 ContentCapturingMode .EVENT_ONLY ,
@@ -737,6 +879,7 @@ def instrumented_generate_content(
737879 model ,
738880 completion_hook ,
739881 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
882+ is_async = False ,
740883 )
741884 request_attributes = _create_request_attributes (
742885 config ,
@@ -774,13 +917,17 @@ def instrumented_generate_content(
774917 finally :
775918 final_attributes = helper .create_final_attributes ()
776919 span .set_attributes (final_attributes )
920+ maybe_tool_definitions = helper ._maybe_get_tool_definitions (
921+ config
922+ )
777923 helper ._maybe_log_completion_details (
778924 extra_attributes ,
779925 request_attributes ,
780926 final_attributes ,
781927 contents ,
782928 candidates ,
783929 config ,
930+ maybe_tool_definitions ,
784931 )
785932 helper ._record_token_usage_metric ()
786933 helper ._record_duration_metric ()
@@ -812,6 +959,7 @@ def instrumented_generate_content_stream(
812959 model ,
813960 completion_hook ,
814961 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
962+ is_async = False ,
815963 )
816964 request_attributes = _create_request_attributes (
817965 config ,
@@ -849,13 +997,17 @@ def instrumented_generate_content_stream(
849997 finally :
850998 final_attributes = helper .create_final_attributes ()
851999 span .set_attributes (final_attributes )
1000+ maybe_tool_definitions = helper ._maybe_get_tool_definitions (
1001+ config
1002+ )
8521003 helper ._maybe_log_completion_details (
8531004 extra_attributes ,
8541005 request_attributes ,
8551006 final_attributes ,
8561007 contents ,
8571008 candidates ,
8581009 config ,
1010+ maybe_tool_definitions ,
8591011 )
8601012 helper ._record_token_usage_metric ()
8611013 helper ._record_duration_metric ()
@@ -886,6 +1038,7 @@ async def instrumented_generate_content(
8861038 model ,
8871039 completion_hook ,
8881040 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
1041+ is_async = True ,
8891042 )
8901043 request_attributes = _create_request_attributes (
8911044 config ,
@@ -923,13 +1076,17 @@ async def instrumented_generate_content(
9231076 finally :
9241077 final_attributes = helper .create_final_attributes ()
9251078 span .set_attributes (final_attributes )
1079+ maybe_tool_definitions = (
1080+ await helper ._maybe_get_tool_definitions_async (config )
1081+ )
9261082 helper ._maybe_log_completion_details (
9271083 extra_attributes ,
9281084 request_attributes ,
9291085 final_attributes ,
9301086 contents ,
9311087 candidates ,
9321088 config ,
1089+ maybe_tool_definitions ,
9331090 )
9341091 helper ._record_token_usage_metric ()
9351092 helper ._record_duration_metric ()
@@ -961,6 +1118,7 @@ async def instrumented_generate_content_stream(
9611118 model ,
9621119 completion_hook ,
9631120 generate_content_config_key_allowlist = generate_content_config_key_allowlist ,
1121+ is_async = True ,
9641122 )
9651123 request_attributes = _create_request_attributes (
9661124 config ,
@@ -991,13 +1149,17 @@ async def instrumented_generate_content_stream(
9911149 helper ._record_token_usage_metric ()
9921150 final_attributes = helper .create_final_attributes ()
9931151 span .set_attributes (final_attributes )
1152+ maybe_tool_definitions = (
1153+ await helper ._maybe_get_tool_definitions_async (config )
1154+ )
9941155 helper ._maybe_log_completion_details (
9951156 extra_attributes ,
9961157 request_attributes ,
9971158 final_attributes ,
9981159 contents ,
9991160 [],
10001161 config ,
1162+ maybe_tool_definitions ,
10011163 )
10021164 helper ._record_duration_metric ()
10031165 with trace .use_span (span , end_on_exit = True ):
@@ -1025,13 +1187,19 @@ async def _response_async_generator_wrapper():
10251187 finally :
10261188 final_attributes = helper .create_final_attributes ()
10271189 span .set_attributes (final_attributes )
1190+ maybe_tool_definitions = (
1191+ await helper ._maybe_get_tool_definitions_async (
1192+ config
1193+ )
1194+ )
10281195 helper ._maybe_log_completion_details (
10291196 extra_attributes ,
10301197 request_attributes ,
10311198 final_attributes ,
10321199 contents ,
10331200 candidates ,
10341201 config ,
1202+ maybe_tool_definitions ,
10351203 )
10361204 helper ._record_token_usage_metric ()
10371205 helper ._record_duration_metric ()
0 commit comments