pydantic-ai

Moss Pydantic AI Cookbook

This cookbook shows how to expose Moss semantic search as a reusable tool inside a Pydantic AI agent.

Overview

Moss is a semantic search platform that delivers sub-10ms retrieval by loading vector indices into local memory. This cookbook provides MossSearchTool, a small wrapper around MossClient that exposes a .tool property for Pydantic AI agents.

Installation

cd examples/cookbook/pydantic-ai
uv sync

Setup

Create a .env file in this directory (see .env.example):

MOSS_PROJECT_ID=your-project-id
MOSS_PROJECT_KEY=your-project-key
MOSS_INDEX_NAME=your-index-name
PYDANTIC_AI_MODEL=openai:gpt-4o
OPENAI_API_KEY=your-openai-api-key

This cookbook eagerly calls await moss.load_index() before the agent runs. On the first run, that preload step may download the local query model used for in-memory search. If MOSS_INDEX_NAME does not exist yet, example.py creates a small demo index automatically before loading it.

Quick Start

Using MossSearchTool

import asyncio
from moss import MossClient
from pydantic_ai import Agent
from moss_pydantic_ai import MossSearchTool

async def main():
    client = MossClient("your-project-id", "your-project-key")
    moss = MossSearchTool(client=client, index_name="my-index")
    await moss.load_index()  # pre-load for fast queries

    agent = Agent("openai:gpt-4o", tools=[moss.tool])
    result = await agent.run("What is the refund policy?")
    print(result.output)

asyncio.run(main())

Configuration

Parameter	Default	Description
client	(required)	A MossClient instance
index_name	(required)	Name of the Moss index to query
tool_name	moss_search	Tool name exposed to the LLM
top_k	5	Number of results to retrieve per query
alpha	0.8	Blend: 1.0 = semantic only, 0.0 = keyword only

Run the Demo

uv run python example.py

The demo creates a MossClient, creates the demo index if needed, loads that index, defines a MossSearchTool, and runs a Pydantic AI agent against it.

Because the demo eagerly preloads the index before agent.run(...), the first run can take longer while Moss fetches the local query model cache.

Run the Tests

uv run python test_integration.py

How It Works

Pydantic AI inspects the tool function's signature and docstring to derive the input schema and description. MossSearchTool._build_tool() creates an async moss_search(query: str) -> str function and wraps it in pydantic_ai.Tool(...), so the parameter schema is auto-generated.

Files

File	Description
moss_pydantic_ai.py	MossSearchTool class
example.py	Runnable cookbook demo using the helper module
test_integration.py	Unit tests (mocked, no credentials required)
pyproject.toml	Package metadata
.env.example	Template for required environment variables

Notes

• This cookbook exposes Moss search because that is the concrete capability exposed by the current Python SDK.
• If Moss later adds first-class workflow or action definitions, the same adapter pattern can be promoted into an official moss.integrations.pydantic_ai module.