Metadata-Version: 2.3
Name: wizit-open-rag
Version: 0.0.2
Summary: AI-powered document transcription and semantic chunking for RAG pipelines
Keywords: rag,retrieval-augmented-generation,llm,chunking,transcription,weaviate,bedrock,langchain,langgraph,pdf,semantic-search
Author: Restebance
Author-email: Restebance <restebance@gmail.com>
License: Apache-2.0
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Dist: boto3>=1.40.23
Requires-Dist: langchain>=1.2.10
Requires-Dist: langchain-aws>=1.3.0
Requires-Dist: langchain-classic>=1.0.7
Requires-Dist: langchain-community>=0.4.1
Requires-Dist: langchain-core>=1.2.16
Requires-Dist: langchain-experimental>=0.4.1
Requires-Dist: langchain-text-splitters>=1.1.1
Requires-Dist: langgraph>=1.0.9
Requires-Dist: pillow>=11.3.0
Requires-Dist: mistralai>=1.0
Requires-Dist: pdfplumber>=0.11
Requires-Dist: pymupdf>=1.27.1
Requires-Dist: anthropic>=0.84.0
Requires-Dist: psycopg2-binary>=2.9.11
Requires-Dist: sqlalchemy[asyncio]>=2.0.43
Requires-Dist: langchain-postgres>=0.0.17
Requires-Dist: weaviate-client>=4.0.0
Requires-Dist: langchain-weaviate>=0.0.3
Requires-Dist: botocore[crt]>=1.43.7
Requires-Python: >=3.12
Project-URL: Bug Tracker, https://github.com/Restebance/open_rag/issues
Project-URL: Repository, https://github.com/Restebance/open_rag
Description-Content-Type: text/markdown

# 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](https://docs.astral.sh/uv/) for dependency management
- AWS credentials configured (standard boto3 credential chain — env vars, `~/.aws/credentials`, or instance profile)
- PostgreSQL with the [pgvector](https://github.com/pgvector/pgvector) extension enabled

---

## Installation

```bash
pip install wizit_open_rag
```

For development (clone + install with dev tools):

```bash
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.

```python
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):**

```python
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:**

```python
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.

```python
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:

```python
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.

```bash
# 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

```bash
uv run pytest
```

### Profiling

```bash
# 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

```bash
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](LICENSE.md).
