# Tech-Spec: Domain-Aware RAG Optimization

**Created:** 2025-12-30
**Status:** Ready for Development

---

## Overview

### Problem Statement

RAGAS scoring for the RAG system based on data chunking is too slow.

**Root Cause Analysis:**
1. **Full Vector Search:** Every query performs exhaustive similarity search across all chunks in the tenant's collection
2. **No Semantic Pre-filtering:** Queries like "What is the EBITDA for Q4?" search through HR documents, technical docs, and all unrelated content
3. **Evaluation Bottleneck:** RAGAS evaluation requires multiple retrieval calls, amplifying the latency issue
4. **Large Chunk Volumes:** Enterprise tenants with 100k+ chunks see significant retrieval latency

**Impact:**
- RAGAS evaluation taking 10-30 seconds per query
- Poor user experience for interactive chatbot queries
- Wasted compute on irrelevant chunks
- Reduced retrieval precision affecting answer quality

### Solution

Implement a two-phase retrieval system using **domain term extraction** to pre-filter chunks before vector similarity search:

1. **Extraction Phase (At Indexing Time):**
   - Extract domain-specific terms from each chunk during the embedding pipeline
   - Store terms as an array field in Milvus with INVERTED index
   - Categorize chunks by high-level domain (finance, HR, engineering, etc.)

2. **Query Phase (At Retrieval Time):**
   - Extract domain terms from user query (keyword matching + LLM semantic expansion)
   - Build filter expression using `array_contains_any()` on domain_terms field
   - Execute filtered vector search (10-100x fewer candidates)

**Expected Improvement:**
- 5-10x faster retrieval for domain-specific queries
- 20-50% improvement in RAGAS precision scores
- Reduced compute costs from filtering before vector search

### Scope (In/Out)

**In Scope:**
- Domain Term Extractor component for text chunks and Excel records
- Updated Milvus schema with `domain_terms` ARRAY and `domain_category` fields
- Query Term Matcher with keyword matching + LLM expansion
- Domain-Aware Retriever with two-phase retrieval
- Airflow DAG integration (`extract_domain_terms_task`)

**Out of Scope:**
- GraphRAG entity relationship extraction (future enhancement)
- Cross-document term linking and knowledge graph construction
- Real-time term taxonomy updates (batch updates only)
- UI for term taxonomy management
- Multi-language term extraction (English only for v1)

---

## Context for Development

### Codebase Patterns

**Current Pipeline (Airflow DAG):**
```
document_extraction_dag
├── validate_event
├── fetch_document
├── branch_by_format
├── parse_{format}_document
├── detect_headers_task (Excel only)
├── chunk_{excel|text}_task
├── generate_embeddings_task
└── [store_chunks_task, store_vectors_task] (parallel)
```

**Key Patterns to Follow:**
- `@task.external_python(python='/opt/airflow/.venv/bin/python')` for isolated execution
- Pickle files for large data transfer (>48KB XCom limit)
- Lazy imports inside task functions
- Graceful degradation with fallbacks
- Dataclass models for textiq-doc-extractions

**Naming Conventions:**
- Classes: PascalCase (`DomainTermExtractor`)
- Functions: snake_case (`extract_from_chunk`)
- Files: snake_case (`domain_term_extractor.py`)
- Constants: UPPER_SNAKE_CASE (`MAX_TERMS_PER_CHUNK`)

### Files to Reference

| File | Line | Pattern to Follow |
|------|------|-------------------|
| `src/extraction_v2/llm_header_detector.py` | 45-80 | LLM call pattern with Azure OpenAI |
| `src/extraction_v2/record_builder.py` | 20-60 | Dataclass usage, metadata handling |
| `src/knowledge/vector_store.py` | 30-90 | Milvus collection operations |
| `airflow/dags/tasks/processing_tasks.py` | 120-180 | External python task structure |
| `airflow/plugins/models/xcom_schemas.py` | 35-50 | Chunk dataclass |

### Technical Decisions

| Decision | Choice | Rationale |
|----------|--------|-----------|
| **Term storage** | Array field with INVERTED index | Chunks can have multiple terms; `array_contains_any()` for efficient filtering |
| **Partition key** | `tenant_id` only | Single-valued per row; domain_terms is multi-valued |
| **Filtering strategy** | Three-level: tenant → category → terms | Progressive narrowing for optimal performance |
| **Term extraction** | Hybrid: rules + LLM | Fast rule-based for most cases; LLM for ambiguous content |
| **Query expansion** | Keyword first, LLM fallback | Minimize LLM calls for common queries |
| **Cache layer** | Redis + local memory | Fast term index lookups during query |

**Three-Level Filtering Strategy:**
```
Level 1: tenant_id (partition key)     → Data isolation, automatic partitioning
Level 2: domain_category (scalar)      → Coarse filter ("finance" vs "hr")
Level 3: domain_terms (array)          → Fine-grained term matching
```

---

## Implementation Plan

### Tasks

#### Story 1: Domain Term Extractor (8h)

- [ ] Create `src/extraction_v2/domain_term_extractor.py`
- [ ] Implement `DomainTermExtractor` class with:
  - `extract_from_chunk(content, metadata)` for text chunks
  - `extract_from_excel_record(headers, row_data, sheet_name)` for Excel
  - `_extract_by_rules()` for rule-based extraction
  - `_match_taxonomy()` for pattern matching
  - `_llm_extract()` for LLM enhancement
- [ ] Create `DomainTaxonomy` dataclass with default categories
- [ ] Add unit tests with mocked LLM (`tests/extraction_v2/test_domain_term_extractor.py`)

#### Story 2: Milvus Schema Update (4h)

- [ ] Create migration script `scripts/migrate_milvus_schema.py`
- [ ] Add fields to schema:
  - `domain_category` VARCHAR(64)
  - `domain_terms` ARRAY<VARCHAR> max_capacity=50
- [ ] Create INVERTED indexes on both new fields
- [ ] Update `VectorStore.index_with_vectors()` to include new fields
- [ ] Test with sample data insertion and querying

#### Story 3: Query Term Matcher (6h)

- [ ] Create `src/knowledge/query_term_matcher.py`
- [ ] Implement `QueryTermMatcher` class with:
  - `extract_query_terms(query, tenant_id)` main method
  - `_keyword_match()` for fast exact/fuzzy matching
  - `_llm_expand()` for semantic expansion
  - `_infer_category_from_terms()` for category detection
- [ ] Create `src/knowledge/term_index.py` for term caching (Redis + local)
- [ ] Add unit tests

#### Story 4: Domain-Aware Retriever (6h)

- [ ] Create `src/knowledge/domain_aware_retriever.py`
- [ ] Implement `DomainAwareRetriever` class with:
  - `retrieve(query, tenant_id, top_k)` two-phase retrieval
  - `_build_filter()` for Milvus filter expression
  - Fallback to full search when no terms match
- [ ] Add integration tests with test Milvus collection
- [ ] Performance benchmarking vs current retriever

#### Story 5: Airflow DAG Integration (4h)

- [ ] Add `extract_domain_terms_task` to `processing_tasks.py`
- [ ] Update DAG flow: `chunk_*_task → extract_domain_terms_task → generate_embeddings_task`
- [ ] Update `store_vectors_task` to pass domain data to VectorStore
- [ ] End-to-end testing with sample documents

### Acceptance Criteria

- [ ] **AC1:** Given a text chunk with finance terms, when processed by DomainTermExtractor, then domain_terms includes ["EBITDA", "revenue", "margin"] and category is "finance"
- [ ] **AC2:** Given an Excel record with headers ["Employee ID", "Salary"], when processed, then domain_terms includes ["employee", "salary"] and category is "hr"
- [ ] **AC3:** Given a query "What is the quarterly revenue?", when matched against term index, then returns terms ["revenue", "quarterly"] with category "finance"
- [ ] **AC4:** Given a domain-specific query, when retrieved with DomainAwareRetriever, then latency is <100ms and results are pre-filtered by domain
- [ ] **AC5:** Given a generic query with no term matches, when retrieved, then full search is used as fallback and returns results
- [ ] **AC6:** Given existing chunks without domain_terms, when queried, then retrieval still works (backward compatibility)
- [ ] **AC7:** Given RAGAS evaluation on test set, when comparing before/after, then precision improves by 20%+

---

## Additional Context

### Dependencies

**New Python Packages:**

| Package | Version | Purpose |
|---------|---------|---------|
| rapidfuzz | 3.x | Fast fuzzy string matching for keyword matching |

**Internal Module Dependencies:**
- `src.knowledge.embeddings.EmbeddingService` - Embedding generation
- `src.knowledge.vector_store.VectorStore` - Milvus operations
- `src.extraction_v2.semantic_chunker.SemanticChunker` - Chunk structure

**External Services:**
- Redis - Term index caching
- Azure OpenAI (gpt-4o) - LLM term expansion
- Milvus 2.6.x - Array field and INVERTED index support

### Testing Strategy

**Unit Tests:**
- `DomainTermExtractor.extract_from_chunk()` with various content types
- `DomainTermExtractor.extract_from_excel_record()` with headers/data
- `QueryTermMatcher._keyword_match()` with edge cases
- `DomainAwareRetriever._build_filter()` with different term counts

**Integration Tests:**
- Full extraction pipeline with test documents
- Query retrieval with pre-populated Milvus collection
- RAGAS evaluation comparison: before vs after optimization

**Performance Tests:**
- Extraction latency per chunk (target: <50ms avg)
- Query latency with filtering (target: <100ms)
- Compare RAGAS scores before/after

**Coverage Target:** 80% for new code

### Notes

**Configuration (Environment Variables):**
```bash
DOMAIN_TERM_MAX_PER_CHUNK=50
DOMAIN_TERM_LLM_THRESHOLD=2
DOMAIN_TERM_CACHE_TTL=3600
QUERY_TERM_MAX_RESULTS=10
QUERY_TERM_LLM_EXPANSION=true
```

**DAG Integration Point:**
```
chunk_excel_task / chunk_text_task
        │
        ▼
extract_domain_terms_task  ← NEW
        │
        ▼
generate_embeddings_task
        │
        ▼
[store_chunks_task, store_vectors_task]
```

**Rollback Plan:**
1. If extraction fails: Set `DOMAIN_TERM_EXTRACTION_ENABLED=false` - chunks without terms use fallback
2. If retrieval fails: DomainAwareRetriever has built-in fallback to full search
3. Full rollback: Revert DAG - domain_terms in Milvus are ignored (harmless)

**Metrics to Track:**
- `domain_terms_extracted_total` - Counter of chunks with terms extracted
- `query_term_matches_total` - Counter of queries with term matches
- `retrieval_filter_ratio` - Ratio of filtered vs full searches
- `retrieval_latency_filtered_ms` vs `retrieval_latency_full_ms`

---

*Generated by BMAD Tech-Spec Workflow*
*Last Updated: 2025-12-30*