# Deep-Dive: Airflow DAG & Extraction V2 Integration

> Generated: 2025-12-30 | Scope: `airflow/` ↔ `src/extraction_v2/`

## Executive Summary

This document provides a comprehensive analysis of how the **Apache Airflow DAG** (`airflow/`) orchestrates document processing using the **Extraction V2** business logic (`src/extraction_v2/`). The Airflow layer acts as the workflow orchestrator while extraction_v2 provides the core document processing capabilities.

**Key Integration Pattern:** Airflow tasks use `@task.external_python` to execute extraction_v2 business logic in an isolated virtual environment (`/opt/airflow/.venv`), enabling heavy processing without blocking the Airflow scheduler.

---

## Architecture Overview

```
┌─────────────────────────────────────────────────────────────────────────────┐
│                           AIRFLOW ORCHESTRATION                              │
│  ┌─────────────┐   ┌─────────────┐   ┌─────────────┐   ┌─────────────┐     │
│  │  Validate   │──▶│   Fetch     │──▶│   Branch    │──▶│   Parse*    │     │
│  │   Event     │   │  Document   │   │  (format)   │   │ (per type)  │     │
│  └─────────────┘   └─────────────┘   └─────────────┘   └─────────────┘     │
│         │                                                     │              │
│         ▼                                                     ▼              │
│  ┌─────────────────────────────────────────────────────────────────────┐   │
│  │                    @task.external_python                             │   │
│  │                    /opt/airflow/.venv/bin/python                     │   │
│  └─────────────────────────────────────────────────────────────────────┘   │
└─────────────────────────────────────────────────────────────────────────────┘
                                      │
                                      ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                         EXTRACTION V2 BUSINESS LOGIC                         │
│  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐             │
│  │ DocumentTo      │  │ HtmlTable       │  │ LargeTable      │             │
│  │ HtmlConverter   │  │ Extractor       │  │ Extractor       │             │
│  └─────────────────┘  └─────────────────┘  └─────────────────┘             │
│  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐             │
│  │ LlmHeader       │  │ RecordBuilder   │  │ SemanticChunker │             │
│  │ Detector        │  │                 │  │                 │             │
│  └─────────────────┘  └─────────────────┘  └─────────────────┘             │
│  ┌─────────────────┐  ┌─────────────────┐                                  │
│  │ BorderTable     │  │ excel_loader    │                                  │
│  │ Detector        │  │                 │                                  │
│  └─────────────────┘  └─────────────────┘                                  │
└─────────────────────────────────────────────────────────────────────────────┘
```

---

## File Inventory

### Airflow Components (`airflow/`)

| File | Purpose | Key Exports |
|------|---------|-------------|
| `dags/config.py` | Shared DAG configuration | `default_args`, `MAX_XCOM_SIZE`, timeout constants |
| `dags/document_extraction_dag.py` | Main DAG definition | DAG with 14 tasks, branching logic |
| `dags/tasks/parse_tasks.py` | Document parsing tasks | `parse_excel_document`, `parse_word_document`, `parse_pdf_document`, `parse_powerpoint_document` |
| `dags/tasks/processing_tasks.py` | Chunking/embedding/storage tasks | `detect_headers_task`, `chunk_excel_task`, `chunk_text_task`, `generate_embeddings_task`, `store_chunks_task`, `store_vectors_task` |
| `plugins/operators/validate_event.py` | Event validation operator | `ValidateEventOperator` |
| `plugins/operators/fetch_document.py` | Document fetch operator | `FetchDocumentOperator` |
| `plugins/models/xcom_schemas.py` | XCom data schemas | `ValidatedEvent`, `FetchedDocument`, `ParsedDocument`, `ContentChunks`, `ChunkEmbeddings` |

### Extraction V2 Components (`src/extraction_v2/`)

| File | Purpose | Key Exports |
|------|---------|-------------|
| `document_converter.py` | Docling-based document conversion | `DocumentToHtmlConverter` |
| `html_table_extractor.py` | HTML table parsing | `HtmlTable`, `HtmlTableExtractor` |
| `large_table_extractor.py` | Large Excel table handling | `LargeTableConfig`, `LargeTableResult`, `LargeTableExtractor` |
| `detect_border.py` | Border-based table detection | `BorderTableDetector` |
| `excel_loader.py` | Excel file loading utilities | `load_excel_file` |
| `llm_header_detector.py` | LLM-based header detection | `LlmHeaderDetector`, `HeaderStructure` |
| `record_builder.py` | Table-to-record conversion | `RecordBuilder`, `ExtractedRecord` |
| `semantic_chunker.py` | Semantic text chunking | `SemanticChunker` |
| `pipeline.py` | Standalone orchestration (non-Airflow) | `ExtractionPipelineV2` |

---

## Dependency Graph

### Airflow → Extraction V2 Import Map

```
airflow/dags/tasks/parse_tasks.py
├── src.extraction_v2.document_converter.DocumentToHtmlConverter
├── src.extraction_v2.html_table_extractor.HtmlTableExtractor
├── src.extraction_v2.html_table_extractor.HtmlTable
├── src.extraction_v2.large_table_extractor.LargeTableExtractor
├── src.extraction_v2.large_table_extractor.LargeTableConfig
└── src.extraction_v2.detect_border.BorderTableDetector

airflow/dags/tasks/processing_tasks.py
├── src.extraction_v2.llm_header_detector.LlmHeaderDetector
├── src.extraction_v2.llm_header_detector.HeaderStructure
├── src.extraction_v2.record_builder.RecordBuilder
├── src.extraction_v2.semantic_chunker.SemanticChunker
├── src.knowledge.embeddings.EmbeddingService
├── src.knowledge.structured_store.StructuredStore
└── src.knowledge.vector_store.VectorStore
```

### Internal Extraction V2 Dependencies

```
pipeline.py (orchestrator)
├── document_converter.DocumentToHtmlConverter
├── html_table_extractor.HtmlTableExtractor
├── large_table_extractor.LargeTableExtractor
├── llm_header_detector.LlmHeaderDetector
├── record_builder.RecordBuilder
└── semantic_chunker.SemanticChunker

large_table_extractor.py
├── detect_border.BorderTableDetector
└── excel_loader.load_excel_file

html_table_extractor.py
└── (standalone, uses BeautifulSoup)

detect_border.py
└── excel_loader.load_excel_file
```

---

## Data Flow Analysis

### Complete Pipeline: Excel Document

```
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 1: VALIDATION & FETCH (Airflow Operators)                              │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  ValidateEventOperator
     │  ├── Input: DAG conf (document_id, object_key, content_type)
     │  └── Output: ValidatedEvent → XCom
     │
     ▼
     │  FetchDocumentOperator
     │  ├── Input: ValidatedEvent from XCom
     │  ├── Action: Download from SeaweedFS/S3 to /tmp/
     │  └── Output: FetchedDocument → XCom
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 2: PARSING (@task.external_python)                                     │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  parse_excel_document()
     │  ├── Input: FetchedDocument dict
     │  │
     │  ├── Step 1: Border Detection
     │  │   └── BorderTableDetector.detect_tables_by_border()
     │  │       ├── Analyzes cell borders in all worksheets
     │  │       └── Returns: list[{sheet, start_row, end_row, start_col, end_col}]
     │  │
     │  ├── Step 2: Table Classification
     │  │   └── BorderTableDetector.classify_tables()
     │  │       ├── Large tables: >1000 rows
     │  │       └── Normal tables: ≤1000 rows
     │  │
     │  ├── Step 3a: Large Table Extraction (if any)
     │  │   └── LargeTableExtractor.extract_table()
     │  │       ├── Chunked processing (100 rows at a time)
     │  │       ├── Header detection from first 10 rows
     │  │       └── Returns: LargeTableResult with records[]
     │  │
     │  ├── Step 3b: Normal Table Extraction
     │  │   └── DocumentToHtmlConverter.convert()
     │  │       ├── Uses Docling for conversion
     │  │       └── Returns: HTML string
     │  │
     │  ├── Step 4: HTML Table Extraction
     │  │   └── HtmlTableExtractor.extract_all_tables()
     │  │       ├── Parses <table> elements
     │  │       ├── Filters: min 2 columns
     │  │       └── Returns: (valid_tables[], raw_text_tables[])
     │  │
     │  └── Output: ParsedDocument → XCom
     │      {
     │        doc_type: 'excel',
     │        tables: [{html, rows, sheet_name}],
     │        raw_text_tables: [{html, rows}],
     │        large_table_records: [ExtractedRecord],
     │        local_path, document_id
     │      }
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 3: HEADER DETECTION (@task.external_python)                            │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  detect_headers_task()
     │  ├── Input: ParsedDocument dict
     │  │
     │  └── For each table:
     │      └── LlmHeaderDetector.detect_headers(html)
     │          ├── Sends table HTML to Azure OpenAI
     │          ├── LLM identifies header rows
     │          └── Returns: HeaderStructure
     │              {
     │                header_rows: [0, 1],  # Row indices
     │                headers: ['Column A', 'Column B']
     │              }
     │
     │  └── Output: header_data[] → XCom
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 4: CHUNKING (@task.external_python)                                    │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  chunk_excel_task()
     │  ├── Input: document_id, ParsedDocument, headers_result
     │  │
     │  ├── Step 1: Include pre-processed large table records
     │  │   └── Already converted to ExtractedRecord in parse phase
     │  │
     │  ├── Step 2: Build records from data tables
     │  │   └── RecordBuilder.build_records_from_table()
     │  │       ├── Uses HeaderStructure to identify headers
     │  │       ├── Each data row → ExtractedRecord
     │  │       └── Record contains: content, headers, metadata
     │  │
     │  ├── Step 3: Build raw text records
     │  │   └── RecordBuilder.build_combined_raw_text_records()
     │  │       ├── Combines text with token limits (8000 tokens)
     │  │       └── Creates semantic-style chunks
     │  │
     │  └── Output: ContentChunks → XCom
     │      {
     │        chunks: [ExtractedRecord dicts],
     │        chunk_count: N,
     │        chunking_strategy: 'table'
     │      }
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 5: EMBEDDING (@task.external_python)                                   │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  generate_embeddings_task()
     │  ├── Input: ContentChunks dict
     │  │
     │  ├── Step 1: Extract text from chunks
     │  │   └── Handle both dict and string formats
     │  │
     │  ├── Step 2: Batch embedding generation
     │  │   └── EmbeddingService.embed_batch()
     │  │       ├── Azure OpenAI text-embedding-3-large
     │  │       ├── Batch size: 16 (max 100)
     │  │       └── Dimensions: 3072
     │  │
     │  ├── Step 3: Save to pickle file (avoid XCom limits)
     │  │   └── Path: /opt/airflow/logs/embeddings_{document_id}.pkl
     │  │
     │  └── Output: {data_file: path, chunk_count, doc_type} → XCom
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 6: STORAGE (@task.external_python, parallel)                           │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     ├── store_chunks_task() [Excel only]
     │   ├── Input: document_id, embedded_result
     │   ├── Load: Pickle file with chunks
     │   └── Store: StructuredStore → PostgreSQL
     │
     └── store_vectors_task()
         ├── Input: document_id, embedded_result
         ├── Load: Pickle file with chunks + embeddings
         ├── Cleanup: Delete pickle file
         └── Store: VectorStore.index_with_vectors() → Milvus
```

### Complete Pipeline: Text Documents (Word/PDF/PowerPoint)

```
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASES 1-2: Same as Excel (Validate → Fetch → Parse)                         │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  parse_word_document() / parse_pdf_document() / parse_powerpoint_document()
     │  ├── Input: FetchedDocument dict
     │  │
     │  ├── Word: Convert .doc/.docm → .docx via LibreOffice
     │  │
     │  └── DocumentToHtmlConverter.convert()
     │      ├── Uses Docling DocumentConverter
     │      └── Returns: markdown_content
     │
     │  └── Output: ParsedDocument → XCom
     │      {
     │        doc_type: 'word'|'pdf'|'powerpoint',
     │        markdown_content: string,
     │        local_path, document_id
     │      }
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 3: SKIP HEADER DETECTION (Text documents bypass this)                  │
└──────────────────────────────────────────────────────────────────────────────┘
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASE 4: SEMANTIC CHUNKING (@task.external_python)                           │
└──────────────────────────────────────────────────────────────────────────────┘
     │
     │  chunk_text_task()
     │  ├── Input: document_id, ParsedDocument, max_tokens=8000
     │  │
     │  ├── Step 1: Re-convert document for full structure
     │  │   └── docling.DocumentConverter().convert()
     │  │       └── Needed because ParsedDocument only has markdown
     │  │
     │  └── Step 2: Semantic chunking
     │      └── SemanticChunker.chunk_document()
     │          ├── Uses Docling HybridChunker
     │          ├── Respects document structure (headings, paragraphs)
     │          ├── Token limit: 8000 (for text-embedding-3-large)
     │          └── Returns: list[dict] with content, metadata
     │
     │  └── Output: ContentChunks → XCom
     │      {
     │        chunks: [chunk dicts],
     │        chunk_count: N,
     │        chunking_strategy: 'semantic'
     │      }
     │
┌──────────────────────────────────────────────────────────────────────────────┐
│ PHASES 5-6: Same as Excel (Embed → Store)                                    │
│ Note: store_chunks_task skips PostgreSQL for text documents                  │
└──────────────────────────────────────────────────────────────────────────────┘
```

---

## Integration Points

### 1. Parse Tasks → Extraction V2

| Airflow Task | Extraction V2 Components | Purpose |
|--------------|--------------------------|---------|
| `parse_excel_document` | `BorderTableDetector`, `LargeTableExtractor`, `DocumentToHtmlConverter`, `HtmlTableExtractor` | Multi-strategy Excel parsing |
| `parse_word_document` | `DocumentToHtmlConverter` | Word → Markdown conversion |
| `parse_pdf_document` | `DocumentToHtmlConverter` | PDF → Markdown conversion |
| `parse_powerpoint_document` | `DocumentToHtmlConverter` | PPT → Markdown conversion |

### 2. Processing Tasks → Extraction V2

| Airflow Task | Extraction V2 Components | Purpose |
|--------------|--------------------------|---------|
| `detect_headers_task` | `LlmHeaderDetector` | AI-powered header identification |
| `chunk_excel_task` | `RecordBuilder`, `HtmlTable`, `HeaderStructure` | Table-to-record conversion |
| `chunk_text_task` | `SemanticChunker` | Structure-aware text splitting |

### 3. Data Transfer Patterns

```python
# Pattern 1: Direct XCom (small data)
@task.external_python(python='/opt/airflow/.venv/bin/python')
def detect_headers_task(parse_result: Dict) -> Dict:
    # Input: Dict from XCom
    # Output: Dict to XCom (small: just header info)
    return {'tables_processed': N, 'header_data': [...]}

# Pattern 2: Pickle file (large data)
@task.external_python(python='/opt/airflow/.venv/bin/python')
def generate_embeddings_task(chunks_result: Dict) -> Dict:
    # Large data → Save to pickle
    data_file = f'/opt/airflow/logs/embeddings_{doc_id}.pkl'
    with open(data_file, 'wb') as f:
        pickle.dump({'chunks': chunks, 'embeddings': embeddings}, f)

    # Return only path reference
    return {'data_file': data_file, 'chunk_count': len(chunks)}
```

---

## Key Design Patterns

### 1. External Python Execution

All heavy processing runs in `/opt/airflow/.venv` to:
- Isolate dependencies (Docling, OpenPyXL, etc.)
- Prevent blocking the Airflow scheduler
- Allow different Python versions if needed

```python
@task.external_python(
    python='/opt/airflow/.venv/bin/python',
    execution_timeout=timedelta(minutes=120)
)
def parse_excel_document(fetch_result: Dict) -> Dict:
    # Heavy imports inside function (not at module level)
    from src.extraction_v2.document_converter import DocumentToHtmlConverter
    ...
```

### 2. Lazy Imports

All extraction_v2 imports happen inside task functions, not at module level:

```python
def parse_excel_document(fetch_result: Dict) -> Dict:
    # Imports inside function
    from src.extraction_v2.large_table_extractor import LargeTableExtractor
    from src.extraction_v2.detect_border import BorderTableDetector

    # Now use the imports
    detector = BorderTableDetector()
```

### 3. Graceful Degradation

Large table detection with fallback:

```python
# In parse_excel_document
detector = BorderTableDetector()
try:
    all_tables = detector.detect_tables_by_border(local_path)
    large_tables, normal_tables = detector.classify_tables(all_tables)
except Exception as e:
    logger.warning(f"Border detection failed: {e}")
    large_tables, normal_tables = [], []
    # Falls back to full Docling processing
```

### 4. Parallel Storage

Storage tasks run in parallel after embedding:

```python
# In DAG definition
embeddings_excel >> [store_chunks_excel, store_vectors_excel]
embeddings_text >> [store_chunks_text, store_vectors_text]
```

---

## Configuration Points

### Shared Configuration (`airflow/dags/config.py`)

```python
default_args = {
    'owner': 'textiq',
    'depends_on_past': False,
    'retries': 3,
    'retry_delay': timedelta(minutes=2),
    'retry_exponential_backoff': True,
}

MAX_XCOM_SIZE = 48 * 1024  # 48KB - triggers pickle fallback
PARSE_TIMEOUT = timedelta(minutes=120)
EMBEDDING_TIMEOUT = timedelta(minutes=60)
STORE_TIMEOUT = timedelta(minutes=30)
```

### Extraction V2 Configuration

```python
# Large table thresholds
LargeTableConfig.max_rows_for_docling = 1000
LargeTableConfig.chunk_size = 100
LargeTableConfig.header_sample_rows = 10

# Border detection
BorderTableDetector.DEFAULT_LARGE_TABLE_THRESHOLD = 1000

# Semantic chunking
SemanticChunker.default_max_tokens = 8000

# Embedding
EmbeddingService.model = 'text-embedding-3-large'
EmbeddingService.dimensions = 3072
```

---

## Error Handling Strategy

| Layer | Strategy | Example |
|-------|----------|---------|
| Airflow Operators | Raise exceptions, let Airflow retry | `ValidateEventOperator` raises `ValueError` for invalid content types |
| Parse Tasks | Log warnings, continue processing | Border detection failure → fallback to full Docling |
| Processing Tasks | Log warnings, skip failed items | Individual table header detection failure → use fallback headers |
| Storage Tasks | Critical failures logged, non-critical as warnings | PostgreSQL failure for text docs is warning (Milvus is primary) |

---

## Comparison: Airflow vs Pipeline.py

| Aspect | Airflow DAG | ExtractionPipelineV2 |
|--------|-------------|---------------------|
| **Use Case** | Production orchestration | Standalone/testing |
| **Execution** | Distributed, scheduler-managed | Single process, async |
| **State** | XCom + pickle files | In-memory |
| **Error Recovery** | Automatic retries, task-level | Manual, exception-based |
| **Monitoring** | Airflow UI, logs | Application logging |
| **Parallelism** | Task-level (workers) | Coroutine-level (asyncio) |

---

## Related Documentation

- [Airflow Component Spec](specs/airflow.md) - Auto-generated API documentation
- [Airflow Migration Tech Spec](feats/airflow-migration-tech-spec.md) - Migration planning
- [Sprint Status](feats/airflow-migration/sprint-status.yaml) - Implementation progress

---

*Generated by BMAD document-project workflow*