How to Use `needle.extract()` for One-Shot Structured Extraction in Python

needle.extract() is a convenience function that performs single-call structured data extraction from text using a Pydantic model or dictionary schema, without requiring a full agent-tool loop.

The needle library provides a lightweight way to extract structured information from unstructured text using local LLMs. The extract() function in needle/__init__.py offers the fastest path from raw text to typed data—ideal for one-shot extraction tasks where you need exactly one structured result.

What needle.extract() Does Under the Hood

The function orchestrates four steps to hide boilerplate from you:

1. Creates a transient Needle agent with your schema as the sole tool
2. Executes a completion against the compiled engine
3. Extracts the first function call from the LLM response
4. Returns a typed result (Pydantic instance or dictionary)

According to the cactus-compute/needle source code, the implementation spans lines 49-79 of needle/__init__.py:

def extract(text, schema, *, system=None, weights=None, max_new_tokens=256):
    # lines 49-53: Build single-tool agent

    agent = Needle(tools=[schema], system=system, weights=weights or _active_weights)
    
    # lines 73-74: Run completion

    response = agent.complete(text, max_new_tokens)
    
    # lines 74-78: Parse first function call

    calls = response.get("function_calls") or []
    if not calls:
        return None
    arguments = calls[0].get("arguments") or {}
    
    # lines 78-79: Return typed or raw result

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

This design means you never manually instantiate Needle, register tools, or parse function_calls yourself.

Method 1: Extract with a Pydantic Model (Type-Safe)

For production code, pass a Pydantic BaseModel as your schema. needle.extract() validates and instantiates the model automatically.

from pydantic import BaseModel
import needle

class Order(BaseModel):
    product: str
    quantity: int
    price: float

text = "I would like to buy 3 notebooks for $7.50 each."

order = needle.extract(text, Order)
print(order)

# → Order(product='notebooks', quantity=3, price=7.5)

The function detects Pydantic models via _is_pydantic_model(schema) and calls schema(**arguments) to construct your typed object.

Method 2: Extract with a Dictionary Schema (Quick & Ad-Hoc)

For prototyping or dynamic schemas, use a plain dictionary with type hints as strings.

import needle

schema = {
    "name": "str",
    "age": "int",
    "city": "str"
}

bio = "Alice is 29 years old and lives in Berlin."
result = needle.extract(bio, schema)
print(result)

# → {'name': 'Alice', 'age': 29, 'city': 'Berlin'}

Dictionary schemas skip Pydantic validation and return raw dictionaries, making them ideal for exploratory extraction.

Method 3: Customize System Prompt and Token Budget

Control extraction behavior with optional keyword arguments:

Parameter	Purpose	Default
system	Custom system prompt for the LLM	None
weights	Specific model weights to load	_active_weights (global)
max_new_tokens	Hard limit on completion length	256

import needle

custom_system = "You are an expert data parser. Extract the fields exactly as requested."
text = "Invoice #12345: total $250, due 2024-09-15."

invoice = needle.extract(
    text,
    {"invoice_id": "str", "total": "float", "due_date": "str"},
    system=custom_system,
    max_new_tokens=128
)
print(invoice)

# → {'invoice_id': '12345', 'total': 250.0, 'due_date': '2024-09-15'}

The system prompt shapes how the LLM interprets your extraction task—use it to enforce format constraints or domain expertise.

Weight Handling and Performance Notes

When you omit weights, needle.extract() reuses the globally active weights (_active_weights). This matters for performance: the underlying compiled engine cannot unload weights, so consecutive calls with different schemas re-initialize the engine but keep the model resident in memory.

From needle/model/architecture.py, this weight management ensures you don't pay full load costs on repeated extractions.

For single extractions with a fresh model, pass explicit weights to trigger a new load.

Summary

• needle.extract() wraps agent creation, tool binding, and call parsing into one function
• Accepts Pydantic models (typed results) or dictionaries (raw results) as schemas
• Returns None when the LLM produces no function_calls
• Reuses global weights by default; customize via the weights parameter
• Implemented in needle/__init__.py lines 49-79, with completion logic in needle/model/run.py

Frequently Asked Questions

What happens if the LLM doesn't return any function calls?

needle.extract() returns None. The source checks response.get("function_calls") or [] and exits early when the list is empty (lines 74-76 of needle/__init__.py). This occurs when the model fails to generate a structured extraction for your input.

Can I use needle.extract() with multiple extraction schemas at once?

No. The function is designed for one-shot extraction with a single schema. Internally it constructs a Needle agent with tools=[schema]—only one tool. For multi-schema extraction, instantiate Needle directly and register multiple tools.

How does needle.extract() handle type conversion?

Type conversion depends on your schema format. With Pydantic models, Pydantic's own validation coerces values (e.g., "7.50" → 7.5). With dictionary schemas, the underlying LLM engine performs extraction and the raw arguments are returned—types match what the engine produces, typically JSON-compatible primitives.

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