How to Use `extract()` for Structured Data Extraction with Pydantic Models in Needle

The extract() method in Needle is a one-shot helper that converts unstructured text into typed Python objects by internally creating a temporary agent, forcing a function-call response, and parsing the arguments into your Pydantic model.

Needle's extract() method eliminates the boilerplate of setting up an agent and managing tool loops when you simply need structured data from raw text. This guide shows you exactly how it works under the hood, with complete code examples for Pydantic models and plain schemas.

How extract() Works Internally

The extract() method follows a streamlined three-step pattern that hides complexity behind a simple API call.

From Public API to Core Logic

The Needle.extract() class method delegates directly to the module-level extract() function in needle/__init__.py:

# needle/__init__.py lines 141-143

@classmethod
def extract(cls, text: str, schema, system: Optional[str] = None):
    return extract(text, schema, system)

This forwarding pattern keeps the class interface consistent while centralizing the implementation.

Temporary Agent Creation

The core extract() function builds a single-purpose agent with your schema as the only available tool:

# needle/__init__.py lines 57-62

def extract(text: str, schema, system: Optional[str] = None):
    agent = Needle(
        tools=[schema],
        system=system,
    )

By registering only your target schema as a tool, the LLM has no alternative but to produce a function call matching that structure.

Response Parsing and Instantiation

After agent.complete() returns, extract() pulls the first function call's arguments and converts them to your desired type:

# needle/__init__.py lines 63-68

calls = result["function_calls"]
arguments = calls[0]["arguments"]

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

• Pydantic models: Instantiated with schema(**arguments) for full validation
• Plain dict schemas: Returned as-is as a dictionary

Using Pydantic Models for Type-Safe Extraction

Pydantic models are the recommended approach for production use. They provide validation, IDE autocomplete, and clear documentation of expected fields.

Defining Your Schema

Create a model with Field descriptions—these become part of the LLM's tool definition:

from pydantic import BaseModel, Field

class WeatherReport(BaseModel):
    """Current weather conditions."""
    temperature: float = Field(..., description="Temperature in Celsius")
    humidity: int = Field(..., ge=0, le=100, description="Relative humidity percent")
    condition: str = Field(..., description="Short weather description")

Running Extraction

Pass raw text and your model class to extract():

from needle import extract

text = """
The weather today is sunny with a temperature of 23.5°C.
Humidity sits at 45%.
"""

report = extract(text, WeatherReport)
print(report)

# → WeatherReport(temperature=23.5, humidity=45, condition='sunny')

print(report.model_dump())

# → {'temperature': 23.5, 'humidity': 45, 'condition': 'sunny'}

The returned object is a fully-instantiated Pydantic model with all validation constraints applied.

Pydantic Schema Conversion in Needle

Needle automatically converts Pydantic models to JSON-Schema tool definitions that LLMs can understand.

Model Detection

The _is_pydantic_model() helper in needle/agent/tools.py identifies Pydantic classes by inspecting the method resolution order:

# needle/agent/tools.py lines 43-46

def _is_pydantic_model(obj) -> bool:
    return (
        inspect.isclass(obj)
        and any(base.__name__ == "BaseModel" for base in obj.__mro__)
    )

JSON-Schema Generation

For validated Pydantic models, pydantic_schema() generates the tool definition:

# needle/agent/tools.py lines 49-60

def pydantic_schema(model: type[BaseModel]) -> CallableTool:
    return {
        "name": model.__name__,
        "description": model.__doc__ or "",
        "parameters": model.model_json_schema(),
    }

The model_json_schema() method produces OpenAPI-compatible JSON Schema, including:

• Type annotations
• Field descriptions
• Validation constraints (ge, le, regex, etc.)

Alternative: Plain Dict Schemas

When you don't need Pydantic's validation layer, you can pass a raw tool definition dictionary:

schema = {
    "name": "extract_contact",
    "parameters": {
        "type": "object",
        "properties": {
            "email": {"type": "string", "description": "User email address"},
            "phone": {"type": "string", "description": "User phone number"},
        },
        "required": ["email"]
    },
    "description": "Extract contact information from the text."
}

text = "Reach me at alice@example.com or call 555-1234."
contact = extract(text, schema)
print(contact)

# → {'email': 'alice@example.com', 'phone': '555-1234'}

The return value is a plain dictionary rather than a typed object.

Adding System Prompts for Guidance

Fine-tune extraction behavior with an optional system prompt:

system_prompt = (
    "You are a precise data extractor. "
    "Only include explicitly stated information. "
    "Never infer or hallucinate values."
)

report = extract(
    "Weather: 28 degrees, partly cloudy",
    WeatherReport,
    system=system_prompt
)

This system prompt becomes the system parameter when the temporary agent is constructed.

Complete Working Example

from pydantic import BaseModel, Field
from needle import extract

# 1. Define schema with clear descriptions

class ProductReview(BaseModel):
    """Extracted product review data."""
    product_name: str = Field(..., description="Name of the reviewed product")
    rating: int = Field(..., ge=1, le=5, description="Star rating from 1-5")
    pros: list[str] = Field(default=[], description="Positive aspects mentioned")
    cons: list[str] = Field(default=[], description="Negative aspects mentioned")

# 2. Raw text to parse

review_text = """
Just tried the AeroBrew Coffee Maker and I'm impressed! 
Easy setup and great taste earn it 5 stars. 
Only downside: a bit noisy in the morning.
"""

# 3. Extract structured data

review = extract(review_text, ProductReview)
assert review.product_name == "AeroBrew Coffee Maker"
assert review.rating == 5
assert "Easy setup" in review.pros
assert "noisy" in str(review.cons).lower()

Summary

• extract() provides one-shot structured data extraction without managing agent state
• Pydantic models are automatically detected via _is_pydantic_model() and converted to JSON-Schema via pydantic_schema() in needle/agent/tools.py
• Temporary agent creation with a single tool forces the LLM to output matching your schema
• Return type depends on input: validated Pydantic instances or plain dictionaries
• System prompts can guide extraction precision without changing the core workflow

Frequently Asked Questions

What happens if the LLM returns malformed data?

Needle relies on Pydantic's validation when using model schemas. If the LLM produces arguments that don't match the schema—wrong types, missing required fields, or out-of-range values—Pydantic raises a ValidationError with specific details about what failed. For plain dict schemas, you receive the raw arguments without validation.

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

No. The extract() method is intentionally designed for single-schema extraction. The implementation in needle/__init__.py passes your schema as the sole item in tools=[schema], forcing a single function call. For multi-step workflows or conditional tool selection, instantiate a full Needle agent with Needle(tools=[schema1, schema2, ...]) and use complete() directly.

How does this compare to using a standard Needle agent?

extract() is a convenience wrapper: it creates a temporary agent, runs one completion, and returns parsed output. Use it for simple, stateless extraction. A standard Needle agent maintains conversation history, supports multiple tool calls across turns, and allows custom handler logic. Choose direct agent instantiation when you need multi-turn dialogue, tool chaining, or persistent context.

Does extract() support async operations?

The current implementation in needle/__init__.py uses synchronous complete() calls. For async workflows, you would need to either wrap extract() in asyncio.to_thread() or manually replicate the pattern: create an agent with your schema and call await agent.complete_async() if available in your Needle version, then parse function_calls from the result.

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