feat! Update session service interface to be async.

Also keep the sync version in the InMemorySessionService as create_session_sync() as a temporary migration option.

PiperOrigin-RevId: 759252188
This commit is contained in:
Google Team Member
2025-05-15 12:23:33 -07:00
committed by Copybara-Service
parent 5b3204c356
commit 1804ca39a6
23 changed files with 268 additions and 264 deletions

View File

@@ -173,7 +173,7 @@ class Runner:
The events generated by the agent.
"""
with tracer.start_as_current_span('invocation'):
session = await self.session_service.get_session(
session = self.session_service.get_session(
app_name=self.app_name, user_id=user_id, session_id=session_id
)
if not session:
@@ -197,7 +197,7 @@ class Runner:
invocation_context.agent = self._find_agent_to_run(session, root_agent)
async for event in invocation_context.agent.run_async(invocation_context):
if not event.partial:
await self.session_service.append_event(session=session, event=event)
self.session_service.append_event(session=session, event=event)
yield event
async def _append_new_message_to_session(
@@ -242,7 +242,7 @@ class Runner:
author='user',
content=new_message,
)
await self.session_service.append_event(session=session, event=event)
self.session_service.append_event(session=session, event=event)
async def run_live(
self,
@@ -324,7 +324,7 @@ class Runner:
)
async for event in invocation_context.agent.run_live(invocation_context):
await self.session_service.append_event(session=session, event=event)
self.session_service.append_event(session=session, event=event)
yield event
async def close_session(self, session: Session):
@@ -335,7 +335,7 @@ class Runner:
"""
if self.memory_service:
await self.memory_service.add_session_to_memory(session)
await self.session_service.close_session(session=session)
self.session_service.close_session(session=session)
def _find_agent_to_run(
self, session: Session, root_agent: BaseAgent