# How to Define Tools Using Pydantic Models in Needle 2

> Learn how to define tools using Pydantic models in Needle 2. Automatically convert Pydantic models to JSON Schema for structured LLM interactions and validation.

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

---

**Needle 2 automatically converts Pydantic models into JSON Schema-compatible tool descriptions when you annotate function arguments or return types, enabling structured LLM interactions with automatic validation.**

When building LLM applications with the `cactus-compute/needle` repository, you can define tools using Pydantic models to ensure type-safe inputs and outputs. This approach allows Needle to generate precise JSON schemas that help the LLM understand exactly what parameters your functions require, including field constraints and validation rules.

## Understanding Pydantic Model Detection

Needle detects Pydantic models through a specialized utility function that inspects type annotations at runtime. In [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), the `_is_pydantic_model` function checks whether a given type inherits from `pydantic.BaseModel` (lines 54-58). This detection triggers the automatic schema generation pipeline, distinguishing Pydantic models from standard Python types that use basic JSON-type inference.

The detection logic ensures that any function argument or return type annotated with a Pydantic class receives special handling. When Needle encounters these models, it bypasses the default type resolution and instead extracts the rich schema definitions that Pydantic provides natively.

## Generating JSON Schema for Tool Definitions

Once Needle identifies a Pydantic model, it generates a compliant JSON Schema using the `pydantic_schema` function defined in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) (lines 60-70). This function calls `model.model_json_schema()` for Pydantic v2 compatibility, falling back to `model.schema()` for older Pydantic versions, then constructs a tool schema containing the model name, field definitions, required properties, and optional descriptions.

The schema injection occurs within `build_schema` (lines 70-73), where Needle replaces the usual JSON-type inference with the Pydantic-generated schema (`pydantic_schema(annotation)["parameters"]`). This seamless substitution ensures that complex validation constraints—such as `gt`, `lt`, `enum`, or regex patterns defined in your Pydantic fields—are automatically reflected in the tool description sent to the LLM.

## Registering Functions with the `@tool` Decorator

To expose a function as a tool, Needle provides the `@tool` decorator, which stores the generated schema directly on the callable. According to the source code in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) (lines 73-76), the decorator attaches the schema dictionary to the function as a `_needle_tool` attribute. This metadata allows the engine to serialize and deserialize Pydantic models correctly when the LLM invokes the tool.

You can also pass functions directly to the `Needle` constructor via the `tools` parameter, as [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) handles the construction of the tool list and ensures all schema metadata is properly registered with the engine.

## Practical Implementation: Weather Tool Example

Here is a complete implementation demonstrating how to define a tool using Pydantic models in Needle 2:

```python

# 1️⃣ Define a Pydantic model that describes the tool's input

import pydantic

class WeatherQuery(pydantic.BaseModel):
    """Ask for weather in a city."""
    city: str                     # required string

    units: str = "metric"         # optional with default

# 2️⃣ Create a function that uses the model as its argument

from needle import tool

@tool                     # registers the function as a Needle tool

def get_weather(query: WeatherQuery) -> str:
    """Return a fake weather report."""
    return f"The weather in {query.city} is sunny ({query.units})"

# 3️⃣ Initialise Needle with the tool (the decorator already attached the schema)

from needle import Needle
needle = Needle(tools=[get_weather])

# 4️⃣ Use the tool via the engine (the LLM will receive the JSON schema)

result = needle.run("What's the weather in Paris?")
print(result)   # → "The weather in Paris is sunny (metric)"

```

The resulting JSON schema that the LLM receives looks like this:

```json
{
  "name": "WeatherQuery",
  "parameters": {
    "type": "object",
    "properties": {
      "city": { "type": "string" },
      "units": { "type": "string", "default": "metric" }
    },
    "required": ["city"]
  },
  "description": "Ask for weather in a city."
}

```

Because the schema is derived directly from the Pydantic model, any field constraints defined in your model automatically appear in the tool description. This enables the LLM to validate inputs before calling the function, reducing runtime errors and improving response quality.

## Summary

- **Automatic Detection**: Needle uses `_is_pydantic_model` in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) to identify Pydantic BaseModel subclasses at runtime.
- **Schema Generation**: The `pydantic_schema` function converts models to JSON Schema using `model_json_schema()`, capturing all field constraints and validation rules.
- **Tool Registration**: The `@tool` decorator stores the generated schema as `_needle_tool` on the function, making it available to the LLM engine.
- **Constraint Preservation**: Field validators like `gt`, `lt`, `enum`, and regex patterns defined in Pydantic models automatically propagate to the LLM's tool description.

## Frequently Asked Questions

### Does Needle support both Pydantic v1 and v2?

Yes. According to the source code in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), the `pydantic_schema` function attempts to call `model.model_json_schema()` first (Pydantic v2), and falls back to `model.schema()` for older Pydantic v1 installations. This dual compatibility ensures your tools work regardless of which Pydantic version is installed in your environment.

### How does Needle handle Pydantic field constraints?

Needle preserves all field constraints automatically. When `pydantic_schema` generates the JSON Schema in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) (lines 60-70), it captures constraints such as `gt`, `lt`, `ge`, `le`, `enum`, and pattern validation directly from the Pydantic model definition. These constraints appear in the parameters object sent to the LLM, allowing the model to validate inputs before attempting to call your function.

### Can I use nested Pydantic models as tool arguments?

Yes. Because Needle relies on Pydantic's native `model_json_schema()` method, nested models are fully supported. When you define a Pydantic model that contains other Pydantic models as fields, the schema generation recursively processes the entire structure, producing a complete JSON Schema with properly nested object definitions that the LLM can understand and populate correctly.

### Where is the tool schema stored after using the `@tool` decorator?

The schema is stored as a `_needle_tool` attribute on the decorated function. As implemented in [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) (lines 73-76), the `@tool` decorator attaches a dictionary containing the generated schema to the callable. When you initialize a `Needle` instance with `tools=[your_function]`, the engine reads this attribute to understand how to serialize arguments and deserialize return values during LLM interactions.