diff --git a/examples/cookbook/google-adk/.env.example b/examples/cookbook/google-adk/.env.example new file mode 100644 index 00000000..987ab00b --- /dev/null +++ b/examples/cookbook/google-adk/.env.example @@ -0,0 +1,7 @@ +# Moss credentials — get these from https://moss.dev +MOSS_PROJECT_ID=your_project_id +MOSS_PROJECT_KEY=your_project_key +MOSS_INDEX_NAME=your_index_name + +# Google API Key for Gemini +GEMINI_API_KEY=your_gemini_api_key diff --git a/examples/cookbook/google-adk/README.md b/examples/cookbook/google-adk/README.md new file mode 100644 index 00000000..9fb592dc --- /dev/null +++ b/examples/cookbook/google-adk/README.md @@ -0,0 +1,84 @@ +# Moss + Google ADK Cookbook + +Use [Moss](https://moss.dev) semantic search as a high-speed retrieval tool for [Google ADK](https://adk.dev) (Agent Development Kit) agents. + +## Why Moss with Google ADK? + +Traditional vector databases add 200–500 ms per retrieval hop. Moss loads the index and model weights directly into your application process, delivering **sub-10 ms** search. Because Google ADK natively supports asynchronous tools, it pairs perfectly with Moss's async-first architecture, allowing you to maximize performance. + +## Installation + +We recommend using [uv](https://docs.astral.sh/uv/) for fast dependency management: + +```bash +uv sync +``` + +Or install dependencies directly: + +```bash +uv pip install "google-adk>=1.10.0" moss python-dotenv +``` + +*Note: This example pins `google-adk>=1.10.0` to ensure optimal parallel execution for async tools.* + +## Setup + +Copy the example env file and fill in your credentials: + +```bash +cp .env.example .env +``` + +Required variables: + +```env +MOSS_PROJECT_ID=your_project_id +MOSS_PROJECT_KEY=your_project_key +MOSS_INDEX_NAME=your_index_name +GEMINI_API_KEY=your_gemini_api_key +``` + +## Running the demo + +```bash +uv run moss_adk_demo.py +``` + +## How it works + +### Loading the index (The Secret Sauce) + +The index must be pulled into local memory **once** before the agent starts. This is the step that switches retrieval from standard cloud-round-trip latency to sub-10ms local speed: + +```python +await client.load_index("my-index") +``` + +Call this in your setup/startup code before invoking the ADK agent. If you skip this, Moss will fall back to querying the cloud API (which works, but is significantly slower). + +### Native Async Tool + +Google ADK natively supports `async def` tool functions. We wrap the `MossClient` in a factory function that returns a fully typed async tool: + +```python +def create_moss_tool(client: MossClient, index_name: str): + async def moss_retrieval(query: str, top_k: int = 5, metadata_filter: dict = None) -> str: + # ... implementation ... + return moss_retrieval +``` + +### Metadata filtering + +Google ADK extracts the tool schema from the function signature and docstrings. We document the Moss filter DSL directly in the docstring so Gemini knows how to use it: + +```python +metadata_filter = { + "$and": [ + {"field": "category", "condition": {"$eq": "refunds"}}, + {"field": "price", "condition": {"$lt": 50}}, + ] +} +``` + +Available operators: `$eq`, `$ne`, `$gt`, `$gte`, `$lt`, `$lte`, `$in`, `$and`, `$or`. diff --git a/examples/cookbook/google-adk/moss_adk_demo.py b/examples/cookbook/google-adk/moss_adk_demo.py new file mode 100644 index 00000000..f9087a97 --- /dev/null +++ b/examples/cookbook/google-adk/moss_adk_demo.py @@ -0,0 +1,74 @@ +import asyncio +import os + +from dotenv import load_dotenv +from google.adk.agents import Agent + +from moss import MossClient +from moss_adk_tool import create_moss_tool + +load_dotenv() + + +async def main(): + project_id = os.getenv("MOSS_PROJECT_ID") + project_key = os.getenv("MOSS_PROJECT_KEY") + index_name = os.getenv("MOSS_INDEX_NAME") + + if not all([project_id, project_key, index_name]): + raise EnvironmentError( + "Please set MOSS_PROJECT_ID, MOSS_PROJECT_KEY, and MOSS_INDEX_NAME " + "in your environment or .env file." + ) + + # Note: Google ADK typically requires GEMINI_API_KEY to be set in the environment. + if not os.getenv("GEMINI_API_KEY"): + raise EnvironmentError( + "Please set GEMINI_API_KEY in your environment or .env file." + ) + + client = MossClient(project_id, project_key) + + # Load the index into local memory before the agent runs. + # This one-time setup is what enables sub-10ms retrieval inside the agent loop. + print(f"Loading index '{index_name}' into local memory...") + await client.load_index(index_name) + print("Index loaded.\n") + + # Create the ADK compatible tool + retrieval_tool = create_moss_tool(client, index_name) + + # Initialize the Google ADK Agent + agent = Agent( + name="moss_assistant", + model="gemini-2.5-flash", + tools=[retrieval_tool], + ) + + from google.adk.runners import InMemoryRunner + from google.genai import types + + runner = InMemoryRunner(agent=agent) + + question = "What is the policy for processing refunds for digital goods?" + print(f"Question: {question}") + print("-" * 50) + + # Format the input message + content = types.Content(role="user", parts=[types.Part(text=question)]) + + print("\n--- Agent Response ---") + + # Run the agent asynchronously via the runner + async for event in runner.run_async( + user_id="user_demo", + session_id="session_demo", + new_message=content + ): + if hasattr(event, "is_final_response") and event.is_final_response(): + print(event.content.parts[0].text) + + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/cookbook/google-adk/moss_adk_tool.py b/examples/cookbook/google-adk/moss_adk_tool.py new file mode 100644 index 00000000..4bd3d528 --- /dev/null +++ b/examples/cookbook/google-adk/moss_adk_tool.py @@ -0,0 +1,38 @@ +from typing import Any, Dict, Optional + +from moss import MossClient, QueryOptions + + +def create_moss_tool(client: MossClient, index_name: str): + """ + Factory function that creates a Google ADK compatible asynchronous tool + for Moss semantic retrieval. + """ + + async def moss_retrieval( + query: str, + top_k: int = 5, + metadata_filter: Optional[Dict[str, Any]] = None, + ) -> str: + """ + Finds relevant information from a knowledge base using semantic search. + Use this when the answer is likely contained in indexed documents. + + Args: + query: The search query string. + top_k: Number of results to return (default: 5). + metadata_filter: Optional filter using the Moss filter DSL. + Example: {'$and': [{'field': 'category', 'condition': {'$eq': 'refunds'}}]} + """ + options = QueryOptions(top_k=top_k, filter=metadata_filter) + results = await client.query(index_name, query, options) + + if not results.docs: + return "No relevant information found." + + return "\n\n".join( + f"--- Result ID: {doc.id} (Score: {doc.score:.3f}) ---\n{doc.text}" + for doc in results.docs + ) + + return moss_retrieval diff --git a/examples/cookbook/google-adk/pyproject.toml b/examples/cookbook/google-adk/pyproject.toml new file mode 100644 index 00000000..3b8a1293 --- /dev/null +++ b/examples/cookbook/google-adk/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "google-adk-moss" +version = "0.1.0" +description = "Google ADK integration for Moss semantic search" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "BSD-2-Clause" } +authors = [ + { name = "InferEdge Inc.", email = "contact@moss.dev" } +] +dependencies = [ + "google-adk>=1.10.0", + "moss>=1.0.0", + "python-dotenv", +] + +[tool.hatch.build.targets.wheel] +packages = ["moss_adk_tool.py"] + +[tool.hatch.build.targets.sdist] +include = [ + "README.md", + "moss_adk_tool.py", + "moss_adk_demo.py", + ".env.example", +] diff --git a/examples/cookbook/google-adk/test_integration.py b/examples/cookbook/google-adk/test_integration.py new file mode 100644 index 00000000..31cab0c5 --- /dev/null +++ b/examples/cookbook/google-adk/test_integration.py @@ -0,0 +1,64 @@ +import asyncio +import unittest +from unittest.mock import AsyncMock, MagicMock + +from moss import QueryOptions +from moss_adk_tool import create_moss_tool + + +class TestMossADKTool(unittest.IsolatedAsyncioTestCase): + async def test_tool_formats_results(self): + mock_client = MagicMock() + mock_docs = [ + MagicMock(id="doc_1", text="First result", score=0.9), + MagicMock(id="doc_2", text="Second result", score=0.8), + ] + + mock_client.query = AsyncMock() + mock_client.query.return_value = MagicMock(docs=mock_docs) + + tool = create_moss_tool(mock_client, "test-index") + + result = await tool("test query", top_k=2) + + self.assertIn("Result ID: doc_1", result) + self.assertIn("First result", result) + self.assertIn("Score: 0.900", result) + self.assertIn("Result ID: doc_2", result) + self.assertIn("Second result", result) + self.assertIn("Score: 0.800", result) + + mock_client.query.assert_called_once() + args, _ = mock_client.query.call_args + self.assertEqual(args[0], "test-index") + self.assertEqual(args[1], "test query") + + async def test_tool_handles_empty_results(self): + mock_client = MagicMock() + mock_client.query = AsyncMock() + mock_client.query.return_value = MagicMock(docs=[]) + + tool = create_moss_tool(mock_client, "test-index") + + result = await tool("empty query") + self.assertEqual(result, "No relevant information found.") + + async def test_tool_passes_metadata_filter(self): + mock_client = MagicMock() + mock_client.query = AsyncMock() + mock_client.query.return_value = MagicMock(docs=[]) + + tool = create_moss_tool(mock_client, "test-index") + filt = {"$and": [{"field": "category", "condition": {"$eq": "refunds"}}]} + + await tool("query", top_k=3, metadata_filter=filt) + + args, kwargs = mock_client.query.call_args + options = args[2] + + self.assertIsInstance(options, QueryOptions) + self.assertEqual(options.top_k, 3) + self.assertEqual(options.filter, filt) + +if __name__ == "__main__": + unittest.main()