google-adk 0.4.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

google/adk/sessions/database_session_service.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 import copy
 from datetime import datetime
+from datetime import timezone
 import json
 import logging
-from typing import Any, Optional
+from typing import Any
+from typing import Optional
 import uuid
 
 from sqlalchemy import Boolean
@@ -49,23 +51,18 @@ from ..events.event import Event
 from . import _session_util
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
-from .base_session_service import ListEventsResponse
 from .base_session_service import ListSessionsResponse
 from .session import Session
 from .state import State
 
-
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("google_adk." + __name__)
 
 DEFAULT_MAX_KEY_LENGTH = 128
 DEFAULT_MAX_VARCHAR_LENGTH = 256
 
 
 class DynamicJSON(TypeDecorator):
-  """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON
-
-  serialization for other databases.
-  """
+  """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON serialization for other databases."""
 
   impl = Text  # Default implementation is TEXT
 
@@ -219,6 +216,7 @@ class StorageAppState(Base):
       DateTime(), default=func.now(), onupdate=func.now()
   )
 
+
 class StorageUserState(Base):
   """Represents a user state stored in the database."""
 
@@ -242,10 +240,7 @@ class DatabaseSessionService(BaseSessionService):
   """A session service that uses a database for storage."""
 
   def __init__(self, db_url: str):
-    """
-    Args:
-        db_url: The database URL to connect to.
-    """
+    """Initializes the database session service with a database URL."""
     # 1. Create DB engine for db connection
     # 2. Create all tables based on schema
     # 3. Initialize all properties
@@ -274,7 +269,7 @@ class DatabaseSessionService(BaseSessionService):
     self.inspector = inspect(self.db_engine)
 
     # DB session factory method
-    self.DatabaseSessionFactory: sessionmaker[DatabaseSessionFactory] = (
+    self.database_session_factory: sessionmaker[DatabaseSessionFactory] = (
         sessionmaker(bind=self.db_engine)
     )
 
@@ -283,7 +278,7 @@ class DatabaseSessionService(BaseSessionService):
     Base.metadata.create_all(self.db_engine)
 
   @override
-  def create_session(
+  async def create_session(
       self,
       *,
       app_name: str,
@@ -297,11 +292,11 @@ class DatabaseSessionService(BaseSessionService):
     # 4. Build the session object with generated id
     # 5. Return the session
 
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
 
       # Fetch app and user states from storage
-      storage_app_state = sessionFactory.get(StorageAppState, (app_name))
-      storage_user_state = sessionFactory.get(
+      storage_app_state = session_factory.get(StorageAppState, (app_name))
+      storage_user_state = session_factory.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -311,12 +306,12 @@ class DatabaseSessionService(BaseSessionService):
       # Create state tables if not exist
       if not storage_app_state:
         storage_app_state = StorageAppState(app_name=app_name, state={})
-        sessionFactory.add(storage_app_state)
+        session_factory.add(storage_app_state)
       if not storage_user_state:
         storage_user_state = StorageUserState(
             app_name=app_name, user_id=user_id, state={}
         )
-        sessionFactory.add(storage_user_state)
+        session_factory.add(storage_user_state)
 
       # Extract state deltas
       app_state_delta, user_state_delta, session_state = _extract_state_delta(
@@ -340,10 +335,10 @@ class DatabaseSessionService(BaseSessionService):
           id=session_id,
           state=session_state,
       )
-      sessionFactory.add(storage_session)
-      sessionFactory.commit()
+      session_factory.add(storage_session)
+      session_factory.commit()
 
-      sessionFactory.refresh(storage_session)
+      session_factory.refresh(storage_session)
 
       # Merge states for response
       merged_state = _merge_state(app_state, user_state, session_state)
@@ -357,7 +352,7 @@ class DatabaseSessionService(BaseSessionService):
       return session
 
   @override
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
@@ -368,29 +363,37 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Get the storage session entry from session table
     # 2. Get all the events based on session id and filtering config
     # 3. Convert and return the session
-    with self.DatabaseSessionFactory() as sessionFactory:
-      storage_session = sessionFactory.get(
+    with self.database_session_factory() as session_factory:
+      storage_session = session_factory.get(
           StorageSession, (app_name, user_id, session_id)
       )
       if storage_session is None:
         return None
 
+      if config and config.after_timestamp:
+        after_dt = datetime.fromtimestamp(
+            config.after_timestamp, tz=timezone.utc
+        )
+        timestamp_filter = StorageEvent.timestamp > after_dt
+      else:
+        timestamp_filter = True
+
       storage_events = (
-          sessionFactory.query(StorageEvent)
+          session_factory.query(StorageEvent)
           .filter(StorageEvent.session_id == storage_session.id)
-          .filter(
-              StorageEvent.timestamp < config.after_timestamp
-              if config
-              else True
-          )
-          .limit(config.num_recent_events if config else None)
+          .filter(timestamp_filter)
           .order_by(StorageEvent.timestamp.asc())
+          .limit(
+              config.num_recent_events
+              if config and config.num_recent_events
+              else None
+          )
           .all()
       )
 
       # Fetch states from storage
-      storage_app_state = sessionFactory.get(StorageAppState, (app_name))
-      storage_user_state = sessionFactory.get(
+      storage_app_state = session_factory.get(StorageAppState, (app_name))
+      storage_user_state = session_factory.get(
           StorageUserState, (app_name, user_id)
       )
 
@@ -431,12 +434,12 @@ class DatabaseSessionService(BaseSessionService):
     return session
 
   @override
-  def list_sessions(
+  async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
       results = (
-          sessionFactory.query(StorageSession)
+          session_factory.query(StorageSession)
           .filter(StorageSession.app_name == app_name)
           .filter(StorageSession.user_id == user_id)
           .all()
@@ -454,20 +457,20 @@ class DatabaseSessionService(BaseSessionService):
       return ListSessionsResponse(sessions=sessions)
 
   @override
-  def delete_session(
+  async def delete_session(
       self, app_name: str, user_id: str, session_id: str
   ) -> None:
-    with self.DatabaseSessionFactory() as sessionFactory:
+    with self.database_session_factory() as session_factory:
       stmt = delete(StorageSession).where(
           StorageSession.app_name == app_name,
           StorageSession.user_id == user_id,
           StorageSession.id == session_id,
       )
-      sessionFactory.execute(stmt)
-      sessionFactory.commit()
+      session_factory.execute(stmt)
+      session_factory.commit()
 
   @override
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     logger.info(f"Append event: {event} to session {session.id}")
 
     if event.partial:
@@ -476,22 +479,25 @@ class DatabaseSessionService(BaseSessionService):
     # 1. Check if timestamp is stale
     # 2. Update session attributes based on event config
     # 3. Store event to table
-    with self.DatabaseSessionFactory() as sessionFactory:
-      storage_session = sessionFactory.get(
+    with self.database_session_factory() as session_factory:
+      storage_session = session_factory.get(
           StorageSession, (session.app_name, session.user_id, session.id)
       )
 
       if storage_session.update_time.timestamp() > session.last_update_time:
         raise ValueError(
-            f"Session last_update_time {session.last_update_time} is later than"
-            f" the upate_time in storage {storage_session.update_time}"
+            "The last_update_time provided in the session object"
+            f" {datetime.fromtimestamp(session.last_update_time):'%Y-%m-%d %H:%M:%S'} is"
+            " earlier than the update_time in the storage_session"
+            f" {storage_session.update_time:'%Y-%m-%d %H:%M:%S'}. Please check"
+            " if it is a stale session."
         )
 
       # Fetch states from storage
-      storage_app_state = sessionFactory.get(
+      storage_app_state = session_factory.get(
           StorageAppState, (session.app_name)
       )
-      storage_user_state = sessionFactory.get(
+      storage_user_state = session_factory.get(
           StorageUserState, (session.app_name, session.user_id)
       )
 
@@ -540,27 +546,18 @@ class DatabaseSessionService(BaseSessionService):
       if event.content:
         storage_event.content = _session_util.encode_content(event.content)
 
-      sessionFactory.add(storage_event)
+      session_factory.add(storage_event)
 
-      sessionFactory.commit()
-      sessionFactory.refresh(storage_session)
+      session_factory.commit()
+      session_factory.refresh(storage_session)
 
       # Update timestamp with commit time
       session.last_update_time = storage_session.update_time.timestamp()
 
     # Also update the in-memory session
-    super().append_event(session=session, event=event)
+    await super().append_event(session=session, event=event)
     return event
 
-  @override
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    raise NotImplementedError()
 
 def convert_event(event: StorageEvent) -> Event:
   """Converts a storage event to an event."""

google/adk/sessions/in_memory_session_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import copy
+import logging
 import time
 from typing import Any
 from typing import Optional
@@ -23,17 +24,19 @@ from typing_extensions import override
 from ..events.event import Event
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
-from .base_session_service import ListEventsResponse
 from .base_session_service import ListSessionsResponse
 from .session import Session
 from .state import State
 
+logger = logging.getLogger('google_adk.' + __name__)
+
 
 class InMemorySessionService(BaseSessionService):
   """An in-memory implementation of the session service."""
 
   def __init__(self):
-    # A map from app name to a map from user ID to a map from session ID to session.
+    # A map from app name to a map from user ID to a map from session ID to
+    # session.
     self.sessions: dict[str, dict[str, dict[str, Session]]] = {}
     # A map from app name to a map from user ID to a map from key to the value.
     self.user_state: dict[str, dict[str, dict[str, Any]]] = {}
@@ -41,7 +44,38 @@ class InMemorySessionService(BaseSessionService):
     self.app_state: dict[str, dict[str, Any]] = {}
 
   @override
-  def create_session(
+  async def create_session(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      state: Optional[dict[str, Any]] = None,
+      session_id: Optional[str] = None,
+  ) -> Session:
+    return self._create_session_impl(
+        app_name=app_name,
+        user_id=user_id,
+        state=state,
+        session_id=session_id,
+    )
+
+  def create_session_sync(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      state: Optional[dict[str, Any]] = None,
+      session_id: Optional[str] = None,
+  ) -> Session:
+    logger.warning('Deprecated. Please migrate to the async method.')
+    return self._create_session_impl(
+        app_name=app_name,
+        user_id=user_id,
+        state=state,
+        session_id=session_id,
+    )
+
+  def _create_session_impl(
       self,
       *,
       app_name: str,
@@ -72,14 +106,45 @@ class InMemorySessionService(BaseSessionService):
     return self._merge_state(app_name, user_id, copied_session)
 
   @override
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
       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,
+        session_id=session_id,
+        config=config,
+    )
+
+  def get_session_sync(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      config: Optional[GetSessionConfig] = None,
+  ) -> Optional[Session]:
+    logger.warning('Deprecated. Please migrate to the async method.')
+    return self._get_session_impl(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        config=config,
+    )
+
+  def _get_session_impl(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      config: Optional[GetSessionConfig] = None,
+  ) -> Optional[Session]:
     if app_name not in self.sessions:
       return None
     if user_id not in self.sessions[app_name]:
@@ -95,18 +160,20 @@ class InMemorySessionService(BaseSessionService):
         copied_session.events = copied_session.events[
             -config.num_recent_events :
         ]
-      elif config.after_timestamp:
-        i = len(session.events) - 1
+      if config.after_timestamp:
+        i = len(copied_session.events) - 1
         while i >= 0:
           if copied_session.events[i].timestamp < config.after_timestamp:
             break
           i -= 1
         if i >= 0:
-          copied_session.events = copied_session.events[i:]
+          copied_session.events = copied_session.events[i + 1 :]
 
     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():
@@ -128,7 +195,18 @@ class InMemorySessionService(BaseSessionService):
     return copied_session
 
   @override
-  def list_sessions(
+  async def list_sessions(
+      self, *, app_name: str, user_id: str
+  ) -> ListSessionsResponse:
+    return self._list_sessions_impl(app_name=app_name, user_id=user_id)
+
+  def list_sessions_sync(
+      self, *, app_name: str, user_id: str
+  ) -> ListSessionsResponse:
+    logger.warning('Deprecated. Please migrate to the async method.')
+    return self._list_sessions_impl(app_name=app_name, user_id=user_id)
+
+  def _list_sessions_impl(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
     empty_response = ListSessionsResponse()
@@ -145,12 +223,26 @@ class InMemorySessionService(BaseSessionService):
       sessions_without_events.append(copied_session)
     return ListSessionsResponse(sessions=sessions_without_events)
 
-  @override
-  def delete_session(
+  async def delete_session(
+      self, *, app_name: str, user_id: str, session_id: str
+  ) -> None:
+    self._delete_session_impl(
+        app_name=app_name, user_id=user_id, session_id=session_id
+    )
+
+  def delete_session_sync(
+      self, *, app_name: str, user_id: str, session_id: str
+  ) -> None:
+    logger.warning('Deprecated. Please migrate to the async method.')
+    self._delete_session_impl(
+        app_name=app_name, user_id=user_id, session_id=session_id
+    )
+
+  def _delete_session_impl(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     if (
-        self.get_session(
+        self._get_session_impl(
             app_name=app_name, user_id=user_id, session_id=session_id
         )
         is None
@@ -160,9 +252,9 @@ class InMemorySessionService(BaseSessionService):
     self.sessions[app_name][user_id].pop(session_id)
 
   @override
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     # Update the in-memory session.
-    super().append_event(session=session, event=event)
+    await super().append_event(session=session, event=event)
     session.last_update_time = event.timestamp
 
     # Update the storage session
@@ -189,18 +281,8 @@ class InMemorySessionService(BaseSessionService):
           ] = event.actions.state_delta[key]
 
     storage_session = self.sessions[app_name][user_id].get(session_id)
-    super().append_event(session=storage_session, event=event)
+    await super().append_event(session=storage_session, event=event)
 
     storage_session.last_update_time = event.timestamp
 
     return event
-
-  @override
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    raise NotImplementedError()

google/adk/sessions/session.py

@@ -14,6 +14,7 @@
 
 from typing import Any
 
+from pydantic import alias_generators
 from pydantic import BaseModel
 from pydantic import ConfigDict
 from pydantic import Field
@@ -37,7 +38,10 @@ class Session(BaseModel):
   model_config = ConfigDict(
       extra='forbid',
       arbitrary_types_allowed=True,
+      alias_generator=alias_generators.to_camel,
+      populate_by_name=True,
   )
+  """The pydantic model config."""
 
   id: str
   """The unique identifier of the session."""

google/adk/sessions/vertex_ai_session_service.py

@@ -14,7 +14,8 @@
 import logging
 import re
 import time
-from typing import Any, Optional
+from typing import Any
+from typing import Optional
 
 from dateutil import parser
 from google import genai
@@ -25,13 +26,11 @@ from ..events.event_actions import EventActions
 from . import _session_util
 from .base_session_service import BaseSessionService
 from .base_session_service import GetSessionConfig
-from .base_session_service import ListEventsResponse
 from .base_session_service import ListSessionsResponse
 from .session import Session
 
-
 isoparse = parser.isoparse
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class VertexAiSessionService(BaseSessionService):
@@ -49,7 +48,7 @@ class VertexAiSessionService(BaseSessionService):
     self.api_client = client._api_client
 
   @override
-  def create_session(
+  async def create_session(
       self,
       *,
       app_name: str,
@@ -57,13 +56,19 @@ class VertexAiSessionService(BaseSessionService):
       state: Optional[dict[str, Any]] = None,
       session_id: Optional[str] = None,
   ) -> Session:
+    if session_id:
+      raise ValueError(
+          'User-provided Session id is not supported for'
+          ' VertexAISessionService.'
+      )
+
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
 
     session_json_dict = {'user_id': user_id}
     if state:
       session_json_dict['session_state'] = state
 
-    api_response = self.api_client.request(
+    api_response = await self.api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions',
         request_dict=session_json_dict,
@@ -75,7 +80,7 @@ class VertexAiSessionService(BaseSessionService):
 
     max_retry_attempt = 5
     while max_retry_attempt >= 0:
-      lro_response = self.api_client.request(
+      lro_response = await self.api_client.async_request(
           http_method='GET',
           path=f'operations/{operation_id}',
           request_dict={},
@@ -88,7 +93,7 @@ class VertexAiSessionService(BaseSessionService):
       max_retry_attempt -= 1
 
     # Get session resource
-    get_session_api_response = self.api_client.request(
+    get_session_api_response = await self.api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
@@ -107,18 +112,18 @@ class VertexAiSessionService(BaseSessionService):
     return session
 
   @override
-  def get_session(
+  async def get_session(
       self,
       *,
       app_name: str,
       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
-    get_session_api_response = self.api_client.request(
+    get_session_api_response = await self.api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
@@ -136,7 +141,7 @@ class VertexAiSessionService(BaseSessionService):
         last_update_time=update_timestamp,
     )
 
-    list_events_api_response = self.api_client.request(
+    list_events_api_response = await self.api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events',
         request_dict={},
@@ -170,7 +175,7 @@ class VertexAiSessionService(BaseSessionService):
     return session
 
   @override
-  def list_sessions(
+  async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
@@ -197,50 +202,23 @@ class VertexAiSessionService(BaseSessionService):
       sessions.append(session)
     return ListSessionsResponse(sessions=sessions)
 
-  def delete_session(
+  async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
     reasoning_engine_id = _parse_reasoning_engine_id(app_name)
-    self.api_client.request(
+    await self.api_client.async_request(
         http_method='DELETE',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
         request_dict={},
     )
 
   @override
-  def list_events(
-      self,
-      *,
-      app_name: str,
-      user_id: str,
-      session_id: str,
-  ) -> ListEventsResponse:
-    reasoning_engine_id = _parse_reasoning_engine_id(app_name)
-    api_response = self.api_client.request(
-        http_method='GET',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}/events',
-        request_dict={},
-    )
-
-    logger.info(f'List events response {api_response}')
-
-    # Handles empty response case
-    if api_response.get('httpHeaders', None):
-      return ListEventsResponse()
-
-    session_events = api_response['sessionEvents']
-
-    return ListEventsResponse(
-        events=[_from_api_event(event) for event in session_events]
-    )
-
-  @override
-  def append_event(self, session: Session, event: Event) -> Event:
+  async def append_event(self, session: Session, event: Event) -> Event:
     # Update the in-memory session.
-    super().append_event(session=session, event=event)
+    await super().append_event(session=session, event=event)
 
     reasoning_engine_id = _parse_reasoning_engine_id(session.app_name)
-    self.api_client.request(
+    await self.api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
         request_dict=_convert_event_to_json(event),

google/adk/telemetry.py

@@ -111,6 +111,9 @@ def trace_call_llm(
   span.set_attribute(
       'gcp.vertex.agent.invocation_id', invocation_context.invocation_id
   )
+  span.set_attribute(
+      'gcp.vertex.agent.session_id', invocation_context.session.id
+  )
   span.set_attribute('gcp.vertex.agent.event_id', event_id)
   # Consider removing once GenAI SDK provides a way to record this info.
   span.set_attribute(