# Story 1.4: Chunking & Embedding Operators

**Epic:** Airflow Infrastructure & Parallel Deployment
**Status:** Done
**Priority:** High
**Estimate:** 6 hours

> **Implementation Evolved:** The original plan below used custom operators (`DetectHeadersOperator`, `ChunkContentOperator`, `GenerateEmbeddingsOperator`). The actual implementation uses `@task.external_python` decorated functions in `airflow/dags/tasks/processing_tasks.py`: `detect_headers_task`, `chunk_excel_task`, `chunk_text_task`, `generate_embeddings_task`. See [`docs/specs/airflow.md`](../../../specs/airflow.md) for current architecture.

---

## User Story

As a **developer**, I want to **implement operators for header detection, content chunking, and embedding generation** so that **documents can be processed into searchable chunks with vector embeddings**.

---

## Acceptance Criteria

- [x] `DetectHeadersOperator` detects headers in Excel tables using LLM
- [x] `ChunkContentOperator` chunks content using table or semantic strategy
- [x] `GenerateEmbeddingsOperator` generates embeddings via Azure OpenAI
- [x] Operators reuse existing business logic from `src/extraction_v2/`
- [x] All operators pass unit tests
- [x] Operators handle rate limiting and retries properly

---

## Technical Details

### Operators to Implement

| Operator | Input | Output | Reuses |
|----------|-------|--------|--------|
| `DetectHeadersOperator` | ParsedDocument (Excel) | DetectedHeaders | `LlmHeaderDetector` |
| `ChunkContentOperator` | ParsedDocument + Headers | ContentChunks | `RecordBuilder`, `SemanticChunker` |
| `GenerateEmbeddingsOperator` | ContentChunks | ChunkEmbeddings | `embeddings.py` |

### Files to Create

| File | Description |
|------|-------------|
| `plugins/operators/detect_headers.py` | Header detection (Excel only) |
| `plugins/operators/chunk_content.py` | Content chunking operator |
| `plugins/operators/generate_embeddings.py` | Embedding generation |

---

## Implementation

### DetectHeadersOperator

```python
# plugins/operators/detect_headers.py
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from typing import Dict

class DetectHeadersOperator(BaseOperator):
    """
    Detect table headers using LLM (Excel documents only).

    Reuses: src.extraction_v2.llm_header_detector.LlmHeaderDetector
    """

    DEFAULT_PARSE_TASK_ID = 'parse_excel'

    @apply_defaults
    def __init__(self, parse_task_id: str = DEFAULT_PARSE_TASK_ID, **kwargs):
        super().__init__(**kwargs)
        self.parse_task_id = parse_task_id

    def execute(self, context) -> Dict:
        from src.extraction_v2.llm_header_detector import LlmHeaderDetector

        parse_result = context['ti'].xcom_pull(task_ids=self.parse_task_id)

        if not parse_result or parse_result.get('doc_type') != 'excel':
            return self._empty_result()

        tables = parse_result.get('tables', [])
        if not tables:
            return self._empty_result()

        detector = LlmHeaderDetector()
        all_headers = []

        for i, table in enumerate(tables):
            try:
                header_info = detector.detect_headers(table.get('html', ''))
                all_headers.append({
                    'table_index': i,
                    'header_rows': header_info.header_rows,
                    'headers': header_info.headers,
                })
            except Exception as e:
                self.log.warning(f"Header detection failed for table {i}: {e}")
                all_headers.append(self._fallback_headers(i))

        return {'tables_processed': len(all_headers), 'header_data': all_headers}
```

### ChunkContentOperator

```python
# plugins/operators/chunk_content.py
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from uuid import UUID
from typing import Dict, Optional

class ChunkContentOperator(BaseOperator):
    """
    Chunk content using appropriate strategy.

    Strategies:
    - 'table': Excel → RecordBuilder (rows become records)
    - 'semantic': Word/PDF/PowerPoint → SemanticChunker

    Reuses:
    - src.extraction_v2.record_builder.RecordBuilder
    - src.extraction_v2.semantic_chunker.SemanticChunker
    """

    VALID_STRATEGIES = ('table', 'semantic')
    DEFAULT_MAX_TOKENS = 8000

    @apply_defaults
    def __init__(self, chunking_strategy: str, max_tokens: int = DEFAULT_MAX_TOKENS, **kwargs):
        super().__init__(**kwargs)
        if chunking_strategy not in self.VALID_STRATEGIES:
            raise ValueError(f"Invalid chunking_strategy. Must be one of: {self.VALID_STRATEGIES}")
        self.chunking_strategy = chunking_strategy
        self.max_tokens = max_tokens

    def execute(self, context) -> Dict:
        validate_result = context['ti'].xcom_pull(task_ids='validate_event')
        document_id = UUID(validate_result['document_id'])

        if self.chunking_strategy == 'table':
            return self._chunk_excel(context, document_id)
        else:
            return self._chunk_text(context, document_id)

    def _chunk_excel(self, context, document_id: UUID) -> Dict:
        """Chunk Excel tables into records using RecordBuilder."""
        from src.extraction_v2.record_builder import RecordBuilder
        from src.extraction_v2.llm_header_detector import HeaderStructure
        from src.extraction_v2.html_table_extractor import HtmlTable

        # ... builds records from tables with detected headers
        return {'chunks': all_records, 'chunk_count': len(all_records), 'chunking_strategy': 'table'}

    def _chunk_text(self, context, document_id: UUID) -> Dict:
        """Chunk text documents semantically using SemanticChunker."""
        from src.extraction_v2.semantic_chunker import SemanticChunker
        from docling.document_converter import DocumentConverter

        # Re-conversion required: SemanticChunker needs DoclingDocument structure
        converter = DocumentConverter()
        result = converter.convert(local_path)

        chunker = SemanticChunker(max_tokens=self.max_tokens)
        chunks = chunker.chunk_document(doc=result.document, file_id=document_id, ...)

        return {'chunks': chunks, 'chunk_count': len(chunks), 'chunking_strategy': 'semantic'}
```

### GenerateEmbeddingsOperator

```python
# plugins/operators/generate_embeddings.py
import asyncio
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from typing import Any, Dict, Optional

class GenerateEmbeddingsOperator(BaseOperator):
    """
    Generate embeddings for chunks using Azure OpenAI.

    Reuses: src.knowledge.embeddings.EmbeddingService
    """

    DEFAULT_BATCH_SIZE = 16
    MAX_API_BATCH_SIZE = 100  # Azure OpenAI API limit
    EMBEDDING_MODEL = 'text-embedding-3-large'
    EMBEDDING_DIMENSIONS = 3072

    @apply_defaults
    def __init__(self, batch_size: int = DEFAULT_BATCH_SIZE, **kwargs):
        super().__init__(**kwargs)
        self.batch_size = min(batch_size, self.MAX_API_BATCH_SIZE)

    def execute(self, context) -> Dict:
        from src.knowledge.embeddings import EmbeddingService

        # Get chunks from both chunking tasks
        excel_chunks = context['ti'].xcom_pull(task_ids='chunk_excel') or {}
        text_chunks = context['ti'].xcom_pull(task_ids='chunk_text') or {}

        all_chunks = excel_chunks.get('chunks', []) + text_chunks.get('chunks', [])
        if not all_chunks:
            return self._empty_result()

        # Extract text and generate embeddings
        texts_to_embed, chunk_ids = self._extract_texts(all_chunks)
        embeddings = self._generate_embeddings_batch(texts_to_embed, chunk_ids)

        return {
            'embeddings': embeddings,
            'embedding_count': len(embeddings),
            'model': self.EMBEDDING_MODEL,
            'chunks': all_chunks,
        }

    def _generate_embeddings_batch(self, texts: list, chunk_ids: list) -> list:
        """Generate embeddings using single event loop for all batches."""
        return asyncio.run(self._generate_embeddings_async(texts, chunk_ids))

    async def _generate_embeddings_async(self, texts: list, chunk_ids: list) -> list:
        """Async batch embedding - single event loop avoids overhead."""
        embedding_service = EmbeddingService()
        embeddings = []

        for i in range(0, len(texts), self.batch_size):
            batch_texts = texts[i:i + self.batch_size]
            batch_ids = chunk_ids[i:i + self.batch_size]

            batch_embeddings = await embedding_service.embed_batch(batch_texts)
            for chunk_id, vector in zip(batch_ids, batch_embeddings):
                embeddings.append({'chunk_id': chunk_id, 'vector': vector})

        return embeddings
```

---

## Tasks

- [x] Create `plugins/operators/detect_headers.py`
- [x] Create `plugins/operators/chunk_content.py`
- [x] Create `plugins/operators/generate_embeddings.py`
- [x] Update `plugins/operators/__init__.py` with exports
- [x] Create unit tests for DetectHeadersOperator
- [x] Create unit tests for ChunkContentOperator
- [x] Create unit tests for GenerateEmbeddingsOperator
- [ ] Integration test with real documents (deferred to Story 1.6 - DAG integration)

---

## Dependencies

- Story 1.3 (Core operators)
- Existing `src/extraction_v2/llm_header_detector.py`
- Existing `src/extraction_v2/record_builder.py`
- Existing `src/extraction_v2/semantic_chunker.py`
- Existing `src/knowledge/embeddings.py`

---

## Notes

- DetectHeadersOperator only runs for Excel documents
- ChunkContentOperator selects strategy based on document type
- GenerateEmbeddingsOperator handles both record and semantic chunk formats
- Batch processing prevents rate limiting issues

---

## File List

| File | Action | Description |
|------|--------|-------------|
| `airflow/plugins/operators/detect_headers.py` | Created | DetectHeadersOperator - detects headers in Excel tables using LLM |
| `airflow/plugins/operators/chunk_content.py` | Created | ChunkContentOperator - chunks content using table or semantic strategy |
| `airflow/plugins/operators/generate_embeddings.py` | Created | GenerateEmbeddingsOperator - generates embeddings via Azure OpenAI |
| `airflow/plugins/operators/__init__.py` | Modified | Added exports for new operators |
| `tests/airflow/test_operators/conftest.py` | Modified | Added Airflow and module mocks for testing |
| `tests/airflow/test_operators/test_detect_headers.py` | Created | 9 unit tests for DetectHeadersOperator |
| `tests/airflow/test_operators/test_chunk_content.py` | Created | 11 unit tests for ChunkContentOperator |
| `tests/airflow/test_operators/test_generate_embeddings.py` | Created | 17 unit tests for GenerateEmbeddingsOperator |

---

## Dev Agent Record

### Implementation Notes

**DetectHeadersOperator:**
- Reuses `LlmHeaderDetector` from `src/extraction_v2/`
- Only processes Excel documents, skips other document types gracefully
- Configurable upstream task ID via `parse_task_id` parameter
- Fallback to first row as header when LLM detection fails
- Returns structured header data: `{tables_processed, header_data[]}`
- Each header entry contains: `table_index`, `header_rows`, `headers`

**ChunkContentOperator:**
- Supports two chunking strategies: 'table' (Excel) and 'semantic' (Word/PDF/PowerPoint)
- Table strategy reuses `RecordBuilder` and `HeaderStructure` from `src/extraction_v2/`
- Semantic strategy reuses `SemanticChunker` with Docling `HybridChunker`
- Handles both data tables and raw text tables for Excel
- Configurable max_tokens (default: 8000 for text-embedding-3-large)
- Returns: `{chunks[], chunk_count, chunking_strategy}`

**GenerateEmbeddingsOperator:**
- Reuses `EmbeddingService` from `src/knowledge/embeddings.py`
- Batch processing with configurable batch_size (default: 16, max: 100)
- Supports both table record format and semantic chunk format
- Extracts text content intelligently from different chunk structures
- Generates unique chunk IDs from source metadata
- Uses `asyncio.run()` to call async embedding service from sync operator
- Returns: `{embeddings[], embedding_count, model, chunks[]}`

### Testing Notes

- 38 unit tests for Story 1.4 operators (9 + 12 + 17)
- All 61 tests pass (including existing operator tests)
- Tests mock Airflow modules via conftest.py for isolated testing
- Comprehensive coverage: happy paths, error conditions, edge cases
- Tests verify XCom data passing between operators

### Validation Results

All acceptance criteria validated:
- ✓ DetectHeadersOperator detects headers in Excel tables using LLM
- ✓ ChunkContentOperator chunks content using table or semantic strategy
- ✓ GenerateEmbeddingsOperator generates embeddings via Azure OpenAI (text-embedding-3-large)
- ✓ Operators reuse existing business logic from src/extraction_v2/ and src/knowledge/
- ✓ All 61 operator tests pass (38 new + 23 existing)
- ✓ Operators handle retries via EmbeddingService (exponential backoff built-in)

### Code Review Fixes (2025-12-09)

Issues found and fixed during adversarial code review:
1. **[HIGH] asyncio.run() in loop** - Now uses single event loop for all batches
2. **[HIGH] Missing strategy validation** - ChunkContentOperator now validates chunking_strategy
3. **[MEDIUM] Magic number for batch size** - Added MAX_API_BATCH_SIZE constant
4. **[LOW] Unused type imports** - Cleaned up imports in all operators
5. **[HIGH] Outdated Implementation section** - Removed misleading example code from story

---

## Change Log

| Date | Change |
|------|--------|
| 2025-12-09 | Created DetectHeadersOperator with LLM header detection |
| 2025-12-09 | Created ChunkContentOperator with table and semantic strategies |
| 2025-12-09 | Created GenerateEmbeddingsOperator with batch processing |
| 2025-12-09 | Updated __init__.py exports for new operators |
| 2025-12-09 | Created comprehensive unit tests (37 tests, all passing) |
| 2025-12-09 | Story marked Ready for Review - all acceptance criteria met |
| 2025-12-09 | **Code Review Fixes Applied:** |
| 2025-12-09 | - Fixed asyncio.run() called multiple times in batch loop (now single event loop) |
| 2025-12-09 | - Added chunking_strategy validation in ChunkContentOperator |
| 2025-12-09 | - Added MAX_API_BATCH_SIZE constant (100) to GenerateEmbeddingsOperator |
| 2025-12-09 | - Removed unused type imports from operators |
| 2025-12-09 | - Removed outdated Implementation section from story file |
| 2025-12-09 | - Added 2 new tests for strategy validation and API limit constant |
| 2025-12-09 | - All 61 tests passing |