""" Processing tasks using @task.external_python decorator. These tasks run in the uv virtual environment at /opt/airflow/.venv which has all the required dependencies. """ from airflow.decorators import task from typing import Dict, Any, List, Optional # Path to the uv virtual environment Python EXTERNAL_PYTHON = '/opt/airflow/.venv/bin/python' @task.external_python(python=EXTERNAL_PYTHON) def detect_headers_task(parse_result: Dict[str, Any]) -> Dict[str, Any]: """ Detect table headers using LLM (Excel documents only). Args: parse_result: Result from parse_excel task (contains data_file path) Returns: DetectedHeaders dict with header info for each table """ from src.extraction_v2.llm_header_detector import LlmHeaderDetector import logging import pickle import os log = logging.getLogger(__name__) # Skip if no parse result if not parse_result: log.info("No parse result provided") return {'tables_processed': 0, 'header_data': []} # Load data from pickle file if data_file path is provided data_file = parse_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading parse data from {data_file}") with open(data_file, 'rb') as f: parse_result = pickle.load(f) if parse_result.get('doc_type') != 'excel': log.info(f"Not an Excel document (type: {parse_result.get('doc_type')})") return {'tables_processed': 0, 'header_data': []} # Get tables from parse result tables = parse_result.get('tables', []) if not tables: log.info("No tables found in parsed document") return {'tables_processed': 0, 'header_data': []} log.info(f"Detecting headers for {len(tables)} tables") # Initialize detector detector = LlmHeaderDetector() # Detect headers for each table all_headers = [] for i, table in enumerate(tables): try: table_html = table.get('html', '') if not table_html: log.warning(f"Table {i} has no HTML content, using fallback") all_headers.append({ 'table_index': i, 'header_rows': [0], 'headers': [], }) continue log.info(f"Detecting headers for table {i}") header_info = detector.detect_headers(table_html) all_headers.append({ 'table_index': i, 'header_rows': header_info.header_rows, 'headers': header_info.headers, }) log.info( f"Table {i}: {len(header_info.header_rows)} header row(s), " f"{len(header_info.headers)} column(s)" ) except Exception as e: log.warning(f"Header detection failed for table {i}: {e}. Using fallback.") all_headers.append({ 'table_index': i, 'header_rows': [0], 'headers': [], }) log.info(f"Header detection complete: {len(all_headers)} tables processed") return { 'tables_processed': len(all_headers), 'header_data': all_headers, } @task.external_python(python=EXTERNAL_PYTHON) def chunk_large_table_records( document_id: str, parse_result: Dict[str, Any], headers_result: Dict[str, Any], tenant_id: str, node_id: str, correlation_id: str, ) -> Dict[str, Any]: """ Chunk large Excel tables into records using RecordBuilder. This task processes LARGE TABLES (>1000 rows) from the NEW pipeline. Args: document_id: Document UUID string parse_result: Result from parse_excel_large_tables task (contains data_file path) headers_result: Result from detect_headers task tenant_id: Tenant UUID for pipeline-event routing node_id: Node UUID for pipeline-event routing correlation_id: Airflow dag_run_id, for event correlation Returns: ContentChunks dict with records """ from uuid import UUID 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 from src.services.pipeline_events import publish_stage_event import logging import pickle import os log = logging.getLogger(__name__) publish_stage_event( "chunking", "started", document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, ) doc_uuid = UUID(document_id) if not parse_result: log.warning("No parse result provided") return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'table'} # Load data from pickle file if data_file path is provided data_file = parse_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading parse data from {data_file}") with open(data_file, 'rb') as f: parse_result = pickle.load(f) # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up parse data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up parse data file {data_file}: {e}") # Get filename from parse result original_filename = parse_result.get('original_filename', '') # Get header data header_data = headers_result.get('header_data', []) if headers_result else [] # Get tables from parse result (NEW pipeline only processes large tables) large_table_records = parse_result.get('large_table_records', []) if not large_table_records: log.info("No large table records found (NEW pipeline only processes large tables)") return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'table'} log.info( f"Chunking {len(large_table_records)} pre-processed large table records" ) builder = RecordBuilder() all_records = [] # NEW PIPELINE: Only process large table records # Normal tables, raw text tables, images, shapes, and comments are handled by DETAILED pipeline log.info(f"Including {len(large_table_records)} pre-processed large table records") all_records.extend(large_table_records) log.info(f"Total: {len(all_records)} records from Excel chunking") # Add filename and col_range to each record's _source (required by ExtractedRecord) for record in all_records: if isinstance(record, dict) and '_source' in record: record['_source']['filename'] = original_filename if 'col_range' not in record['_source']: record['_source']['col_range'] = '' # Save large data to pickle file to avoid XCom size limits chunk_data = { 'chunks': all_records, 'chunk_count': len(all_records), 'chunking_strategy': 'table', 'doc_type': 'excel', 'original_filename': original_filename, } data_file = f"/opt/airflow/logs/chunks_excel_{document_id}.pkl" with open(data_file, 'wb') as f: pickle.dump(chunk_data, f) log.info(f"Saved chunk result to {data_file}") # Return only small metadata via XCom return { 'data_file': data_file, 'chunk_count': len(all_records), 'chunking_strategy': 'table', 'doc_type': 'excel', 'original_filename': original_filename, 'document_id': document_id, } @task.external_python(python=EXTERNAL_PYTHON) def convert_large_tables_to_natural_language( chunks_result: Dict[str, Any] ) -> Dict[str, Any]: """ Convert large table records to natural language using rule-based formatting. Takes flat record dictionaries from chunk_excel_tables and converts them to natural language descriptions with proper structure for downstream tasks. Args: chunks_result: Result from chunk_excel_tables (contains data_file path) Returns: Dict with data_file path containing structured records with natural language """ import logging import pickle import os from uuid import uuid4 log = logging.getLogger(__name__) if not chunks_result: log.warning("No chunks result provided") return {'chunks': [], 'chunk_count': 0, 'doc_type': 'excel'} # Load data from pickle file data_file = chunks_result.get('data_file') if not data_file or not os.path.exists(data_file): log.warning("No data file found") return chunks_result log.info(f"Loading chunk data from {data_file}") with open(data_file, 'rb') as f: chunk_data = pickle.load(f) # Clean up input file try: os.remove(data_file) log.info(f"Cleaned up input file: {data_file}") except Exception as e: log.warning(f"Failed to clean up input file {data_file}: {e}") chunks = chunk_data.get('chunks', []) if not chunks: log.info("No chunks to convert") return chunks_result log.info(f"Converting {len(chunks)} large table records to natural language") # Get original filename for metadata original_filename = chunk_data.get('original_filename', '') from src.extraction_v2.record_to_natural_language import RecordToNaturalLanguage converter = RecordToNaturalLanguage() # Prepare batch data records = [] sheet_names = [] row_numbers = [] for idx, chunk in enumerate(chunks): if isinstance(chunk, dict): sheet_name = chunk.get('_sheet_name', 'Table') row_number = chunk.get('_row_number', idx + 1) records.append(chunk) sheet_names.append(sheet_name) row_numbers.append(row_number) # Convert to natural language (rule-based, no LLM) natural_language_texts = converter.convert_batch( records, sheet_names, row_numbers ) # Group 3-4 records into one chunk, respecting 8k token limit for embedding MAX_TOKENS = 8000 TARGET_GROUP_SIZE = 4 from src.llm.factory import get_token_counter token_len = get_token_counter() from langchain_text_splitters import RecursiveCharacterTextSplitter text_splitter = RecursiveCharacterTextSplitter( chunk_size=MAX_TOKENS, chunk_overlap=200, length_function=token_len, separators=["\n\n", "\n", " ", ""], ) def make_chunk(text, sheet, row): return { 'content': {'text': text}, 'text': text, '_source': { 'filename': original_filename, 'sheet': sheet, 'row': row, }, 'headers': ['text'], } structured_chunks = [] i = 0 while i < len(natural_language_texts): group_texts = [] group_rows = [] group_sheet = sheet_names[i] while i < len(natural_language_texts) and len(group_texts) < TARGET_GROUP_SIZE: nl_text = natural_language_texts[i] if not nl_text: i += 1 continue # Check token count of the joined result (including separators) candidate = "\n\n".join(group_texts + [nl_text]) if token_len(candidate) > MAX_TOKENS and group_texts: break group_texts.append(nl_text) group_rows.append(row_numbers[i]) i += 1 if not group_texts: continue grouped_text = "\n\n".join(group_texts) # If a single record exceeds the limit, split it into multiple chunks if token_len(grouped_text) > MAX_TOKENS: splits = text_splitter.split_text(grouped_text) for split_text in splits: structured_chunks.append(make_chunk(split_text, group_sheet, group_rows[0])) else: structured_chunks.append(make_chunk(grouped_text, group_sheet, group_rows[0])) log.info( f"Grouped {len(natural_language_texts)} records into {len(structured_chunks)} chunks " f"(target {TARGET_GROUP_SIZE} records/chunk, max {MAX_TOKENS} tokens)" ) # Save to new pickle file output_data = { 'chunks': structured_chunks, 'chunk_count': len(structured_chunks), 'chunking_strategy': chunk_data.get('chunking_strategy', 'table'), 'doc_type': 'excel', 'original_filename': chunk_data.get('original_filename', ''), } nl_doc_id = chunks_result.get('document_id', uuid4().hex) output_file = f"/opt/airflow/logs/nl_chunks_{nl_doc_id}_{uuid4().hex[:8]}.pkl" with open(output_file, 'wb') as f: pickle.dump(output_data, f) log.info(f"Saved structured chunks to {output_file}") result = { 'data_file': output_file, 'chunk_count': len(structured_chunks), 'chunking_strategy': 'table', 'doc_type': 'excel', 'original_filename': chunk_data.get('original_filename', ''), } # Propagate document_id for downstream tasks and scoped cleanup if chunks_result.get('document_id'): result['document_id'] = chunks_result['document_id'] return result @task.external_python(python=EXTERNAL_PYTHON) def hierarchical_chunk_text_task( document_id: str, parse_result: Dict[str, Any] = None, chunk_sizes: List[int] = None, fetch_result_md: Dict[str, Any] = None, ) -> Dict[str, Any]: """ Hierarchically chunk text documents using LlamaIndex HierarchicalNodeParser. Creates 2-level hierarchy (MarkdownNodeParser -> SentenceSplitter(512)) with parent-child relationships for Word/PDF/PPTX documents. For markdown (when ``fetch_result_md`` is provided), reads the .md file directly and uses the same mermaid-aware Excel parser map (MarkdownNodeParser -> MermaidAwareLineSplitter), so ```mermaid fences stay inline with their surrounding content as atomic line-groups. A leaf containing a fence may grow up to ``chunker_max = embed_max - EMBED_SAFETY_BUFFER`` so the diagram stays adjacent to its explanatory prose; oversize single fences fall through to ``split_oversize_mermaid``. Args: document_id: Document UUID string parse_result: Result from parse task (contains data_file path with DoclingDocument). Required for non-markdown formats; ignored when ``fetch_result_md`` is set. chunk_sizes: Legacy kwarg, kept for backward compatibility. Ignored by both branches (the parser map is fixed: markdown -> 512). fetch_result_md: Result from fetch task for the .md file. When set, the task takes the markdown branch; otherwise uses the existing parse_result path. Returns: Dict with nodes_file path and metadata """ import logging import pickle import os from uuid import UUID from llama_index.core import Document from llama_index.core.node_parser import HierarchicalNodeParser from src.llm.factory import get_token_counter, get_tokenizer log = logging.getLogger(__name__) # Legacy kwarg, retained for backward compatibility; the parser map is fixed. if chunk_sizes is None: chunk_sizes = [512] doc_uuid = UUID(document_id) # Provider-aware tokenizer (Azure -> tiktoken cl100k_base; vLLM -> HF tokenizer) tokenizer = get_tokenizer() token_count = get_token_counter() # ===== Markdown branch ===== # The DAG always wires fetch_result_md (it can't conditionally omit kwargs at # XCom resolution time), so we gate on the actual filename extension here. _md_filename = (fetch_result_md or {}).get('original_filename') or '' if fetch_result_md is not None and _md_filename.lower().endswith('.md'): from llama_index.core.node_parser import get_leaf_nodes from llama_index.core.schema import NodeRelationship, RelatedNodeInfo from src.core.config import Settings from src.extraction_v2.mermaid_chunking import ( TEXT_PARSER_IDS, build_excel_parser_map, prune_header_only_nodes, ) embed_max_tokens = Settings().embedding_model_context_tokens # Match the Excel branch: leave headroom against the embed ceiling for # retained metadata (~100 tok template + filename/document_id) plus # Azure-tokenizer variance vs local cl100k_base counts. EMBED_SAFETY_BUFFER = 400 chunker_max = embed_max_tokens - EMBED_SAFETY_BUFFER local_path = fetch_result_md.get('local_path') original_filename = _md_filename or 'document.md' log.info( f"Markdown chunking: {original_filename} (path={local_path}, " f"embed_max={embed_max_tokens} chunker_max={chunker_max} buffer={EMBED_SAFETY_BUFFER})" ) with open(local_path, 'rb') as fh: content = fh.read().decode('utf-8') # let UnicodeDecodeError propagate if not content.strip(): raise ValueError("Markdown file is empty") # Mermaid-aware 2-level hierarchy (same parser map as the Excel branch): # MarkdownNodeParser carves sections by ATX headers, then # MermaidAwareLineSplitter produces leaves where any ```mermaid fence # stays atomic and sits inline with its surrounding prose if the leaf # fits under chunker_max. Oversize single fences fall through to # split_oversize_mermaid (body-line split only — fence boundary preserved). md_node_parser_map = build_excel_parser_map( tokenize_fn=token_count, embed_max_tokens=chunker_max, ) md_parser = HierarchicalNodeParser( node_parser_ids=TEXT_PARSER_IDS, node_parser_map=md_node_parser_map, ) doc = Document( doc_id=str(doc_uuid), text=content, metadata={ "document_id": str(doc_uuid), "filename": original_filename, "document_type": "markdown", "page_number": 1, }, ) nodes = md_parser.get_nodes_from_documents([doc]) nodes = list(prune_header_only_nodes(nodes)) leaf_nodes = get_leaf_nodes(nodes) mermaid_bearing = sum(1 for n in leaf_nodes if n.metadata.get("has_mermaid")) log.info( f"Markdown: {len(nodes)} total nodes, {len(leaf_nodes)} leaves " f"({mermaid_bearing} mermaid-bearing)" ) # Common post-processing: ref_doc_id, token_count metadata, save pickle. doc_id_str = str(doc_uuid) for node in nodes: if not node.ref_doc_id: node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str) node.metadata["token_count"] = token_count(node.text) node.metadata["char_count"] = len(node.text) nodes_dir = '/opt/airflow/logs/hierarchical_nodes' os.makedirs(nodes_dir, exist_ok=True) nodes_file = os.path.join(nodes_dir, f'nodes_{document_id}.pkl') nodes_data = { 'nodes': nodes, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': 'markdown', 'original_filename': original_filename, 'document_id': document_id, 'chunk_sizes': [512], 'mermaid_bearing_leaves': mermaid_bearing, } with open(nodes_file, 'wb') as f: pickle.dump(nodes_data, f) log.info(f"Saved {len(nodes)} markdown nodes to {nodes_file}") return { 'nodes_file': nodes_file, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': 'markdown', 'original_filename': original_filename, 'document_id': document_id, } # ===== End markdown branch ===== if not parse_result: log.warning("No parse result provided") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} # Load data from pickle file if data_file path is provided data_file = parse_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading parse data from {data_file}") with open(data_file, 'rb') as f: parse_result = pickle.load(f) # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up parse data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up parse data file {data_file}: {e}") original_filename = parse_result.get('original_filename', 'document') doc_type = parse_result.get('doc_type', 'word') log.info(f"Hierarchical chunking document: {original_filename} (type: {doc_type})") # Non-md branch uses the unified text parser map: # MarkdownNodeParser -> SentenceSplitter(512, "\n\n") # `chunk_sizes` is a legacy kwarg ignored by this branch. from src.extraction_v2.mermaid_chunking import ( TEXT_PARSER_IDS, build_text_parser_map, prune_header_only_nodes, ) if chunk_sizes != [512]: log.info( f"chunk_sizes={chunk_sizes} ignored for non-md branch; " "using unified text parser map (markdown, 512)." ) node_parser_map = build_text_parser_map(tokenizer) node_parser = HierarchicalNodeParser( node_parser_ids=TEXT_PARSER_IDS, node_parser_map=node_parser_map, ) # Create LlamaIndex Documents from pre-serialized per-page markdown documents = [] pages = parse_result.get('pages', []) if pages: total_pages = len(pages) for page_data in pages: page_number = page_data.get('page_number', 1) markdown = page_data.get('markdown', '') if not markdown or not markdown.strip(): continue doc = Document( doc_id=str(doc_uuid), text=markdown, metadata={ "document_id": str(doc_uuid), "filename": original_filename, "document_type": doc_type, "page_number": page_number, "total_pages": total_pages, }, ) documents.append(doc) else: # Legacy fallback: use raw text content log.warning("No pages found, using fallback text extraction") text_content = parse_result.get('markdown_content', '') or parse_result.get('text', '') or parse_result.get('content', '') if text_content: doc = Document( doc_id=str(doc_uuid), text=text_content, metadata={ "document_id": str(doc_uuid), "filename": original_filename, "document_type": doc_type, "page_number": 1, }, ) documents.append(doc) if not documents: log.warning("No content to chunk") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} log.info(f"Created {len(documents)} LlamaIndex Documents") # Apply hierarchical parsing nodes = node_parser.get_nodes_from_documents(documents) log.info(f"Created {len(nodes)} total nodes (before header-only prune)") nodes = list(prune_header_only_nodes(nodes)) log.info(f"After header-only prune: {len(nodes)} nodes") # Count leaf nodes from llama_index.core.node_parser import get_leaf_nodes leaf_nodes = get_leaf_nodes(nodes) log.info(f"Leaf nodes: {len(leaf_nodes)}") # Ensure ALL nodes have ref_doc_id set (used as doc_id in Milvus) from llama_index.core.schema import NodeRelationship, RelatedNodeInfo doc_id_str = str(doc_uuid) for node in nodes: if not node.ref_doc_id: node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str) # Add token count to each node's metadata for node in nodes: node.metadata["token_count"] = token_count(node.text) node.metadata["char_count"] = len(node.text) # Save nodes to pickle file (they can be large) nodes_dir = '/opt/airflow/logs/hierarchical_nodes' os.makedirs(nodes_dir, exist_ok=True) nodes_file = os.path.join(nodes_dir, f'nodes_{document_id}.pkl') nodes_data = { 'nodes': nodes, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': doc_type, 'original_filename': original_filename, 'document_id': document_id, 'parser_ids': TEXT_PARSER_IDS, } with open(nodes_file, 'wb') as f: pickle.dump(nodes_data, f) log.info(f"Saved {len(nodes)} nodes to {nodes_file}") return { 'nodes_file': nodes_file, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': doc_type, 'original_filename': original_filename, 'document_id': document_id, } @task.external_python(python=EXTERNAL_PYTHON) def chunk_text_task( document_id: str, parse_result: Dict[str, Any], max_tokens: int = 8000 ) -> Dict[str, Any]: """ Chunk text documents semantically using SemanticChunker. Args: document_id: Document UUID string parse_result: Result from parse task (word/pdf/powerpoint) - contains data_file path max_tokens: Maximum tokens per chunk Returns: ContentChunks dict with semantic chunks """ from uuid import UUID from src.extraction_v2.semantic_chunker import SemanticChunker from src.extraction_v2.document_converter import DocumentToHtmlConverter import logging import pickle import os log = logging.getLogger(__name__) doc_uuid = UUID(document_id) if not parse_result: log.warning("No parse result provided") return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'semantic'} # Load data from pickle file if data_file path is provided data_file = parse_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading parse data from {data_file}") with open(data_file, 'rb') as f: parse_result = pickle.load(f) # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up parse data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up parse data file {data_file}: {e}") local_path = parse_result.get('local_path') if not local_path: log.warning("No local_path in parse result") return {'chunks': [], 'chunk_count': 0, 'chunking_strategy': 'semantic'} original_filename = parse_result.get('original_filename', 'document') doc_type = parse_result.get('doc_type', 'word') log.info(f"Semantic chunking document: {original_filename} (type: {doc_type})") # Initialize variables that may be used in finally block converted_path = None converted_path_from_parse = None try: # USE STORED DOCLING DOCUMENT - NO RE-CONVERSION! docling_doc = parse_result.get('docling_document') if docling_doc is None: # Fallback for backward compatibility (if parse task didn't store DoclingDocument) log.warning("No stored DoclingDocument, falling back to re-conversion") from pathlib import Path import os from docling.document_converter import DocumentConverter converted_path_from_parse = parse_result.get('converted_path', '') file_path = Path(local_path) converted_path = None file_to_convert = local_path log.info(f"Original file path: {local_path}") log.info(f"Original file exists: {os.path.exists(local_path)}") log.info(f"Converted path from parse: {converted_path_from_parse}") if converted_path_from_parse: log.info(f"Converted file exists: {os.path.exists(converted_path_from_parse)}") # Use pre-converted file if available (avoids double conversion) if converted_path_from_parse and os.path.exists(converted_path_from_parse): log.info(f"Using pre-converted .docx from parse task: {converted_path_from_parse}") file_to_convert = converted_path_from_parse converted_path = converted_path_from_parse # Otherwise, check if we need to convert legacy Word formats elif file_path.suffix.lower() in ['.doc', '.docm']: log.info(f"Converting {file_path.suffix} to .docx before chunking...") log.info(f"File path before conversion: {file_path}") log.info(f"File exists before conversion: {os.path.exists(file_path)}") if not os.path.exists(file_path): raise ValueError(f"Original file not found: {file_path}. The temporary directory may have been cleaned up.") html_converter = DocumentToHtmlConverter() converted_path = html_converter._convert_doc_to_docx(file_path) if converted_path is None: raise ValueError(f"Failed to convert {file_path.suffix} to .docx: LibreOffice conversion failed") file_to_convert = converted_path log.info(f"Converted to .docx: {file_to_convert}") # Re-convert document for full structure converter = DocumentConverter() result = converter.convert(file_to_convert) docling_doc = result.document # Initialize chunker and chunk document chunker = SemanticChunker(max_tokens=max_tokens) chunks = chunker.chunk_document( doc=docling_doc, file_id=doc_uuid, original_filename=original_filename, ) # Add doc_type to each chunk for downstream tasks for chunk in chunks: chunk['doc_type'] = doc_type chunk['original_filename'] = original_filename log.info(f"Created {len(chunks)} semantic chunks") # Save large data to pickle file to avoid XCom size limits chunk_data = { 'chunks': chunks, 'chunk_count': len(chunks), 'chunking_strategy': 'semantic', 'doc_type': doc_type, 'original_filename': original_filename, } data_file = f"/opt/airflow/logs/chunks_text_{document_id}.pkl" with open(data_file, 'wb') as f: pickle.dump(chunk_data, f) log.info(f"Saved chunk result to {data_file}") # Return only small metadata via XCom return { 'data_file': data_file, 'chunk_count': len(chunks), 'chunking_strategy': 'semantic', 'doc_type': doc_type, 'original_filename': original_filename, 'document_id': document_id, } except Exception as e: log.error(f"Semantic chunking failed: {e}") raise finally: # Clean up temporary converted file ONLY if we created it # Don't delete pre-converted files from parse task if converted_path and converted_path != converted_path_from_parse and os.path.exists(converted_path): try: import shutil temp_dir = os.path.dirname(converted_path) if temp_dir and os.path.exists(temp_dir) and 'tmp' in temp_dir: shutil.rmtree(temp_dir) log.debug(f"Cleaned up temporary directory: {temp_dir}") except Exception as cleanup_error: log.warning(f"Failed to clean up temp file {converted_path}: {cleanup_error}") @task.external_python(python=EXTERNAL_PYTHON) def extract_domain_terms_task(chunks_result: Dict[str, Any], tenant_id: str, node_id: str) -> Dict[str, Any]: """ Extract domain terms from chunks for domain-aware RAG optimization. Analyzes chunk content and extracts domain-specific terms and categories to enable efficient pre-filtering during vector search. Args: chunks_result: Result from chunk task (contains data_file path) Returns: Dict with file path to chunk data enriched with domain terms """ from src.extraction_v2.domain_term_extractor import DomainTermExtractor import asyncio import logging import pickle import os from uuid import uuid4 log = logging.getLogger(__name__) if not chunks_result: log.warning("No chunks result provided") return { 'chunk_count': 0, 'doc_type': 'excel', 'original_filename': '', } # Load data from pickle file if data_file path is provided # Preserve original metadata before loading pickle original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'} data_file = chunks_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading chunk data from {data_file}") with open(data_file, 'rb') as f: chunks_result = pickle.load(f) # Restore important metadata from XCom (like 'adapted' flag) for key in ['adapted', 'document_id', 'filename']: if key in original_metadata and key not in chunks_result: chunks_result[key] = original_metadata[key] # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up chunk data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up chunk data file {data_file}: {e}") # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records') chunks = chunks_result.get('chunks') or chunks_result.get('records', []) if not chunks: log.info("No chunks to process") return { 'chunk_count': 0, 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } log.info(f"Extracting domain terms from {len(chunks)} chunks") # Initialize extractor extractor = DomainTermExtractor() # Extract domain terms concurrently doc_type = chunks_result.get('doc_type', 'excel') domain_data = asyncio.run( extractor.extract_from_chunks_async(chunks, doc_type=doc_type) ) log.info(f"Domain term extraction complete: {len(domain_data)} chunks processed") # Store enriched data to temporary file (include document_id for scoped cleanup) data_dir = '/opt/airflow/logs/domain_data' os.makedirs(data_dir, exist_ok=True) doc_id = chunks_result.get('document_id', uuid4().hex) output_file = os.path.join(data_dir, f'domain_terms_{doc_id}_{uuid4().hex[:8]}.pkl') # Determine the correct key name: 'records' for OLD pipeline, 'chunks' for NEW pipeline # Also preserve embeddings if they exist (from OLD pipeline) chunk_key = 'records' if chunks_result.get('adapted') or chunks_result.get('records') else 'chunks' enriched_data = { chunk_key: chunks, 'domain_data': domain_data, 'chunk_count': len(chunks), 'doc_type': doc_type, 'original_filename': chunks_result.get('original_filename', ''), } # Preserve embeddings if they exist (from OLD pipeline which already has them) if 'embeddings' in chunks_result: enriched_data['embeddings'] = chunks_result['embeddings'] with open(output_file, 'wb') as f: pickle.dump(enriched_data, f) log.info(f"Saved enriched data to {output_file}") # Persist new terms to domain dictionary all_terms = set() term_confidences: dict[str, float] = {} for data in domain_data: all_terms.update(data['domain_terms']) # Max across chunks: same term appearing in multiple chunks keeps the # strongest confidence. Rule-only chunks emit 0.95; LLM chunks emit # exp(min(token_logprobs)) or 0.70 on alignment failure. for t, c in data.get('domain_confidences', {}).items(): term_confidences[t] = max(term_confidences.get(t, 0.0), c) # Infer dominant category from chunk categories from collections import Counter category_counts_list = Counter(data['domain_category'] for data in domain_data) dominant_category = category_counts_list.most_common(1)[0][0] if category_counts_list else 'general' # Get document_id from chunks_result metadata (passed via XCom) document_id = chunks_result.get('document_id') # May be None for some pipelines if document_id: from uuid import UUID document_id = UUID(document_id) if isinstance(document_id, str) else document_id # Persist new terms to dictionary (skip duplicates, translate via LLM) if all_terms: try: new_count = extractor.persist_new_terms( terms=list(all_terms), domain_category=dominant_category, document_id=document_id, tenant_id=tenant_id, node_id=node_id, confidence_by_term=term_confidences, ) log.info(f"Persisted {new_count} new terms to domain dictionary") except Exception as e: log.warning(f"Failed to persist terms to dictionary: {e}. Continuing with pipeline.") # Return file path via XCom (preserve 'adapted' flag and document_id) result = { 'data_file': output_file, 'chunk_count': len(chunks), 'doc_type': doc_type, 'original_filename': chunks_result.get('original_filename', ''), } # Preserve 'adapted' flag from OLD pipeline if present if chunks_result.get('adapted'): result['adapted'] = True # Propagate document_id for downstream tasks and scoped cleanup if chunks_result.get('document_id'): result['document_id'] = chunks_result['document_id'] return result @task.external_python(python=EXTERNAL_PYTHON) def tag_records_task(chunks_result: Dict[str, Any]) -> Dict[str, Any]: """ Assign semantic tags to records using SemanticTagger. This task runs AFTER domain_terms extraction and BEFORE embedding generation. Tags are added to each record for downstream filtering in Milvus. Tags include: - req:functional, req:non_functional, req:constraint, etc. (ISO 29148) - risk:safety, risk:security, risk:assumption - info:background, info:scope - topic:{keyword} (LLM-extracted domain topics) Args: chunks_result: Result from domain_terms task (contains data_file path) Returns: Dict with file path to tagged data and metadata """ from src.extraction_v2.semantic_tagger import SemanticTagger import asyncio import logging import pickle import os from uuid import uuid4 log = logging.getLogger(__name__) if not chunks_result: log.warning("No chunks result provided") return { 'chunk_count': 0, 'doc_type': 'excel', 'original_filename': '', } # Preserve original metadata before loading pickle original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'} data_file = chunks_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading chunk data from {data_file}") with open(data_file, 'rb') as f: chunks_result = pickle.load(f) # Restore important metadata from XCom for key in ['adapted', 'document_id', 'filename']: if key in original_metadata and key not in chunks_result: chunks_result[key] = original_metadata[key] # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up chunk data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up chunk data file {data_file}: {e}") # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records') chunks = chunks_result.get('chunks') or chunks_result.get('records', []) if not chunks: log.info("No chunks to tag") return { 'chunk_count': 0, 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } log.info(f"Tagging {len(chunks)} chunks with semantic tags") try: # Initialize tagger tagger = SemanticTagger() # Prepare records for tagging (tagger expects specific format) records_for_tagging = [] for idx, chunk in enumerate(chunks): record = { 'id': idx, 'text': '', 'parent_header': '', 'sheet_name': '', } # Extract text content if isinstance(chunk, dict): content = chunk.get('content') or chunk.get('text', '') if isinstance(content, dict): # Excel record: format as key: value pairs record['text'] = ' | '.join(f"{k}: {v}" for k, v in content.items() if v) else: record['text'] = str(content) if content else '' # Extract context from metadata source = chunk.get('_source', {}) metadata = chunk.get('metadata', {}) record['parent_header'] = metadata.get('parent_header', '') record['sheet_name'] = source.get('sheet', '') else: record['text'] = str(chunk) if record['text'].strip(): records_for_tagging.append(record) if not records_for_tagging: log.info("No valid text content to tag") else: # Tag records concurrently (modifies in-place) asyncio.run(tagger.tag_records_async(records_for_tagging)) # Copy tags back to original chunks for record in records_for_tagging: idx = record['id'] tags = record.get('tags', []) # TEXTIQ-569: SemanticTagger also emits verified_tag_names # post-inheritance; carry it forward to Milvus. verified_tag_names = record.get('verified_tag_names', []) if idx < len(chunks): if isinstance(chunks[idx], dict): chunks[idx]['tags'] = tags chunks[idx]['verified_tag_names'] = verified_tag_names elif hasattr(chunks[idx], 'tags'): chunks[idx].tags = tags if hasattr(chunks[idx], 'verified_tag_names'): chunks[idx].verified_tag_names = verified_tag_names # Count tags assigned total_tags = sum(len(r.get('tags', [])) for r in records_for_tagging) log.info(f"Assigned {total_tags} total tags to {len(records_for_tagging)} chunks") except Exception as e: log.warning(f"Tagging failed: {e}. Continuing with empty tags.") # On failure, ensure all chunks have empty tags list for chunk in chunks: if isinstance(chunk, dict) and 'tags' not in chunk: chunk['tags'] = [] # Save tagged data to temporary file (include document_id for scoped cleanup) data_dir = '/opt/airflow/logs/tag_data' os.makedirs(data_dir, exist_ok=True) doc_id = chunks_result.get('document_id', uuid4().hex) output_file = os.path.join(data_dir, f'tagged_{doc_id}_{uuid4().hex[:8]}.pkl') # Determine the correct key name chunk_key = 'records' if chunks_result.get('adapted') or chunks_result.get('records') else 'chunks' tagged_data = { chunk_key: chunks, 'domain_data': chunks_result.get('domain_data', []), 'embeddings': chunks_result.get('embeddings', []), 'chunk_count': len(chunks), 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } with open(output_file, 'wb') as f: pickle.dump(tagged_data, f) log.info(f"Saved tagged data to {output_file}") # Return file path via XCom result = { 'data_file': output_file, 'chunk_count': len(chunks), 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } # Preserve 'adapted' flag if present if chunks_result.get('adapted'): result['adapted'] = True # Propagate document_id for downstream tasks and scoped cleanup if chunks_result.get('document_id'): result['document_id'] = chunks_result['document_id'] return result @task.external_python(python=EXTERNAL_PYTHON) def extract_domain_terms_nodes_task( nodes_result: Dict[str, Any], tenant_id: str, node_id: str, ) -> Dict[str, Any]: """ Extract domain terms from LlamaIndex nodes for domain-aware RAG. Works with hierarchical nodes from hierarchical_chunk_*_task. Enriches node.metadata with domain_category and domain_terms. Args: nodes_result: Result from hierarchical_chunk task (contains nodes_file path) Returns: Dict with enriched nodes_file path and metadata """ from src.extraction_v2.domain_term_extractor import DomainTermExtractor import asyncio import logging import pickle import os from uuid import UUID from collections import Counter log = logging.getLogger(__name__) if not nodes_result: log.warning("No nodes result provided") return {'nodes_file': '', 'node_count': 0} # Load nodes from pickle file nodes_file = nodes_result.get('nodes_file') if not nodes_file or not os.path.exists(nodes_file): log.warning(f"Nodes file not found: {nodes_file}") return {'nodes_file': '', 'node_count': 0} log.info(f"Loading nodes from {nodes_file}") with open(nodes_file, 'rb') as f: nodes_data = pickle.load(f) nodes = nodes_data.get('nodes', []) if not nodes: log.info("No nodes to process") return nodes_result log.info(f"Extracting domain terms from {len(nodes)} nodes concurrently") # Initialize extractor and run async extraction extractor = DomainTermExtractor() doc_type = nodes_data.get('doc_type', 'word') all_terms, category_counts, term_confidences = asyncio.run( extractor.extract_from_nodes_async(nodes, doc_type=doc_type) ) log.info(f"Domain term extraction complete: {len(nodes)} nodes processed") log.info(f"Category distribution: {category_counts}") # Persist new terms to dictionary dominant_category = Counter(category_counts).most_common(1)[0][0] if category_counts else 'general' document_id = nodes_data.get('document_id') if document_id: document_id = UUID(document_id) if isinstance(document_id, str) else document_id if all_terms: try: new_count = extractor.persist_new_terms( terms=list(all_terms), domain_category=dominant_category, document_id=document_id, tenant_id=tenant_id, node_id=node_id, confidence_by_term=term_confidences, ) log.info(f"Persisted {new_count} new terms to domain dictionary") except Exception as e: log.warning(f"Failed to persist terms: {e}") # Save enriched nodes back nodes_data['nodes'] = nodes with open(nodes_file, 'wb') as f: pickle.dump(nodes_data, f) log.info(f"Saved enriched nodes to {nodes_file}") return { 'nodes_file': nodes_file, 'node_count': len(nodes), 'leaf_count': nodes_data.get('leaf_count', 0), 'doc_type': doc_type, 'original_filename': nodes_data.get('original_filename', ''), 'document_id': nodes_data.get('document_id', ''), } @task.external_python(python=EXTERNAL_PYTHON) def tag_records_nodes_task( nodes_result: Dict[str, Any], ) -> Dict[str, Any]: """ Assign semantic tags to LlamaIndex nodes using SemanticTagger. Works with hierarchical nodes from hierarchical_chunk_*_task. Enriches node.metadata with tags list. Tags include: - req:functional, req:non_functional, req:constraint, etc. (ISO 29148) - risk:safety, risk:security, risk:assumption - info:background, info:scope - topic:{keyword} (LLM-extracted domain topics) Args: nodes_result: Result from extract_domain_terms_nodes_task Returns: Dict with tagged nodes_file path and metadata """ from src.extraction_v2.semantic_tagger import SemanticTagger import asyncio import logging import pickle import os log = logging.getLogger(__name__) if not nodes_result: log.warning("No nodes result provided") return {'nodes_file': '', 'node_count': 0} # Load nodes from pickle file nodes_file = nodes_result.get('nodes_file') if not nodes_file or not os.path.exists(nodes_file): log.warning(f"Nodes file not found: {nodes_file}") return {'nodes_file': '', 'node_count': 0} log.info(f"Loading nodes from {nodes_file}") with open(nodes_file, 'rb') as f: nodes_data = pickle.load(f) nodes = nodes_data.get('nodes', []) if not nodes: log.info("No nodes to tag") return nodes_result log.info(f"Tagging {len(nodes)} nodes with semantic tags concurrently") try: # Initialize tagger tagger = SemanticTagger() # Prepare records for tagging records_for_tagging = [] for idx, node in enumerate(nodes): text = node.text if hasattr(node, 'text') else str(node) metadata = node.metadata if hasattr(node, 'metadata') else {} if text.strip(): records_for_tagging.append({ 'id': idx, 'text': text, 'parent_header': metadata.get('parent_header', ''), 'sheet_name': metadata.get('sheet_name', ''), }) if records_for_tagging: # Tag records concurrently (modifies in-place) asyncio.run(tagger.tag_records_async(records_for_tagging)) # Copy tags back to nodes (convert TagAssignment objects to dicts for JSON serialization) for record in records_for_tagging: idx = record['id'] tags = record.get('tags', []) if idx < len(nodes): # TagAssignment is a Pydantic model - convert to dict serialized_tags = [ t.model_dump() if hasattr(t, 'model_dump') else t for t in tags ] if tags else [] nodes[idx].metadata['tags'] = serialized_tags # TEXTIQ-569: SemanticTagger emits `verified_tag_names` on # each record post-inheritance; carry it onto the node so # LlamaIndex writes it into the Milvus dynamic envelope. nodes[idx].metadata['verified_tag_names'] = record.get( 'verified_tag_names', [] ) total_tags = sum(len(r.get('tags', [])) for r in records_for_tagging) log.info(f"Assigned {total_tags} total tags to {len(records_for_tagging)} nodes") else: log.info("No valid text content to tag") except Exception as e: log.warning(f"Tagging failed: {e}. Continuing with empty tags.") for node in nodes: if 'tags' not in node.metadata: node.metadata['tags'] = [] if 'verified_tag_names' not in node.metadata: node.metadata['verified_tag_names'] = [] # Save tagged nodes back nodes_data['nodes'] = nodes with open(nodes_file, 'wb') as f: pickle.dump(nodes_data, f) log.info(f"Saved tagged nodes to {nodes_file}") return { 'nodes_file': nodes_file, 'node_count': len(nodes), 'leaf_count': nodes_data.get('leaf_count', 0), 'doc_type': nodes_data.get('doc_type', 'word'), 'original_filename': nodes_data.get('original_filename', ''), 'document_id': nodes_data.get('document_id', ''), } @task.external_python(python=EXTERNAL_PYTHON) def generate_embeddings_task( chunks_result: Dict[str, Any], batch_size: int = 16 ) -> Dict[str, Any]: """ Generate embeddings for content chunks. Stores large data (chunks + embeddings + domain_data) to a temporary file to avoid XCom size limits. Returns file path via XCom instead of the actual data. Args: chunks_result: Result from extract_domain_terms task (contains data_file path) batch_size: Number of chunks to process in each batch (max 100 for Azure OpenAI) Returns: Dict with file path to embedded data and metadata """ from src.knowledge.embeddings import EmbeddingService import asyncio import logging import pickle import tempfile import os from uuid import uuid4 log = logging.getLogger(__name__) if not chunks_result: log.warning("No chunks result provided") return { 'chunk_count': 0, 'doc_type': 'excel', 'original_filename': '', } # Load data from pickle file if data_file path is provided # Preserve original metadata before loading pickle original_metadata = {k: v for k, v in chunks_result.items() if k != 'data_file'} data_file = chunks_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading chunk data from {data_file}") with open(data_file, 'rb') as f: chunks_result = pickle.load(f) # Restore important metadata from XCom (like 'adapted' flag) for key in ['adapted', 'document_id', 'filename']: if key in original_metadata and key not in chunks_result: chunks_result[key] = original_metadata[key] # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up chunk data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up chunk data file {data_file}: {e}") # Handle both NEW pipeline format ('chunks') and OLD pipeline format ('records') chunks = chunks_result.get('chunks') or chunks_result.get('records', []) domain_data = chunks_result.get('domain_data', []) # Load domain data if not chunks: log.info("No chunks to embed") return { 'chunk_count': 0, 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } log.info(f"Generating embeddings for {len(chunks)} chunks (batch_size={batch_size})") # Extract text from chunks for embedding texts = [] valid_chunks = [] for idx, chunk in enumerate(chunks): if isinstance(chunk, dict): text = chunk.get('text') or chunk.get('content') or chunk.get('raw_content', '') text = str(text) if text else '' if text.strip(): texts.append(text) valid_chunks.append(chunk) else: log.warning(f"Skipping chunk {idx} with no valid text content") else: text = str(chunk) if text.strip(): texts.append(text) valid_chunks.append(chunk) chunks = valid_chunks if not texts: log.warning("No valid texts to embed after filtering empty content") return { 'chunk_count': 0, 'doc_type': chunks_result.get('doc_type', 'excel'), 'original_filename': chunks_result.get('original_filename', ''), } skipped_count = len(chunks_result.get('chunks', [])) - len(texts) if skipped_count > 0: log.info(f"Filtered out {skipped_count} chunks with empty content, {len(texts)} remaining") # Generate unique run ID for checkpointing run_id = uuid4().hex[:8] checkpoint_file = f'/opt/airflow/logs/embedding_checkpoint_{run_id}.pkl' async def embed_with_checkpoints(): # Initialize embedding service # Note: EmbeddingService creates and closes clients per batch internally service = EmbeddingService() all_embeddings = [] start_batch = 0 # Check for existing checkpoint (retry scenario) if os.path.exists(checkpoint_file): log.info(f"Found checkpoint, resuming from {checkpoint_file}...") with open(checkpoint_file, 'rb') as f: checkpoint = pickle.load(f) all_embeddings = checkpoint['embeddings'] start_batch = checkpoint['next_batch'] log.info(f"Resuming from batch {start_batch}, {len(all_embeddings)} embeddings recovered") # Generate embeddings in batches (max 100 per Azure OpenAI call) effective_batch_size = min(batch_size, 100) total_batches = (len(texts) - 1) // effective_batch_size + 1 for i in range(start_batch * effective_batch_size, len(texts), effective_batch_size): batch = texts[i:i + effective_batch_size] batch_num = i // effective_batch_size + 1 log.info(f"Processing batch {batch_num}/{total_batches}") batch_embeddings = await service.embed_batch(batch) all_embeddings.extend(batch_embeddings) # Save checkpoint after each batch checkpoint = { 'embeddings': all_embeddings, 'next_batch': batch_num, 'total_batches': total_batches, } with open(checkpoint_file, 'wb') as f: pickle.dump(checkpoint, f) log.info(f"Checkpoint saved: {len(all_embeddings)}/{len(texts)} embeddings") # Clean up checkpoint on success if os.path.exists(checkpoint_file): os.remove(checkpoint_file) return all_embeddings all_embeddings = asyncio.run(embed_with_checkpoints()) log.info(f"Generated {len(all_embeddings)} embeddings") # Store large data to temporary file to avoid XCom size limits # Use /opt/airflow/logs as shared storage (mounted volume) data_dir = '/opt/airflow/logs/embedding_data' os.makedirs(data_dir, exist_ok=True) doc_id = chunks_result.get('document_id', uuid4().hex) data_file = os.path.join(data_dir, f'embeddings_{doc_id}_{uuid4().hex[:8]}.pkl') embedded_data = { 'chunks': chunks, 'embeddings': all_embeddings, 'domain_data': domain_data, # Pass through domain data 'chunk_count': len(chunks), 'doc_type': chunks_result.get('doc_type'), 'original_filename': chunks_result.get('original_filename', ''), } with open(data_file, 'wb') as f: pickle.dump(embedded_data, f) file_size_mb = os.path.getsize(data_file) / (1024 * 1024) log.info(f"Saved embedded data to {data_file} ({file_size_mb:.2f} MB)") # Return only metadata and file path via XCom (small payload) result = { 'data_file': data_file, 'chunk_count': len(chunks), 'doc_type': chunks_result.get('doc_type'), 'original_filename': chunks_result.get('original_filename', ''), } # Propagate document_id for downstream tasks and scoped cleanup if chunks_result.get('document_id'): result['document_id'] = chunks_result['document_id'] return result @task.external_python(python=EXTERNAL_PYTHON) def generate_embeddings_chunks( detailed_pipeline_result: Dict[str, Any] ) -> Dict[str, Any]: """ Generate embeddings for DETAILED pipeline chunks. Converts detailed extraction chunks (images, shapes, flowcharts, text) to ExtractedRecord structure and generates embeddings using LegacyToMilvusAdapter. Args: detailed_pipeline_result: Result from parse_excel_detailed task (contains data_file path) Returns: Dict with embedded chunks data_file path and metadata """ # All imports must be inside for external_python import sys sys.path.insert(0, '/opt/airflow') from src.adapters.legacy_to_milvus_adapter import LegacyToMilvusAdapter from uuid import UUID import logging import pickle import os import asyncio log = logging.getLogger(__name__) if not detailed_pipeline_result: log.warning("No DETAILED pipeline result provided") return {'record_count': 0, 'adapted': False} # Load chunks from pickle file data_file = detailed_pipeline_result.get('data_file') if not data_file or not os.path.exists(data_file): log.error(f"DETAILED pipeline data file not found: {data_file}") return {'record_count': 0, 'adapted': False} log.info(f"Loading DETAILED pipeline chunks from {data_file}") with open(data_file, 'rb') as f: chunks = pickle.load(f) # Clean up DETAILED pipeline file after loading try: os.remove(data_file) log.info(f"Cleaned up DETAILED pipeline data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up DETAILED pipeline file {data_file}: {e}") if not chunks: log.info("No chunks to adapt") return {'record_count': 0, 'adapted': True} log.info(f"Adapting {len(chunks)} DETAILED pipeline chunks to schema") # Get metadata from result document_id = UUID(detailed_pipeline_result.get('document_id')) filename = detailed_pipeline_result.get('filename', 'document.xlsx') # Convert chunks to dicts if they're dataclass instances chunks_as_dicts = [] for chunk in chunks: if hasattr(chunk, '__dict__'): # Dataclass instance - convert to dict chunk_dict = { 'element_type': getattr(chunk, 'element_type', 'Text'), 'text': getattr(chunk, 'text', ''), 'keywords_text': getattr(chunk, 'keywords_text', ''), 'created_at': getattr(chunk, 'created_at', ''), 'metadata': getattr(chunk, 'metadata', {}), } chunks_as_dicts.append(chunk_dict) elif isinstance(chunk, dict): chunks_as_dicts.append(chunk) else: log.warning(f"Skipping unknown chunk type: {type(chunk)}") # Initialize adapter adapter = LegacyToMilvusAdapter() # Adapt chunks (async) async def adapt(): return await adapter.adapt_chunks( legacy_chunks=chunks_as_dicts, document_id=document_id, filename=filename, ) adapted_result = asyncio.run(adapt()) log.info(f"Adapted {adapted_result['record_count']} records with embeddings") # Save adapted result to pickle file adapted_file = f"/opt/airflow/logs/adapted_{document_id}.pkl" with open(adapted_file, 'wb') as f: pickle.dump(adapted_result, f) log.info(f"Saved adapted result to {adapted_file}") # Return metadata via XCom return { 'data_file': adapted_file, 'document_id': str(document_id), 'filename': filename, 'record_count': adapted_result['record_count'], 'adapted': True, } @task.external_python(python=EXTERNAL_PYTHON) def merge_pipeline_results( new_pipeline_result: Dict[str, Any], adapted_old_result: Dict[str, Any] ) -> Dict[str, Any]: """ Merge NEW pipeline results with DETAILED pipeline results. Combines embeddings and records from both pipelines into a single dataset ready for Milvus storage. Args: new_pipeline_result: Result from generate_embeddings_large_tables (NEW pipeline) adapted_old_result: Result from generate_embeddings_chunks (DETAILED pipeline) Returns: Dict with merged data_file path and combined metadata """ import logging import pickle import os log = logging.getLogger(__name__) # Load NEW pipeline data new_chunks = [] new_embeddings = [] new_domain_data = [] if new_pipeline_result and new_pipeline_result.get('chunk_count', 0) > 0: new_data_file = new_pipeline_result.get('data_file', '') if new_data_file and os.path.exists(new_data_file): log.info(f"Loading NEW pipeline data from {new_data_file}") with open(new_data_file, 'rb') as f: new_data = pickle.load(f) new_chunks = new_data.get('chunks', []) new_embeddings = new_data.get('embeddings', []) new_domain_data = new_data.get('domain_data', []) # Clean up temp file try: os.remove(new_data_file) log.info(f"Cleaned up NEW pipeline file: {new_data_file}") except Exception as e: log.warning(f"Failed to clean up NEW pipeline file: {e}") else: log.info("NEW pipeline returned 0 chunks (no large tables found)") else: log.info("NEW pipeline skipped or returned no results (no large tables found)") # Load adapted DETAILED pipeline data detailed_records = [] detailed_embeddings = [] detailed_domain_data = [] if adapted_old_result and adapted_old_result.get('adapted'): detailed_data_file = adapted_old_result.get('data_file') if detailed_data_file and os.path.exists(detailed_data_file): log.info(f"Loading adapted DETAILED pipeline data from {detailed_data_file}") with open(detailed_data_file, 'rb') as f: detailed_data = pickle.load(f) detailed_records = detailed_data.get('records', []) detailed_embeddings = detailed_data.get('embeddings', []) detailed_domain_data = detailed_data.get('domain_data', []) log.info(f"DEBUG merge: DETAILED pipeline loaded - records={len(detailed_records)}, embeddings={len(detailed_embeddings)}, domain_data={len(detailed_domain_data)}") # Clean up temp file try: os.remove(detailed_data_file) log.info(f"Cleaned up adapted DETAILED pipeline file: {detailed_data_file}") except Exception as e: log.warning(f"Failed to clean up adapted file: {e}") log.info( f"Merging pipelines: NEW ({len(new_chunks)} records) + " f"DETAILED ({len(detailed_records)} records)" ) # Merge records, embeddings, and domain data merged_chunks = new_chunks + detailed_records merged_embeddings = new_embeddings + detailed_embeddings merged_domain_data = new_domain_data + detailed_domain_data if not merged_chunks: log.warning("No records to merge from either pipeline") return { 'chunk_count': 0, 'merged': False, 'doc_type': 'excel', 'original_filename': '', } log.info(f"Merged total: {len(merged_chunks)} records with embeddings and domain data") # Get metadata from whichever result is available doc_type = (new_pipeline_result.get('doc_type') if new_pipeline_result else 'excel') original_filename = (new_pipeline_result.get('original_filename') if new_pipeline_result else adapted_old_result.get('filename', '')) # Save merged data to pickle file log.info(f"DEBUG merge: About to save merged_data with domain_data length = {len(merged_domain_data)}") merged_data = { 'chunks': merged_chunks, 'embeddings': merged_embeddings, 'domain_data': merged_domain_data, 'chunk_count': len(merged_chunks), 'doc_type': doc_type, 'original_filename': original_filename, } # Use document_id from either result document_id = (new_pipeline_result.get('document_id') if new_pipeline_result and 'document_id' in new_pipeline_result else adapted_old_result.get('document_id', 'merged')) merged_file = f"/opt/airflow/logs/merged_{document_id}.pkl" with open(merged_file, 'wb') as f: pickle.dump(merged_data, f) log.info(f"Saved merged result to {merged_file}") # Return metadata via XCom return { 'data_file': merged_file, 'chunk_count': len(merged_chunks), 'doc_type': doc_type, 'original_filename': original_filename, 'merged': True, 'new_count': len(new_chunks), 'detailed_count': len(detailed_records), } @task.external_python(python=EXTERNAL_PYTHON) def upload_images_to_seaweedfs( document_id: str, detailed_pipeline_result: Dict[str, Any], tenant_id: str, node_id: str, ) -> Dict[str, Any]: """ Upload images from DETAILED pipeline to SeaweedFS. Processes chunks with image_base64 data, uploads to SeaweedFS, and replaces base64 with image_url and image_s3_key. Args: document_id: Document UUID string detailed_pipeline_result: Result from parse_excel_detailed task (contains data_file path) tenant_id: Tenant UUID string for S3 key scoping node_id: Node UUID string for S3 key scoping Returns: Updated result dict with images uploaded and metadata updated """ from uuid import UUID from src.services.seaweedfs_service import SeaweedFSService import logging import pickle import os log = logging.getLogger(__name__) doc_uuid = UUID(document_id) if not detailed_pipeline_result: log.info("No detailed pipeline result provided") return { 'uploaded_count': 0, 'success': True, 'data_file': None, 'document_id': document_id, 'filename': '', 'chunk_count': 0, } # Load data from pickle file if data_file path is provided data_file = detailed_pipeline_result.get('data_file') if not data_file or not os.path.exists(data_file): log.warning("No data_file found in detailed_pipeline_result") return { 'uploaded_count': 0, 'success': True, 'data_file': None, 'document_id': detailed_pipeline_result.get('document_id'), 'filename': detailed_pipeline_result.get('filename', ''), 'chunk_count': 0, } log.info(f"Loading detailed parse result from {data_file}") with open(data_file, 'rb') as f: parse_result = pickle.load(f) # parse_to_markdown returns {sheets, image_chunks, flowchart_chunks, color_map, cell_values_map} image_chunks = parse_result.get('image_chunks', []) if not image_chunks: log.info("No image chunks to upload") return { 'uploaded_count': 0, 'success': True, 'data_file': data_file, 'document_id': detailed_pipeline_result.get('document_id'), 'filename': detailed_pipeline_result.get('filename', ''), } # Access metadata via attribute (ChunkedElementDTO) or dict key def get_meta(chunk): if hasattr(chunk, 'metadata'): return chunk.metadata elif isinstance(chunk, dict): return chunk.get('metadata', {}) return {} # Count images with base64 data image_count = sum( 1 for chunk in image_chunks if 'image_base64' in get_meta(chunk) ) log.info(f"Found {image_count} images to upload") if image_count == 0: log.info("No images with base64 data found") return { 'uploaded_count': 0, 'success': True, 'data_file': data_file, 'document_id': detailed_pipeline_result.get('document_id'), 'filename': detailed_pipeline_result.get('filename', ''), } # Upload images one at a time, freeing base64 data immediately after upload log.info("Uploading images to SeaweedFS...") service = SeaweedFSService() uploaded_count = 0 for chunk in image_chunks: metadata = get_meta(chunk) if 'image_base64' not in metadata: continue image_number = metadata.get('image_number', 1) sheet_name = metadata.get('sheet_name', 'Sheet1') mime_type = metadata.get('mime_type', 'image/png') result = service.upload_image( image_base64=metadata['image_base64'], document_id=doc_uuid, image_number=image_number, sheet_name=sheet_name, mime_type=mime_type, tenant_id=tenant_id, node_id=node_id, ) if result['success']: metadata['image_s3_key'] = result['image_s3_key'] metadata['image_url'] = result['image_url'] # Remove base64 immediately to free memory del metadata['image_base64'] uploaded_count += 1 log.info(f"Uploaded image {image_number}: {result['image_s3_key']}") else: # Still remove base64 to avoid OOM on large files metadata.pop('image_base64', None) log.warning(f"Failed to upload image {image_number}: {result.get('error', 'Unknown')}") log.info(f"Successfully uploaded {uploaded_count}/{image_count} images to SeaweedFS") # Save updated parse result back (image base64 removed, URLs added) with open(data_file, 'wb') as f: pickle.dump(parse_result, f) log.info(f"Updated parse result saved to {data_file}") return { 'uploaded_count': uploaded_count, 'total_images': image_count, 'success': True, 'data_file': data_file, 'document_id': detailed_pipeline_result.get('document_id'), 'filename': detailed_pipeline_result.get('filename'), } @task.external_python(python=EXTERNAL_PYTHON) def store_vectors_task( document_id: str, embedded_result: Dict[str, Any], pg_records_result: Dict[str, Any] = None, tenant_id: str = "", node_id: str = "", ) -> Dict[str, Any]: """ Store vectors in Milvus. Args: document_id: Document UUID string embedded_result: Result from generate_embeddings task (may contain data_file path) pg_records_result: Result from save_extracted_records_task with record_ids_file. If provided, uses PostgreSQL record IDs for Milvus linking. If None, generates new UUIDs (legacy behavior, no PostgreSQL linking). tenant_id: Tenant UUID string for Milvus data isolation node_id: Node UUID string for Milvus data isolation Returns: Storage result dict """ from uuid import UUID, uuid4 from src.knowledge.vector_store import VectorStore from src.resilience.breakers import async_breaker_call, milvus_breaker import asyncio import logging import pickle import os log = logging.getLogger(__name__) doc_uuid = UUID(document_id) if not embedded_result: log.warning("No embedded result provided") return {'stored_count': 0, 'success': False} # Load data from file if data_file path is provided (large data case) data_file = embedded_result.get('data_file', '') if data_file and os.path.exists(data_file): log.info(f"Loading embedded data from {data_file}") with open(data_file, 'rb') as f: embedded_result = pickle.load(f) # Clean up the temp file after loading try: os.remove(data_file) log.info(f"Cleaned up temp file: {data_file}") except Exception as e: log.warning(f"Failed to clean up temp file {data_file}: {e}") chunks = embedded_result.get('chunks', []) embeddings = embedded_result.get('embeddings', []) domain_data = embedded_result.get('domain_data', []) # Extract domain data # DEBUG: Log domain_data status log.info(f"DEBUG store_vectors: embedded_result keys = {list(embedded_result.keys())}") log.info(f"DEBUG store_vectors: domain_data type = {type(domain_data)}, len = {len(domain_data) if domain_data else 0}") if not chunks or not embeddings: log.info("No chunks or embeddings to store") return {'stored_count': 0, 'success': True} if len(chunks) != len(embeddings): log.error(f"Mismatch: {len(chunks)} chunks vs {len(embeddings)} embeddings") return {'stored_count': 0, 'success': False} # Log domain data statistics if domain_data: category_counts = {} for data in domain_data: category = data.get('domain_category', 'general') category_counts[category] = category_counts.get(category, 0) + 1 log.info(f"Domain category distribution: {category_counts}") log.info(f"Storing {len(chunks)} vectors to Milvus") # Load PostgreSQL record IDs if available (for proper PostgreSQL-Milvus linking) record_ids = None record_ids_file = pg_records_result.get('record_ids_file', '') if pg_records_result else '' if record_ids_file and os.path.exists(record_ids_file): log.info(f"Loading PostgreSQL record IDs from {record_ids_file}") with open(record_ids_file, 'rb') as f: record_ids = pickle.load(f) # Clean up the record IDs file try: os.remove(record_ids_file) log.info(f"Cleaned up record IDs file: {record_ids_file}") except Exception as e: log.warning(f"Failed to clean up record IDs file {record_ids_file}: {e}") if len(record_ids) != len(chunks): log.error( f"Mismatch: {len(record_ids)} PostgreSQL record IDs vs {len(chunks)} chunks. " "Falling back to generated UUIDs (no PostgreSQL linking)." ) record_ids = None else: log.info(f"Using {len(record_ids)} PostgreSQL record IDs for Milvus linking") # Fall back to generating new UUIDs if no PostgreSQL IDs available if record_ids is None: log.warning("No PostgreSQL record IDs available - generating new UUIDs (no PostgreSQL linking)") record_ids = [uuid4() for _ in chunks] # Determine document type from embedded_result (passed through from chunk task) doc_type = embedded_result.get('doc_type') or 'text' original_filename = embedded_result.get('original_filename', '') # Convert Excel dict chunks to ExtractedRecord objects for VectorStore compatibility if doc_type == 'excel': from src.extraction.models import ExtractedRecord records = [] for chunk in chunks: if isinstance(chunk, dict): content = chunk.get('content', {}) source = chunk.get('_source', {}) # Ensure all required _source fields are present complete_source = { 'document_id': source.get('document_id', str(doc_uuid)), 'filename': source.get('filename', ''), 'sheet': source.get('sheet', 'Sheet1'), 'table_id': source.get('table_id', 0), 'row': source.get('row', 0), 'col_range': source.get('col_range', ''), } # Skip if content is empty if not content: log.debug(f"Skipping chunk with empty content") continue record = ExtractedRecord( content=content, headers=chunk.get('headers', list(content.keys())), _source=complete_source, _merge_metadata=chunk.get('_merge_metadata'), resolved_content=chunk.get('resolved_content'), _legacy_metadata=chunk.get('_legacy_metadata'), tags=chunk.get('tags', []), ) records.append(record) else: records.append(chunk) else: records = chunks async def store(): store = VectorStore() count = await async_breaker_call( milvus_breaker, store.index_with_vectors, records=records, vectors=embeddings, document_id=doc_uuid, record_ids=record_ids, document_type=doc_type, filename=original_filename, domain_data=domain_data if domain_data else None, tenant_id=tenant_id, node_id=node_id, ) return count stored_count = asyncio.run(store()) log.info(f"Successfully stored {stored_count} vectors") return { 'stored_count': stored_count, 'success': True, } @task.external_python(python=EXTERNAL_PYTHON) def save_document_metadata_task( document_id: str, fetch_result: Dict[str, Any], validation_result: Dict[str, Any], tenant_id: str, node_id: str, correlation_id: str, ) -> Dict[str, Any]: """ Create Document record in PostgreSQL. Args: document_id: Document UUID string fetch_result: Result from fetch_document task with file info validation_result: Result from validate_event task tenant_id: Tenant UUID string for data isolation node_id: Node UUID string for data isolation correlation_id: Airflow dag_run_id, for pipeline-event correlation Returns: Dict with document_id and success status """ import asyncio import logging from uuid import UUID from datetime import datetime from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession from sqlalchemy.orm import sessionmaker from src.db.models import Document from src.core.config import settings from src.services.pipeline_events import publish_stage_event log = logging.getLogger(__name__) publish_stage_event( "pending", "started", document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, ) # Extract file info filename = fetch_result.get('filename', validation_result.get('object_key', 'unknown')) file_path = validation_result.get('object_key', '') file_size = fetch_result.get('size', 0) log.info(f"Creating document record: {document_id}, filename: {filename}") async def save_document(): # Create async engine engine = create_async_engine( settings.database_url, echo=False, pool_pre_ping=True, ) async_session = sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, ) async with async_session() as session: try: # Set RLS tenant context for row-level security policy from sqlalchemy import text as sa_text if tenant_id: from uuid import UUID as _UUID _tid = str(_UUID(tenant_id)) # validate UUID format await session.execute(sa_text(f"SET LOCAL app.tenant_id = '{_tid}'")) # Upsert: if upload endpoint already created the record (status=pending), # update it to processing. Otherwise insert a new record (legacy CLI flow). from sqlalchemy.dialects.postgresql import insert as pg_insert stmt = pg_insert(Document).values( id=UUID(document_id), filename=filename, file_path=file_path, file_size_bytes=file_size, status='processing', created_at=datetime.utcnow(), tenant_id=UUID(tenant_id) if tenant_id else None, node_id=UUID(node_id) if node_id else None, ).on_conflict_do_update( index_elements=['id'], set_={ 'status': 'processing', }, ) await session.execute(stmt) await session.commit() log.info(f"Document record upserted successfully: {document_id}") return True except Exception as e: await session.rollback() log.error(f"Failed to upsert document record: {e}") raise finally: await engine.dispose() success = asyncio.run(save_document()) if success: publish_stage_event( "processing", "started", document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, ) return { 'document_id': document_id, 'success': success, } @task.external_python(python=EXTERNAL_PYTHON) def save_extracted_records_task( document_id: str, chunks_result: Dict[str, Any], ) -> Dict[str, Any]: """ Save ExtractedRecords to PostgreSQL using StructuredStore. This task MUST run BEFORE store_vectors_task so that record IDs exist in PostgreSQL before Milvus references them. Args: document_id: Document UUID string chunks_result: Result containing chunks with embeddings Returns: Dict with record_ids_file (path to pickled UUIDs) and record_count. The record_ids are the PostgreSQL primary keys that must be used by store_vectors_task to link Milvus vectors to PostgreSQL records. """ import asyncio import logging import pickle import os from uuid import UUID from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession from sqlalchemy.orm import sessionmaker from src.knowledge.structured_store import StructuredStore from src.extraction.models import ExtractedRecord from src.core.config import settings log = logging.getLogger(__name__) # Load chunks from pickle if data_file is provided data_file = chunks_result.get('data_file') if data_file and os.path.exists(data_file): log.info(f"Loading chunks from {data_file}") with open(data_file, 'rb') as f: chunks_result = pickle.load(f) chunks = chunks_result.get('chunks', []) if not chunks: log.warning(f"No chunks to save for document {document_id}") return { 'record_ids_file': '', # Empty string for Airflow 3.x XCom compatibility 'record_count': 0, 'success': True, 'document_id': document_id, } log.info(f"Saving {len(chunks)} records to PostgreSQL for document {document_id}") async def save_records(): # Create async engine engine = create_async_engine( settings.database_url, echo=False, pool_pre_ping=True, ) async_session = sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, ) async with async_session() as session: try: # Convert chunks to ExtractedRecord objects records = [] for chunk_idx, chunk in enumerate(chunks): if isinstance(chunk, dict): # Handle text documents (Word/PDF/PowerPoint) which have 'text' field # vs Excel documents which have 'content' dict if 'text' in chunk and 'content' not in chunk: # Text document chunk - convert to ExtractedRecord format text_content = chunk.get('text', '') if not text_content.strip(): log.debug(f"Skipping empty text chunk at index {chunk_idx}") continue content = {'text': text_content} headers = ['text'] # Use metadata from text chunk metadata = chunk.get('metadata', {}) source = { 'document_id': document_id, 'filename': metadata.get('source_filename', chunks_result.get('original_filename', 'unknown')), 'sheet': 'Page', # Use 'Page' for text documents 'table_id': metadata.get('page_number', 1) or 1, 'row': chunk_idx, 'col_range': '', } else: # Excel chunk - use existing format content = chunk.get('content', {}) headers = chunk.get('headers', list(content.keys()) if content else []) source = chunk.get('_source', {}) # Ensure source has required fields if 'document_id' not in source: source['document_id'] = document_id if 'filename' not in source: source['filename'] = chunks_result.get('original_filename', 'unknown') if 'sheet' not in source: source['sheet'] = 'Sheet1' if 'table_id' not in source: source['table_id'] = 1 if 'row' not in source: source['row'] = chunk_idx if 'col_range' not in source: source['col_range'] = '' # Skip if content is empty if not content: log.debug(f"Skipping chunk {chunk_idx} with empty content") continue record = ExtractedRecord( content=content, headers=headers, _source=source, resolved_content=chunk.get('resolved_content'), tags=chunk.get('tags', []), ) records.append(record) elif isinstance(chunk, ExtractedRecord): records.append(chunk) # Use StructuredStore to save records - now returns StoreResult with IDs store = StructuredStore(session) result = await store.store_records( records=records, document_id=UUID(document_id), ) log.info(f"Saved {result.count} records to PostgreSQL with {len(result.record_ids)} IDs") return result.record_ids, result.count except Exception as e: await session.rollback() log.error(f"Failed to save records: {e}") raise finally: await engine.dispose() record_ids, saved_count = asyncio.run(save_records()) # Save record_ids to pickle file (can be large for many records) # These IDs link PostgreSQL records to Milvus vectors record_ids_file = '' if record_ids: record_ids_dir = '/opt/airflow/logs/record_ids' os.makedirs(record_ids_dir, exist_ok=True) record_ids_file = os.path.join(record_ids_dir, f'record_ids_{document_id}.pkl') with open(record_ids_file, 'wb') as f: pickle.dump(record_ids, f) log.info(f"Saved {len(record_ids)} record IDs to {record_ids_file}") return { 'record_ids_file': record_ids_file, 'record_count': saved_count, 'success': True, 'document_id': document_id, } @task.external_python(python=EXTERNAL_PYTHON) def update_document_status_task( document_id: str, outcome: Dict[str, Any], record_count: int = 0, tenant_id: str = '', node_id: str = '', correlation_id: str = '', ) -> Dict[str, Any]: """ Write terminal document status to PostgreSQL and emit the terminal pipeline event. Status is resolved upstream by ``resolve_terminal_outcome`` (regular @task in the DAG file, which has access to Airflow's task-instance state via context). Args: document_id: Document UUID string outcome: Dict from resolve_terminal_outcome with keys ``status`` ('completed'|'failed') and ``error`` (str | None). record_count: Number of records extracted tenant_id: Tenant UUID string for RLS context + event routing node_id: Node UUID string for event routing correlation_id: Airflow dag_run_id, for event correlation Returns: Dict with document_id, status, success """ import asyncio import logging from uuid import UUID from datetime import datetime from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession from sqlalchemy.orm import sessionmaker from sqlalchemy import update, text as sa_text from src.db.models import Document from src.core.config import settings from src.services.pipeline_events import publish_stage_event log = logging.getLogger(__name__) outcome = outcome or {} status = outcome.get('status', 'failed') error_msg = outcome.get('error') if status != 'completed' else None if status != 'completed' and not error_msg: error_msg = 'terminal outcome unavailable' log.info(f"Updating document {document_id} status to {status} (error={error_msg})") async def update_document(): engine = create_async_engine( settings.database_url, echo=False, pool_pre_ping=True, ) async_session = sessionmaker( engine, class_=AsyncSession, expire_on_commit=False, ) async with async_session() as session: try: if tenant_id: _tid = str(UUID(tenant_id)) # validate UUID format await session.execute(sa_text(f"SET LOCAL app.tenant_id = '{_tid}'")) stmt = ( update(Document) .where(Document.id == UUID(document_id)) .values( status=status, record_count=record_count, processed_at=datetime.utcnow() if status == 'completed' else None, ) ) await session.execute(stmt) await session.commit() log.info(f"Document {document_id} updated successfully with status={status}") return True except Exception as e: await session.rollback() log.error(f"Failed to update document: {e}") raise finally: await engine.dispose() success = asyncio.run(update_document()) publish_stage_event( 'completed' if status == 'completed' else 'failed', 'succeeded' if status == 'completed' else 'failed', document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, **({'error': error_msg} if error_msg else {}), ) return { 'document_id': document_id, 'status': status, 'success': success, } @task.external_python(python=EXTERNAL_PYTHON) def hierarchical_chunk_excel_detailed_task( document_id: str, parse_result: Dict[str, Any], chunk_sizes: List[int] = None, ) -> Dict[str, Any]: """ Hierarchically chunk Excel Detailed pipeline output using LlamaIndex. Accepts per-sheet markdown from parse_to_markdown. Creates one LlamaIndex Document per sheet (preserving sheet_name in metadata), then applies a 2-level hierarchy via the mermaid-aware Excel parser map: MarkdownNodeParser -> sections (parents, by ATX headers) MermaidAwareLineSplitter(512)-> leaves (line-based, mermaid atomic) Mermaid fences stay in-context (no orphan leaves); pure-prose leaves stay <= 512 tokens, mermaid-bearing leaves may grow up to the embedding model context. Header-only nodes are pruned post-parse. Image and flowchart content remain flat leaf-level TextNodes. Args: document_id: Document UUID string parse_result: Result from upload_images_to_seaweedfs (contains data_file path to pickle with {sheets, image_chunks, flowchart_chunks, color_map, cell_values_map}) chunk_sizes: DEPRECATED -- ignored. Excel branch now uses a fixed 2-level mermaid-aware hierarchy. Kept for back-compat with callers that pass it. Returns: Dict with nodes_file path and metadata """ import logging import pickle import os from uuid import UUID from llama_index.core import Document from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes from src.core.config import Settings from src.extraction_v2.mermaid_chunking import ( EXCEL_PARSER_IDS, build_excel_parser_map, prune_header_only_nodes, ) from src.llm.factory import get_token_counter, get_tokenizer log = logging.getLogger(__name__) if chunk_sizes is not None: log.warning( f"chunk_sizes={chunk_sizes} is deprecated for the Excel branch; " "ignoring and using the fixed mermaid-aware 2-level hierarchy." ) doc_uuid = UUID(document_id) if not parse_result: log.warning("No parse result provided") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} # Load data from pickle file data_file = parse_result.get('data_file') if not data_file or not os.path.exists(data_file): log.warning(f"Data file not found: {data_file}") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} log.info(f"Loading Excel Detailed parse result from {data_file}") with open(data_file, 'rb') as f: parsed_data = pickle.load(f) sheets = parsed_data.get('sheets', []) image_chunks = parsed_data.get('image_chunks', []) flowchart_chunks = parsed_data.get('flowchart_chunks', []) color_map = parsed_data.get('color_map', {}) cell_values_map = parsed_data.get('cell_values_map', {}) original_filename = parse_result.get('filename', 'document.xlsx') if not sheets and not image_chunks and not flowchart_chunks: log.warning("No content to chunk") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} log.info(f"Hierarchical chunking Excel Detailed: {original_filename}") log.info(f"Input: {len(sheets)} sheets, {len(image_chunks)} images, {len(flowchart_chunks)} flowcharts") tokenizer = get_tokenizer() token_count = get_token_counter() embed_max_tokens = Settings().embedding_model_context_tokens # Headroom against the 8192-token embed ceiling: covers retained # metadata (~100 tok template + filename/sheet_name/document_id) # plus Azure-tokenizer variance vs local cl100k_base counts. EMBED_SAFETY_BUFFER = 400 chunker_max = embed_max_tokens - EMBED_SAFETY_BUFFER log.info( f"Excel hierarchy: 2-level [markdown, chunk_size_512], " f"embed_max={embed_max_tokens} chunker_max={chunker_max} buffer={EMBED_SAFETY_BUFFER}" ) node_parser_map = build_excel_parser_map( tokenize_fn=token_count, embed_max_tokens=chunker_max, ) node_parser = HierarchicalNodeParser( node_parser_ids=EXCEL_PARSER_IDS, node_parser_map=node_parser_map, ) # Create one LlamaIndex Document per sheet from markdown documents = [] for sheet_info in sheets: sheet_name = sheet_info['sheet_name'] markdown = sheet_info['markdown'] if not markdown or not markdown.strip(): continue doc_metadata = { "document_id": str(doc_uuid), "filename": original_filename, "document_type": "excel", "sheet_name": sheet_name, } doc = Document( doc_id=str(doc_uuid), text=markdown, metadata=doc_metadata, ) documents.append(doc) log.info(f"Sheet '{sheet_name}': {token_count(markdown)} tokens") # Add image descriptions as separate Documents (leaf-level, not hierarchically chunked) def _dto_to_dict(chunk): if hasattr(chunk, 'model_dump'): return chunk.model_dump() elif hasattr(chunk, 'dict'): return chunk.dict() elif hasattr(chunk, '__dict__'): return {'text': getattr(chunk, 'text', ''), 'metadata': getattr(chunk, 'metadata', {}), 'element_type': getattr(chunk, 'element_type', 'Image')} return chunk image_documents = [] for img_chunk in image_chunks: img = _dto_to_dict(img_chunk) text = img.get('text', '') if not text or not text.strip(): continue meta = img.get('metadata', {}) doc_metadata = { "document_id": str(doc_uuid), "filename": original_filename, "document_type": "excel", "sheet_name": meta.get('sheet_name', 'Unknown'), "element_type": "Image", } for key in ['image_url', 'image_s3_key', 'image_number', 'position_key']: if key in meta: doc_metadata[key] = meta[key] image_documents.append(Document(doc_id=str(doc_uuid), text=text, metadata=doc_metadata)) flowchart_documents = [] for fc_chunk in flowchart_chunks: fc = _dto_to_dict(fc_chunk) text = fc.get('text', '') if not text or not text.strip(): continue meta = fc.get('metadata', {}) doc_metadata = { "document_id": str(doc_uuid), "filename": original_filename, "document_type": "excel", "sheet_name": meta.get('sheet_name', 'Unknown'), "element_type": "Flowchart", } flowchart_documents.append(Document(doc_id=str(doc_uuid), text=text, metadata=doc_metadata)) if not documents and not image_documents and not flowchart_documents: log.warning("No content to chunk after processing") return {'nodes_file': '', 'node_count': 0, 'leaf_count': 0} log.info(f"Created {len(documents)} sheet Documents for hierarchical chunking") # Apply hierarchical parsing to sheet documents nodes = [] if documents: nodes = node_parser.get_nodes_from_documents(documents) log.info(f"Created {len(nodes)} hierarchical nodes from {len(documents)} sheets") nodes = list(prune_header_only_nodes(nodes)) log.info(f"After header-only prune: {len(nodes)} nodes") # Add image/flowchart as flat TextNodes (no hierarchy needed for these) from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo doc_id_str = str(doc_uuid) for doc in image_documents + flowchart_documents: leaf_node = TextNode( text=doc.text, metadata=doc.metadata, ) leaf_node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str) nodes.append(leaf_node) # Count leaf nodes leaf_nodes = get_leaf_nodes(nodes) mermaid_bearing_leaves = sum( 1 for n in leaf_nodes if n.metadata.get("has_mermaid") ) log.info( f"Total nodes: {len(nodes)}, leaf nodes: {len(leaf_nodes)}, " f"mermaid-bearing leaves: {mermaid_bearing_leaves}" ) leaf_ids = {l.node_id for l in leaf_nodes} # Ensure ALL nodes have ref_doc_id set (used as doc_id in Milvus) for node in nodes: if not node.ref_doc_id: node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(node_id=doc_id_str) node.metadata["token_count"] = token_count(node.text) node.metadata["char_count"] = len(node.text) # cell_colors is only attached to non-leaf (parent / sheet-level) # nodes. Leaves go to Milvus, where the dynamic field is capped at # 64KB; a single highlighted sheet can hold 10k+ colored cells # (~200KB JSON) and would blow that cap on every leaf. Parents go # to PostgresDocumentStore with no such limit, and the # AutoMergingRetriever pulls them at query time so highlight-aware # consumers still see cell_colors via the parent chain. # cell_values is dropped entirely: it's a much larger duplicate # of the markdown content and nothing in this repo reads it back. if node.node_id not in leaf_ids: sheet_name = node.metadata.get("sheet_name") if sheet_name and sheet_name in color_map: node.metadata["cell_colors"] = color_map[sheet_name] # Per-leaf token audit -- including the metadata that actually ends up # in the embed text -- so any leaf that exceeds the embed ceiling is # visible in the chunking log instead of only at the storage stage. from llama_index.core.schema import MetadataMode embed_max = Settings().embedding_model_context_tokens leaf_sizes = [] for idx, leaf in enumerate(leaf_nodes): embed_text = leaf.get_content(metadata_mode=MetadataMode.EMBED) tot = token_count(embed_text) text_tok = token_count(leaf.text) leaf_sizes.append((idx, tot, text_tok)) if tot > embed_max: log.warning( f"leaf[{idx}] OVER BUDGET total={tot} text_tok={text_tok} " f"excluded_keys={leaf.excluded_embed_metadata_keys} " f"sheet={leaf.metadata.get('sheet_name')!r}" ) leaf_sizes.sort(key=lambda x: x[1], reverse=True) log.info( f"Leaf embed-text token sizes (top 5): {leaf_sizes[:5]} " f"(embed_max={embed_max})" ) # Save nodes to pickle file nodes_dir = '/opt/airflow/logs/hierarchical_nodes' os.makedirs(nodes_dir, exist_ok=True) nodes_file = os.path.join(nodes_dir, f'nodes_excel_detailed_{document_id}.pkl') nodes_data = { 'nodes': nodes, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': 'excel', 'original_filename': original_filename, 'document_id': document_id, 'parser_ids': EXCEL_PARSER_IDS, 'mermaid_bearing_leaves': mermaid_bearing_leaves, } with open(nodes_file, 'wb') as f: pickle.dump(nodes_data, f) log.info(f"Saved {len(nodes)} nodes to {nodes_file}") return { 'nodes_file': nodes_file, 'node_count': len(nodes), 'leaf_count': len(leaf_nodes), 'doc_type': 'excel', 'original_filename': original_filename, 'document_id': document_id, } @task.external_python(python=EXTERNAL_PYTHON) def store_to_llamaindex_task( document_id: str, nodes_result: Dict[str, Any], tenant_id: str, node_id: str, correlation_id: str, ) -> Dict[str, Any]: """ Store hierarchical nodes to LlamaIndex MilvusVectorStore + PostgresDocumentStore. Pattern: ALL nodes → PostgresDocumentStore, LEAF nodes only → Milvus. This enables AutoMergingRetriever to expand context at query time. Args: document_id: Document UUID string nodes_result: Result from hierarchical_chunk task (contains nodes_file path) tenant_id: Tenant UUID for pipeline-event routing node_id: Node UUID for pipeline-event routing correlation_id: Airflow dag_run_id, for event correlation Returns: Dict with stored counts and success status """ import asyncio import logging import pickle import os import sys from uuid import UUID from llama_index.core import VectorStoreIndex, StorageContext, Settings from llama_index.core.node_parser import get_leaf_nodes from llama_index.vector_stores.milvus import MilvusVectorStore from llama_index.vector_stores.milvus.utils import BM25BuiltInFunction from llama_index.storage.docstore.postgres import PostgresDocumentStore from src.core.config import settings from src.llm import get_embed_model from src.resilience.breakers import breaker_call, milvus_breaker from src.services.pipeline_events import publish_stage_event if '/opt/airflow/dags' not in sys.path: sys.path.insert(0, '/opt/airflow/dags') from tasks.chunk_schema import CHUNK_SCHEMA_VERSION log = logging.getLogger(__name__) publish_stage_event( "indexing", "started", document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, ) doc_uuid = UUID(document_id) if not nodes_result: log.warning("No nodes result provided") return {'stored_count': 0, 'leaf_count': 0, 'success': False} # Load nodes from pickle file nodes_file = nodes_result.get('nodes_file') if not nodes_file or not os.path.exists(nodes_file): log.warning(f"Nodes file not found: {nodes_file}") return {'stored_count': 0, 'leaf_count': 0, 'success': False} log.info(f"Loading nodes from {nodes_file}") with open(nodes_file, 'rb') as f: nodes_data = pickle.load(f) nodes = nodes_data.get('nodes', []) if not nodes: log.warning("No nodes to store") return {'stored_count': 0, 'leaf_count': 0, 'success': True} # Inject tenant/node context + chunk schema version into all node metadata for node in nodes: node.metadata['tenant_id'] = tenant_id node.metadata['node_id'] = node_id node.metadata['schema_version'] = CHUNK_SCHEMA_VERSION # TEXTIQ-569: ensure key is always present so the Milvus dynamic # envelope carries it on every new chunk (tag-cascade filter on # missing key would silently no-match). tag_records_task already # populates it; this is the safety default for paths that skip # tagging (e.g., tagging failure, empty text). node.metadata.setdefault('verified_tag_names', []) # Constrain the per-node embed text to semantic-only metadata. Excel # leaves can be ~8k tokens already; serialising bulk metadata such as # cell_colors, cell_values, the tag list, or routing IDs blows past # the 8192-token embedding ceiling. Whitelist the few keys that # genuinely help retrieval, then exclude everything else. EMBED_INCLUDE_KEYS = { "filename", "sheet_name", "document_type", "element_type", } for node in nodes: excl = [k for k in node.metadata.keys() if k not in EMBED_INCLUDE_KEYS] node.excluded_embed_metadata_keys = excl node.excluded_llm_metadata_keys = excl log.info(f"Storing {len(nodes)} nodes to LlamaIndex stores") async def store_nodes(): """Async wrapper for LlamaIndex operations (MilvusVectorStore uses async internally).""" # Configure embedding model via provider factory Settings.embed_model = get_embed_model() # Initialize MilvusVectorStore with hybrid search (dense + BM25 sparse) milvus_uri = f"http://{settings.milvus_host}:{settings.milvus_port}" log.info(f"Connecting to Milvus: {milvus_uri}") vector_store = MilvusVectorStore( uri=milvus_uri, collection_name=settings.milvus_collection, dim=settings.embedding_dimensions, overwrite=False, # Don't drop existing data enable_sparse=True, sparse_embedding_function=BM25BuiltInFunction(), hybrid_ranker="RRFRanker", hybrid_ranker_params={"k": 60}, ) # Initialize PostgresDocumentStore log.info("Connecting to PostgreSQL docstore") # Create engines with RLS tenant context set on every connection from uuid import UUID as _UUID from sqlalchemy import create_engine as _create_engine, event as sa_event from sqlalchemy.ext.asyncio import create_async_engine as _create_async_engine _tid = str(_UUID(tenant_id)) # validate UUID _nid = str(_UUID(node_id)) # validate UUID db_url = settings.database_url db_url_sync = db_url.replace("+asyncpg", "") # asyncpg: use server_settings to set tenant context at connection time async_engine = _create_async_engine( db_url, connect_args={"server_settings": {"app.tenant_id": _tid, "app.node_id": _nid}}, ) # psycopg2: use connect event to set tenant context sync_engine = _create_engine(db_url_sync) @sa_event.listens_for(sync_engine, "connect") def _set_tenant(dbapi_conn, connection_record): cursor = dbapi_conn.cursor() cursor.execute(f"SET app.tenant_id = '{_tid}'; SET app.node_id = '{_nid}'") cursor.close() from llama_index.storage.kvstore.postgres import PostgresKVStore kvstore = PostgresKVStore( table_name=settings.docstore_table, engine=sync_engine, async_engine=async_engine, perform_setup=True, use_jsonb=True, ) docstore = PostgresDocumentStore(postgres_kvstore=kvstore, namespace=settings.docstore_namespace) # Create storage context storage_context = StorageContext.from_defaults( vector_store=vector_store, docstore=docstore, ) # Add ALL nodes to docstore (for parent lookup during retrieval) log.info(f"Adding {len(nodes)} nodes to PostgresDocumentStore") storage_context.docstore.add_documents(nodes) # Get leaf nodes for vector indexing leaf_nodes = get_leaf_nodes(nodes) log.info(f"Indexing {len(leaf_nodes)} leaf nodes to Milvus") # Create index with only leaf nodes (vectors stored in Milvus). # Wrapped with milvus_breaker — VectorStoreIndex(...) executes the # underlying Milvus inserts during construction. index = breaker_call( milvus_breaker, VectorStoreIndex, leaf_nodes, storage_context=storage_context, show_progress=True, ) log.info(f"Successfully stored {len(nodes)} nodes ({len(leaf_nodes)} leaf nodes)") return len(nodes), len(leaf_nodes) # Run async operations stored_count, leaf_count = asyncio.run(store_nodes()) # Clean up nodes file try: os.remove(nodes_file) log.info(f"Cleaned up nodes file: {nodes_file}") except Exception as e: log.warning(f"Failed to clean up nodes file: {e}") return { 'stored_count': stored_count, 'leaf_count': leaf_count, 'document_id': document_id, 'success': True, } @task.external_python(python=EXTERNAL_PYTHON) def store_large_table_to_llamaindex_task( document_id: str, nl_result: Dict[str, Any], tenant_id: str, node_id: str, correlation_id: str, ) -> Dict[str, Any]: """ Store Large Table records to LlamaIndex MilvusVectorStore. For Large Table pipeline: 1 record = 1 Document (no hierarchical parsing). Text = NL description, metadata.record_json = original JSON. Args: document_id: Document UUID string nl_result: Result from convert_large_tables_to_natural_language task tenant_id: Tenant UUID for pipeline-event routing node_id: Node UUID for pipeline-event routing correlation_id: Airflow dag_run_id, for event correlation Returns: Dict with stored count and success status """ import asyncio import logging import pickle import os import sys from uuid import UUID from llama_index.core import Document, VectorStoreIndex, StorageContext, Settings from llama_index.vector_stores.milvus import MilvusVectorStore from llama_index.vector_stores.milvus.utils import BM25BuiltInFunction from src.core.config import settings from src.llm import get_embed_model from src.resilience.breakers import breaker_call, milvus_breaker from src.services.pipeline_events import publish_stage_event if '/opt/airflow/dags' not in sys.path: sys.path.insert(0, '/opt/airflow/dags') from tasks.chunk_schema import CHUNK_SCHEMA_VERSION log = logging.getLogger(__name__) publish_stage_event( "indexing", "started", document_id=document_id, tenant_id=tenant_id, node_id=node_id, correlation_id=correlation_id, ) doc_uuid = UUID(document_id) if not nl_result: log.warning("No NL result provided") return {'stored_count': 0, 'success': False} # Load data from pickle file data_file = nl_result.get('data_file') if not data_file or not os.path.exists(data_file): log.warning(f"Data file not found: {data_file}") return {'stored_count': 0, 'success': False} log.info(f"Loading NL records from {data_file}") with open(data_file, 'rb') as f: nl_data = pickle.load(f) chunks = nl_data.get('chunks', []) if not chunks: log.warning("No records to store") return {'stored_count': 0, 'success': True} original_filename = nl_data.get('original_filename', '') log.info(f"Storing {len(chunks)} Large Table records to LlamaIndex") # Convert records to LlamaIndex Documents (before async block) documents = [] for record in chunks: nl_text = record.get('text', '') if not nl_text or not nl_text.strip(): continue source = record.get('_source', {}) doc = Document( doc_id=str(doc_uuid), # Set doc_id to our document_id text=nl_text, metadata={ "document_id": str(doc_uuid), "tenant_id": tenant_id, "node_id": node_id, "schema_version": CHUNK_SCHEMA_VERSION, "filename": source.get('filename', original_filename), "document_type": "excel", "sheet_name": source.get('sheet', ''), "row_number": source.get('row', 0), "record_json": record.get('content', {}), "domain_category": record.get('domain_category', 'general'), "domain_terms": record.get('domain_terms', []), "tags": record.get('tags', []), # TEXTIQ-569: emitted by SemanticTagger.tag_records_async # post-inheritance. Stored under the Milvus dynamic envelope # via LlamaIndex's node_to_metadata_dict; indexed via the # JSON-path INVERTED index for tag-cascade filtering. "verified_tag_names": record.get('verified_tag_names', []), }, ) documents.append(doc) if not documents: log.warning("No valid documents to store") return {'stored_count': 0, 'success': True} log.info(f"Created {len(documents)} LlamaIndex Documents") async def store_documents(): """Async wrapper for LlamaIndex operations (MilvusVectorStore uses async internally).""" # Configure embedding model via provider factory Settings.embed_model = get_embed_model() # Initialize MilvusVectorStore with hybrid search (dense + BM25 sparse) milvus_uri = f"http://{settings.milvus_host}:{settings.milvus_port}" log.info(f"Connecting to Milvus: {milvus_uri}") vector_store = MilvusVectorStore( uri=milvus_uri, collection_name=settings.milvus_collection, dim=settings.embedding_dimensions, overwrite=False, enable_sparse=True, sparse_embedding_function=BM25BuiltInFunction(), hybrid_ranker="RRFRanker", hybrid_ranker_params={"k": 60}, ) storage_context = StorageContext.from_defaults(vector_store=vector_store) # Store directly via VectorStoreIndex (no hierarchy). # Wrapped with milvus_breaker — from_documents executes Milvus inserts. index = breaker_call( milvus_breaker, VectorStoreIndex.from_documents, documents, storage_context=storage_context, show_progress=True, ) return len(documents) stored_count = asyncio.run(store_documents()) log.info(f"Successfully stored {stored_count} Large Table records") # Clean up data file try: os.remove(data_file) log.info(f"Cleaned up data file: {data_file}") except Exception as e: log.warning(f"Failed to clean up data file: {e}") return { 'stored_count': stored_count, 'document_id': document_id, 'success': True, }