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

> Learn one-shot structured extraction with Pydantic in Needle 2. Convert free-form text into type-safe Pydantic objects using the extract helper and temporary agents.

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

---

**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`](https://github.com/cactus-compute/needle/blob/main/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.

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

```

*(Source: [[`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/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.

```python
response = agent.complete(text, max_new_tokens)

```

*(Source: [[`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/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.

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

```

*(Source: [[`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/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.

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

```

*(Source: [[`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/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:

```python
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:

```python
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:

```python
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`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py)) normalizes model outputs across Pydantic v1 and v2 APIs by calling `model_dump` (v2) or `dict` (v1):

```python
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`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py), while the test suite in [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) (core logic), [`needle/agent/tools.py`](https://github.com/cactus-compute/needle/blob/main/needle/agent/tools.py) (tool infrastructure), and [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/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`](https://github.com/cactus-compute/needle/blob/main/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.