How to Use Pydantic Models for Structured Extraction with Needle

Needle automatically converts Pydantic BaseModel classes into OpenAI-compatible JSON schemas, enabling type-safe structured extraction from LLM outputs.

The Needle framework makes Pydantic models for structured extraction effortless. When you annotate a function parameter with a Pydantic model, Needle detects the type, generates a precise JSON schema, and exposes it to the LLM as a tool definition. This eliminates manual schema construction and guarantees that extracted data matches your Python types.

How Needle Detects Pydantic Models

Needle identifies Pydantic models through _is_pydantic_model in needle/agent/tools.py. This utility walks the class method resolution order (MRO) searching for a base class whose module starts with pydantic and whose name is BaseModel【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L43-L46】:

def _is_pydantic_model(cls) -> bool:
    # Simplified logic: checks MRO for pydantic.BaseModel

    for base in cls.__mro__:
        if base.__module__.startswith("pydantic") and base.__name__ == "BaseModel":
            return True
    return False

This detection works with Pydantic v1 and v2, ensuring broad compatibility without explicit version checks.

Converting Pydantic Models to LLM Schemas

Once detected, pydantic_schema transforms the model into an OpenAI-compatible parameter description【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L49-L60】:

• Properties: Extracted from the model's JSON schema via model_json_schema() (Pydantic v2) or model.schema() (Pydantic v1)
• Required fields: Pulled from the schema's "required" array
• Documentation: The model's docstring becomes the parameter description【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L56-L59】

The build_schema routine orchestrates this: whenever it encounters a Pydantic-annotated parameter, it delegates to pydantic_schema【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L21-L23】.

Complete Example: Weather Extraction Tool

Here's a practical Pydantic structured extraction workflow with Needle:

from pydantic import BaseModel
from needle import tool

# 1. Define your extraction schema as a Pydantic model

class WeatherQuery(BaseModel):
    """Parameters for fetching weather data."""
    city: str                 # required field

    units: str = "metric"     # optional with default

# 2. Create a Needle tool using the model

@tool
def get_weather(query: WeatherQuery) -> str:
    """Fetch and return weather information for a city."""
    # LLM populates `query` with validated WeatherQuery data

    temp = 22 if query.units == "metric" else 72
    unit_symbol = "°C" if query.units == "metric" else "°F"
    return f"The weather in {query.city} is {temp}{unit_symbol}."

# 3. Inspect the generated schema

import json
print(json.dumps(get_weather._needle_tool, indent=2))

Generated schema output:

{
  "name": "get_weather",
  "parameters": {
    "type": "object",
    "properties": {
      "query": {
        "type": "object",
        "properties": {
          "city": {"type": "string"},
          "units": {"type": "string", "default": "metric"}
        },
        "required": ["city"],
        "description": "Parameters for fetching weather data."
      }
    },
    "required": ["query"]
  },
  "description": "Fetch and return weather information for a city."
}

Notice how Needle preserves:

• Type information (string for both fields)
• Required vs. optional (city required, units optional)
• Default values ("metric" annotation visible)
• Documentation (model docstring propagated)

Decorating Functions with @tool

The @tool decorator in needle/agent/tools.py attaches the generated schema to your function via fn._needle_tool【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L62-L65】. This simple mechanism makes the schema available to Needle's LLM integration layer without modifying function behavior.

You can also use build_schema directly for advanced scenarios:

from needle import build_schema

def extract_person(name: str, details: PersonDetails) -> dict:
    """Extract person information."""
    return {"name": name, "details": details.model_dump()}

schema = build_schema(extract_person)

# Use schema manually or attach to custom tool implementations

Key Files and Implementation Details

File	Purpose
needle/agent/tools.py	Core implementation: _is_pydantic_model, pydantic_schema, build_schema, @tool decorator【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py】
needle/__init__.py	Public API exports (tool, build_schema, pydantic_schema)【https://github.com/cactus-compute/needle/blob/main/needle/__init__.py】
tests/test_tools.py	Validation suite for Pydantic detection and schema generation【https://github.com/cactus-compute/needle/blob/main/tests/test_tools.py】

Advanced: Nested Models and Complex Types

Needle's Pydantic structured extraction supports arbitrarily nested models. Define related schemas and compose them:

from pydantic import BaseModel
from typing import List
from needle import tool

class Address(BaseModel):
    """A physical address."""
    city: str
    country: str = "US"

class Person(BaseModel):
    """A person with multiple addresses."""
    name: str
    age: int | None = None
    addresses: List[Address]

@tool
def register_user(data: Person) -> str:
    """Register a new user in the system."""
    return f"Registered {data.name} with {len(data.addresses)} address(es)."

The generated schema recursively includes Address properties under data.addresses, giving the LLM a complete, typed contract for complex extractions.

Summary

• Automatic detection: _is_pydantic_model in needle/agent/tools.py recognizes Pydantic classes by MRO inspection
• Schema generation: pydantic_schema converts models to OpenAI-compatible JSON with properties, required fields, and documentation
• Seamless integration: @tool decorator attaches schemas via fn._needle_tool for LLM consumption
• Type safety: Runtime validation ensures extracted data matches your Pydantic definitions

Frequently Asked Questions

What Pydantic versions does Needle support?

Needle supports both Pydantic v1 and v2. The pydantic_schema function adapts automatically: it uses model_json_schema() for v2 models and falls back to model.schema() for v1【https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py#L51-L54】. No version-specific configuration is required.

Can I use Pydantic models for return types as well?

Yes, though the primary use case is parameter annotation for tool inputs. Needle's schema builder processes type annotations wherever they appear. For return types, the schema helps document expected outputs, but LLM tool calling typically focuses on input parameters.

How does Needle handle optional fields and defaults?

Optional fields with default values are marked accordingly in the generated schema. The "required" array in the JSON schema only includes fields without defaults, and default values are preserved in the property definitions. The LLM receives accurate guidance on which fields it must provide versus which are optional.

What happens if my Pydantic model uses complex validators?

Needle extracts the static JSON schema from your model, which includes field types, descriptions, and constraints defined via Field(). Custom validators run during Pydantic instantiation after the LLM returns data—Needle doesn't need to represent validators in the schema, only the structural contract.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}