diff --git a/contributing/samples/mongodb_service/README.md b/contributing/samples/mongodb_service/README.md new file mode 100644 index 00000000..6cda362a --- /dev/null +++ b/contributing/samples/mongodb_service/README.md @@ -0,0 +1,142 @@ +# MongoDB Session Service Sample + +This sample shows how to persist ADK sessions and state in MongoDB using the +community `MongoSessionService`. + +## Prerequisites + +- Python 3.10+ (Python 3.11+ recommended) +- A running MongoDB instance (local or Atlas) and a connection string with + create/read/write permissions +- ADK and ADK Community installed +- Google API key for the sample agent (Gemini), set as `GOOGLE_API_KEY` + +## Setup + +### 1. Install dependencies + +```bash + pip install "google-adk-community[mongo]" google-adk python-dotenv +``` + +### 2. Configure environment variables + +Create a `.env` in this directory: + +```bash +# Required: Google API key for the agent +GOOGLE_API_KEY=your-google-api-key + +# Recommended: Mongo connection string (Atlas or local) +MONGODB_URI=mongodb+srv://:@/ +``` + +**Note:** Keep your Mongo credentials out of source control. The sample loads the connection string from the `MONGODB_URI` environment variable, which is loaded from the `.env` file at runtime. + +### 3. Pick a database name + +By default the sample uses `adk_sessions_db`. Collections are created +automatically if they do not exist. + +## Usage + +### Option 1: Run the included sample + +```bash +python main.py +``` + +`main.py`: +- Creates a `MongoSessionService` with a connection string +- Creates a session for the demo user +- Runs the `financial_advisor_agent` with `Runner.run_async` +- Prints the agent's final response + +### Option 2: Use `MongoSessionService` with your own runner + +```python +import os +from google.adk.runners import Runner +from google.genai import types +from google.adk_community.sessions import MongoSessionService + +session_service = MongoSessionService( + connection_string=os.environ.get("MONGODB_URI") +) + +await session_service.create_session( + app_name="my_app", user_id="user1", session_id="demo" +) + +runner = Runner(app_name="my_app", agent=your_agent, session_service=session_service) +query = "Hello, can you help me with my account?" +content = types.Content(role="user", parts=[types.Part(text=query)]) + +async for event in runner.run_async( + user_id="user1", + session_id="demo", + new_message=content, +): + if event.is_final_response(): + print(event.content.parts[0].text) +``` + +If you already have an `AsyncMongoClient`, pass it instead of a connection +string: + +```python +from pymongo import AsyncMongoClient + +client = AsyncMongoClient(host="localhost", port=27017) +session_service = MongoSessionService(client=client) +``` + +## Collections and indexing + +`MongoSessionService` writes to two collections (configurable): +- `sessions`: conversation history and session-level state +- `session_state`: shared app/user state across sessions + +Indexes are created on first use: +- Unique session identity: `(app_name, user_id, id)` +- Last update for recency queries: `(app_name, user_id, last_update_time)` + +## Sample structure + +``` +mongodb_service/ +├── main.py # Runs the sample with Mongo-backed sessions +├── mongo_service_agent/ +│ ├── __init__.py # Agent package init +│ └── agent.py # Financial advisor agent with two tools +└── README.md # This file +``` + +## Sample agent + +The agent (`mongo_service_agent/agent.py`) includes: +- `get_invoice_status(service)` tool for simple invoice lookups +- `calculate_service_tax(amount)` tool for tax calculations +- Gemini model (`gemini-2.0-flash`) with instructions to route to the tools + +## Sample query + +``` +What is the status of my university invoice? Also, calculate the tax for a service amount of 9500 MXN. +``` + +## Configuration options (`MongoSessionService`) + +- `database_name` (str, default `adk_sessions_db`): Mongo database to store session data. +- `connection_string` (str, optional): Mongo URI (mutually exclusive with `client`) +- `client` (AsyncMongoClient, optional): Provide your own client/connection pool +- `session_collection` (str, default `sessions`): Collection for session docs +- `state_collection` (str, default `session_state`): Collection for shared state +- `default_app_name` (str, optional): Fallback app name when not provided per call (defaults to `adk-mongo-session-service`) + +## Tips + +- Use `runner.run_async` (as in `main.py`) to keep the Mongo client on the same + event loop and avoid loop-bound client errors. +- For production, prefer environment variables or secrets managers for the + connection string and database credentials. diff --git a/contributing/samples/mongodb_service/main.py b/contributing/samples/mongodb_service/main.py new file mode 100644 index 00000000..40efe020 --- /dev/null +++ b/contributing/samples/mongodb_service/main.py @@ -0,0 +1,77 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Example of using MongoDB for Session Service.""" + +import asyncio +import os + +from google.adk.errors.already_exists_error import AlreadyExistsError +from google.adk.runners import Runner +from google.genai import types +from mongo_service_agent import root_agent + +from google.adk_community.sessions import MongoSessionService + +APP_NAME = "financial_advisor_agent" +USER_ID = "demo_user" +SESSION_ID = "demo_session" + + +async def main(): + """Main function to run the agent asynchronously.""" + + # You can create the MongoSessionService in two ways: + # 1. With an existing AsyncMongoClient instance + # 2. By providing a connection string directly + # from pymongo import AsyncMongoClient + # client = AsyncMongoClient(host="localhost", port=27017) + # session_service = MongoSessionService(client=client) + + connection_string = os.environ.get("MONGODB_URI") + if not connection_string: + raise ValueError( + "MONGODB_URI environment variable not set. See README.md for setup." + ) + session_service = MongoSessionService(connection_string=connection_string) + try: + await session_service.create_session( + app_name=APP_NAME, user_id=USER_ID, session_id=SESSION_ID + ) + except AlreadyExistsError: + # Session already exists, which is fine for this example. + pass + + runner = Runner( + agent=root_agent, app_name=APP_NAME, session_service=session_service + ) + + query = ( + "What is the status of my university invoice? Also, calculate the tax for" + " a service amount of 9500 MXN." + ) + print(f"User Query -> {query}") + content = types.Content(role="user", parts=[types.Part(text=query)]) + + async for event in runner.run_async( + user_id=USER_ID, + session_id=SESSION_ID, + new_message=content, + ): + if event.is_final_response(): + print(f"Agent Response -> {event.content.parts[0].text}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/contributing/samples/mongodb_service/mongo_service_agent/__init__.py b/contributing/samples/mongodb_service/mongo_service_agent/__init__.py new file mode 100644 index 00000000..97557b3a --- /dev/null +++ b/contributing/samples/mongodb_service/mongo_service_agent/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .agent import root_agent diff --git a/contributing/samples/mongodb_service/mongo_service_agent/agent.py b/contributing/samples/mongodb_service/mongo_service_agent/agent.py new file mode 100644 index 00000000..b78ab52e --- /dev/null +++ b/contributing/samples/mongodb_service/mongo_service_agent/agent.py @@ -0,0 +1,91 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from dotenv import load_dotenv +from google.adk.agents import Agent +from google.adk.tools import FunctionTool + +_TAX_RATE = 0.16 +load_dotenv() + + +# Tool 1 +def get_invoice_status(service: str) -> dict: + """Return the invoice status details for the requested service. + + Args: + service: Business category to query (for example 'gas' or 'restaurant'). + + Returns: + dict: Contains a ``status`` entry with the invoice state and a ``report`` + entry that gives the relevant context for that status. + """ + if service == "university": + return {"status": "success", "report": "All TEC invoices are paid."} + elif service == "gas": + return {"status": "pending", "report": "Gas invoice due in 5 days."} + elif service == "restaurant": + return { + "status": "overdue", + "report": "Restaurant invoice is overdue by 10 days.", + } + else: + return {"status": "error", "report": "Service not recognized."} + + +invoice_tool = FunctionTool(func=get_invoice_status) + + +# Tool 2 +def calculate_service_tax(amount: float) -> dict: + """Calculate tax for a service amount using a fixed rate. + + Args: + amount: Untaxed amount that needs a tax calculation. + + Returns: + dict: Keys ``amount``, ``tax_amount``, and ``total_amount`` capturing the + original value, the computed tax, and the amount plus tax respectively. + """ + tax_amount = amount * _TAX_RATE + total_amount = amount + tax_amount + return { + "amount": amount, + "tax_amount": tax_amount, + "total_amount": total_amount, + } + + +tax_tool = FunctionTool(func=calculate_service_tax) + +# Agent +root_agent = Agent( + model="gemini-2.5-flash", + name="financial_advisor_agent", + description=( + "Financial advisor agent for managing invoices and calculating service" + " taxes." + ), + instruction=( + "You are an AI agent designed to assist users with financial inquiries" + " related to invoices and service tax calculations.\n**Available" + " Tools:**\n1. get_invoice_status(service): Retrieves the status of" + " invoices\n2. calculate_service_tax(amount): Calculates the tax for a" + " given amount\nUse these tools to assist users with their financial" + " inquiries.\nIf the user asks about other financial topics, respond" + " politely that you can only assist with invoice status and tax" + " calculations.\n" + ), + tools=[invoice_tool, tax_tool], +) diff --git a/pyproject.toml b/pyproject.toml index a03bdcab..ce61c461 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ changelog = "https://github.com/google/adk-python-community/blob/main/CHANGELOG. documentation = "https://google.github.io/adk-docs/" [project.optional-dependencies] +mongo = ["pymongo>=4.9,<5.0.0"] s3 = [ "aioboto3>=13.0.0", # For S3ArtifactService ] diff --git a/src/google/adk_community/sessions/__init__.py b/src/google/adk_community/sessions/__init__.py index 90bf28d7..ad24db4e 100644 --- a/src/google/adk_community/sessions/__init__.py +++ b/src/google/adk_community/sessions/__init__.py @@ -14,6 +14,7 @@ """Community session services for ADK.""" +from .mongo_session_service import MongoSessionService from .redis_session_service import RedisSessionService -__all__ = ["RedisSessionService"] +__all__ = ["MongoSessionService", "RedisSessionService"] diff --git a/src/google/adk_community/sessions/mongo_session_service.py b/src/google/adk_community/sessions/mongo_session_service.py new file mode 100644 index 00000000..6ca6576f --- /dev/null +++ b/src/google/adk_community/sessions/mongo_session_service.py @@ -0,0 +1,443 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from __future__ import annotations + +import asyncio +import logging +import time +from typing import Any +from typing import Optional +from typing import TYPE_CHECKING +import uuid + +from google.adk.errors.already_exists_error import AlreadyExistsError +from google.adk.events import Event +from google.adk.sessions import _session_util +from google.adk.sessions import Session +from google.adk.sessions import State +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 +import orjson +from typing_extensions import override + +from .utils import _json_serializer + +if TYPE_CHECKING: + from pymongo import AsyncMongoClient + from pymongo.asynchronous.collection import AsyncCollection + +logger = logging.getLogger("google_adk." + __name__) + + +def _normalize_value(value: Any) -> Any: + """Convert a value through the shared JSON serializer into BSON types.""" + return orjson.loads(orjson.dumps(value, default=_json_serializer)) + + +class MongoKeys: + """Helper to generate composite keys for Mongo-backed storage.""" + + @staticmethod + def session(app_name: str, user_id: str, session_id: str) -> str: + return f"session:{app_name}:{user_id}:{session_id}" + + @staticmethod + def app_state(app_name: str) -> str: + return f"{State.APP_PREFIX}{app_name}" + + @staticmethod + def user_state(app_name: str, user_id: str) -> str: + return f"{State.USER_PREFIX}{app_name}:{user_id}" + + +class MongoSessionService(BaseSessionService): + """Session service backed by MongoDB. + + The MongoDB client is bound to the event loop where it is constructed. When + loop affinity matters, construct the client on the target loop and pass it + through ``client=``. + """ + + def __init__( + self, + client: Optional[AsyncMongoClient] = None, + connection_string: Optional[str] = None, + database_name: str = "adk_sessions_db", + session_collection: str = "sessions", + state_collection: str = "session_state", + default_app_name: Optional[str] = "adk-mongo-session-service", + ) -> None: + try: + import pymongo + except ImportError as exc: + raise ImportError( + "pymongo is required to use MongoSessionService. " + "Install it with: pip install google-adk-community[mongo]" + ) from exc + + if bool(connection_string) == bool(client): + raise ValueError( + "Provide either 'connection_string' or 'client', but not both." + ) + self._client = client or pymongo.AsyncMongoClient(connection_string) + db = self._client[database_name] + self._sessions: AsyncCollection = db[session_collection] + self._kv: AsyncCollection = db[state_collection] + self._ascending = pymongo.ASCENDING + self._descending = pymongo.DESCENDING + self._duplicate_key_error = pymongo.errors.DuplicateKeyError + self._default_app_name = default_app_name + self._indexes_built = False + self._indexes_lock = asyncio.Lock() + + @override + async def create_session( + self, + *, + app_name: str, + user_id: str, + state: Optional[dict[str, Any]] = None, + session_id: Optional[str] = None, + ) -> Session: + app_name = self._resolve_app_name(app_name) + await self._ensure_indexes() + + session_id = (session_id or "").strip() or str(uuid.uuid4()) + doc_id = MongoKeys.session(app_name, user_id, session_id) + + state_deltas = _session_util.extract_state_delta(state or {}) + await self._apply_state_delta( + app_name, user_id, state_deltas.get("app"), state_deltas.get("user") + ) + + session_doc = { + "_id": doc_id, + "app_name": app_name, + "user_id": user_id, + "id": session_id, + "state": { + key: _normalize_value(value) + for key, value in state_deltas.get("session", {}).items() + }, + "events": [], + "last_update_time": time.time(), + } + + try: + await self._sessions.insert_one(session_doc) + except self._duplicate_key_error as exc: + raise AlreadyExistsError( + f"Session with id {session_id} already exists." + ) from exc + + session = self._doc_to_session(session_doc) + app_doc, user_doc = await self._fetch_state( + session.app_name, session.user_id + ) + return self._merge_state(session, app_doc, user_doc) + + @override + async def get_session( + self, + *, + app_name: str, + user_id: str, + session_id: str, + config: Optional[GetSessionConfig] = None, + ) -> Optional[Session]: + app_name = self._resolve_app_name(app_name) + await self._ensure_indexes() + + doc = await self._sessions.find_one( + {"_id": MongoKeys.session(app_name, user_id, session_id)} + ) + if not doc: + return None + + session = self._doc_to_session(doc) + session = self._apply_event_filters(session, config) + app_doc, user_doc = await self._fetch_state( + session.app_name, session.user_id + ) + return self._merge_state(session, app_doc, user_doc) + + @override + async def list_sessions( + self, + *, + app_name: str, + user_id: Optional[str] = None, + ) -> ListSessionsResponse: + """List all matching sessions, materialized because no pagination is exposed.""" + app_name = self._resolve_app_name(app_name) + await self._ensure_indexes() + + filters: dict[str, Any] = {"app_name": app_name} + if user_id is not None: + filters["user_id"] = user_id + + cursor = self._sessions.find(filters, projection={"events": False}) + # NOTE: BaseSessionService expects the full list to be returned, so we have + # to materialize the entire cursor which may load many sessions into memory. + docs = await cursor.to_list(length=None) + if len(docs) > 1000: + logger.warning( + "Loading a large number of sessions (%d) into memory for app '%s'.", + len(docs), + app_name, + ) + + sessions: list[Session] = [self._doc_to_session(doc) for doc in docs] + app_doc = await self._kv.find_one({"_id": MongoKeys.app_state(app_name)}) + if user_id is not None: + user_docs = {} + user_doc = await self._kv.find_one( + {"_id": MongoKeys.user_state(app_name, user_id)} + ) + if user_doc: + user_docs[user_doc["_id"]] = user_doc + else: + user_ids = {session.user_id for session in sessions} + user_docs = {} + if user_ids: + cursor = self._kv.find({ + "_id": { + "$in": [ + MongoKeys.user_state(app_name, current_user_id) + for current_user_id in user_ids + ] + } + }) + user_docs = { + doc["_id"]: doc for doc in await cursor.to_list(length=None) + } + + sessions = [ + self._merge_state( + session, + app_doc, + user_docs.get(MongoKeys.user_state(app_name, session.user_id)), + ) + for session in sessions + ] + + return ListSessionsResponse(sessions=sessions) + + @override + async def delete_session( + self, + *, + app_name: str, + user_id: str, + session_id: str, + ) -> None: + app_name = self._resolve_app_name(app_name) + await self._ensure_indexes() + + await self._sessions.delete_one( + {"_id": MongoKeys.session(app_name, user_id, session_id)} + ) + + @override + async def append_event(self, session: Session, event: Event) -> Event: + if event.partial: + return event + + await self._ensure_indexes() + + event = await super().append_event(session, event) + session.last_update_time = event.timestamp + for key in list(session.state): + if key.startswith(State.TEMP_PREFIX): + del session.state[key] + + state_delta = event.actions.state_delta if event.actions else None + state_deltas = _session_util.extract_state_delta(state_delta or {}) + + await self._apply_state_delta( + session.app_name, + session.user_id, + state_deltas.get("app"), + state_deltas.get("user"), + ) + + updates: dict[str, Any] = { + "$push": { + "events": _normalize_value(event.model_dump(exclude_none=True)) + }, + "$set": {"last_update_time": event.timestamp}, + } + + session_state_set = { + f"state.{key}": _normalize_value(value) + for key, value in state_deltas["session"].items() + if value is not None + } + session_state_unset = { + f"state.{key}": "" + for key, value in state_deltas["session"].items() + if value is None + } + + if session_state_set: + updates.setdefault("$set", {}).update(session_state_set) + if session_state_unset: + updates["$unset"] = session_state_unset + + result = await self._sessions.update_one( + { + "_id": MongoKeys.session( + session.app_name, session.user_id, session.id + ) + }, + updates, + ) + if result.matched_count == 0: + logger.warning( + "Failed to append event: session %s/%s/%s not found in storage", + session.app_name, + session.user_id, + session.id, + ) + + return event + + async def _ensure_indexes(self) -> None: + if self._indexes_built: + return + async with self._indexes_lock: + if self._indexes_built: + return + await self._sessions.create_index( + [ + ("app_name", self._ascending), + ("user_id", self._ascending), + ("id", self._ascending), + ], + unique=True, + name="session_identity_idx", + ) + await self._sessions.create_index( + [ + ("app_name", self._ascending), + ("user_id", self._ascending), + ("last_update_time", self._descending), + ], + name="session_last_update_idx", + ) + self._indexes_built = True + + async def _fetch_state( + self, app_name: str, user_id: str + ) -> tuple[Optional[dict[str, Any]], Optional[dict[str, Any]]]: + app_doc, user_doc = await asyncio.gather( + self._kv.find_one({"_id": MongoKeys.app_state(app_name)}), + self._kv.find_one({"_id": MongoKeys.user_state(app_name, user_id)}), + ) + return app_doc, user_doc + + def _merge_state( + self, + session: Session, + app_doc: Optional[dict[str, Any]], + user_doc: Optional[dict[str, Any]], + ) -> Session: + merged_state = dict(session.state) + if app_doc and app_doc.get("state"): + for key, value in app_doc["state"].items(): + merged_state[State.APP_PREFIX + key] = value + if user_doc and user_doc.get("state"): + for key, value in user_doc["state"].items(): + merged_state[State.USER_PREFIX + key] = value + + return session.model_copy(update={"state": merged_state}) + + async def _apply_state_delta( + self, + app_name: str, + user_id: str, + app_state_delta: Optional[dict[str, Any]], + user_state_delta: Optional[dict[str, Any]], + ) -> None: + tasks = [] + if app_state_delta: + tasks.append( + self._update_state_document( + MongoKeys.app_state(app_name), app_state_delta + ) + ) + if user_state_delta: + tasks.append( + self._update_state_document( + MongoKeys.user_state(app_name, user_id), user_state_delta + ) + ) + if tasks: + await asyncio.gather(*tasks) + + async def _update_state_document( + self, + key: str, + delta: dict[str, Any], + ) -> None: + set_ops = { + f"state.{k}": _normalize_value(v) + for k, v in delta.items() + if v is not None + } + unset_ops = {f"state.{k}": "" for k, v in delta.items() if v is None} + + update: dict[str, Any] = {} + if set_ops: + update["$set"] = set_ops + if unset_ops: + update["$unset"] = unset_ops + if not update: + return + + await self._kv.update_one({"_id": key}, update, upsert=True) + + def _apply_event_filters( + self, session: Session, config: Optional[GetSessionConfig] + ) -> Session: + if not config: + return session + events = session.events + + if config.after_timestamp is not None: + events = [e for e in events if e.timestamp > config.after_timestamp] + if config.num_recent_events is not None: + events = events[-config.num_recent_events :] + + return session.model_copy(update={"events": events}) + + def _doc_to_session(self, doc: dict[str, Any]) -> Session: + events = [Event.model_validate(e) for e in doc.get("events", [])] + return Session( + id=doc.get("id"), + app_name=doc.get("app_name"), + user_id=doc.get("user_id"), + state=doc.get("state", {}), + events=events, + last_update_time=doc.get("last_update_time", 0.0), + ) + + def _resolve_app_name(self, app_name: Optional[str]) -> str: + resolved = app_name or self._default_app_name + if not resolved: + raise ValueError( + "app_name must be provided either in the call or in default_app_name." + ) + return resolved diff --git a/tests/unittests/sessions/test_mongo_session_service.py b/tests/unittests/sessions/test_mongo_session_service.py new file mode 100644 index 00000000..379a1d42 --- /dev/null +++ b/tests/unittests/sessions/test_mongo_session_service.py @@ -0,0 +1,420 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from datetime import datetime +from datetime import timezone +from decimal import Decimal +from unittest.mock import AsyncMock +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("pymongo") + +from google.adk.events import Event +from google.adk.events.event_actions import EventActions +from google.adk.sessions import State +from google.adk.sessions.base_session_service import GetSessionConfig +import pytest_asyncio + +from google.adk_community.sessions.mongo_session_service import MongoKeys +from google.adk_community.sessions.mongo_session_service import MongoSessionService + + +class TestMongoSessionService: + """Tests for MongoSessionService mirroring Redis coverage.""" + + @pytest_asyncio.fixture + async def mongo_service(self): + """Create a Mongo session service with mocked collections.""" + sessions_collection = AsyncMock() + sessions_collection.create_index = AsyncMock() + sessions_collection.insert_one = AsyncMock() + sessions_collection.find_one = AsyncMock(return_value=None) + sessions_collection.find = MagicMock() + sessions_collection.delete_one = AsyncMock() + sessions_collection.update_one = AsyncMock( + return_value=MagicMock(matched_count=1) + ) + + kv_collection = AsyncMock() + kv_collection.find_one = AsyncMock(return_value=None) + kv_collection.find = MagicMock() + kv_collection.update_one = AsyncMock() + + db = MagicMock() + + def _get_collection(name: str): + if name == "sessions": + return sessions_collection + if name == "session_state": + return kv_collection + raise KeyError(name) + + db.__getitem__.side_effect = _get_collection + + client = MagicMock() + client.__getitem__.return_value = db + + service = MongoSessionService(database_name="test_db", client=client) + return service, sessions_collection, kv_collection + + @pytest.mark.asyncio + async def test_get_empty_session(self, mongo_service): + """get_session should return None for missing sessions.""" + service, sessions_collection, _ = mongo_service + sessions_collection.find_one.return_value = None + + session = await service.get_session( + app_name="test_app", user_id="test_user", session_id="missing" + ) + + assert session is None + sessions_collection.find_one.assert_awaited_once_with({ + "_id": MongoKeys.session("test_app", "test_user", "missing"), + }) + assert service._indexes_built is True + + @pytest.mark.asyncio + async def test_create_and_get_session(self, mongo_service): + """create_session persists and get_session retrieves a session.""" + service, sessions_collection, kv_collection = mongo_service + kv_collection.find_one.return_value = None + + session = await service.create_session( + app_name="test_app", + user_id="test_user", + session_id="session-1", + state={"key": "value"}, + ) + + inserted_doc = sessions_collection.insert_one.await_args.args[0] + assert inserted_doc["_id"] == MongoKeys.session( + "test_app", "test_user", "session-1" + ) + assert inserted_doc["state"] == {"key": "value"} + + sessions_collection.find_one.return_value = dict(inserted_doc) + + retrieved = await service.get_session( + app_name="test_app", user_id="test_user", session_id="session-1" + ) + + assert retrieved is not None + assert retrieved.id == "session-1" + assert retrieved.app_name == session.app_name + assert retrieved.user_id == session.user_id + assert retrieved.state["key"] == "value" + + # Indexes should only be created once even when called from multiple methods. + assert sessions_collection.create_index.await_count == 2 + + @pytest.mark.asyncio + async def test_list_sessions_merges_state_and_strips_events( + self, mongo_service + ): + """list_sessions returns sessions without events but with merged state.""" + service, sessions_collection, kv_collection = mongo_service + app_name = "test_app" + user_id = "user1" + + docs = [ + { + "_id": MongoKeys.session(app_name, user_id, "s1"), + "app_name": app_name, + "user_id": user_id, + "id": "s1", + "state": {"session_key": "v1"}, + }, + { + "_id": MongoKeys.session(app_name, user_id, "s2"), + "app_name": app_name, + "user_id": user_id, + "id": "s2", + "state": {"session_key": "v2"}, + }, + ] + + cursor = MagicMock() + cursor.to_list = AsyncMock(return_value=docs) + sessions_collection.find.return_value = cursor + + async def _kv_find_one(query): + if query["_id"] == MongoKeys.app_state(app_name): + return {"_id": query["_id"], "state": {"theme": "dark"}} + if query["_id"] == MongoKeys.user_state(app_name, user_id): + return {"_id": query["_id"], "state": {"pref": "value"}} + return None + + kv_collection.find_one.side_effect = _kv_find_one + + response = await service.list_sessions(app_name=app_name, user_id=user_id) + + sessions_collection.find.assert_called_once_with( + {"app_name": app_name, "user_id": user_id}, projection={"events": False} + ) + + assert len(response.sessions) == 2 + for doc, sess in zip(docs, response.sessions): + assert sess.id == doc["id"] + assert sess.events == [] + assert sess.state["session_key"] == doc["state"]["session_key"] + assert sess.state[State.APP_PREFIX + "theme"] == "dark" + assert sess.state[State.USER_PREFIX + "pref"] == "value" + + @pytest.mark.asyncio + async def test_session_state_management_on_append_event(self, mongo_service): + """append_event should persist state deltas and mutate in-memory session.""" + service, sessions_collection, kv_collection = mongo_service + app_name = "test_app" + user_id = "test_user" + session_id = "session-123" + kv_collection.find_one.return_value = None + + session = await service.create_session( + app_name=app_name, + user_id=user_id, + session_id=session_id, + state={"initial_key": "initial_value"}, + ) + + event = Event( + invocation_id="invocation", + author="user", + timestamp=datetime.now().astimezone(timezone.utc).timestamp(), + actions=EventActions( + state_delta={ + f"{State.APP_PREFIX}key": "app_value", + f"{State.USER_PREFIX}key1": "user_value", + "temp:key": "temp_value", + "initial_key": "updated_value", + } + ), + ) + + await service.append_event(session=session, event=event) + + assert session.state[State.APP_PREFIX + "key"] == "app_value" + assert session.state[State.USER_PREFIX + "key1"] == "user_value" + assert session.state["initial_key"] == "updated_value" + assert session.state.get("temp:key") is None + + # App and user deltas are stored in the kv collection. + assert kv_collection.update_one.await_count == 2 + kv_update_filters = [ + call.args[0] for call in kv_collection.update_one.await_args_list + ] + kv_updates = [ + call.args[1] for call in kv_collection.update_one.await_args_list + ] + assert {"_id": MongoKeys.app_state(app_name)} in kv_update_filters + assert {"_id": MongoKeys.user_state(app_name, user_id)} in kv_update_filters + assert any( + update.get("$set", {}).get("state.key") == "app_value" + for update in kv_updates + ) + assert any( + update.get("$set", {}).get("state.key1") == "user_value" + for update in kv_updates + ) + + update_filter, update_doc = sessions_collection.update_one.await_args.args + assert update_filter == { + "_id": MongoKeys.session(app_name, user_id, session_id) + } + assert update_doc["$set"]["state.initial_key"] == "updated_value" + # Temp state should not be persisted. + assert "state.temp:key" not in update_doc.get("$set", {}) + assert "state.temp:key" not in update_doc.get("$unset", {}) + + @pytest.mark.asyncio + async def test_get_session_with_config(self, mongo_service): + """get_session applies after_timestamp and num_recent_events filters.""" + service, sessions_collection, kv_collection = mongo_service + kv_collection.find_one.return_value = None + + events = [ + Event(author="user", timestamp=float(i)).model_dump( + mode="json", exclude_none=True + ) + for i in range(1, 6) + ] + doc = { + "_id": MongoKeys.session("app", "user", "session"), + "app_name": "app", + "user_id": "user", + "id": "session", + "events": events, + "state": {}, + } + + sessions_collection.find_one.return_value = doc + + config = GetSessionConfig(num_recent_events=3) + filtered = await service.get_session( + app_name="app", user_id="user", session_id="session", config=config + ) + assert [e.timestamp for e in filtered.events] == [3.0, 4.0, 5.0] + + config = GetSessionConfig(after_timestamp=3.0) + filtered = await service.get_session( + app_name="app", user_id="user", session_id="session", config=config + ) + assert [e.timestamp for e in filtered.events] == [4.0, 5.0] + + @pytest.mark.asyncio + async def test_delete_session(self, mongo_service): + """delete_session removes session documents.""" + service, sessions_collection, _ = mongo_service + + await service.delete_session( + app_name="test_app", user_id="user", session_id="session-1" + ) + + sessions_collection.delete_one.assert_awaited_once_with( + {"_id": MongoKeys.session("test_app", "user", "session-1")} + ) + + @pytest.mark.asyncio + async def test_non_bson_state_values_are_normalized(self, mongo_service): + """Session, app, and user state use the shared JSON serializer.""" + service, sessions_collection, kv_collection = mongo_service + kv_collection.find_one.return_value = None + state = { + "session_set": {"value"}, + "session_decimal": Decimal("1.25"), + "session_datetime": datetime(2025, 1, 1, tzinfo=timezone.utc), + "session_bytes": b"bytes", + f"{State.APP_PREFIX}app_set": {"value"}, + f"{State.USER_PREFIX}user_set": {"value"}, + } + + await service.create_session( + app_name="app", user_id="user", session_id="session", state=state + ) + + inserted = sessions_collection.insert_one.await_args.args[0] + assert inserted["state"] == { + "session_set": ["value"], + "session_decimal": 1.25, + "session_datetime": "2025-01-01T00:00:00+00:00", + "session_bytes": "Ynl0ZXM=", + } + updates = [ + call.args[1] for call in kv_collection.update_one.await_args_list + ] + set_values = { + key: value + for update in updates + for key, value in update.get("$set", {}).items() + } + assert set_values["state.app_set"] == ["value"] + assert set_values["state.user_set"] == ["value"] + + @pytest.mark.asyncio + async def test_event_is_stored_as_nested_document(self, mongo_service): + """Events remain addressable Mongo documents after normalization.""" + service, sessions_collection, kv_collection = mongo_service + kv_collection.find_one.return_value = None + session = await service.create_session( + app_name="app", user_id="user", session_id="session" + ) + + event = Event(author="user", timestamp=1.0) + await service.append_event(session, event) + + update = sessions_collection.update_one.await_args.args[1] + stored_event = update["$push"]["events"] + assert isinstance(stored_event, dict) + assert stored_event["author"] == "user" + assert not isinstance(stored_event, bytes) + + @pytest.mark.asyncio + async def test_list_sessions_fetches_shared_state_once(self, mongo_service): + """Listing sessions for one user does not issue per-session state reads.""" + service, sessions_collection, kv_collection = mongo_service + app_name = "app" + user_id = "user" + docs = [ + { + "_id": MongoKeys.session(app_name, user_id, f"s{i}"), + "app_name": app_name, + "user_id": user_id, + "id": f"s{i}", + "state": {}, + } + for i in range(3) + ] + session_cursor = MagicMock() + session_cursor.to_list = AsyncMock(return_value=docs) + sessions_collection.find.return_value = session_cursor + kv_collection.find_one.side_effect = [ + {"_id": MongoKeys.app_state(app_name), "state": {"a": 1}}, + {"_id": MongoKeys.user_state(app_name, user_id), "state": {"u": 2}}, + ] + + response = await service.list_sessions(app_name=app_name, user_id=user_id) + + assert len(response.sessions) == 3 + assert kv_collection.find_one.await_count == 2 + assert all( + session.state[State.APP_PREFIX + "a"] == 1 + for session in response.sessions + ) + assert all( + session.state[State.USER_PREFIX + "u"] == 2 + for session in response.sessions + ) + + @pytest.mark.asyncio + async def test_list_sessions_bulk_fetches_users(self, mongo_service): + """Listing all users uses one state query and attributes state correctly.""" + service, sessions_collection, kv_collection = mongo_service + app_name = "app" + docs = [ + { + "_id": MongoKeys.session(app_name, user_id, session_id), + "app_name": app_name, + "user_id": user_id, + "id": session_id, + "state": {}, + } + for user_id, session_id in (("user1", "s1"), ("user2", "s2")) + ] + session_cursor = MagicMock() + session_cursor.to_list = AsyncMock(return_value=docs) + sessions_collection.find.return_value = session_cursor + user_cursor = MagicMock() + user_cursor.to_list = AsyncMock( + return_value=[ + { + "_id": MongoKeys.user_state(app_name, user_id), + "state": {"value": user_id}, + } + for user_id in ("user1", "user2") + ] + ) + kv_collection.find.return_value = user_cursor + kv_collection.find_one.return_value = None + + response = await service.list_sessions(app_name=app_name) + + assert kv_collection.find.call_count == 1 + assert [ + session.state[State.USER_PREFIX + "value"] + for session in response.sessions + ] == [ + "user1", + "user2", + ] diff --git a/tests/unittests/sessions/test_mongo_session_service_optional.py b/tests/unittests/sessions/test_mongo_session_service_optional.py new file mode 100644 index 00000000..6ad43cad --- /dev/null +++ b/tests/unittests/sessions/test_mongo_session_service_optional.py @@ -0,0 +1,32 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from unittest.mock import MagicMock + +import pytest + +from google.adk_community.sessions.mongo_session_service import MongoSessionService + + +def test_constructor_without_pymongo_names_mongo_extra(monkeypatch): + real_import = __import__ + + def _import(name, *args, **kwargs): + if name == "pymongo": + raise ModuleNotFoundError("No module named 'pymongo'") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr("builtins.__import__", _import) + with pytest.raises(ImportError, match=r"google-adk-community\[mongo\]"): + MongoSessionService(client=MagicMock())