google-adk-mongo-session-service 0.0.1__py3-none-any.whl

google_adk_mongo_session_service/__init__.py

@@ -0,0 +1,3 @@
+from .mongo_session_service import MongoSessionService
+
+__all__ = ["MongoSessionService"]

google_adk_mongo_session_service/models.py

@@ -0,0 +1,185 @@
+"""DAL models and _id helpers for MongoSessionService.
+
+Document _id is computed at read/write time from PK fields (not stored in model).
+Collections: sessions, events, app_states, user_states, adk_internal_metadata.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from datetime import timezone
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+
+from google.adk.events.event import Event
+from google.adk.sessions.session import Session
+
+
+# --- _id helpers (deterministic, PK-derived) ---
+
+
+def session_doc_id(app_name: str, user_id: str, id: str) -> str:
+    return f"{app_name}_{user_id}_{id}"
+
+
+def event_doc_id(event_id: str, app_name: str, user_id: str, session_id: str) -> str:
+    return f"{event_id}_{app_name}_{user_id}_{session_id}"
+
+
+def app_state_doc_id(app_name: str) -> str:
+    return app_name
+
+
+def user_state_doc_id(app_name: str, user_id: str) -> str:
+    return f"{app_name}_{user_id}"
+
+
+def metadata_doc_id(key: str) -> str:
+    return key
+
+
+# --- DAL models (no _id field; _id added at serialization time) ---
+
+
+class MongoSession(BaseModel):
+    """Document model for sessions collection."""
+
+    app_name: str
+    user_id: str
+    id: str
+    state: dict[str, Any] = Field(default_factory=dict)
+    create_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    update_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def to_doc(self) -> dict[str, Any]:
+        doc = self.model_dump()
+        doc["_id"] = session_doc_id(self.app_name, self.user_id, self.id)
+        return doc
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> MongoSession:
+        data = {k: v for k, v in doc.items() if k != "_id"}
+        return cls.model_validate(data)
+
+    def to_session(
+        self,
+        state: dict[str, Any] | None = None,
+        events: list[Event] | None = None,
+    ) -> Session:
+        if state is None:
+            state = {}
+        if events is None:
+            events = []
+        update_ts = self.update_time.timestamp()
+        if self.update_time.tzinfo is None:
+            update_ts = self.update_time.replace(tzinfo=timezone.utc).timestamp()
+        return Session(
+            app_name=self.app_name,
+            user_id=self.user_id,
+            id=self.id,
+            state=state,
+            events=events,
+            last_update_time=update_ts,
+        )
+
+
+class MongoEvent(BaseModel):
+    """Document model for events collection."""
+
+    id: str
+    app_name: str
+    user_id: str
+    session_id: str
+    invocation_id: str = ""
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    event_data: dict[str, Any] | None = None
+
+    def to_doc(self) -> dict[str, Any]:
+        doc = self.model_dump()
+        doc["_id"] = event_doc_id(self.id, self.app_name, self.user_id, self.session_id)
+        return doc
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> MongoEvent:
+        data = {k: v for k, v in doc.items() if k != "_id"}
+        return cls.model_validate(data)
+
+    def to_event(self) -> Event:
+        ts = self.timestamp.timestamp()
+        if self.timestamp.tzinfo is None:
+            ts = self.timestamp.replace(tzinfo=timezone.utc).timestamp()
+        return Event.model_validate({
+            **(self.event_data or {}),
+            "id": self.id,
+            "invocation_id": self.invocation_id,
+            "timestamp": ts,
+        })
+
+    @classmethod
+    def from_event(cls, session: Session, event: Event) -> MongoEvent:
+        ts = datetime.fromtimestamp(event.timestamp, tz=timezone.utc)
+        return cls(
+            id=event.id,
+            invocation_id=event.invocation_id,
+            session_id=session.id,
+            app_name=session.app_name,
+            user_id=session.user_id,
+            timestamp=ts,
+            event_data=event.model_dump(exclude_none=True, mode="json"),
+        )
+
+
+class MongoAppState(BaseModel):
+    """Document model for app_states collection."""
+
+    app_name: str
+    state: dict[str, Any] = Field(default_factory=dict)
+    update_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def to_doc(self) -> dict[str, Any]:
+        doc = self.model_dump()
+        doc["_id"] = app_state_doc_id(self.app_name)
+        return doc
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> MongoAppState:
+        data = {k: v for k, v in doc.items() if k != "_id"}
+        return cls.model_validate(data)
+
+
+class MongoUserState(BaseModel):
+    """Document model for user_states collection."""
+
+    app_name: str
+    user_id: str
+    state: dict[str, Any] = Field(default_factory=dict)
+    update_time: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def to_doc(self) -> dict[str, Any]:
+        doc = self.model_dump()
+        doc["_id"] = user_state_doc_id(self.app_name, self.user_id)
+        return doc
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> MongoUserState:
+        data = {k: v for k, v in doc.items() if k != "_id"}
+        return cls.model_validate(data)
+
+
+class MongoMetadata(BaseModel):
+    """Document model for adk_internal_metadata collection."""
+
+    key: str
+    value: str
+
+    def to_doc(self) -> dict[str, Any]:
+        doc = self.model_dump()
+        doc["_id"] = metadata_doc_id(self.key)
+        return doc
+
+    @classmethod
+    def from_doc(cls, doc: dict[str, Any]) -> MongoMetadata:
+        data = {k: v for k, v in doc.items() if k != "_id"}
+        return cls.model_validate(data)

google_adk_mongo_session_service/mongo_session_service.py

@@ -0,0 +1,339 @@
+"""MongoSessionService: BaseSessionService implementation using MongoDB."""
+
+from __future__ import annotations
+
+import copy
+import uuid
+from datetime import datetime
+from datetime import timezone
+from typing import Any
+from typing import Optional
+
+import pymongo
+from typing_extensions import override
+
+from google.adk.errors.already_exists_error import AlreadyExistsError
+from google.adk.events.event import Event
+from google.adk.sessions import _session_util
+from google.adk.sessions.base_session_service import BaseSessionService
+from google.adk.sessions.base_session_service import GetSessionConfig
+from google.adk.sessions.base_session_service import ListSessionsResponse
+from google.adk.sessions.session import Session
+from google.adk.sessions.state import State
+
+from .models import (
+    MongoAppState,
+    MongoEvent,
+    MongoSession,
+    MongoUserState,
+    app_state_doc_id,
+    session_doc_id,
+    user_state_doc_id,
+)
+
+# Collection names (fixed, not configurable)
+COLLECTION_SESSIONS = "sessions"
+COLLECTION_EVENTS = "events"
+COLLECTION_APP_STATES = "app_states"
+COLLECTION_USER_STATES = "user_states"
+
+
+def _merge_state(
+    app_state: dict[str, Any],
+    user_state: dict[str, Any],
+    session_state: dict[str, Any],
+) -> dict[str, Any]:
+    """Merge app, user, and session states into a single state dictionary."""
+    merged_state = copy.deepcopy(session_state)
+    for key in app_state.keys():
+        merged_state[State.APP_PREFIX + key] = app_state[key]
+    for key in user_state.keys():
+        merged_state[State.USER_PREFIX + key] = user_state[key]
+    return merged_state
+
+
+def _update_timestamp_tz(update_time: datetime) -> float:
+    """Return update_time as UTC timestamp (float)."""
+    if update_time.tzinfo is None:
+        return update_time.replace(tzinfo=timezone.utc).timestamp()
+    return update_time.timestamp()
+
+
+class MongoSessionService(BaseSessionService):
+    """Session service that uses MongoDB for storage."""
+
+    def __init__(self, conn_string: str, db_name: str):
+        self._mongo_client: pymongo.AsyncMongoClient = pymongo.AsyncMongoClient(
+            conn_string
+        )
+        self._database = self._mongo_client[db_name]
+
+    def _sessions(self):
+        return self._database[COLLECTION_SESSIONS]
+
+    def _events(self):
+        return self._database[COLLECTION_EVENTS]
+
+    def _app_states(self):
+        return self._database[COLLECTION_APP_STATES]
+
+    def _user_states(self):
+        return self._database[COLLECTION_USER_STATES]
+
+    @override
+    async def create_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        state: Optional[dict[str, Any]] = None,
+        session_id: Optional[str] = None,
+    ) -> Session:
+        state = state or {}
+        if session_id is not None:
+            sid = session_doc_id(app_name, user_id, session_id)
+            existing = await self._sessions().find_one({"_id": sid})
+            if existing is not None:
+                raise AlreadyExistsError(
+                    f"Session with id {session_id} already exists."
+                )
+
+        # Get or create app_state and user_state
+        app_state_id = app_state_doc_id(app_name)
+        app_doc = await self._app_states().find_one({"_id": app_state_id})
+        if app_doc is None:
+            app_state = MongoAppState(app_name=app_name, state={})
+            await self._app_states().insert_one(app_state.to_doc())
+        else:
+            app_state = MongoAppState.from_doc(app_doc)
+
+        user_state_id = user_state_doc_id(app_name, user_id)
+        user_doc = await self._user_states().find_one({"_id": user_state_id})
+        if user_doc is None:
+            user_state = MongoUserState(
+                app_name=app_name, user_id=user_id, state={}
+            )
+            await self._user_states().insert_one(user_state.to_doc())
+        else:
+            user_state = MongoUserState.from_doc(user_doc)
+
+        # Extract and apply state deltas
+        state_deltas = _session_util.extract_state_delta(state)
+        app_state_delta = state_deltas["app"]
+        user_state_delta = state_deltas["user"]
+        session_state = state_deltas["session"]
+        if app_state_delta:
+            app_state.state = app_state.state | app_state_delta
+            await self._app_states().replace_one(
+                {"_id": app_state_id}, app_state.to_doc()
+            )
+        if user_state_delta:
+            user_state.state = user_state.state | user_state_delta
+            await self._user_states().replace_one(
+                {"_id": user_state_id}, user_state.to_doc()
+            )
+
+        # Generate session id if not provided
+        if session_id is None:
+            session_id = str(uuid.uuid4())
+
+        now = datetime.now(timezone.utc)
+        mongo_session = MongoSession(
+            app_name=app_name,
+            user_id=user_id,
+            id=session_id,
+            state=session_state,
+            create_time=now,
+            update_time=now,
+        )
+        await self._sessions().insert_one(mongo_session.to_doc())
+
+        merged_state = _merge_state(
+            app_state.state, user_state.state, session_state
+        )
+        return mongo_session.to_session(state=merged_state)
+
+    @override
+    async def get_session(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        config: Optional[GetSessionConfig] = None,
+    ) -> Optional[Session]:
+        sid = session_doc_id(app_name, user_id, session_id)
+        session_doc = await self._sessions().find_one({"_id": sid})
+        if session_doc is None:
+            return None
+
+        mongo_session = MongoSession.from_doc(session_doc)
+
+        # Query events: app_name, user_id, session_id; filter/sort/limit
+        filter_query: dict[str, Any] = {
+            "app_name": app_name,
+            "user_id": user_id,
+            "session_id": session_id,
+        }
+        if config and config.after_timestamp is not None:
+            after_dt = datetime.fromtimestamp(
+                config.after_timestamp, tz=timezone.utc
+            )
+            filter_query["timestamp"] = {"$gte": after_dt}
+
+        cursor = (
+            self._events()
+            .find(filter_query)
+            .sort("timestamp", -1)
+        )
+        if config and config.num_recent_events is not None:
+            cursor = cursor.limit(config.num_recent_events)
+        event_docs = await cursor.to_list(length=None)
+
+        events = [
+            MongoEvent.from_doc(d).to_event()
+            for d in reversed(event_docs)
+        ]
+
+        # Load app_state and user_state
+        app_doc = await self._app_states().find_one(
+            {"_id": app_state_doc_id(app_name)}
+        )
+        user_doc = await self._user_states().find_one(
+            {"_id": user_state_doc_id(app_name, user_id)}
+        )
+        app_state = MongoAppState.from_doc(app_doc).state if app_doc else {}
+        user_state = MongoUserState.from_doc(user_doc).state if user_doc else {}
+        merged_state = _merge_state(
+            app_state, user_state, mongo_session.state
+        )
+        return mongo_session.to_session(state=merged_state, events=events)
+
+    @override
+    async def list_sessions(
+        self, *, app_name: str, user_id: Optional[str] = None
+    ) -> ListSessionsResponse:
+        filter_query: dict[str, Any] = {"app_name": app_name}
+        if user_id is not None:
+            filter_query["user_id"] = user_id
+        cursor = self._sessions().find(filter_query)
+        session_docs = await cursor.to_list(length=None)
+
+        app_doc = await self._app_states().find_one(
+            {"_id": app_state_doc_id(app_name)}
+        )
+        app_state = MongoAppState.from_doc(app_doc).state if app_doc else {}
+
+        user_states_map: dict[str, dict[str, Any]] = {}
+        if user_id is not None:
+            user_doc = await self._user_states().find_one(
+                {"_id": user_state_doc_id(app_name, user_id)}
+            )
+            if user_doc:
+                user_states_map[user_id] = MongoUserState.from_doc(
+                    user_doc
+                ).state
+        else:
+            cursor_user = self._user_states().find({"app_name": app_name})
+            user_docs = await cursor_user.to_list(length=None)
+            for ud in user_docs:
+                u = MongoUserState.from_doc(ud)
+                user_states_map[u.user_id] = u.state
+
+        sessions = []
+        for doc in session_docs:
+            mongo_session = MongoSession.from_doc(doc)
+            user_state = user_states_map.get(mongo_session.user_id, {})
+            merged_state = _merge_state(
+                app_state, user_state, mongo_session.state
+            )
+            sessions.append(mongo_session.to_session(state=merged_state))
+        return ListSessionsResponse(sessions=sessions)
+
+    @override
+    async def delete_session(
+        self, *, app_name: str, user_id: str, session_id: str
+    ) -> None:
+        await self._events().delete_many(
+            {
+                "app_name": app_name,
+                "user_id": user_id,
+                "session_id": session_id,
+            }
+        )
+        sid = session_doc_id(app_name, user_id, session_id)
+        await self._sessions().delete_one({"_id": sid})
+
+    @override
+    async def append_event(self, session: Session, event: Event) -> Event:
+        if event.partial:
+            return event
+        event = self._trim_temp_delta_state(event)
+
+        sid = session_doc_id(session.app_name, session.user_id, session.id)
+        session_doc = await self._sessions().find_one({"_id": sid})
+        if session_doc is None:
+            raise ValueError(f"Session {session.id} not found.")
+        mongo_session = MongoSession.from_doc(session_doc)
+        stored_update_ts = _update_timestamp_tz(mongo_session.update_time)
+        if stored_update_ts > session.last_update_time:
+            raise ValueError(
+                "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" {datetime.fromtimestamp(stored_update_ts):'%Y-%m-%d %H:%M:%S'}."
+                " Please check if it is a stale session."
+            )
+
+        app_state_id = app_state_doc_id(session.app_name)
+        user_state_id = user_state_doc_id(session.app_name, session.user_id)
+        app_doc = await self._app_states().find_one({"_id": app_state_id})
+        user_doc = await self._user_states().find_one({"_id": user_state_id})
+        app_state = (
+            MongoAppState.from_doc(app_doc)
+            if app_doc
+            else MongoAppState(app_name=session.app_name, state={})
+        )
+        user_state = (
+            MongoUserState.from_doc(user_doc)
+            if user_doc
+            else MongoUserState(
+                app_name=session.app_name, user_id=session.user_id, state={}
+            )
+        )
+
+        if event.actions and event.actions.state_delta:
+            state_deltas = _session_util.extract_state_delta(
+                event.actions.state_delta
+            )
+            if state_deltas["app"]:
+                app_state.state = app_state.state | state_deltas["app"]
+                await self._app_states().replace_one(
+                    {"_id": app_state_id},
+                    app_state.to_doc(),
+                    upsert=True,
+                )
+            if state_deltas["user"]:
+                user_state.state = user_state.state | state_deltas["user"]
+                await self._user_states().replace_one(
+                    {"_id": user_state_id},
+                    user_state.to_doc(),
+                    upsert=True,
+                )
+            if state_deltas["session"]:
+                mongo_session.state = mongo_session.state | state_deltas[
+                    "session"
+                ]
+
+        update_time = datetime.fromtimestamp(event.timestamp, tz=timezone.utc)
+        mongo_session.update_time = update_time
+        await self._sessions().replace_one(
+            {"_id": sid}, mongo_session.to_doc()
+        )
+
+        mongo_event = MongoEvent.from_event(session, event)
+        await self._events().insert_one(mongo_event.to_doc())
+
+        session.last_update_time = _update_timestamp_tz(mongo_session.update_time)
+        await super().append_event(session=session, event=event)
+        return event

google_adk_mongo_session_service-0.0.1.dist-info/METADATA

@@ -0,0 +1,57 @@
+Metadata-Version: 2.4
+Name: google-adk-mongo-session-service
+Version: 0.0.1
+Summary: Implementation of `google.adk.sessions.base_session_service.BaseSessionService` for storing session information in MongoDB.
+Project-URL: Homepage, https://github.com/kirlut/google-adk-mongo-session-service
+Project-URL: Repository, https://github.com/kirlut/google-adk-mongo-session-service
+Project-URL: Issues, https://github.com/kirlut/google-adk-mongo-session-service/issues
+Author: Kirill Lutcenko
+License: MIT License
+        
+        Copyright (c) 2026 Kirill Lutcenko
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Requires-Python: >=3.13
+Requires-Dist: google-adk>=1.23.0
+Requires-Dist: pymongo>=4.16.0
+Description-Content-Type: text/markdown
+
+[![PyPI version](https://img.shields.io/pypi/v/google-adk-mongo-session-service.svg)](https://pypi.org/project/google-adk-mongo-session-service/)
+
+**Disclaimer:** This code was written with the help of an AI agent and has never been tested in production. Please test it carefully before using it in production.
+
+Implementation of `google.adk.sessions.base_session_service.BaseSessionService` for storing session information in MongoDB. It plugs into [Google ADK](https://github.com/google/adk) agents so sessions can be backed by MongoDB instead of in-memory/default storage.
+
+## Installation
+
+```bash
+pip install google-adk-mongo-session-service
+```
+
+## Usage
+
+```python
+from google_adk_mongo_session_service import MongoSessionService
+
+session_service = MongoSessionService(conn_string="mongodb://...", db_name="mydb")
+# Pass session_service to your Agent / Runner as the session backend
+```
+
+**Full example:** [example/](https://github.com/kirlut/google-adk-mongo-session-service/tree/main/example/) — runnable weather agent with MongoSessionService; see [example/README.md](https://github.com/kirlut/google-adk-mongo-session-service/blob/main/example/README.md) for setup and run instructions.

google_adk_mongo_session_service-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+google_adk_mongo_session_service/__init__.py,sha256=deiRYAwHnNBW4R-ra8g5gJgsWejxK4eca3rQbhJatP0,90
+google_adk_mongo_session_service/models.py,sha256=1XhC6wssMsddRJaAQQXSzIkPOVWWS__CYtVJxelUaNI,5559
+google_adk_mongo_session_service/mongo_session_service.py,sha256=Fk-Pur-8k8OFZeHY5pG2rlrhRwV3Uoijab8g9tvWrpA,12454
+google_adk_mongo_session_service-0.0.1.dist-info/METADATA,sha256=cA-PV2kkz-5ZMcdF-It7ew3pKEm7wHFHo6otyIkDSDI,3079
+google_adk_mongo_session_service-0.0.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+google_adk_mongo_session_service-0.0.1.dist-info/licenses/LICENSE,sha256=Xwko7DhJQtVO6Uk5C-pTVTj5HXoTg3ZTyyjrJMKw-WU,1072
+google_adk_mongo_session_service-0.0.1.dist-info/RECORD,,

google_adk_mongo_session_service-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

google_adk_mongo_session_service-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Kirill Lutcenko
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.