diff --git a/pyproject.toml b/pyproject.toml index 67824f41f8f..9f311e0bbd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -122,6 +122,7 @@ optional-dependencies.all = [ "pandas>=2.2.3", "protobuf>=6", "pyarrow>=14", + "pymongo>=4.9,<5", "pypika>=0.50", "python-dateutil>=2.9.0.post0,<3", "redis>=4.2", @@ -218,6 +219,7 @@ optional-dependencies.extensions = [ "lxml>=5.3", "nltk!=3.10.1", # Transitive via llama-index-core; 3.10.1's import hook breaks any venv living inside the working directory (reverted upstream in nltk/nltk#3732). "openai>=2.20,<3", + "pymongo>=4.9,<5", # For MongoDB tools. "pypika>=0.50", "toolbox-adk>=1,<2", ] @@ -256,6 +258,7 @@ optional-dependencies.mcp = [ "anyio>=4.9,<5", "mcp>=1.24,<3", ] +optional-dependencies.mongodb = [ "pymongo>=4.9,<5" ] optional-dependencies.oci = [ "oci>=2.126", # OCI Generative AI native SDK (OCIGenAILlm) ] diff --git a/src/google/adk/features/_feature_registry.py b/src/google/adk/features/_feature_registry.py index 720189a5548..60f0edb3f83 100644 --- a/src/google/adk/features/_feature_registry.py +++ b/src/google/adk/features/_feature_registry.py @@ -57,6 +57,8 @@ class FeatureName(str, Enum): # enum member by name. Keeping it private avoids a backward-compat # obligation for what is intended as a temporary, internal kill-switch. _MCP_GRACEFUL_ERROR_HANDLING = "MCP_GRACEFUL_ERROR_HANDLING" + MONGODB_TOOLSET = "MONGODB_TOOLSET" + MONGODB_TOOL_SETTINGS = "MONGODB_TOOL_SETTINGS" PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING" PUBSUB_TOOL_CONFIG = "PUBSUB_TOOL_CONFIG" PUBSUB_TOOLSET = "PUBSUB_TOOLSET" @@ -180,6 +182,12 @@ class FeatureConfig: FeatureName._MCP_GRACEFUL_ERROR_HANDLING: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=True ), + FeatureName.MONGODB_TOOLSET: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), + FeatureName.MONGODB_TOOL_SETTINGS: FeatureConfig( + FeatureStage.EXPERIMENTAL, default_on=True + ), FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig( FeatureStage.EXPERIMENTAL, default_on=True ), diff --git a/src/google/adk/integrations/mongodb/__init__.py b/src/google/adk/integrations/mongodb/__init__.py new file mode 100644 index 00000000000..77ced2325a1 --- /dev/null +++ b/src/google/adk/integrations/mongodb/__init__.py @@ -0,0 +1,48 @@ +# Copyright 2026 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. + +"""MongoDB Integration (Experimental). + +MongoDB tools for vector search and hybrid search against collections on +MongoDB Atlas or MongoDB 8.0+ deployments. +""" + +from __future__ import annotations + +import typing + +if typing.TYPE_CHECKING: + from ._mongodb_toolset import MongoDbToolset + from ._settings import MongoDbToolSettings + +# Map attribute names to relative module paths +_lazy_imports = { + "MongoDbToolset": "._mongodb_toolset", + "MongoDbToolSettings": "._settings", +} + + +def __getattr__(name: str) -> typing.Any: + if name in _lazy_imports: + import importlib + + module_path = _lazy_imports[name] + # __name__ is 'google.adk.integrations.mongodb' + module = importlib.import_module(module_path, __name__) + return getattr(module, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return list(_lazy_imports.keys()) diff --git a/src/google/adk/integrations/mongodb/_client.py b/src/google/adk/integrations/mongodb/_client.py new file mode 100644 index 00000000000..e0dec7189c7 --- /dev/null +++ b/src/google/adk/integrations/mongodb/_client.py @@ -0,0 +1,53 @@ +# Copyright 2026 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 + +from typing import TYPE_CHECKING + +from ... import version + +if TYPE_CHECKING: + from pymongo import MongoClient + + +def get_mongo_client(connection_string: str) -> MongoClient: + """Creates a PyMongo client for the given connection string. + + PyMongo connects lazily, so the returned client does not perform any + network I/O until the first operation. + + Args: + connection_string: The MongoDB connection string (URI), e.g. + "mongodb+srv://user:password@cluster.mongodb.net/". + + Returns: + A PyMongo client configured with ADK driver metadata. + + Raises: + ImportError: If the `pymongo` package is not installed. + """ + try: + from pymongo import MongoClient # pylint: disable=import-outside-toplevel + from pymongo.driver_info import DriverInfo # pylint: disable=import-outside-toplevel + except ImportError as exc: + raise ImportError( + "MongoDB tools require the 'pymongo' package. " + "Please install it using `pip install google-adk[mongodb]`." + ) from exc + + return MongoClient( + connection_string, + driver=DriverInfo(name="adk-mongodb-tool", version=version.__version__), + ) diff --git a/src/google/adk/integrations/mongodb/_mongodb_toolset.py b/src/google/adk/integrations/mongodb/_mongodb_toolset.py new file mode 100644 index 00000000000..6cf032327cb --- /dev/null +++ b/src/google/adk/integrations/mongodb/_mongodb_toolset.py @@ -0,0 +1,172 @@ +# Copyright 2026 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 + +"""MongoDB toolset for vector search and hybrid search operations.""" + +import asyncio +import inspect +from typing import Any +from typing import Callable +from typing import TYPE_CHECKING + +from typing_extensions import override + +from . import _client +from . import _search_tool +from ...agents.readonly_context import ReadonlyContext +from ...features import experimental +from ...features import FeatureName +from ...tools.base_tool import BaseTool +from ...tools.base_toolset import BaseToolset +from ...tools.base_toolset import ToolPredicate +from ...tools.function_tool import FunctionTool +from ...tools.tool_context import ToolContext +from ._settings import MongoDbToolSettings + +if TYPE_CHECKING: + from pymongo import MongoClient + + +class _MongoDbTool(FunctionTool): + """FunctionTool that injects the bound MongoDB client, database and settings. + + The `client`, `database_name` and `settings` parameters are configured on + the toolset and hidden from the LLM, so the model only sees the search + parameters of each tool. + """ + + def __init__( + self, + func: Callable[..., Any], + *, + client: MongoClient, + database_name: str, + settings: MongoDbToolSettings, + ): + super().__init__(func=func) + self._ignore_params.append("client") + self._ignore_params.append("database_name") + self._ignore_params.append("settings") + self._client = client + self._database_name = database_name + self._settings = settings + + @override + async def run_async( + self, *, args: dict[str, Any], tool_context: ToolContext + ) -> Any: + args_to_call = args.copy() + signature = inspect.signature(self.func) + if "client" in signature.parameters: + args_to_call["client"] = self._client + if "database_name" in signature.parameters: + args_to_call["database_name"] = self._database_name + if "settings" in signature.parameters: + args_to_call["settings"] = self._settings + return await super().run_async(args=args_to_call, tool_context=tool_context) + + +DEFAULT_MONGODB_TOOL_NAME_PREFIX = "mongodb" + + +@experimental(FeatureName.MONGODB_TOOLSET) +class MongoDbToolset(BaseToolset): + """MongoDB Toolset contains tools for vector search and hybrid search. + + The tool names are: + - mongodb_vector_search + - mongodb_hybrid_search + + Example: + ```python + toolset = MongoDbToolset( + connection_string="mongodb+srv://user:pass@cluster.mongodb.net/", + database_name="products_db", + ) + agent = Agent(model="gemini-2.5-flash", tools=[toolset]) + ``` + """ + + def __init__( + self, + *, + database_name: str, + connection_string: str | None = None, + mongo_client: MongoClient | None = None, + tool_filter: ToolPredicate | list[str] | None = None, + settings: MongoDbToolSettings | None = None, + ): + """Initializes the MongoDbToolset. + + Args: + database_name: The MongoDB database the search tools operate on. + connection_string: The MongoDB connection string (URI) used to create a + client owned by this toolset. Requires the `pymongo` package + (`pip install google-adk[mongodb]`). + mongo_client: An existing PyMongo client to use instead of creating one + from `connection_string`. The caller keeps ownership of the client. + tool_filter: Filter to apply to tools. + settings: The settings for the MongoDB tools. + """ + super().__init__( + tool_filter=tool_filter, + tool_name_prefix=DEFAULT_MONGODB_TOOL_NAME_PREFIX, + ) + if mongo_client is not None and connection_string is not None: + raise ValueError( + "Only one of `connection_string` and `mongo_client` may be provided." + ) + if mongo_client is not None: + self._client = mongo_client + self._owns_client = False + elif connection_string is not None: + self._client = _client.get_mongo_client(connection_string) + self._owns_client = True + else: + raise ValueError( + "Either `connection_string` or `mongo_client` must be provided." + ) + self._database_name = database_name + self._settings = settings if settings else MongoDbToolSettings() + + @override + async def get_tools( + self, readonly_context: ReadonlyContext | None = None + ) -> list[BaseTool]: + """Get tools from the toolset.""" + all_tools = [ + _MongoDbTool( + func=func, + client=self._client, + database_name=self._database_name, + settings=self._settings, + ) + for func in [ + _search_tool.vector_search, + _search_tool.hybrid_search, + ] + ] + return [ + tool + for tool in all_tools + if self._is_tool_selected(tool, readonly_context) + ] + + @override + async def close(self) -> None: + """Closes the MongoDB client if it was created by this toolset.""" + if self._owns_client: + await asyncio.to_thread(self._client.close) diff --git a/src/google/adk/integrations/mongodb/_search_tool.py b/src/google/adk/integrations/mongodb/_search_tool.py new file mode 100644 index 00000000000..edb9622b900 --- /dev/null +++ b/src/google/adk/integrations/mongodb/_search_tool.py @@ -0,0 +1,341 @@ +# Copyright 2026 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 + +"""Tools to run vector and hybrid search against MongoDB collections.""" + +import asyncio +import json +import logging +from typing import Any + +from ._settings import MongoDbToolSettings + +logger = logging.getLogger("google_adk." + __name__) + +_SEARCH_SCORE_ALIAS = "search_score" +_VECTOR_PIPELINE_NAME = "vector" +_FULL_TEXT_PIPELINE_NAME = "full_text" + + +def _json_safe(value: Any) -> Any: + """Returns the value unchanged if JSON-serializable, else its string form.""" + try: + json.dumps(value) + return value + except (TypeError, ValueError, OverflowError): + return str(value) + + +def _resolve_limits( + limit: int | None, num_candidates: int | None, settings: MongoDbToolSettings +) -> tuple[int, int]: + """Resolves the result limit and the candidate count for a search operation.""" + resolved_limit = min(limit or settings.default_limit, settings.max_results) + if num_candidates is None: + resolved_num_candidates = max( + resolved_limit * 10, settings.default_num_candidates + ) + else: + # $vectorSearch requires numCandidates to be at least the limit. + resolved_num_candidates = max(num_candidates, resolved_limit) + return resolved_limit, resolved_num_candidates + + +def _build_result_projection( + embedding_field: str, output_fields: list[str] | None, score_meta: str +) -> dict[str, Any]: + """Builds the projection stage applied to search results. + + The raw embedding vector is excluded by default to keep results compact; + callers can opt into exact fields via `output_fields`. The search score is + always added under the `search_score` field. + """ + if output_fields: + projection: dict[str, Any] = {field: 1 for field in output_fields} + else: + projection = {embedding_field: 0} + projection[_SEARCH_SCORE_ALIAS] = {"$meta": score_meta} + return {"$project": projection} + + +def _aggregate_documents( + client: Any, # pymongo.MongoClient; kept as Any so pymongo stays optional. + database_name: str, + collection_name: str, + pipeline: list[dict[str, Any]], +) -> dict[str, Any]: + """Runs an aggregation pipeline and returns JSON-safe rows.""" + cursor = client[database_name][collection_name].aggregate(pipeline) + rows = [ + {key: _json_safe(value) for key, value in document.items()} + for document in cursor + ] + return {"status": "SUCCESS", "rows": rows} + + +async def vector_search( + collection_name: str, + query_embedding: list[float], + client: Any, # pymongo.MongoClient; kept as Any so pymongo stays optional. + database_name: str, + settings: MongoDbToolSettings, + filter: dict[str, Any] | None = None, + limit: int | None = None, + num_candidates: int | None = None, + index_name: str | None = None, + embedding_field: str | None = None, + output_fields: list[str] | None = None, +) -> dict[str, Any]: + """Runs an Atlas Vector Search query against a MongoDB collection. + + Finds documents whose embedding is most similar to `query_embedding` using + the `$vectorSearch` aggregation stage. Requires a vector search index on the + collection (available on MongoDB Atlas and MongoDB 8.0+). + + Args: + collection_name (str): The name of the collection to search. + query_embedding (list[float]): The embedding vector of the query, e.g. + produced by an embedding model for the user's text. + filter (dict): An optional MongoDB query filter to pre-filter documents + before searching, e.g. {"category": "kitchen"}. Only fields indexed as + filter fields in the vector search index can be used. + limit (int): The maximum number of documents to return. Capped by the + toolset settings. + num_candidates (int): The number of nearest neighbors to consider during + the search. Higher values improve recall at the cost of latency. + index_name (str): The name of the vector search index to query. Defaults + to the toolset settings. + embedding_field (str): The document field that stores the embedding + vectors. Defaults to the toolset settings. + output_fields (list[str]): The document fields to return in the results. + By default all fields except the embedding vector are returned. The + `_id` and the `search_score` are always included. + + Returns: + dict: A dictionary with the search results. + On success: {"status": "SUCCESS", "rows": [...]}, where each row is a + matching document with a "search_score" field. + On error: {"status": "ERROR", "error_details": "..."}. + + Examples: + Find the two products most similar to a query embedding, restricted to + a category: + >>> await vector_search( + ... collection_name="products", + ... query_embedding=[0.12, -0.03, ...], + ... filter={"category": "kitchen"}, + ... limit=2, + ... ) + { + "status": "SUCCESS", + "rows": [ + {"_id": "...", "name": "Robot Vacuum", "search_score": 0.93}, + {"_id": "...", "name": "Steam Mop", "search_score": 0.88}, + ], + } + """ + try: + resolved_index_name = index_name or settings.default_vector_index_name + resolved_embedding_field = ( + embedding_field or settings.default_embedding_field + ) + resolved_limit, resolved_num_candidates = _resolve_limits( + limit, num_candidates, settings + ) + + vector_search_stage: dict[str, Any] = { + "index": resolved_index_name, + "path": resolved_embedding_field, + "queryVector": query_embedding, + "numCandidates": resolved_num_candidates, + "limit": resolved_limit, + } + if filter: + vector_search_stage["filter"] = filter + + pipeline = [ + {"$vectorSearch": vector_search_stage}, + _build_result_projection( + resolved_embedding_field, output_fields, "vectorSearchScore" + ), + ] + + return await asyncio.to_thread( + _aggregate_documents, client, database_name, collection_name, pipeline + ) + except Exception as ex: + logger.exception("MongoDB vector search failed") + return { + "status": "ERROR", + "error_details": str(ex), + } + + +async def hybrid_search( + collection_name: str, + query: str, + query_embedding: list[float], + text_search_field: str, + client: Any, # pymongo.MongoClient; kept as Any so pymongo stays optional. + database_name: str, + settings: MongoDbToolSettings, + filter: dict[str, Any] | None = None, + limit: int | None = None, + num_candidates: int | None = None, + vector_index_name: str | None = None, + search_index_name: str | None = None, + embedding_field: str | None = None, + vector_weight: float | None = None, + text_weight: float | None = None, + output_fields: list[str] | None = None, +) -> dict[str, Any]: + """Runs a hybrid (full-text + vector) search against a MongoDB collection. + + Combines full-text search and vector search with reciprocal rank fusion + using the `$rankFusion` aggregation stage, so documents matching either the + text query or the embedding similarity are ranked together. Requires a + full-text search index and a vector search index on the collection, and a + deployment that supports `$rankFusion` (MongoDB 8.0+, or MongoDB Atlas). + + Args: + collection_name (str): The name of the collection to search. + query (str): The text query for full-text search. + query_embedding (list[float]): The embedding vector of the query, e.g. + produced by an embedding model for the user's text. + text_search_field (str): The document field to run the full-text search + against. + filter (dict): An optional MongoDB query filter to pre-filter documents + before the vector search, e.g. {"category": "kitchen"}. Only fields + indexed as filter fields in the vector search index can be used. + limit (int): The maximum number of documents to return. Capped by the + toolset settings. + num_candidates (int): The number of nearest neighbors to consider during + the vector search. Higher values improve recall at the cost of + latency. + vector_index_name (str): The name of the vector search index to query. + Defaults to the toolset settings. + search_index_name (str): The name of the full-text search index to + query. Defaults to the toolset settings. + embedding_field (str): The document field that stores the embedding + vectors. Defaults to the toolset settings. + vector_weight (float): The weight of the vector search ranking in the + fused score. Defaults to 1.0. + text_weight (float): The weight of the full-text search ranking in the + fused score. Defaults to 1.0. + output_fields (list[str]): The document fields to return in the results. + By default all fields except the embedding vector are returned. The + `_id` and the `search_score` are always included. + + Returns: + dict: A dictionary with the search results. + On success: {"status": "SUCCESS", "rows": [...]}, where each row is a + matching document with a "search_score" field holding the fused + reciprocal rank fusion score. + On error: {"status": "ERROR", "error_details": "..."}. + + Examples: + Find products relevant to "cordless vacuum for pet hair", weighing + vector matches twice as much as text matches: + >>> await hybrid_search( + ... collection_name="products", + ... query="cordless vacuum for pet hair", + ... query_embedding=[0.12, -0.03, ...], + ... text_search_field="description", + ... vector_weight=2.0, + ... limit=3, + ... ) + { + "status": "SUCCESS", + "rows": [ + {"_id": "...", "name": "Pet Hair Vacuum", "search_score": 0.032}, + ... + ], + } + """ + try: + resolved_vector_index_name = ( + vector_index_name or settings.default_vector_index_name + ) + resolved_search_index_name = ( + search_index_name or settings.default_search_index_name + ) + resolved_embedding_field = ( + embedding_field or settings.default_embedding_field + ) + resolved_limit, resolved_num_candidates = _resolve_limits( + limit, num_candidates, settings + ) + + vector_search_stage: dict[str, Any] = { + "index": resolved_vector_index_name, + "path": resolved_embedding_field, + "queryVector": query_embedding, + "numCandidates": resolved_num_candidates, + "limit": resolved_num_candidates, + } + if filter: + vector_search_stage["filter"] = filter + + pipeline = [ + { + "$rankFusion": { + "input": { + "pipelines": { + _VECTOR_PIPELINE_NAME: [ + {"$vectorSearch": vector_search_stage} + ], + _FULL_TEXT_PIPELINE_NAME: [ + { + "$search": { + "index": resolved_search_index_name, + "text": { + "query": query, + "path": text_search_field, + }, + } + }, + {"$limit": resolved_num_candidates}, + ], + } + }, + "combination": { + "weights": { + _VECTOR_PIPELINE_NAME: ( + vector_weight if vector_weight is not None else 1.0 + ), + _FULL_TEXT_PIPELINE_NAME: ( + text_weight if text_weight is not None else 1.0 + ), + } + }, + "scoreDetails": False, + } + }, + {"$limit": resolved_limit}, + _build_result_projection( + resolved_embedding_field, output_fields, "score" + ), + ] + + return await asyncio.to_thread( + _aggregate_documents, client, database_name, collection_name, pipeline + ) + except Exception as ex: + logger.exception("MongoDB hybrid search failed") + return { + "status": "ERROR", + "error_details": str(ex), + } diff --git a/src/google/adk/integrations/mongodb/_settings.py b/src/google/adk/integrations/mongodb/_settings.py new file mode 100644 index 00000000000..0bc1ac77704 --- /dev/null +++ b/src/google/adk/integrations/mongodb/_settings.py @@ -0,0 +1,44 @@ +# Copyright 2026 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 + +from pydantic import BaseModel +from pydantic import Field + +from ...features import experimental +from ...features import FeatureName + + +@experimental(FeatureName.MONGODB_TOOL_SETTINGS) +class MongoDbToolSettings(BaseModel): + """Settings for MongoDB tools.""" + + default_vector_index_name: str = "vector_index" + """Default name of the vector search index to query.""" + + default_search_index_name: str = "default" + """Default name of the full-text search index used by hybrid search.""" + + default_embedding_field: str = "embedding" + """Default document field that stores embedding vectors.""" + + default_limit: int = Field(default=4, gt=0) + """Default number of documents returned by a search operation.""" + + max_results: int = Field(default=50, gt=0) + """Maximum number of documents a search operation may return.""" + + default_num_candidates: int = Field(default=100, gt=0) + """Default number of nearest neighbors considered by vector search.""" diff --git a/tests/unittests/integrations/mongodb/test_client.py b/tests/unittests/integrations/mongodb/test_client.py new file mode 100644 index 00000000000..ebcb730d789 --- /dev/null +++ b/tests/unittests/integrations/mongodb/test_client.py @@ -0,0 +1,53 @@ +# Copyright 2026 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. + +"""Tests for the MongoDB client factory.""" + +import sys +from types import ModuleType +from unittest import mock + +from google.adk.integrations.mongodb._client import get_mongo_client +import pytest + + +def test_get_mongo_client_raises_import_error_without_pymongo(monkeypatch): + """get_mongo_client raises a helpful ImportError when pymongo is missing.""" + monkeypatch.setitem(sys.modules, "pymongo", None) + + with pytest.raises(ImportError, match=r"google-adk\[mongodb\]"): + get_mongo_client("mongodb://localhost:27017") + + +def test_get_mongo_client_creates_client_with_driver_metadata(monkeypatch): + """get_mongo_client builds a MongoClient from the connection string.""" + fake_pymongo = ModuleType("pymongo") + fake_driver_info = ModuleType("pymongo.driver_info") + mongo_client_cls = mock.MagicMock() + driver_info_cls = mock.MagicMock() + fake_pymongo.MongoClient = mongo_client_cls + fake_driver_info.DriverInfo = driver_info_cls + monkeypatch.setitem(sys.modules, "pymongo", fake_pymongo) + monkeypatch.setitem(sys.modules, "pymongo.driver_info", fake_driver_info) + + result = get_mongo_client("mongodb://localhost:27017") + + assert result is mongo_client_cls.return_value + mongo_client_cls.assert_called_once() + assert mongo_client_cls.call_args.args[0] == "mongodb://localhost:27017" + assert ( + mongo_client_cls.call_args.kwargs["driver"] + is driver_info_cls.return_value + ) + assert driver_info_cls.call_args.kwargs["name"] == "adk-mongodb-tool" diff --git a/tests/unittests/integrations/mongodb/test_mongodb_toolset.py b/tests/unittests/integrations/mongodb/test_mongodb_toolset.py new file mode 100644 index 00000000000..9e4ee3287f9 --- /dev/null +++ b/tests/unittests/integrations/mongodb/test_mongodb_toolset.py @@ -0,0 +1,167 @@ +# Copyright 2026 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. + +"""Tests for MongoDbToolset. + +Verifies that the toolset exposes prefixed, filterable search tools and +injects the bound client, database and settings at run time. +""" + +from unittest import mock + +from google.adk.integrations.mongodb import MongoDbToolset +from google.adk.integrations.mongodb import MongoDbToolSettings +from google.adk.integrations.mongodb._mongodb_toolset import DEFAULT_MONGODB_TOOL_NAME_PREFIX +import pytest + + +def _make_toolset(**kwargs): + return MongoDbToolset( + database_name="test_db", mongo_client=mock.MagicMock(), **kwargs + ) + + +def test_mongodb_toolset_name_prefix(): + """MongoDbToolset prefixes its tool names with 'mongodb'.""" + toolset = _make_toolset() + assert toolset.tool_name_prefix == DEFAULT_MONGODB_TOOL_NAME_PREFIX + + +async def test_mongodb_toolset_tools_default(): + """The default toolset exposes the vector and hybrid search tools.""" + toolset = _make_toolset() + + tools = await toolset.get_tools() + + assert set([tool.name for tool in tools]) == { + "vector_search", + "hybrid_search", + } + + +async def test_mongodb_toolset_tools_prefixed(): + """Tools are returned with the 'mongodb' name prefix applied.""" + toolset = _make_toolset() + + tools = await toolset.get_tools_with_prefix() + + assert set([tool.name for tool in tools]) == { + "mongodb_vector_search", + "mongodb_hybrid_search", + } + + +async def test_mongodb_toolset_tools_selective(): + """tool_filter restricts the exposed tools to the listed names.""" + toolset = _make_toolset(tool_filter=["vector_search"]) + + tools = await toolset.get_tools() + + assert [tool.name for tool in tools] == ["vector_search"] + + +async def test_mongodb_toolset_unknown_tool_filtered_out(): + """Unknown names in tool_filter yield no tools.""" + toolset = _make_toolset(tool_filter=["unknown"]) + + tools = await toolset.get_tools() + + assert tools == [] + + +def test_mongodb_toolset_requires_client_or_connection_string(): + """Constructing without a client or connection string raises ValueError.""" + with pytest.raises(ValueError, match="must be provided"): + MongoDbToolset(database_name="test_db") + + +def test_mongodb_toolset_rejects_client_and_connection_string(): + """Constructing with both a client and a connection string raises ValueError.""" + with pytest.raises(ValueError, match="Only one of"): + MongoDbToolset( + database_name="test_db", + connection_string="mongodb://localhost:27017", + mongo_client=mock.MagicMock(), + ) + + +async def test_mongodb_tool_injects_client_database_and_settings(): + """Running a tool injects the bound client, database and settings.""" + client = mock.MagicMock() + client["test_db"]["test_coll"].aggregate.return_value = iter( + [{"_id": 1, "title": "Doc"}] + ) + toolset = MongoDbToolset( + database_name="test_db", + mongo_client=client, + settings=MongoDbToolSettings(default_limit=9), + ) + tools = await toolset.get_tools() + tool = next(tool for tool in tools if tool.name == "vector_search") + + result = await tool.run_async( + args={"collection_name": "test_coll", "query_embedding": [0.1]}, + tool_context=mock.MagicMock(), + ) + + assert result["status"] == "SUCCESS" + assert result["rows"] == [{"_id": 1, "title": "Doc"}] + pipeline = client["test_db"]["test_coll"].aggregate.call_args[0][0] + assert pipeline[0]["$vectorSearch"]["queryVector"] == [0.1] + # The custom settings flow through to the tool. + assert pipeline[0]["$vectorSearch"]["limit"] == 9 + + +async def test_mongodb_tool_declaration_hides_injected_parameters(): + """The generated function schema only exposes search parameters to the model.""" + toolset = _make_toolset() + tools = await toolset.get_tools() + + declaration = next( + tool for tool in tools if tool.name == "vector_search" + )._get_declaration() + + properties = declaration.parameters_json_schema["properties"] + assert "collection_name" in properties + assert "query_embedding" in properties + for injected in ("client", "database_name", "settings"): + assert injected not in properties + + +async def test_close_does_not_close_injected_client(): + """close() leaves a caller-owned client open.""" + injected_client = mock.MagicMock() + toolset = MongoDbToolset( + database_name="test_db", mongo_client=injected_client + ) + + await toolset.close() + + injected_client.close.assert_not_called() + + +async def test_close_closes_client_created_from_connection_string(monkeypatch): + """close() closes the client the toolset created from a connection string.""" + created_client = mock.MagicMock() + monkeypatch.setattr( + "google.adk.integrations.mongodb._mongodb_toolset._client.get_mongo_client", + lambda connection_string: created_client, + ) + toolset = MongoDbToolset( + database_name="test_db", connection_string="mongodb://localhost:27017" + ) + + await toolset.close() + + created_client.close.assert_called_once() diff --git a/tests/unittests/integrations/mongodb/test_search_tool.py b/tests/unittests/integrations/mongodb/test_search_tool.py new file mode 100644 index 00000000000..dfa733732c5 --- /dev/null +++ b/tests/unittests/integrations/mongodb/test_search_tool.py @@ -0,0 +1,315 @@ +# Copyright 2026 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. + +"""Tests for MongoDB search tools. + +Verifies that vector_search and hybrid_search build the expected MongoDB +aggregation pipelines and return JSON-safe results. +""" + +from unittest import mock + +from google.adk.integrations.mongodb import _search_tool +from google.adk.integrations.mongodb import MongoDbToolSettings + +_EMBEDDING = [0.1, 0.2, 0.3] + + +def _make_client(documents=None): + """Returns a mock MongoClient whose aggregate() yields the given documents.""" + client = mock.MagicMock() + client["test_db"]["test_coll"].aggregate.return_value = iter(documents or []) + return client + + +def _aggregate_pipeline(client): + """Returns the pipeline passed to aggregate() on the mock client.""" + return client["test_db"]["test_coll"].aggregate.call_args[0][0] + + +async def test_vector_search_uses_settings_defaults(): + """Vector search queries the collection with index, field and limits from settings.""" + client = _make_client() + + result = await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + assert result == {"status": "SUCCESS", "rows": []} + pipeline = _aggregate_pipeline(client) + assert pipeline[0]["$vectorSearch"] == { + "index": "vector_index", + "path": "embedding", + "queryVector": _EMBEDDING, + "numCandidates": 100, + "limit": 4, + } + + +async def test_vector_search_applies_explicit_arguments(): + """Explicit index, field, filter and limits override the settings defaults.""" + client = _make_client() + + await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + filter={"category": "kitchen"}, + limit=7, + num_candidates=42, + index_name="my_index", + embedding_field="text_embedding", + ) + + pipeline = _aggregate_pipeline(client) + assert pipeline[0]["$vectorSearch"] == { + "index": "my_index", + "path": "text_embedding", + "queryVector": _EMBEDDING, + "filter": {"category": "kitchen"}, + "numCandidates": 42, + "limit": 7, + } + + +async def test_vector_search_caps_limit_at_max_results(): + """A limit above settings.max_results is capped.""" + client = _make_client() + + await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(max_results=10), + limit=50, + ) + + pipeline = _aggregate_pipeline(client) + assert pipeline[0]["$vectorSearch"]["limit"] == 10 + + +async def test_vector_search_raises_num_candidates_to_limit(): + """numCandidates below the limit is raised, as $vectorSearch requires it.""" + client = _make_client() + + await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + limit=8, + num_candidates=5, + ) + + pipeline = _aggregate_pipeline(client) + assert pipeline[0]["$vectorSearch"]["limit"] == 8 + assert pipeline[0]["$vectorSearch"]["numCandidates"] == 8 + + +async def test_vector_search_excludes_embedding_field_from_results(): + """The default projection hides the raw embedding vector and adds the score.""" + client = _make_client() + + await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + pipeline = _aggregate_pipeline(client) + assert pipeline[1] == { + "$project": { + "embedding": 0, + "search_score": {"$meta": "vectorSearchScore"}, + } + } + + +async def test_vector_search_projects_output_fields_when_given(): + """output_fields switches the projection to inclusion mode.""" + client = _make_client() + + await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + output_fields=["title", "price"], + ) + + pipeline = _aggregate_pipeline(client) + assert pipeline[1] == { + "$project": { + "title": 1, + "price": 1, + "search_score": {"$meta": "vectorSearchScore"}, + } + } + + +async def test_vector_search_returns_json_safe_rows(): + """Non-JSON-serializable values in result documents are converted to strings.""" + object_id = object() + client = _make_client( + [{"_id": object_id, "title": "Doc", "search_score": 0.9}] + ) + + result = await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + assert result["status"] == "SUCCESS" + assert result["rows"] == [ + {"_id": str(object_id), "title": "Doc", "search_score": 0.9} + ] + + +async def test_vector_search_returns_error_on_failure(): + """A failing aggregation returns an ERROR result instead of raising.""" + client = _make_client() + client["test_db"]["test_coll"].aggregate.side_effect = RuntimeError("boom") + + result = await _search_tool.vector_search( + collection_name="test_coll", + query_embedding=_EMBEDDING, + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + assert result == {"status": "ERROR", "error_details": "boom"} + + +async def test_hybrid_search_builds_rank_fusion_pipeline(): + """Hybrid search fuses vector and full-text rankings via $rankFusion.""" + client = _make_client() + + result = await _search_tool.hybrid_search( + collection_name="test_coll", + query="cordless vacuum", + query_embedding=_EMBEDDING, + text_search_field="description", + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + assert result == {"status": "SUCCESS", "rows": []} + pipeline = _aggregate_pipeline(client) + rank_fusion = pipeline[0]["$rankFusion"] + pipelines = rank_fusion["input"]["pipelines"] + assert pipelines["vector"] == [{ + "$vectorSearch": { + "index": "vector_index", + "path": "embedding", + "queryVector": _EMBEDDING, + "numCandidates": 100, + "limit": 100, + } + }] + assert pipelines["full_text"] == [ + { + "$search": { + "index": "default", + "text": {"query": "cordless vacuum", "path": "description"}, + } + }, + {"$limit": 100}, + ] + assert rank_fusion["combination"]["weights"] == { + "vector": 1.0, + "full_text": 1.0, + } + assert rank_fusion["scoreDetails"] is False + assert pipeline[1] == {"$limit": 4} + assert pipeline[2] == { + "$project": {"embedding": 0, "search_score": {"$meta": "score"}} + } + + +async def test_hybrid_search_applies_weights_filter_and_index_names(): + """Explicit weights, filter and index names are applied to the pipeline.""" + client = _make_client() + + await _search_tool.hybrid_search( + collection_name="test_coll", + query="cordless vacuum", + query_embedding=_EMBEDDING, + text_search_field="description", + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + filter={"in_stock": True}, + limit=5, + num_candidates=25, + vector_index_name="v_idx", + search_index_name="s_idx", + embedding_field="vec", + vector_weight=2.0, + text_weight=0.5, + ) + + pipeline = _aggregate_pipeline(client) + rank_fusion = pipeline[0]["$rankFusion"] + assert rank_fusion["input"]["pipelines"]["vector"] == [{ + "$vectorSearch": { + "index": "v_idx", + "path": "vec", + "queryVector": _EMBEDDING, + "filter": {"in_stock": True}, + "numCandidates": 25, + "limit": 25, + } + }] + full_text = rank_fusion["input"]["pipelines"]["full_text"] + assert full_text[0]["$search"]["index"] == "s_idx" + assert full_text[1] == {"$limit": 25} + assert rank_fusion["combination"]["weights"] == { + "vector": 2.0, + "full_text": 0.5, + } + assert pipeline[1] == {"$limit": 5} + + +async def test_hybrid_search_returns_error_on_failure(): + """A failing aggregation returns an ERROR result instead of raising.""" + client = _make_client() + client["test_db"]["test_coll"].aggregate.side_effect = RuntimeError("boom") + + result = await _search_tool.hybrid_search( + collection_name="test_coll", + query="cordless vacuum", + query_embedding=_EMBEDDING, + text_search_field="description", + client=client, + database_name="test_db", + settings=MongoDbToolSettings(), + ) + + assert result == {"status": "ERROR", "error_details": "boom"}