How to Implement Structured Knowledge Extraction from Documents: A Complete Two-Stage Pipeline

Structured knowledge extraction from documents requires a two-stage pipeline that first discovers a modular schema from corpus data, then extracts structured JSON fact cards using LLM prompts tailored to that schema.

This guide explains the implementation found in the bojieli/ai-agent-book repository, which demonstrates how to transform unstructured legal texts (or any document collection) into a queryable structured knowledge base. The approach separates schema generation from per-document processing, enabling efficient caching and clean integration with downstream retrieval systems.

Understanding the Two-Stage Extraction Architecture

The repository implements a decoupled pipeline that distinguishes between schema creation and document processing. This architecture, found in chapter3/structured-knowledge-extraction/, solves the cold-start problem of structured extraction by letting the LLM discover what factors matter before attempting to extract them.

• Stage 1 (Factor Discovery): Analyzes corpus batches to generate a reusable schema.json containing modular factor definitions.
• Stage 2 (Structured Extraction): Applies the cached schema to individual documents, outputting consistent JSON objects that downstream agents can consume.

Stage 1: Factor Discovery and Schema Generation

The discovery.py module implements bottom-up schema generation that avoids hard-coded ontologies. Instead of manually defining extraction fields, the system prompts an LLM to identify relevant factors from sample documents.

How Discovery Works

1. Batch Processing: Raw case facts are processed in batches (default size: 12) via the discover_schema() function.
2. Factor Listing: Each batch prompt asks the model to list every factor that could affect a legal sentence.
3. Schema Consolidation: A second LLM call consolidates raw factor lists across all batches into a modular schema containing:
   • core – Factors common to all document types (e.g., "self-surrender", "compensation").
   • extensions – Type-specific factors (e.g., "amount involved" for theft cases).
4. Disk Caching: The resulting schema is cached at data/schema.json to avoid redundant API calls across experiments.

Key functions in chapter3/structured-knowledge-extraction/discovery.py include discover_schema() for generation and load_schema() for retrieval, alongside factors_for_charge() and all_factors() for schema querying.

# Generate schema from corpus

from chapter3.structured_knowledge_extraction.discovery import discover_schema, load_schema

cases = [
    {"id": "c1", "charge": "盗窃罪", "fact": "..."},
    {"id": "c2", "charge": "故意伤害罪", "fact": "..."},
]

# Build and cache schema

schema = discover_schema(cases, batch_size=12, verbose=True)

# Inspect results

print(schema["core"][:3])                       # Universal factors

print(schema["extensions"]["盗窃罪"][:2])       # Theft-specific factors

Stage 2: Per-Document Structured Extraction

The extractor.py module handles the runtime extraction phase, converting raw text into structured JSON using the previously generated schema.

Extraction Workflow

1. Charge Detection: If not provided, extract_one() uses a lightweight LLM prompt to determine the document category (e.g., crime type).
2. Factor Retrieval: The system calls factors_for_charge() to retrieve only the relevant schema subset for that document type.
3. Structured Generation: A second LLM prompt outputs a JSON object where each factor key contains the extracted value or null when information is missing.
4. Normalization: The internal _normalize() method cleans raw JSON output, handling numeric conversion, boolean normalization, and null consistency to guarantee a predictable data shape.

# Extract from single document

from chapter3.structured_knowledge_extraction.extractor import extract_one
from chapter3.structured_knowledge_extraction.discovery import load_schema

schema = load_schema()

fact_text = """
被告人于2023年5月1日盗窃价值人民币12,000元的手机...
"""

structured = extract_one(fact_text, schema=schema)
print(structured)

Typical output:

{
  "charge": "盗窃罪",
  "amount": 12000,
  "self_surrender": null,
  "compensation": true
}

Bulk Processing with Caching

For production datasets, extract_dataset() in chapter3/structured-knowledge-extraction/extractor.py processes entire corpora with intelligent caching to extracted.jsonl. This dramatically reduces API usage after the first pass by reusing cached extractions.

# Bulk extraction

from chapter3.structured_knowledge_extraction.extractor import extract_dataset

schema = load_schema()
results = extract_dataset(schema, use_cache=True, verbose=True)

# Returns list of dicts: original fields + `extracted` JSON

Configuration and Client Management

The config.py module centralizes LLM configuration, reading API keys from environment variables and instantiating OpenAI clients with retry logic. As implemented in chapter3/structured-knowledge-extraction/config.py, the default model is gpt-5.6-luna, though the architecture supports OpenAI or OpenRouter backends.

# Configuration pattern from config.py

from chapter3.structured_knowledge_extraction.config import get_client

client = get_client()  # Handles retries, timeouts, and model selection

Integrating with Downstream Retrieval Systems

The extraction pipeline feeds directly into RAG architectures through schema utilities. The all_factors() function returns a flattened list of factor definitions suitable for vector indexing.

# Integration with GraphRAG

from chapter3.structured_knowledge_extraction.discovery import load_schema, all_factors
from chapter3.structured_index.graphrag_indexer import GraphRAGIndexer

schema = load_schema()
vectors = all_factors(schema)  # Structured entities for indexing

indexer = GraphRAGIndexer(vectors=vectors)

According to the bojieli/ai-agent-book source code, this integration enables the GraphRAGIndexer (found in chapter3/structured-index/graphrag_indexer.py) to reason over structured factor embeddings rather than raw text chunks, improving retrieval precision for agent queries.

Summary

• Use discovery.py for bottom-up schema generation that separates core (universal) from extensions (type-specific) factors, caching results to data/schema.json.
• Use extractor.py for runtime JSON extraction, leveraging extract_one() for single documents or extract_dataset() for bulk processing with extracted.jsonl caching.
• Reference config.py for centralized LLM client management with environment-based API key handling and default model gpt-5.6-luna.
• Leverage load_schema() and all_factors() to bridge extraction pipelines with downstream RAG or knowledge-graph indexers.

Frequently Asked Questions

How does the two-stage pipeline reduce API costs compared to single-pass extraction?

The Factor Discovery stage runs once per corpus to generate a reusable schema.json, while Structured Extraction caches results to extracted.jsonl. According to the implementation in chapter3/structured-knowledge-extraction/extractor.py, subsequent runs with use_cache=True skip LLM calls entirely for previously processed documents, eliminating redundant API usage across experiments.

What happens when a document contains factors not present in the cached schema?

The factors_for_charge() function retrieves only the factors defined for the detected document type. If the schema lacks coverage for specific edge cases, the extraction prompt in extractor.py instructs the model to return null for missing values rather than hallucinating new fields. To handle new factor types, you must re-run discover_schema() on an updated corpus batch.

Can this pipeline work with non-legal documents?

Yes. While the bojieli/ai-agent-book repository uses legal cases as the demonstration corpus, the schema generation logic in discovery.py is domain-agnostic. Simply provide a list of documents with charge (category) and fact (text) fields, and the pipeline will discover relevant factors for your specific document type.

How is the extracted JSON normalized for downstream consistency?

The extractor.py module includes a private _normalize() method that post-processes raw LLM output to ensure numeric fields are converted to proper numbers, boolean values are standardized, and missing fields are explicitly set to null. This guarantees that all entries in extracted.jsonl conform to the schema contract expected by downstream agents.

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