How to Perform One-Shot Structured Extraction with Pydantic Models in Needle 2

Use the extract helper in Needle 2 to convert free-form text into type-safe Pydantic objects by treating your schema as the only available tool in a temporary agent.

Needle 2 provides a streamlined workflow for structured data extraction through its extract function. This helper eliminates boilerplate by automatically handling tool registration, prompt completion, and result parsing when working with Pydantic models or plain dictionaries. According to the cactus-compute/needle source code, the implementation leverages a temporary agent pattern that forces the LLM to use your schema as the sole callable function.

How the extract Function Works

The extract helper in needle/__init__.py orchestrates four distinct operations to deliver one-shot structured extraction:

1. Creating a Temporary Agent

The function instantiates a Needle agent with your Pydantic model provided as the exclusive tool. This constraint forces the language model to map the input text onto your schema structure.

agent = Needle(tools=[schema], system=system)

(Source: [needle/__init__.py line 49](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L49))

2. Running Completion

The helper sends your text prompt to the shared LLM engine with an expanded token budget to accommodate the structured output.

response = agent.complete(text, max_new_tokens)

(Source: [needle/__init__.py line 50](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L50))

3. Extracting Function Calls

Needle returns a JSON-compatible structure containing any generated function calls. The helper parses the first call's arguments from the response payload.

calls = response.get("function_calls") or []
arguments = calls[0].get("arguments") or {}

(Source: [needle/__init__.py lines 51-55](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L51-L55))

4. Instantiating the Model

If the supplied schema is a Pydantic model (verified via _is_pydantic_model), the function returns a populated instance; otherwise, it returns a plain dictionary.

return schema(**arguments) if _is_pydantic_model(schema) else arguments

(Source: [needle/__init__.py line 55](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L55))

Usage Examples

Basic Pydantic Model Extraction

Define your data structure using Pydantic BaseModel and pass it to extract along with the source text:

from pydantic import BaseModel, Field
from needle import extract

class WeatherReport(BaseModel):
    location: str = Field(..., description="City name")
    temperature_c: float = Field(..., description="Temperature in Celsius")
    condition: str = Field(..., description="Weather condition (e.g., sunny, rainy)")

prompt = """
The forecast for Paris today: 18.5°C and partly cloudy.
"""

report = extract(prompt, WeatherReport)

print(report)

# → WeatherReport(location='Paris', temperature_c=18.5, condition='partly cloudy')

Dictionary Schema Extraction

For lightweight scenarios, pass a type-mapping dictionary instead of a Pydantic class:

schema = {"location": str, "temperature_c": float, "condition": str}
result = extract(prompt, schema)
print(result)

# → {'location': 'Paris', 'temperature_c': 18.5, 'condition': 'partly cloudy'}

Custom System Messages

Override the default system prompt and adjust token limits for complex extraction tasks:

system_msg = "You are a helpful assistant that extracts weather data."
report = extract(prompt, WeatherReport, system=system_msg, max_new_tokens=128)

Internal Implementation Details

The extraction pipeline relies on auxiliary utilities to handle Pydantic serialization. The _jsonable function (lines 37-42 in needle/__init__.py) normalizes model outputs across Pydantic v1 and v2 APIs by calling model_dump (v2) or dict (v1):

def _jsonable(obj):
    if hasattr(obj, "model_dump"):  # Pydantic v2

        return obj.model_dump()
    if hasattr(obj, "dict"):      # Pydantic v1

        return obj.dict()
    return obj

(Source: [needle/__init__.py lines 37-42](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py#L37-L42))

The tool-call infrastructure that enables this functionality resides in needle/agent/tools.py, while the test suite in tests/test_inference.py validates the end-to-end extraction workflow against actual LLM responses.

Summary

• One-shot extraction in Needle 2 uses the extract helper to transform text into structured data via temporary agents.
• Schema flexibility allows both Pydantic models and plain dictionaries as extraction targets.
• Automatic instantiation returns fully-typed objects when using Pydantic, or raw dictionaries for type mappings.
• Source locations include needle/__init__.py (core logic), needle/agent/tools.py (tool infrastructure), and tests/test_inference.py (validation).

Frequently Asked Questions

Can I use Needle 2 extraction without Pydantic?

Yes. While Pydantic models provide type safety and validation, the extract function accepts plain dictionaries mapping field names to Python types. When using dictionaries, the function returns a standard Python dict instead of a model instance.

How does Needle handle Pydantic v1 versus v2 compatibility?

The library includes an internal _jsonable utility that detects the available API by checking for model_dump (v2) or dict (v1) methods. This ensures consistent JSON serialization regardless of which Pydantic version is installed in your environment.

What happens if the LLM returns multiple function calls?

The extraction helper specifically selects the first function call from the response array using calls[0].get("arguments"). For one-shot extraction workflows, design your schema to capture all required data in a single structured output rather than relying on multiple calls.

Where is the extraction logic tested?

The one-shot structured extraction functionality is verified in tests/test_inference.py, which contains test cases confirming that the extract helper correctly parses LLM outputs into the expected Pydantic models and handles edge cases like missing fields or invalid JSON.

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