Using Pydantic Models with Needle for Type-Safe Extraction

Needle automatically converts Pydantic models into OpenAI-compatible function schemas, enabling strict type validation and structured data extraction from LLM responses without manual JSON Schema authoring.

The cactus-compute/needle library streamlines LLM function calling by generating OpenAI-compatible tool schemas directly from Python type hints. When you use Pydantic models with Needle for type-safe extraction, the framework inspects your model definitions, builds JSON Schemas automatically, and validates incoming LLM payloads against your Python types at runtime.

How Needle Detects Pydantic Models

During schema construction in needle/agent/tools.py, the build_schema function inspects each parameter's type annotation. If the annotation satisfies the internal _is_pydantic_model check, Needle recognizes the argument as a structured Pydantic object rather than a primitive type.

The detection logic (lines 143–147) inspects the type hint to determine if it inherits from Pydantic's BaseModel:

if _is_pydantic_model(annotation):
    return pydantic_schema(annotation)["parameters"]

This check triggers the specialized conversion pipeline, ensuring that your model's fields are properly exposed to the LLM as structured inputs.

Automatic Schema Generation

Once a Pydantic model is detected, the pydantic_schema function (lines 149–159) extracts the model's JSON schema using model_json_schema() (with a legacy schema() fallback for older Pydantic versions). It reshapes this output to conform to the OpenAI tool specification, propagating:

• Properties and required fields from your model definition
• Nested definitions for complex types
• The model's docstring as the schema description

Because the schema derives directly from the model, any change to your Pydantic fields—new constraints, type updates, or validation rules—is instantly reflected in the tool contract. This eliminates mismatches between the LLM's expected input format and your Python implementation.

Practical Implementation Examples

Define your data structure as a Pydantic model and expose it via the @tool decorator. Needle stores the generated schema on the function's _needle_tool attribute.

from needle import tool, Field
import pydantic

# Define the extraction target

class Weather(pydantic.BaseModel):
    """Current weather conditions."""
    city: str = Field(description="Name of the city")
    temperature_c: float = Field(description="Temperature in Celsius")
    cloudy: bool = Field(default=False, description="Is it cloudy?")

# Expose to the LLM

@tool
def get_weather() -> Weather:
    """Ask the model to provide weather information."""
    # Body is never executed; Needle uses the signature for schema generation

    ...

When the LLM calls get_weather, Needle validates the returned JSON against the Weather model before your code executes.

Combining Primitives and Models

You can mix standard arguments with Pydantic models for partial structure:

@tool
def schedule_meeting(topic: str, when: pydantic.BaseModel) -> str:
    """Schedule a meeting with a structured time description."""
    ...

In this case, when is processed by _is_pydantic_model and the generated schema embeds the nested model's fields alongside the primitive topic string.

Runtime Validation and Type Safety

When the LLM invokes a tool, Needle uses the stored schema to validate the payload against your Pydantic model before passing control to your function. This guarantees that:

1. All required fields are present
2. Types match your annotations (e.g., temperature_c must be a float)
3. Validation rules defined in your model (ranges, regex patterns, etc.) are enforced

If validation fails, the error is raised before your business logic runs, preventing type errors downstream. The test suites in tests/test_tools.py and tests/test_inference.py demonstrate this end-to-end validation behavior.

Summary

• Automatic detection: The _is_pydantic_model function in needle/agent/tools.py identifies Pydantic types during schema construction.
• Zero-config schemas: The pydantic_schema function converts models to OpenAI-compatible JSON Schema using model_json_schema().
• Runtime guarantees: Needle validates LLM outputs against your models before execution, ensuring type safety.
• Immediate synchronization: Changes to Pydantic fields instantly update the tool contract without manual schema editing.

Frequently Asked Questions

What versions of Pydantic does Needle support?

Needle supports both Pydantic v1 and v2. The pydantic_schema implementation attempts model_json_schema() first (v2) and falls back to schema() for v1 compatibility, ensuring broad version coverage.

Can I use nested Pydantic models with Needle?

Yes. Needle recursively processes nested Pydantic models through pydantic_schema, embedding the full JSON Schema definition in the definitions or $defs section of the tool specification. The LLM receives the complete structure required to populate nested objects correctly.

How does Needle handle Pydantic validation errors?

If the LLM returns data that violates your model constraints (missing fields, wrong types, or failed validators), Needle raises the standard Pydantic validation error before your tool function executes. This prevents invalid data from reaching your application logic and allows you to handle extraction failures explicitly.

Where are the Pydantic integration tests located?

The primary validation tests reside in tests/test_tools.py, which verifies that _is_pydantic_model correctly identifies models and that pydantic_schema generates expected OpenAI-compatible output. End-to-end inference tests using Pydantic-based tools are available in tests/test_inference.py.

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