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

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, 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 and exposes the following signature:

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, 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 and demonstrate practical extraction scenarios.

Extracting into a Pydantic Model

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

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

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, 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.
• 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.

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