fix: fix return type of get_session.

PiperOrigin-RevId: 760638917
This commit is contained in:
Shangjie Chen 2025-05-19 08:46:19 -07:00 committed by Copybara-Service
parent 3b5232c14f
commit 0497fc578c
3 changed files with 8 additions and 6 deletions

View File

@ -99,7 +99,7 @@ class BaseSessionService(abc.ABC):
session.events.append(event)
return event
def __update_session_state(self, session: Session, event: Event):
def __update_session_state(self, session: Session, event: Event) -> None:
"""Updates the session state based on the event."""
if not event.actions or not event.actions.state_delta:
return

View File

@ -113,7 +113,7 @@ class InMemorySessionService(BaseSessionService):
user_id: str,
session_id: str,
config: Optional[GetSessionConfig] = None,
) -> Session:
) -> Optional[Session]:
return self._get_session_impl(
app_name=app_name,
user_id=user_id,
@ -128,7 +128,7 @@ class InMemorySessionService(BaseSessionService):
user_id: str,
session_id: str,
config: Optional[GetSessionConfig] = None,
) -> Session:
) -> Optional[Session]:
logger.warning('Deprecated. Please migrate to the async method.')
return self._get_session_impl(
app_name=app_name,
@ -144,7 +144,7 @@ class InMemorySessionService(BaseSessionService):
user_id: str,
session_id: str,
config: Optional[GetSessionConfig] = None,
) -> Session:
) -> Optional[Session]:
if app_name not in self.sessions:
return None
if user_id not in self.sessions[app_name]:
@ -171,7 +171,9 @@ class InMemorySessionService(BaseSessionService):
return self._merge_state(app_name, user_id, copied_session)
def _merge_state(self, app_name: str, user_id: str, copied_session: Session):
def _merge_state(
self, app_name: str, user_id: str, copied_session: Session
) -> Session:
# Merge app state
if app_name in self.app_state:
for key in self.app_state[app_name].keys():

View File

@ -119,7 +119,7 @@ class VertexAiSessionService(BaseSessionService):
user_id: str,
session_id: str,
config: Optional[GetSessionConfig] = None,
) -> Session:
) -> Optional[Session]:
reasoning_engine_id = _parse_reasoning_engine_id(app_name)
# Get session resource