feat: add Moss retrieval cookbook example for Google ADK (#93)

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

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`.

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())

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

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 = "[email protected]" }
+]
+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",
+]

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()