fix:Remove unnecessary Aclosing

PiperOrigin-RevId: 822496695
This commit is contained in:
Alexey Guseynov
2025-10-22 02:40:07 -07:00
committed by Copybara-Service
parent 4df926388b
commit 2f4f5611bd
+23 -33
View File
@@ -282,27 +282,22 @@ class BaseAgent(BaseModel):
Event: the events generated by the agent.
"""
async def _run_with_trace() -> AsyncGenerator[Event, None]:
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
ctx = self._create_invocation_context(parent_context)
tracing.trace_agent_invocation(span, self, ctx)
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
ctx = self._create_invocation_context(parent_context)
tracing.trace_agent_invocation(span, self, ctx)
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
async with Aclosing(self._run_async_impl(ctx)) as agen:
async for event in agen:
yield event
if ctx.end_invocation:
return
if event := await self._handle_after_agent_callback(ctx):
async with Aclosing(self._run_async_impl(ctx)) as agen:
async for event in agen:
yield event
async with Aclosing(_run_with_trace()) as agen:
async for event in agen:
if ctx.end_invocation:
return
if event := await self._handle_after_agent_callback(ctx):
yield event
@final
@@ -320,24 +315,19 @@ class BaseAgent(BaseModel):
Event: the events generated by the agent.
"""
async def _run_with_trace() -> AsyncGenerator[Event, None]:
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
ctx = self._create_invocation_context(parent_context)
tracing.trace_agent_invocation(span, self, ctx)
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
with tracer.start_as_current_span(f'invoke_agent {self.name}') as span:
ctx = self._create_invocation_context(parent_context)
tracing.trace_agent_invocation(span, self, ctx)
if event := await self._handle_before_agent_callback(ctx):
yield event
if ctx.end_invocation:
return
async with Aclosing(self._run_live_impl(ctx)) as agen:
async for event in agen:
yield event
if event := await self._handle_after_agent_callback(ctx):
async with Aclosing(self._run_live_impl(ctx)) as agen:
async for event in agen:
yield event
async with Aclosing(_run_with_trace()) as agen:
async for event in agen:
if event := await self._handle_after_agent_callback(ctx):
yield event
async def _run_async_impl(