# Component Tech Spec: airflow

> Generated: 2025-12-17 | Source: airflow/

## Overview

The `airflow` component implements the document extraction DAG (Directed Acyclic Graph) for TextIQ's document processing pipeline using Apache Airflow 3.x. This component orchestrates multi-format document extraction, chunking, embedding generation, and storage through a series of tasks that branch based on file type (Excel, Word, PDF, PowerPoint).

**Key Capabilities:**
- Event-driven document processing triggered via API or SeaweedFS webhooks
- File-type-based branching (Excel uses table-based chunking; Word/PDF/PPT use semantic chunking)
- Heavy processing tasks run in an external Python virtual environment (`/opt/airflow/.venv`) via `@task.external_python`
- Dual storage: chunks to PostgreSQL, vectors to Milvus
- XCom file-based data transfer for large payloads to avoid size limits

## Architecture

### DAG Structure

```
validate_event → fetch_document → detect_format (branch)
                                        │
        ┌───────────────┬───────────────┼───────────────┬───────────────┐
        ↓               ↓               ↓               ↓               ↓
   parse_excel    parse_word      parse_pdf     parse_powerpoint  unsupported
        │               │               │               │
        ↓               └───────────────┴───────────────┘
  detect_headers                        │
        │                               ↓
        ↓                        select_parse_result
   chunk_excel                          │
        │                               ↓
        ↓                          chunk_text
generate_embeddings_excel               │
        │                               ↓
   ┌────┴────┐              generate_embeddings_text
   ↓         ↓                          │
store_chunks store_vectors         ┌────┴────┐
 (excel)      (excel)              ↓         ↓
                              store_chunks store_vectors
                                (text)      (text)
```

### File Structure

```
airflow/
├── config/                    # Airflow configuration
├── dags/
│   ├── __init__.py
│   ├── config.py              # Shared DAG configuration
│   ├── document_extraction_dag.py  # Main DAG definition
│   └── tasks/
│       ├── __init__.py
│       ├── parse_tasks.py     # @task.external_python parse functions
│       └── processing_tasks.py # @task.external_python processing functions
├── logs/                      # Task logs and embedding data storage
└── plugins/
    ├── __init__.py            # TextIQPlugin registration
    ├── hooks/
    │   └── __init__.py        # Future: SeaweedFSHook, MilvusHook
    ├── models/
    │   ├── __init__.py
    │   └── xcom_schemas.py    # XCom data schema definitions
    └── operators/
        ├── __init__.py
        ├── validate_event.py  # ValidateEventOperator
        └── fetch_document.py  # FetchDocumentOperator
```

## API / Interface

### DAG Configuration

#### `document_extraction_dag`

**DAG ID:** `document_extraction_dag`
**Schedule:** `None` (event-triggered only)
**Max Active Runs:** 1
**Tags:** `['extraction', 'v2', 'docling']`

**Parameters:**
- `document_id` (string): Unique document identifier (UUID)
- `object_key` (string): Path to the document in SeaweedFS storage
- `content_type` (string): MIME type of the document

**Triggering via API:**
```json
{
    "document_id": "uuid",
    "object_key": "path/to/file.xlsx",
    "content_type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
}
```

### Operators

#### Class: `ValidateEventOperator`

Validates incoming webhook/API events and extracts document metadata.

**Supported Event Formats:**
1. SeaweedFS webhook format (S3-compatible event notifications)
2. Direct API trigger format

**Class Constants:**
- `EXTENSION_TO_CONTENT_TYPE`: Mapping of file extensions to MIME types
- `ALLOWED_CONTENT_TYPES`: Set of supported MIME types
- `MAX_FILE_SIZE`: 100MB limit
- `DEFAULT_BUCKET`: `'documents'`

**Method: `execute(context) -> Dict`**

Validates and extracts event metadata.

**Returns:** ValidatedEvent dict containing:
- `bucket`: Storage bucket name
- `object_key`: File path/key in storage
- `content_type`: MIME type of the document
- `size_bytes`: File size in bytes
- `document_id`: Unique document identifier
- `event_id`: Event identifier

**Raises:**
- `ValueError`: If content type is unsupported or file is too large
- `KeyError`: If required fields are missing

---

#### Class: `FetchDocumentOperator`

Downloads documents from SeaweedFS/S3 storage to local temporary storage.

**Constructor:**
- `validate_event_task_id` (str, default: `'validate_event'`): Task ID to pull XCom from

**Method: `execute(context) -> Dict`**

Downloads document from storage.

**Returns:** FetchedDocument dict containing:
- `local_path`: Local filesystem path to downloaded file
- `content_type`: MIME type of the document
- `original_filename`: Original filename from storage
- `file_extension`: File extension (e.g., 'xlsx', 'pdf')
- `document_id`: Document identifier

**Environment Variables Used:**
- `AWS_ACCESS_KEY_ID`
- `AWS_SECRET_ACCESS_KEY`
- `AWS_ENDPOINT_URL`

### TaskFlow Functions (DAG Tasks)

#### `detect_format(event: dict) -> str`

Branch task that routes to appropriate parser based on file extension.

**Supported Extensions:**
- Excel: `xlsx`, `xls`, `xlsm` → `'parse_excel'`
- Word: `docx`, `doc`, `docm` → `'parse_word'`
- PDF: `pdf` → `'parse_pdf'`
- PowerPoint: `pptx`, `ppt`, `pptm` → `'parse_powerpoint'`
- Other: → `'unsupported_format'`

---

#### `get_document_id(event: dict) -> str`

Extracts `document_id` from validated event result.

---

#### `select_parse_result(word_result, pdf_result, ppt_result) -> dict`

Selects non-None parse result from text document branches.
**Trigger Rule:** `none_failed_min_one_success`

### Parse Tasks (`parse_tasks.py`)

All parse tasks use `@task.external_python(python='/opt/airflow/.venv/bin/python')`.

#### `parse_excel_document(fetch_result: Dict) -> Dict`

Parses Excel documents using Docling and custom extractors.

**Processing Steps:**
1. Detect tables via border detection (`LargeTableExtractor`)
2. Classify tables as large (>1000 rows) or normal
3. Process large tables with chunked extraction
4. Process normal tables with Docling HTML conversion
5. Extract tables from HTML (`HtmlTableExtractor`)

**Returns:** ParsedDocument dict:
- `doc_type`: `'excel'`
- `html_content`: HTML representation
- `markdown_content`: Empty for Excel
- `tables`: List of table data with HTML, rows, sheet_name
- `raw_text_tables`: Tables without structured data
- `large_table_records`: Pre-processed records from large tables
- `local_path`, `document_id`, `original_filename`

---

#### `parse_word_document(fetch_result: Dict) -> Dict`

Parses Word documents, handling `.doc`/`.docm` conversion to `.docx`.

**Processing:**
1. Convert `.doc`/`.docm` to `.docx` using LibreOffice if needed
2. Save converted file for downstream task reuse
3. Convert to markdown using `DocumentToHtmlConverter`

**Returns:** ParsedDocument dict with `doc_type: 'word'`, `markdown_content`, `converted_path`

---

#### `parse_pdf_document(fetch_result: Dict) -> Dict`

Parses PDF documents using Docling.

**Returns:** ParsedDocument dict with `doc_type: 'pdf'`, `markdown_content`

---

#### `parse_powerpoint_document(fetch_result: Dict) -> Dict`

Parses PowerPoint documents using Docling.

**Returns:** ParsedDocument dict with `doc_type: 'powerpoint'`, `markdown_content`

### Processing Tasks (`processing_tasks.py`)

#### `detect_headers_task(parse_result: Dict) -> Dict`

Detects table headers using LLM (`LlmHeaderDetector`) for Excel documents only.

**Returns:**
- `tables_processed`: Count of tables analyzed
- `header_data`: List of header info per table (header_rows, headers)

---

#### `chunk_excel_task(document_id: str, parse_result: Dict, headers_result: Dict) -> Dict`

Chunks Excel tables into records using `RecordBuilder`.

**Processing:**
1. Include pre-processed large table records
2. Build records from data tables with detected headers
3. Process raw text tables using `build_combined_raw_text_records` with token limits

**Returns:** ContentChunks dict:
- `chunks`: List of ExtractedRecord dicts
- `chunk_count`: Total records
- `chunking_strategy`: `'table'`
- `doc_type`: `'excel'`

---

#### `chunk_text_task(document_id: str, parse_result: Dict, max_tokens: int = 8000) -> Dict`

Chunks text documents semantically using `SemanticChunker`.

**Processing:**
1. Re-convert document using `docling.DocumentConverter` for full structure
2. Apply semantic chunking with token limits

**Returns:** ContentChunks dict:
- `chunks`: List of semantic chunk dicts
- `chunk_count`: Total chunks
- `chunking_strategy`: `'semantic'`
- `doc_type`: Document type from parse result

---

#### `generate_embeddings_task(chunks_result: Dict, batch_size: int = 16) -> Dict`

Generates embeddings for content chunks using `EmbeddingService`.

**Processing:**
1. Extract text from chunks (handles both dict and string formats)
2. Filter empty texts
3. Generate embeddings in batches (max 100 per Azure OpenAI call)
4. Store large data to pickle file to avoid XCom size limits

**Returns:**
- `data_file`: Path to pickle file with chunks and embeddings
- `chunk_count`: Number of chunks processed
- `doc_type`, `original_filename`

---

#### `store_chunks_task(document_id: str, embedded_result: Dict) -> Dict`

Stores chunks in PostgreSQL.

**Processing:**
- Excel: Uses `StructuredStore` to store `ExtractedRecord` objects
- Text: Skips PostgreSQL (stored directly in Milvus)

**Returns:**
- `stored_count`: Records stored
- `success`: Boolean
- `record_ids`: List of record IDs

---

#### `store_vectors_task(document_id: str, embedded_result: Dict) -> Dict`

Stores vectors in Milvus using `VectorStore.index_with_vectors()`.

**Processing:**
1. Load embedded data from pickle file
2. Clean up temp file after loading
3. Convert Excel chunks to `ExtractedRecord` objects
4. Index vectors with metadata

**Returns:**
- `stored_count`: Vectors stored
- `success`: Boolean

### XCom Data Schemas (`xcom_schemas.py`)

```python
@dataclass
class ValidatedEvent:
    document_id: UUID
    file_path: str
    content_type: str
    size_bytes: int
    event_id: str
    bucket: str
    object_key: str

@dataclass
class FetchedDocument:
    local_path: str
    content_type: str
    original_filename: str
    file_extension: str

@dataclass
class ParsedDocument:
    doc_type: str
    html_content: Optional[str]
    markdown_content: Optional[str]
    tables: list[dict]
    local_path: str

@dataclass
class ContentChunks:
    chunks: list[dict]
    chunk_count: int
    chunking_strategy: str

@dataclass
class ChunkEmbeddings:
    embeddings: list[dict]
    embedding_count: int
    model: str
    cached_count: int
```

## Dependencies

### Internal

| Module | Usage |
|--------|-------|
| `src.extraction_v2.document_converter.DocumentToHtmlConverter` | HTML/Markdown conversion via Docling |
| `src.extraction_v2.html_table_extractor.HtmlTableExtractor` | Table extraction from HTML |
| `src.extraction_v2.large_table_extractor.LargeTableExtractor` | Large Excel table detection and extraction |
| `src.extraction_v2.llm_header_detector.LlmHeaderDetector` | AI-powered header detection |
| `src.extraction_v2.record_builder.RecordBuilder` | Build ExtractedRecord from table data |
| `src.extraction_v2.semantic_chunker.SemanticChunker` | Semantic text chunking |
| `src.knowledge.embeddings.EmbeddingService` | Azure OpenAI embedding generation |
| `src.knowledge.structured_store.StructuredStore` | PostgreSQL record storage |
| `src.knowledge.vector_store.VectorStore` | Milvus vector indexing |
| `src.db.session.async_session_maker` | Database session factory |
| `src.extraction.models.ExtractedRecord` | Record data model |

### External

| Package | Usage |
|---------|-------|
| `airflow` | DAG definition, operators, task decorators |
| `airflow.decorators.task` | TaskFlow API (`@task`, `@task.branch`, `@task.external_python`) |
| `airflow.models.BaseOperator` | Custom operator base class |
| `airflow.models.param.Param` | DAG parameter definitions |
| `airflow.providers.standard.operators.empty.EmptyOperator` | Placeholder tasks |
| `boto3` | S3-compatible storage access (SeaweedFS) |
| `docling.document_converter.DocumentConverter` | Document conversion engine |
| `pickle` | Serialization for large XCom data |
| `tempfile` | Temporary file/directory management |

### Services

| Service | Purpose |
|---------|---------|
| SeaweedFS | S3-compatible document storage |
| PostgreSQL | Structured record storage, Airflow metadata |
| Milvus | Vector database for similarity search |
| Azure OpenAI | Embedding generation (`text-embedding-3-large`) |
| LibreOffice | Legacy Word format conversion (`.doc`, `.docm`) |

## Data Flow / State

### Inputs

1. **API Trigger Event:**
   ```json
   {
     "document_id": "UUID",
     "object_key": "path/to/document.xlsx",
     "content_type": "MIME type"
   }
   ```

2. **SeaweedFS Webhook Event:**
   ```json
   {
     "Records": [{
       "s3": {
         "bucket": {"name": "documents"},
         "object": {"key": "path/to/file", "contentType": "...", "size": 1234}
       }
     }],
     "document_id": "UUID"
   }
   ```

### Processing Flow

1. **Validate:** Extract and validate event metadata, infer content type from extension if needed
2. **Fetch:** Download document from SeaweedFS to local temp directory
3. **Branch:** Route to appropriate parser based on file extension
4. **Parse:** Convert document to structured format (HTML/Markdown + tables)
5. **Detect Headers (Excel only):** Use LLM to identify table headers
6. **Chunk:**
   - Excel: Table-based chunking with `RecordBuilder`
   - Text: Semantic chunking with `SemanticChunker`
7. **Embed:** Generate vectors using Azure OpenAI (batch processing)
8. **Store:** Save to PostgreSQL (chunks) and Milvus (vectors) in parallel

### Outputs

- **PostgreSQL:** `ExtractedRecord` objects with content, headers, source metadata
- **Milvus:** Vector embeddings indexed by document_id and record_id

### State Management

- **Stateless Tasks:** Each task is independent, data passed via XCom
- **Large Data Handling:** Embeddings stored to pickle files, paths passed via XCom
- **Temp File Cleanup:** Embedding pickle files cleaned after Milvus storage

## Configuration

### Timeouts (Environment Variables)

| Variable | Default | Description |
|----------|---------|-------------|
| `AIRFLOW_PARSE_TIMEOUT_MINUTES` | 120 | Parse task timeout |
| `AIRFLOW_EMBEDDING_TIMEOUT_MINUTES` | 60 | Embedding generation timeout |
| `AIRFLOW_STORE_TIMEOUT_MINUTES` | 30 | Storage task timeout |

### DAG Default Args

```python
default_args = {
    'owner': 'textiq',
    'depends_on_past': False,
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 0,
}
```

### Shared Config (`config.py`)

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

PARSE_TIMEOUT = timedelta(minutes=30)
EMBEDDING_TIMEOUT = timedelta(minutes=60)
STORE_TIMEOUT = timedelta(minutes=15)
MAX_XCOM_SIZE = 48 * 1024  # 48KB
```

## Usage Examples

### Trigger DAG via Airflow REST API

```bash
curl -X POST "http://localhost:8081/api/v1/dags/document_extraction_dag/dagRuns" \
  -H "Content-Type: application/json" \
  -d '{
    "conf": {
      "document_id": "550e8400-e29b-41d4-a716-446655440000",
      "object_key": "uploads/2025/12/report.xlsx",
      "content_type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
    }
  }'
```

### Trigger via Python

```python
from airflow.api.client.local_client import Client

client = Client(api_base_url="http://localhost:8081")
client.trigger_dag(
    dag_id="document_extraction_dag",
    conf={
        "document_id": "550e8400-e29b-41d4-a716-446655440000",
        "object_key": "uploads/report.pdf",
        "content_type": "application/pdf"
    }
)
```

### SeaweedFS Webhook Configuration

Configure SeaweedFS to send S3-compatible event notifications:

```json
{
  "EventName": "s3:ObjectCreated:*",
  "Records": [{
    "s3": {
      "bucket": {"name": "documents"},
      "object": {
        "key": "uploads/document.docx",
        "contentType": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "size": 102400
      }
    }
  }],
  "document_id": "550e8400-e29b-41d4-a716-446655440000"
}
```

## Error Handling

### Operator-Level

- **ValidateEventOperator:** Raises `KeyError` for missing fields, `ValueError` for invalid content types or file sizes
- **FetchDocumentOperator:** Cleans up temp directory on download failure

### Task-Level

- **Parse Tasks:** Raise `ValueError` for missing input or conversion failures
- **Chunk Tasks:** Log warnings for individual table failures, continue processing
- **Embedding Task:** Filter empty texts, log skip counts
- **Store Tasks:** Non-critical PostgreSQL failures logged as warnings (vector storage continues)

### Trigger Rules

- `select_parse_result`: `none_failed_min_one_success` - Succeeds if any text parser completes
- `chunk_text`: `none_failed_min_one_success` - Runs if any text parse branch succeeds
- `embeddings_text`: `none_failed_min_one_success` - Same as above
- `store_chunks_text` / `store_vectors_text`: `none_failed_min_one_success`

---

*This spec was auto-generated. Review and update as needed.*