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

> Implement structured knowledge extraction from documents with a two-stage pipeline. Discover schemas and extract JSON fact cards using tailored LLM prompts.

- Repository: [Bojie Li/ai-agent-book](https://github.com/bojieli/ai-agent-book)
- Tags: how-to-guide
- Published: 2026-07-26

---

**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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/data/schema.json) to avoid redundant API calls across experiments.

Key functions in [`chapter3/structured-knowledge-extraction/discovery.py`](https://github.com/bojieli/ai-agent-book/blob/main/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.

```python

# 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`](https://github.com/bojieli/ai-agent-book/blob/main/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.

```python

# 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:**

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

```

### Bulk Processing with Caching

For production datasets, `extract_dataset()` in [`chapter3/structured-knowledge-extraction/extractor.py`](https://github.com/bojieli/ai-agent-book/blob/main/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.

```python

# 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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/chapter3/structured-knowledge-extraction/config.py), the default model is `gpt-5.6-luna`, though the architecture supports OpenAI or OpenRouter backends.

```python

# 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.

```python

# 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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/discovery.py)** for bottom-up schema generation that separates **core** (universal) from **extensions** (type-specific) factors, caching results to [`data/schema.json`](https://github.com/bojieli/ai-agent-book/blob/main/data/schema.json).
- **Use [`extractor.py`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/schema.json), while **Structured Extraction** caches results to `extracted.jsonl`. According to the implementation in [`chapter3/structured-knowledge-extraction/extractor.py`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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`](https://github.com/bojieli/ai-agent-book/blob/main/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.