# How Needle 2 Performs Structured Extraction: A Deep Dive into the `extract` API

> Discover how Needle 2 performs structured extraction using its extract API. Learn how Pydantic models and function calling enable precise data parsing.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-08-24

---

**Needle 2 performs structured extraction by temporarily treating a Pydantic model or dictionary schema as the sole available tool, forcing the underlying language model to return data via function calling that is then parsed into the desired structure.**

Structured extraction transforms unstructured text into type-safe objects without writing complex parsing logic. In Needle 2, this capability centers on the `extract` function defined in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), which provides a one-shot API that leverages the framework's tool-calling infrastructure to guarantee valid output formats.

## The Core Mechanism: Tool-Based Forcing

The `extract` function works by hijacking Needle’s tool-use pipeline. Instead of allowing the language model to choose from multiple functions, it registers your schema as the **only** available tool, compelling the model to structure its response as a function call matching that schema.

### Function Signature and Parameters

The implementation resides at lines 66–78 of [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py) and exposes the following signature:

```python
def extract(
    text: str,
    schema: type | dict,
    system: str | None = None,
    max_new_tokens: int = 256,
    weights: str | None = None
) -> object

```

- **`text`**: The raw input string to parse.
- **`schema`**: Either a Pydantic `BaseModel` class or a plain dictionary describing the desired structure.
- **`system`**: Optional system prompt to guide extraction behavior.
- **`weights`**: Specific model weights to use; defaults to `_active_weights` if omitted.

### Engine Re-initialization Strategy

When invoked, `extract` constructs a temporary `Needle` agent instance with `tools=[schema]`. This isolated agent reinitializes the underlying engine with your schema as the exclusive callable tool. According to the source code, this ensures the language model "sees" only your extraction target, eliminating hallucinated or irrelevant tool selections.

## Parsing and Type Handling

After the temporary agent generates a completion, the function inspects the response for `function_calls`. If none are present, it returns `None`. Otherwise, it extracts the arguments from the first function call and processes them based on the schema type.

### Pydantic Model Instantiation

When `schema` is a Pydantic model, Needle unpacks the function arguments using `schema(**arguments)`. This provides runtime validation and type coercion, ensuring the returned object conforms exactly to your model definition.

### Raw Dictionary Output

If you supply a plain dictionary description rather than a Pydantic class, the function returns the raw arguments dictionary directly. This offers flexibility for quick prototyping without defining formal model classes.

## Isolation from Existing Tools

A critical implementation detail is that the temporary agent creation does not pollute existing `Needle` instances. Because `extract` builds a separate internal agent, any tools previously registered to a user-defined instance remain untouched. This behavior is verified by the test `test_extract_keeps_agent_tools` in [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_inference.py), which confirms that after calling `extract`, the original agent retains full access to its own tool registry.

## Complete Usage Examples

The following patterns mirror the test suite in [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_inference.py) and demonstrate practical extraction scenarios.

### Extracting into a Pydantic Model

```python
import pydantic
from needle import extract

class Contact(pydantic.BaseModel):
    name: str
    email: str

result = extract(
    "John Doe can be reached at john@doe.com",
    Contact,
)
print(result)

# Output: Contact(name='John Doe', email='john@doe.com')

```

### Extracting into a Plain Dictionary

```python
schema = {"name": "str", "email": "str"}

result = extract(
    "Alice Smith, alice@example.org",
    schema,
)
print(result)

# Output: {'name': 'Alice Smith', 'email': 'alice@example.org'}

```

### Custom System Prompts and Weights

```python
result = extract(
    "Bob <bob@company.com>",
    Contact,
    system="You are an extraction assistant. Return only JSON.",
    weights="my-finetuned-needle-weights",
)
print(result)

# Output: Contact(name='Bob', email='bob@company.com')

```

## Summary

- **Single-tool forcing**: Needle 2 achieves structured extraction by temporarily registering your schema as the only available tool in [`needle/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/__init__.py), compelling function-call output.
- **Flexible schemas**: Accepts both Pydantic models for type-safe objects and raw dictionaries for simple key-value extraction.
- **Zero side effects**: The implementation creates isolated temporary agents, leaving existing tool configurations unaffected as proven by [`tests/test_inference.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_inference.py).
- **Unified engine**: Falls back to active weights (`_active_weights`) when no specific model is requested, maintaining consistency with the standard `Needle` API.

## Frequently Asked Questions

### What happens if the model fails to generate a function call?

If the language model does not produce any `function_calls` in its response, the `extract` function returns `None`. This typically occurs when the input text contains no data matching the schema or when the model confidence is low.

### Can I use structured extraction with custom fine-tuned models?

Yes. Pass your custom weights identifier to the `weights` parameter in `extract`. The function will use these specific weights instead of the default `_active_weights`, allowing extraction behavior tailored to your fine-tuned Needle engine.

### Does using `extract` modify my existing Needle agent's tools?

No. The function creates an internal temporary agent specifically for the extraction task. As demonstrated in `test_extract_keeps_agent_tools`, your original `Needle` instance retains all its previously registered tools and state after `extract` completes.

### Is Pydantic required for structured extraction?

No. While Pydantic models provide validation and IDE support, you can pass a plain dictionary describing the desired fields. The function will return a raw dictionary containing the extracted values, offering a lightweight alternative for simple use cases.