wizit-open-rag 0.0.2

AI-powered document transcription and semantic chunking for RAG pipelines

wizit_open_rag

A Python library for AI-powered document transcription and semantic chunking with RAG (Retrieval-Augmented Generation). It processes PDFs through a cost-aware tiered pipeline — plain-text extraction first, OCR second, LLM last — then chunks the resulting Markdown semantically, enriches each chunk with surrounding context, and returns ready-to-index Document objects for PostgreSQL pgvector.

Version: 0.0.1 | Python: >=3.12 | Build: uv

---

Features

• Cost-aware tiered transcription: pdfplumber (free) → OCR (AWS Textract or Mistral Document AI) → Claude Haiku (LLM fallback). Each page escalates only when the previous tier scores below the quality threshold.
• LangGraph-based transcription workflow with configurable retry logic and accuracy thresholds
• Semantic chunking with 85th-percentile breakpoints (plus recursive and Markdown-header strategies)
• Per-chunk context enrichment — each chunk is wrapped with <context> and <content> tags
• Pluggable storage backends: local filesystem or AWS S3
• Vector indexing into PostgreSQL pgvector via LangChain PGVectorStore
• LangSmith tracing support

---

Prerequisites

• Python 3.12 or higher
• uv for dependency management
• AWS credentials configured (standard boto3 credential chain — env vars, ~/.aws/credentials, or instance profile)
• PostgreSQL with the pgvector extension enabled

---

Installation

pip install wizit_open_rag

For development (clone + install with dev tools):

git clone https://github.com/Restebance/open_rag.git
cd open_rag
uv sync --group dev
cp example.env .env

---

Usage

Document Transcription — default (LLM every page)

OpenRagTranscriber accepts the raw bytes of a single PDF page and returns a ParsedDocPage with the Markdown result.

import asyncio
import fitz  # PyMuPDF
from wizit_open_rag import OpenRagTranscriber

transcriber = OpenRagTranscriber(
    langsmith_project_name="my-project",   # required
    langsmith_api_key="lsv2_...",          # required
    llm_model_id="global.anthropic.claude-sonnet-4-6",
    target_language="es-CO",
)

# Extract a single-page PDF from a multi-page document
with fitz.open("document.pdf") as doc:
    single = fitz.open()
    single.insert_pdf(doc, from_page=0, to_page=0)
    page_bytes = single.tobytes()

result = asyncio.run(transcriber.transcribe_document(page_number=1, page_content=page_bytes))
print(result.page_text)  # Markdown string

Document Transcription — cost-aware tiered pipeline

Enable the tiered pipeline with use_tiered_transcription=True. Each page is processed by the cheapest method that meets the quality threshold (transcription_accuracy_threshold).

Page bytes
  ↓
Tier 1: pdfplumber          (free — digital text + tables)
  ↓ score < threshold
Tier 2: AWS Textract         (OCR — tables, forms, key-value pairs)
      OR Mistral Document AI  (swappable via tier2_ocr)
  ↓ score < threshold
Tier 3: Claude Haiku         (LLM fallback — always returns a result)

With AWS Textract (default):

transcriber = OpenRagTranscriber(
    langsmith_project_name="my-project",
    langsmith_api_key="lsv2_...",
    use_tiered_transcription=True,
    tier2_ocr="textract",                                      # default
    tier3_model_id="global.anthropic.claude-haiku-4-5-20251001",
    transcription_accuracy_threshold=0.90,
)

AWS credentials must be available in the environment (boto3 credential chain). No extra API key needed.

With Mistral Document AI:

transcriber = OpenRagTranscriber(
    langsmith_project_name="my-project",
    langsmith_api_key="lsv2_...",
    use_tiered_transcription=True,
    tier2_ocr="mistral",
    mistral_api_key="...",   # or set MISTRAL_API_KEY env var
    tier3_model_id="global.anthropic.claude-haiku-4-5-20251001",
)

Both options share the same transcribe_document call — the tier selection only affects what happens internally.

Semantic Chunking with Context

ChunksManager takes a pre-loaded Markdown string and returns a list of LangChain Document objects, each enriched with a contextual summary.

import asyncio
from wizit_open_rag import ChunksManager

manager = ChunksManager(
    langsmith_project_name="my-project",   # required
    langsmith_api_key="lsv2_...",          # required
)

with open("document.md") as f:
    markdown_content = f.read()

docs = asyncio.run(manager.gen_context_chunks(
    file_key="document.md",
    file_markdown_content=markdown_content,
    file_tags={"category": "hr", "department": "onboarding"},
))

# docs is a List[Document]; index to pgvector as needed
for doc in docs:
    print(doc.page_content)

gen_context_chunks does not load files from storage — the caller must pass the content as a string. Indexing to pgvector is the caller's responsibility.

Full Pipeline Example

Transcribe every page of a PDF with the tiered pipeline and collect the Markdown:

import asyncio
import fitz
from wizit_open_rag import OpenRagTranscriber

async def transcribe_pdf(pdf_path: str) -> str:
    transcriber = OpenRagTranscriber(
        langsmith_project_name="my-project",
        langsmith_api_key="lsv2_...",
        use_tiered_transcription=True,
        tier2_ocr="textract",
    )
    pages_text = []
    with fitz.open(pdf_path) as doc:
        for i in range(len(doc)):
            single = fitz.open()
            single.insert_pdf(doc, from_page=i, to_page=i)
            result = await transcriber.transcribe_document(
                page_number=i + 1,
                page_content=single.tobytes(),
            )
            pages_text.append(result.page_text or "")
    return "\n\n".join(pages_text)

markdown = asyncio.run(transcribe_pdf("document.pdf"))

---

Configuration Reference

All parameters are set on OpenRagTranscriber.__init__:

Parameter	Default	Description
langsmith_project_name	required	LangSmith project name for run tracing
langsmith_api_key	required	LangSmith API key
llm_model_id	"global.anthropic.claude-sonnet-4-6"	Bedrock model used when use_tiered_transcription=False
target_language	"es"	BCP-47 language tag for the expected output
transcription_accuracy_threshold	0.90	Minimum quality score [0.0, 0.95] to accept a tier's output
max_transcription_retries	2	Retry attempts [1, 3] within the LLM tier's LangGraph loop
use_tiered_transcription	False	Enable cost-aware tiered pipeline
tier2_ocr	"textract"	Tier 2 OCR backend: "textract" or "mistral"
tier3_model_id	"global.anthropic.claude-haiku-4-5-20251001"	Bedrock model for the LLM fallback tier
mistral_api_key	None	Mistral API key for tier2_ocr="mistral"; falls back to MISTRAL_API_KEY env var

---

Environment Variables

Variable	Purpose
LANGSMITH_API_KEY	LangSmith API key for tracing
LANGCHAIN_PROJECT	LangSmith project name
LANGSMITH_TRACING	Enable LangSmith tracing (true / false)
MISTRAL_API_KEY	Mistral API key (only needed for tier2_ocr="mistral")
VECTOR_STORE_CONNECTION	PostgreSQL connection string (pgvector)
VECTOR_STORE_TABLE	pgvector table name
SUPABASE_KEY / SUPABASE_URL	Supabase credentials (optional)

AWS credentials are read from the standard boto3 credential chain and are not set via environment variables in this library.

---

Architecture

wizit_open_rag/
├── transcription.py              # Public API — OpenRagTranscriber
├── chunks.py                     # Public API — ChunksManager
├── domain/                       # Core data models (PageToTranscribe, ParsedDocPage, ParsedDoc)
├── application/
│   ├── interfaces.py             # ABCs for all swappable services + PageTranscriptionTier
│   ├── transcription_app.py      # LangGraph transcription orchestrator
│   ├── tiered_transcription_app.py  # Cost-aware tier sequencer
│   └── context_chunk_app.py      # Per-chunk context enrichment
├── data/                         # Shared enums and prompt strings
├── infra/
│   ├── llms/                     # AWS Bedrock chat (ChatBedrockConverse)
│   ├── embeddings/               # AWS Bedrock embeddings (BedrockEmbeddings)
│   ├── transcription/
│   │   ├── pdfplumber_tier.py    # Tier 1 — digital text + tables, no external calls
│   │   ├── textract_tier.py      # Tier 2a — AWS Textract OCR
│   │   ├── mistral_ocr_tier.py   # Tier 2b — Mistral Document AI OCR
│   │   └── llm_tier.py           # Tier 3 — LLM wrapper (adapts TranscriptionApp)
│   ├── persistence/              # Local filesystem, AWS S3, PostgreSQL managers
│   ├── rag/                      # SemanticChunks, RecursiveChunks, MarkdownHeadersChunks, pgvector, Weaviate
│   └── secrets/                  # AWS Secrets Manager helper
├── utils/
└── workflows/                    # LangGraph state machines (transcription + context)

---

Development

Smoke tests

test_tiered.py exercises the full tiered pipeline against a local PDF. Credentials are read from environment variables.

# Tier 1 only — no credentials needed
uv run python test_tiered.py --skip-llm --pages 2

# Tier 1 + Textract score check + full tiered pipeline + LLM baseline
uv run python test_tiered.py --tier2 textract --pages 2

# Same with Mistral OCR as Tier 2
uv run python test_tiered.py --tier2 mistral --pages 2

# Different PDF
uv run python test_tiered.py --pdf data/TBBC-2025.pdf --pages 5 --tier2 textract

Flag	Default	Description
--pdf	data/GenAI-TBBC.pdf	Path to a local PDF
--pages	2	Number of pages to process
--tier2	textract	OCR backend: textract or mistral
--skip-llm	off	Skip tests that call Bedrock (Tests 2 and 3)

Unit tests

uv run pytest

Profiling

# CPU
uv run pyinstrument test_tiered.py --skip-llm --pages 5

# Memory
uv run python -m memray run test_tiered.py --skip-llm --pages 5

Building the package

uv build

---

Gotchas

• SemanticChunks calls AWS Bedrock at construction time — make sure credentials are available before instantiating ChunksManager.
• Both transcribe_document and gen_context_chunks are async; wrap them in asyncio.run(...) from synchronous code.
• OpenRagTranscriber and ChunksManager require langsmith_project_name and langsmith_api_key as constructor arguments — they are not read from environment variables.
• AWS Bedrock cross-region model IDs use the global. prefix (e.g. global.anthropic.claude-sonnet-4-6).
• transcribe_document takes page_number (1-based int) and page_content (raw bytes of a single-page PDF). Use PyMuPDF to split pages before calling it.
• When use_tiered_transcription=True, each OpenRagTranscriber instance holds two compiled LangGraph workflows (Sonnet for the default path, Haiku for Tier 3). Instantiate once and reuse across pages.

---

License

Licensed under the Apache License 2.0.