# How to Use Pydantic Models for Tool Definitions in Needle 2

> Learn to define agent tools declaratively with Pydantic models in Needle 2. Automatically generate JSON schemas and enable runtime type validation for LLM prompts.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-21

---

**Needle 2 leverages Pydantic models in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) to define agent tools declaratively, automatically generating JSON schemas for LLM prompts while providing runtime type validation and typed execution through the `run` method.**

Needle 2 is a lightweight LLM agent framework that uses Pydantic models to create type-safe tool interfaces. By inheriting from `BaseTool` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), developers define tool arguments using standard Pydantic field types, enabling automatic schema generation for LLM consumption and robust input validation. This approach eliminates manual JSON parsing and ensures that agents receive structured, validated data before executing business logic.

## The BaseTool Architecture

In [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), the framework defines `BaseTool` as a subclass of `pydantic.BaseModel`. This design allows every tool to inherit Pydantic's validation machinery while maintaining a consistent interface for the agent runtime. The module maintains a `TOOL_REGISTRY: Dict[str, Type[BaseTool]]` dictionary that maps string identifiers to their corresponding model classes, enabling dynamic discovery and instantiation during agent execution.

## Defining Your First Tool

To create a tool, subclass `BaseTool` and declare parameters as typed fields. The class docstring becomes the tool's description for the LLM.

```python
from needle.agent.tools import BaseTool, TOOL_REGISTRY

class SearchTool(BaseTool):
    """Search the web for a given query."""
    query: str
    num_results: int = 5
    
    def run(self) -> str:
        # Access validated fields as typed attributes

        return f"Searching for '{self.query}' with {self.num_results} results"

# Register to make discoverable by the agent

TOOL_REGISTRY["search"] = SearchTool

```

### Registration Requirements

Every tool must be added to the `TOOL_REGISTRY` dictionary with a unique string key. This registry located in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) acts as the single source of truth for the agent when determining available capabilities and routing tool calls.

## Schema Generation for LLM Prompts

When constructing system prompts, Needle extracts the JSON schema from each tool model. The framework calls Pydantic's `model_json_schema()` method (or equivalent schema helpers) to generate a complete JSON Schema object describing required fields, types, and descriptions directly from your Python type hints.

```python
def build_system_prompt():
    schemas = []
    for name, tool_cls in TOOL_REGISTRY.items():
        # Generate schema from Pydantic model

        schema = tool_cls().model_json_schema()
        schemas.append(f"Tool: {name}\nSchema: {schema}")
    return "\n\n".join(schemas)

```

## Runtime Validation and Execution Flow

When the LLM returns a tool call payload, Needle validates the JSON against the Pydantic model before invoking business logic.

```python
import json
from pydantic import ValidationError

def execute_tool_call(tool_name: str, json_payload: str):
    tool_cls = TOOL_REGISTRY[tool_name]
    
    try:
        args = json.loads(json_payload)
        # Validation occurs during instantiation

        tool_instance = tool_cls(**args)
    except (json.JSONDecodeError, ValidationError) as e:
        raise RuntimeError(f"Tool '{tool_name}' received invalid input: {e}")
    
    # Execute with fully typed instance

    return tool_instance.run()

```

The `run` method receives `self` as a validated Pydantic model instance, allowing direct access to properly typed attributes without manual casting.

## Leveraging Advanced Pydantic Features

Because `BaseTool` extends `pydantic.BaseModel`, you utilize the full Pydantic v2 feature set for sophisticated tool definitions.

- **Field validation**: Use `Field(..., min_length=1)` or `@field_validator` decorators for fine-grained constraints
- **Default values**: Provide defaults for optional parameters the LLM may omit
- **Nested models**: Compose complex arguments using nested `BaseModel` subclasses for structured data

```python
from pydantic import Field, field_validator

class CalculateTool(BaseTool):
    """Perform safe mathematical calculations."""
    expression: str = Field(..., min_length=1, description="Math expression to evaluate")
    precision: int = Field(default=2, ge=0, le=10, description="Decimal places")
    
    @field_validator('expression')
    @classmethod
    def validate_safety(cls, v: str):
        if any(op in v for op in ['import', 'exec', 'eval']):
            raise ValueError("Unsafe expression detected")
        return v
    
    def run(self) -> float:
        # self.expression is guaranteed to be a string matching constraints

        result = eval(self.expression)  # Simplified - use safeeval in production

        return round(result, self.precision)

```

## Summary

- **BaseTool** in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) inherits from Pydantic's `BaseModel`, providing automatic type validation and serialization
- Tools register themselves in **TOOL_REGISTRY** to become discoverable by the agent runtime
- **Schema generation** uses native Pydantic methods to create JSON schemas that inform LLM tool selection
- The **run** method executes on a fully validated instance with typed attributes, eliminating defensive coding
- Advanced features like **Field constraints**, **validators**, and **nested models** work natively without framework modifications

## Frequently Asked Questions

### What Pydantic version does Needle 2 require?

Needle 2 supports Pydantic v2, utilizing `model_json_schema()` and modern validation patterns. The `BaseTool` class in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) inherits from `pydantic.BaseModel` using v2 configuration semantics, ensuring compatibility with current Pydantic features.

### How do I handle optional parameters in tool definitions?

Define optional fields by providing default values or using `Optional[Type] = None` annotations. Pydantic marks these as non-required in the generated JSON schema, allowing the LLM to omit them while your `run` method receives the default value or `None`.

### Can I use nested Pydantic models for complex tool arguments?

Yes. Define separate `BaseModel` subclasses for nested structures and use them as field types in your `BaseTool` subclass. The generated JSON schema in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) will include these as nested objects, and Pydantic automatically instantiates the entire object hierarchy during validation in the execution flow.

### Where does tool validation occur in the source code?

Validation occurs in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) when the tool class constructor (inherited from Pydantic) receives the JSON payload from the LLM. The `BaseTool` instantiation parses and validates the input dictionary, raising `pydantic.ValidationError` if required fields are missing or types mismatch, which the agent catches before calling `run`.