How to Use Pydantic Models as Tools in Needle: A Complete Guide

Yes, Needle natively supports Pydantic BaseModel classes as first-class tools, automatically converting them to JSON Schema through built-in detection logic in needle/agent/tools.py.

Needle's tool system lets you expose Python functions to LLMs with structured inputs. When you use Pydantic models as type hints, Needle detects them automatically and generates proper JSON Schema definitions—no manual schema writing required. This integration lives in the cactus-compute/needle repository and eliminates boilerplate when building complex tool interfaces.

How Needle Detects Pydantic Models

The detection mechanism lives in needle/agent/tools.py. The _is_pydantic_model function (lines 49-53) inspects a class's method resolution order (MRO) to identify Pydantic BaseModel subclasses:

def _is_pydantic_model(annotation) -> bool:
    # Checks if annotation's MRO contains pydantic.BaseModel

    for base in getattr(annotation, "__mro__", []):
        if (
            getattr(base, "__module__", "").startswith("pydantic")
            and base.__name__ == "BaseModel"
        ):
            return True
    return False

This detection triggers whenever Needle processes function type hints during schema building.

Converting Pydantic Models to JSON Schema

Once detected, the pydantic_schema function (lines 55-66) handles conversion. It calls Pydantic's native model_json_schema() (or schema() for v1), then reshapes the output for Needle's tool system:

def pydantic_schema(model) -> dict:
    """Convert a Pydantic model to Needle's JSON-Schema format."""
    # Get native Pydantic schema

    if hasattr(model, "model_json_schema"):  # Pydantic v2

        raw_schema = model.model_json_schema()
    else:  # Pydantic v1 fallback

        raw_schema = model.schema()
    
    # Reshape for tool compatibility

    return {
        "name": model.__name__,
        "description": model.__doc__ or "",
        "parameters": raw_schema
    }

The build_schema function (lines 70-72) then injects this schema when encountering Pydantic-typed parameters, using pydantic_schema(annotation)["parameters"] as the parameter definition.

Practical Example: Weather Tool with Pydantic

Here's a complete, runnable example showing Pydantic model tools in action:

import pydantic
from needle import tool

class WeatherQuery(pydantic.BaseModel):
    """Parameters for a weather lookup."""
    city: str               # Required field

    units: str = "metric"   # Optional with default

@tool
def get_weather(request: WeatherQuery) -> str:
    """Fetch current weather conditions for a specified city."""
    # Implementation would call weather API here

    return f"Weather in {request.city}: 72°F, sunny"

# Inspect the generated schema

print(get_weather._needle_tool)

The resulting schema structure:

{
    "name": "get_weather",
    "description": "Fetch current weather conditions for a specified city.",
    "parameters": {
        "type": "object",
        "properties": {
            "request": {
                "type": "object",
                "properties": {
                    "city": {"type": "string"},
                    "units": {"type": "string", "default": "metric"}
                },
                "required": ["city"]
            }
        },
        "required": ["request"]
    }
}

Needle automatically nests the Pydantic model's fields under the request parameter, preserving all validation rules and defaults.

Direct Schema Generation Without Decorators

You can also generate schemas programmatically using the public API:

from needle.agent.tools import pydantic_schema

schema = pydantic_schema(WeatherQuery)

# Returns: {"name": "WeatherQuery", "description": "...", "parameters": {...}}

This is exported through needle/__init__.py alongside tool and Field for convenient access.

Test Coverage and Validation

The implementation is validated in tests/test_tools.py (lines 18-31) via test_pydantic_model_schema. This test confirms:

• Correct Pydantic model detection via _is_pydantic_model
• Accurate JSON-Schema conversion through pydantic_schema
• Proper integration into final tool schemas

Run the test suite to verify behavior: pytest tests/test_tools.py::test_pydantic_model_schema -v

Summary

• Pydantic models work as tools in Needle through automatic detection and conversion
• _is_pydantic_model in needle/agent/tools.py identifies Pydantic classes via MRO inspection
• pydantic_schema converts models to JSON-Schema using native Pydantic methods
• build_schema injects model schemas as parameter definitions during tool registration
• Tests in tests/test_tools.py validate the complete pipeline

Frequently Asked Questions

When should I use Pydantic models versus plain type hints in Needle tools?

Use Pydantic models when you need complex nested structures, field-level validation, or self-documenting schemas. For simple tools with 1-2 primitive parameters, plain type hints like def tool(x: int, y: str) suffice. Models shine when your LLM needs to supply structured data with constraints.

Does Needle support both Pydantic v1 and v2?

Yes. The pydantic_schema function checks for model_json_schema() (v2) first, falling back to schema() for v1 compatibility. Both versions generate valid tool schemas.

Can I nest Pydantic models inside other Pydantic models for tools?

Absolutely. Since pydantic_schema calls Pydantic's native JSON-Schema generator recursively, nested models automatically produce correct $ref structures. The LLM receives the complete, flattened parameter schema.

What happens if my Pydantic model has complex validators?

Pydantic's validators run before Needle sees the data—they execute during model instantiation in your tool function. The JSON-Schema exposure only describes the input shape; validation occurs at runtime when the decorated function receives the model instance.

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