mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: aclose all async generators to fix OTel tracing context
See https://github.com/google/adk-python/issues/1670#issuecomment-3115891100 PiperOrigin-RevId: 794659547
This commit is contained in:
Kacper Jawoszek
committed by
Copybara-Service
Copybara-Service
parent
c5af44cfc0
commit
a30c63c593
@@ -21,6 +21,7 @@ from typing import Union
|
||||
from google.genai import live
|
||||
from google.genai import types
|
||||
|
||||
from ..utils.context_utils import Aclosing
|
||||
from .base_llm_connection import BaseLlmConnection
|
||||
from .llm_response import LlmResponse
|
||||
|
||||
@@ -142,90 +143,92 @@ class GeminiLlmConnection(BaseLlmConnection):
|
||||
"""
|
||||
|
||||
text = ''
|
||||
async for message in self._gemini_session.receive():
|
||||
logger.debug('Got LLM Live message: %s', message)
|
||||
if message.server_content:
|
||||
content = message.server_content.model_turn
|
||||
if content and content.parts:
|
||||
llm_response = LlmResponse(
|
||||
content=content, interrupted=message.server_content.interrupted
|
||||
)
|
||||
if content.parts[0].text:
|
||||
text += content.parts[0].text
|
||||
llm_response.partial = True
|
||||
# don't yield the merged text event when receiving audio data
|
||||
elif text and not content.parts[0].inline_data:
|
||||
async with Aclosing(self._gemini_session.receive()) as agen:
|
||||
async for message in agen:
|
||||
logger.debug('Got LLM Live message: %s', message)
|
||||
if message.server_content:
|
||||
content = message.server_content.model_turn
|
||||
if content and content.parts:
|
||||
llm_response = LlmResponse(
|
||||
content=content, interrupted=message.server_content.interrupted
|
||||
)
|
||||
if content.parts[0].text:
|
||||
text += content.parts[0].text
|
||||
llm_response.partial = True
|
||||
# don't yield the merged text event when receiving audio data
|
||||
elif text and not content.parts[0].inline_data:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
yield llm_response
|
||||
if (
|
||||
message.server_content.input_transcription
|
||||
and message.server_content.input_transcription.text
|
||||
):
|
||||
user_text = message.server_content.input_transcription.text
|
||||
parts = [
|
||||
types.Part.from_text(
|
||||
text=user_text,
|
||||
)
|
||||
]
|
||||
llm_response = LlmResponse(
|
||||
content=types.Content(role='user', parts=parts)
|
||||
)
|
||||
yield llm_response
|
||||
if (
|
||||
message.server_content.output_transcription
|
||||
and message.server_content.output_transcription.text
|
||||
):
|
||||
# TODO: Right now, we just support output_transcription without
|
||||
# changing interface and data protocol. Later, we can consider to
|
||||
# support output_transcription as a separate field in LlmResponse.
|
||||
|
||||
# Transcription is always considered as partial event
|
||||
# We rely on other control signals to determine when to yield the
|
||||
# full text response(turn_complete, interrupted, or tool_call).
|
||||
text += message.server_content.output_transcription.text
|
||||
parts = [
|
||||
types.Part.from_text(
|
||||
text=message.server_content.output_transcription.text
|
||||
)
|
||||
]
|
||||
llm_response = LlmResponse(
|
||||
content=types.Content(role='model', parts=parts), partial=True
|
||||
)
|
||||
yield llm_response
|
||||
|
||||
if message.server_content.turn_complete:
|
||||
if text:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
yield LlmResponse(
|
||||
turn_complete=True,
|
||||
interrupted=message.server_content.interrupted,
|
||||
)
|
||||
break
|
||||
# in case of empty content or parts, we sill surface it
|
||||
# in case it's an interrupted message, we merge the previous partial
|
||||
# text. Other we don't merge. because content can be none when model
|
||||
# safety threshold is triggered
|
||||
if message.server_content.interrupted and text:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
yield llm_response
|
||||
if (
|
||||
message.server_content.input_transcription
|
||||
and message.server_content.input_transcription.text
|
||||
):
|
||||
user_text = message.server_content.input_transcription.text
|
||||
parts = [
|
||||
types.Part.from_text(
|
||||
text=user_text,
|
||||
)
|
||||
]
|
||||
llm_response = LlmResponse(
|
||||
content=types.Content(role='user', parts=parts)
|
||||
)
|
||||
yield llm_response
|
||||
if (
|
||||
message.server_content.output_transcription
|
||||
and message.server_content.output_transcription.text
|
||||
):
|
||||
# TODO: Right now, we just support output_transcription without
|
||||
# changing interface and data protocol. Later, we can consider to
|
||||
# support output_transcription as a separate field in LlmResponse.
|
||||
|
||||
# Transcription is always considered as partial event
|
||||
# We rely on other control signals to determine when to yield the
|
||||
# full text response(turn_complete, interrupted, or tool_call).
|
||||
text += message.server_content.output_transcription.text
|
||||
parts = [
|
||||
types.Part.from_text(
|
||||
text=message.server_content.output_transcription.text
|
||||
)
|
||||
]
|
||||
llm_response = LlmResponse(
|
||||
content=types.Content(role='model', parts=parts), partial=True
|
||||
)
|
||||
yield llm_response
|
||||
|
||||
if message.server_content.turn_complete:
|
||||
yield LlmResponse(interrupted=message.server_content.interrupted)
|
||||
if message.tool_call:
|
||||
if text:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
yield LlmResponse(
|
||||
turn_complete=True, interrupted=message.server_content.interrupted
|
||||
parts = [
|
||||
types.Part(function_call=function_call)
|
||||
for function_call in message.tool_call.function_calls
|
||||
]
|
||||
yield LlmResponse(content=types.Content(role='model', parts=parts))
|
||||
if message.session_resumption_update:
|
||||
logger.info('Redeived session reassumption message: %s', message)
|
||||
yield (
|
||||
LlmResponse(
|
||||
live_session_resumption_update=message.session_resumption_update
|
||||
)
|
||||
)
|
||||
break
|
||||
# in case of empty content or parts, we sill surface it
|
||||
# in case it's an interrupted message, we merge the previous partial
|
||||
# text. Other we don't merge. because content can be none when model
|
||||
# safety threshold is triggered
|
||||
if message.server_content.interrupted and text:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
yield LlmResponse(interrupted=message.server_content.interrupted)
|
||||
if message.tool_call:
|
||||
if text:
|
||||
yield self.__build_full_text_response(text)
|
||||
text = ''
|
||||
parts = [
|
||||
types.Part(function_call=function_call)
|
||||
for function_call in message.tool_call.function_calls
|
||||
]
|
||||
yield LlmResponse(content=types.Content(role='model', parts=parts))
|
||||
if message.session_resumption_update:
|
||||
logger.info('Redeived session reassumption message: %s', message)
|
||||
yield (
|
||||
LlmResponse(
|
||||
live_session_resumption_update=message.session_resumption_update
|
||||
)
|
||||
)
|
||||
|
||||
async def close(self):
|
||||
"""Closes the llm server connection."""
|
||||
|
||||
@@ -32,6 +32,7 @@ from google.genai.types import FinishReason
|
||||
from typing_extensions import override
|
||||
|
||||
from .. import version
|
||||
from ..utils.context_utils import Aclosing
|
||||
from ..utils.variant_utils import GoogleLLMVariant
|
||||
from .base_llm import BaseLlm
|
||||
from .base_llm_connection import BaseLlmConnection
|
||||
@@ -141,39 +142,40 @@ class Gemini(BaseLlm):
|
||||
# contents are sent, we send an accumulated event which contains all the
|
||||
# previous partial content. The only difference is bidi rely on
|
||||
# complete_turn flag to detect end while sse depends on finish_reason.
|
||||
async for response in responses:
|
||||
logger.debug(_build_response_log(response))
|
||||
llm_response = LlmResponse.create(response)
|
||||
usage_metadata = llm_response.usage_metadata
|
||||
if (
|
||||
llm_response.content
|
||||
and llm_response.content.parts
|
||||
and llm_response.content.parts[0].text
|
||||
):
|
||||
part0 = llm_response.content.parts[0]
|
||||
if part0.thought:
|
||||
thought_text += part0.text
|
||||
else:
|
||||
text += part0.text
|
||||
llm_response.partial = True
|
||||
elif (thought_text or text) and (
|
||||
not llm_response.content
|
||||
or not llm_response.content.parts
|
||||
# don't yield the merged text event when receiving audio data
|
||||
or not llm_response.content.parts[0].inline_data
|
||||
):
|
||||
parts = []
|
||||
if thought_text:
|
||||
parts.append(types.Part(text=thought_text, thought=True))
|
||||
if text:
|
||||
parts.append(types.Part.from_text(text=text))
|
||||
yield LlmResponse(
|
||||
content=types.ModelContent(parts=parts),
|
||||
usage_metadata=llm_response.usage_metadata,
|
||||
)
|
||||
thought_text = ''
|
||||
text = ''
|
||||
yield llm_response
|
||||
async with Aclosing(responses) as agen:
|
||||
async for response in agen:
|
||||
logger.debug(_build_response_log(response))
|
||||
llm_response = LlmResponse.create(response)
|
||||
usage_metadata = llm_response.usage_metadata
|
||||
if (
|
||||
llm_response.content
|
||||
and llm_response.content.parts
|
||||
and llm_response.content.parts[0].text
|
||||
):
|
||||
part0 = llm_response.content.parts[0]
|
||||
if part0.thought:
|
||||
thought_text += part0.text
|
||||
else:
|
||||
text += part0.text
|
||||
llm_response.partial = True
|
||||
elif (thought_text or text) and (
|
||||
not llm_response.content
|
||||
or not llm_response.content.parts
|
||||
# don't yield the merged text event when receiving audio data
|
||||
or not llm_response.content.parts[0].inline_data
|
||||
):
|
||||
parts = []
|
||||
if thought_text:
|
||||
parts.append(types.Part(text=thought_text, thought=True))
|
||||
if text:
|
||||
parts.append(types.Part.from_text(text=text))
|
||||
yield LlmResponse(
|
||||
content=types.ModelContent(parts=parts),
|
||||
usage_metadata=llm_response.usage_metadata,
|
||||
)
|
||||
thought_text = ''
|
||||
text = ''
|
||||
yield llm_response
|
||||
|
||||
# generate an aggregated content at the end regardless the
|
||||
# response.candidates[0].finish_reason
|
||||
|
||||
Reference in New Issue
Block a user