How to Perform Structured Extraction with Needle 2: A Complete Guide

Needle 2 provides a one-shot extract function that treats a Pydantic model or dictionary schema as the sole available tool, enabling direct structured data extraction from plain text without managing a full agent loop.

The cactus-compute/needle repository offers a streamlined approach to structured data extraction through its top-level API. The extract function defined in needle/__init__.py handles the complexity of tool registration and language model completion internally, allowing developers to convert unstructured text into typed objects using minimal code. This implementation leverages the same underlying engine as the full Needle agent class while isolating extraction operations to prevent side effects on existing agent configurations.

The Core extract Function Signature

The extract function is implemented at lines 66-78 of needle/__init__.py with the following signature:

extract(
    text: str,
    schema: type | dict,
    system: str | None = None,
    max_new_tokens: int = 256,
    weights: str | None = None
) → object

This function accepts raw text alongside a schema definition—either a Pydantic model class or a plain dictionary describing the desired structure. When invoked, the function initializes a temporary Needle agent internally, registering the schema as the only available tool via the tools=[schema] parameter. This forces the underlying language model to treat the schema as the sole callable function, ensuring the completion output conforms to the specified structure.

Extracting with Pydantic Models

For type-safe structured extraction with Needle 2, pass a Pydantic BaseModel subclass as the schema:

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)  # → Contact(name='John Doe', email='john@doe.com')

When the schema is a Pydantic model, the function automatically instantiates the class with the extracted arguments (schema(**arguments)), returning a fully typed object rather than a raw dictionary.

Dictionary-Based Schema Extraction

If you prefer schema flexibility without defining Pydantic classes, pass a dictionary describing the field types:

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

result = extract(
    "Alice Smith, alice@example.org",
    schema,
)
print(result)  # → {'name': 'Alice Smith', 'email': 'alice@example.org'}

In this mode, extract returns the raw dictionary of extracted arguments directly, bypassing Pydantic instantiation while maintaining the same validation guarantees from the underlying model.

Advanced Configuration Options

The extract function exposes parameters to customize the extraction behavior without modifying global state:

• system: Inject a custom system prompt to guide the extraction logic
• max_new_tokens: Control the generation length (defaults to 256)
• weights: Specify alternative model weights for the operation

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

If no weights argument is provided, the function falls back to _active_weights, ensuring consistency with the currently loaded model.

Implementation Architecture

Under the hood, the extraction mechanism operates by temporarily reconfiguring the language model engine. According to the source code in needle/__init__.py, the function constructs an isolated Needle agent instance that exists solely for the duration of the extraction call.

Tool Preservation: Because this agent is created separately from any user-defined instances, existing tool configurations remain unaffected. The test suite verifies this behavior in tests/test_inference.py through the test_extract_keeps_agent_tools test (lines 72-90), which confirms that agents retain their original tool sets after extract invocations.

Result Handling: After triggering completion, the function inspects the response for function_calls. If no function calls are present, it returns None. Otherwise, it extracts the first call's arguments and either builds a Pydantic instance or returns the raw dict, depending on whether schema is a class or dictionary.

Engine Efficiency: The implementation reuses the shared engine that powers the standard Needle API, avoiding the overhead of loading separate model instances while maintaining thread safety through temporary agent isolation.

Summary

• Location: The extract function lives in needle/__init__.py at lines 66-78
• Schemas: Supports both Pydantic models and plain dictionaries for structure definition
• Isolation: Creates temporary agents to prevent side effects on existing tool configurations
• Fallback: Uses _active_weights when no explicit weights parameter is specified
• Return: Returns typed Pydantic instances, raw dictionaries, or None based on extraction success

Frequently Asked Questions

What is the exact function signature of needle.extract?

The function signature is extract(text: str, schema: type | dict, system: str | None = None, max_new_tokens: int = 256, weights: str | None = None) → object. It is defined at lines 66-78 of needle/__init__.py in the cactus-compute/needle repository.

Does calling extract affect existing Needle agent tools?

No. The function creates a temporary, isolated Needle agent instance for each extraction operation. As verified by the test_extract_keeps_agent_tools test in tests/test_inference.py, existing agent tool configurations remain intact after calling extract.

Can I use a dictionary instead of a Pydantic model for the schema?

Yes. The schema parameter accepts either a Pydantic model class or a dictionary describing field types. When using a dictionary, extract returns the raw extracted arguments as a Python dict rather than instantiating a class.

How does extract handle cases where no structured data is found?

If the language model completion contains no function_calls, the function returns None. This occurs when the input text does not contain data matching the provided schema or when the model fails to identify extractable entities.

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