ontology_builder.py

#!/usr/bin/env python3
"""Ontology Builder - Acronym-Ontology Builder/Disambiguator"""
from __future__ import annotations
import argparse, asyncio, datetime as _dt, json, logging, pathlib as _pl, sys, random, warnings, re
from typing import Dict, List, Any, Tuple
from google import genai
from prompts import PROMPTS
from google.genai import types
from langchain.text_splitter import TokenTextSplitter
import tiktoken, tqdm

warnings.filterwarnings("ignore")

def extract_acronyms_from_content(content: str) -> Dict[str, int]:
    """Extract acronyms from text content using regex patterns"""
    # Use separate patterns to avoid group issues
    patterns = [
        r'\([A-Z][A-Z0-9\.\-]*[A-Z0-9]\)',  # Parentheses pattern
        r'\b[A-Z][A-Z0-9\.\-]*[A-Z0-9]\b',  # Word boundary pattern  
        r'[A-Z]\.[A-Z]\.(?:[A-Z]\.)*'       # Dotted acronyms (fixed)
    ]
    
    all_matches = []
    for pattern in patterns:
        matches = re.findall(pattern, content)
        all_matches.extend(matches)
    
    # Clean matches (remove parentheses)
    cleaned_matches = []
    for match in all_matches:
        text = str(match).strip('()')
        
        # Filter: must start and end with alphanumeric, be mostly uppercase
        if text and re.match(r'^[A-Z].*[A-Z0-9]$', text) and len(text) >= 2:
            cleaned_matches.append(text)
    
    # Count occurrences
    acronym_counts = {}
    for acronym in cleaned_matches:
        acronym_counts[acronym] = acronym_counts.get(acronym, 0) + 1
    
    print(f"🔍 Found {len(all_matches)} raw matches, {len(cleaned_matches)} valid acronyms")
    
    return acronym_counts

def create_acronyms_file(input_path: str, run_dir: _pl.Path) -> _pl.Path:
    """Create primitive ontology file from input content"""
    primitive_ontology_file = run_dir / "primitive_ontology.json"
    
    if primitive_ontology_file.exists():
        print(f"📄 Using existing primitive ontology: {primitive_ontology_file}")
        return primitive_ontology_file
    
    print("🔍 Extracting primitive ontology from input content...")
    
    # Load content from JSONL chunks file
    input_path_obj = _pl.Path(input_path)
    full_text = ""
    chunk_count = 0
    
    if input_path_obj.is_file() and input_path_obj.suffix == '.jsonl':
        # Handle JSONL chunks file
        with open(input_path_obj, 'r', encoding='utf-8') as f:
            for line in f:
                try:
                    chunk = json.loads(line.strip())
                    if isinstance(chunk, dict) and 'text' in chunk:
                        full_text += " " + chunk['text']
                        chunk_count += 1
                except:
                    continue
        print(f"📚 Processed {chunk_count} text chunks")
    else:
        # Handle directories or other content for direct processing
        content = load_content(input_path, False)
        full_text = "\n".join(content)
    
    print(f"📊 Total text length: {len(full_text):,} characters")
    
    # Extract acronyms
    acronym_counts = extract_acronyms_from_content(full_text)
    
    # Sort by frequency (descending)
    sorted_acronyms = dict(sorted(acronym_counts.items(), key=lambda x: x[1], reverse=True))
    
    # Save to file
    primitive_ontology_file.write_text(json.dumps(sorted_acronyms, indent=2, ensure_ascii=False))
    
    print(f"✅ Extracted {len(sorted_acronyms)} unique acronyms")
    print(f"📄 Saved primitive ontology to: {primitive_ontology_file}")
    
    # Show top 10 most frequent
    top_10 = list(sorted_acronyms.items())[:10]
    if top_10:
        print("🔝 Top 10 most frequent acronyms:")
        for acronym, count in top_10:
            print(f"   {acronym}: {count}")
    
    return primitive_ontology_file

def main():
    script_name = _pl.Path(sys.argv[0]).stem
    
    parser = argparse.ArgumentParser(description='Acronym Ontology Builder')
    parser.add_argument('--input', '-i', type=str, help='Input chunks.jsonl file or directory (for extract) or ontology JSON (for disambiguate)')
    parser.add_argument('--acronyms', '-a', type=str, help='Primitive ontology file (auto-created if not provided)')
    parser.add_argument('--output', '-o', type=str, default=f'./{script_name}', help='Output directory')
    parser.add_argument('--context-window', '-w', type=int, default=128000, help='Context window size')
    parser.add_argument('--parallel', '-p', type=int, default=2, help='Parallel workers')
    parser.add_argument('--api-key', required=True, help='Google API key')
    parser.add_argument('--json-chunks', action='store_true', help='Force treat input as JSON chunks')
    parser.add_argument('--enable-thinking', action='store_true', help='Enable AI thinking')
    parser.add_argument('--disable-streaming', action='store_true', help='Disable streaming')
    parser.add_argument('--skip-extraction', action='store_true', help='Skip extraction step (auto-finds most recent ontology)')
    parser.add_argument('--skip-filtering', action='store_true', help='Skip filtering step')
    parser.add_argument('--skip-disambiguation', action='store_true', help='Skip disambiguation step')
    parser.add_argument('--continue-run', '-c', type=str, help='Continue existing run (provide run directory name, e.g., run_20250610_002020)')
    args = parser.parse_args()
    
    if not args.skip_extraction and not args.input:
        print("Input file/directory required for extraction")
        return
    
    output_dir = _pl.Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # Auto-detect most recent ontology when skipping extraction (BEFORE creating new run dir)
    input_path = args.input
    auto_continue_run = None
    if args.skip_extraction:
        if not args.input:
            # Find most recent ontology.json in output directory
            run_dirs = [d for d in output_dir.iterdir() if d.is_dir() and d.name.startswith('run_')]
            if run_dirs:
                most_recent_run = sorted(run_dirs, key=lambda x: x.name)[-1]
                ontology_file = most_recent_run / "ontology.json"
                if ontology_file.exists():
                    input_path = str(ontology_file)
                    auto_continue_run = most_recent_run.name
                    print(f"Auto-continuing run: {most_recent_run.name}")
                else:
                    print(f"No ontology.json found in most recent run: {most_recent_run.name}")
                    return
            else:
                print("No previous runs found, ontology input required when skipping extraction")
                return
        else:
            # Check if user provided chunks.jsonl instead of ontology.json
            input_file = _pl.Path(args.input)
            if input_file.name == "chunks.jsonl":
                print("⚠️  chunks.jsonl provided but --skip-extraction requires ontology.json")
                print("💡 To run full pipeline: remove --skip-extraction flag")
                print("💡 To skip extraction: provide path to ontology.json from previous run")
                return
            input_path = args.input
    
    # Handle continue-run or create new run directory (AFTER input detection)
    if args.continue_run:
        run_dir = output_dir / args.continue_run
        if not run_dir.exists():
            print(f"Run directory {args.continue_run} not found, creating new run")
            timestamp = _dt.datetime.now().strftime('%Y%m%d_%H%M%S')
            run_dir = output_dir / f"run_{timestamp}"
            run_dir.mkdir(parents=True, exist_ok=True)
        else:
            print(f"Continuing run: {args.continue_run}")
    elif auto_continue_run:
        # Auto-continue the run where we found the ontology
        run_dir = output_dir / auto_continue_run
        print(f"Continuing run: {auto_continue_run}")
    else:
        timestamp = _dt.datetime.now().strftime('%Y%m%d_%H%M%S')
        run_dir = output_dir / f"run_{timestamp}"
        run_dir.mkdir(parents=True, exist_ok=True)
    
    # Initialize client
    client = genai.Client(api_key=args.api_key)
    
    # Auto-create acronyms file if not provided and not skipping extraction
    acronyms_path = args.acronyms
    if not args.skip_extraction and not acronyms_path:
        acronyms_path = str(create_acronyms_file(args.input, run_dir))
    
    asyncio.run(run_full_pipeline(
        input_path, acronyms_path, run_dir, args.json_chunks, client,
        args.context_window, args.parallel, args.enable_thinking, args.disable_streaming,
        args.skip_extraction, args.skip_disambiguation, args.skip_filtering
    ))

# Configuration
MODEL_ID = "gemini-2.5-flash-preview-05-20"
MODEL_PRO = "gemini-2.5-pro-preview-03-25"
MODEL_FLASH = "gemini-2.5-flash-preview-05-20"
MAX_WINDOW, PROMPT_TOK, ANSWER_TOK = 128_000, 200, 8_192
OVERLAP_TOK, MAX_CTX_LEN = 64, 120
DEFAULT_PARALLEL, MAX_RETRIES = 2, 10
BASE_DELAY, MAX_DELAY = 2.0, 120.0
ENC = tiktoken.get_encoding("cl100k_base")
tlen = lambda s: len(ENC.encode(s))

# Logging
logging.basicConfig(level=logging.INFO, format="%(message)s", stream=sys.stderr)
for logger in ["google.api_core.http", "httpx"]:
    logging.getLogger(logger).setLevel(logging.WARNING)

# Global configuration
SCHEMAS = {
    "extract": {"type": "array", "items": {"type": "object", "properties": {
        "acronym": {"type": "string"}, "definition": {"type": "string"},
        "context": {"type": "string"}, "confidence": {"type": "string", "enum": ["high", "medium", "low"]}
    }, "required": ["acronym", "definition", "confidence"]}},
    
    "disamb": {"type": "array", "items": {"type": "object", "properties": {
        "acronym": {"type": "string"}, "definition": {"type": "string"}
    }, "required": ["acronym", "definition"]}},
    
    "filter": {"type": "array", "items": {"type": "object", "properties": {
        "acronym": {"type": "string"}, "definition": {"type": "string"},
        "action": {"type": "string", "enum": ["keep", "discard", "merge"]},
        "reason": {"type": "string"}, "merge_with": {"type": "string"}
    }, "required": ["acronym", "definition", "action", "reason"]}}
}

async def retry_api_call(func, *args, **kwargs):
    """Generic retry wrapper for API calls"""
    # Extract pbar and chunk_idx from kwargs if available
    pbar = kwargs.pop('pbar', None)
    chunk_idx = kwargs.pop('chunk_idx', None)
    last_error = None
    
    for attempt in range(MAX_RETRIES):
        if attempt > 0:
            delay = min(BASE_DELAY * (2 ** (attempt - 1)) + random.random() * 0.2, MAX_DELAY)
            if pbar and chunk_idx is not None and last_error:
                pbar.write(f"⏳ [{chunk_idx:3d}] Retry {attempt}: waiting {delay:.1f}s - {str(last_error)[:40]}")
            await asyncio.sleep(delay)
        try:
            # Pass pbar to generate_content if it's the function being called
            if func.__name__ == 'generate_content' and pbar:
                return await func(*args, **kwargs, pbar=pbar)
            else:
                return await func(*args, **kwargs)
        except Exception as e:
            last_error = e
            if attempt == MAX_RETRIES - 1:
                raise
            if not (pbar and chunk_idx is not None):
                logging.info(f"⏳ Retry {attempt + 1}: {str(e)[:50]}")

async def retry_api_call_with_temp(func, *args, **kwargs):
    """Retry wrapper that passes attempt number to function for temperature adjustment"""
    # Extract pbar and chunk_idx from kwargs if available
    pbar = kwargs.pop('pbar', None)
    chunk_idx = kwargs.pop('chunk_idx', None)
    last_error = None
    
    for attempt in range(MAX_RETRIES):
        if attempt > 0:
            delay = min(BASE_DELAY * (2 ** (attempt - 1)) + random.random() * 0.2, MAX_DELAY)
            if pbar and chunk_idx is not None and last_error:
                error_type = "JSON" if "json" in str(last_error).lower() else "API"
                pbar.write(f"⏳ [{chunk_idx:3d}] {error_type} retry {attempt}: waiting {delay:.1f}s - {str(last_error)[:40]}")
            await asyncio.sleep(delay)
        try:
            return await func(attempt, *args, **kwargs)
        except Exception as e:
            last_error = e
            if attempt == MAX_RETRIES - 1:
                raise
            if not (pbar and chunk_idx is not None):
                logging.info(f"⏳ Retry {attempt + 1}: {str(e)[:50]}")

def api_config(temp: float, schema: dict = None, mime: str = None, enable_thinking: bool = False):
    """Build API config"""
    cfg = {"temperature": temp, "automatic_function_calling": types.AutomaticFunctionCallingConfig(disable=True)}
    if schema: cfg.update({"response_mime_type": mime or "application/json", "response_schema": schema})
    if not enable_thinking: cfg["thinking_config"] = types.ThinkingConfig(thinking_budget=0)
    return types.GenerateContentConfig(**cfg)

async def generate_content(client, model: str, prompt: str, config: types.GenerateContentConfig, disable_streaming: bool = False, pbar=None):
    """Generate content with streaming support"""
    if disable_streaming:
        resp = await client.aio.models.generate_content(model=model, contents=prompt, config=config)
        response_text = resp.text or ""
        response_tokens = tlen(response_text)
        if pbar:
            pbar.write(f"🔢 Number of tokens received from LLM: {response_tokens:,}")
        else:
            print(f"🔢 Number of tokens received from LLM: {response_tokens:,}")
        return response_text, getattr(resp.candidates[0], "finish_reason", " ") if resp.candidates else " "
    else:
        stream = await client.aio.models.generate_content_stream(model=model, contents=prompt, config=config)
        text, finish = "", " "
        async for part in stream:
            if part.text: text += part.text
            if getattr(part, "candidates", None):
                finish = getattr(part.candidates[0], "finish_reason", " ")
        response_tokens = tlen(text)
        if pbar:
            pbar.write(f"🔢 Number of tokens received from LLM: {response_tokens:,}")
        else:
            print(f"🔢 Number of tokens received from LLM: {response_tokens:,}")
        return text, finish


def load_acronyms(path: str) -> Dict[str, int]:
    """Load acronyms from file"""
    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        return {k.upper(): int(v) for k, v in data.items()} if isinstance(data, dict) else {}
    except:
        d = {}
        with open(path, encoding="utf-8", errors="ignore") as f:
            for ln in f:
                p = ln.strip().split()
                if p: d[(p[1] if p[0].isdigit() else p[0]).upper()] = int(p[0]) if p[0].isdigit() else 1
        return d

def load_content(path: str, json_chunks: bool) -> List[str]:
    """Load content from various sources"""
    path = _pl.Path(path)
    
    if path.is_dir():
        texts = []
        md_files = list(path.glob("**/*.md"))
        logging.info(f"📁 Found {len(md_files)} .md files")
        for f in md_files:
            try:
                content = f.read_text(encoding="utf-8", errors="ignore")
                if content:
                    texts.append(content)
                    logging.info(f"   📄 {f.name}: {len(content):,} chars")
            except: pass
        return ["\n\n".join(texts)]
    
    if json_chunks:
        chunks = []
        content = path.read_text(encoding="utf-8")
        
        # Handle both JSON array and newline-delimited JSON
        if content.startswith('['):
            try:
                data = json.loads(content)
                chunks = [item["content"] for item in data if isinstance(item, dict) and "content" in item]
            except: pass
        else:
            for line in content.split('\n'):
                if line.strip():
                    try:
                        item = json.loads(line.strip())
                        if isinstance(item, dict) and "content" in item:
                            chunks.append(item["content"])
                    except: pass
        
        logging.info(f"📦 Loaded {len(chunks)} JSON chunks")
        return chunks
    
    return [path.read_text(encoding="utf-8", errors="ignore")]

async def call_chunk(client, idx: int, chunk: str, acr_list: str, sem: asyncio.Semaphore, enable_thinking: bool, disable_streaming: bool, pbar=None):
    """Process a single chunk"""
    prompt = PROMPTS["extract"].format(max_ctx_len=MAX_CTX_LEN, acr_list=acr_list, passage=chunk)
    prompt_tokens = tlen(prompt)
    
    async with sem:
        try:
            # First, get the response from API with standard retry (for API errors)
            text, finish = await retry_api_call(
                generate_content, client, MODEL_ID, prompt, 
                api_config(0.1, SCHEMAS["extract"], enable_thinking=enable_thinking), 
                disable_streaming, pbar=pbar, chunk_idx=idx
            )
            
            # Then handle JSON parsing with temperature retries (for JSON errors only)
            data = None
            for attempt in range(MAX_RETRIES):
                try:
                    data = json.loads(text)
                    break
                except json.JSONDecodeError as json_err:
                    if attempt == MAX_RETRIES - 1:
                        raise json_err
                    
                    # Use random temperature for JSON retry attempts
                    temp = random.uniform(0.0, 0.8)
                    if pbar:
                        pbar.write(f"🎲 [{idx:3d}] JSON retry {attempt + 1} with temp={temp:.2f}")
                    
                    # Get new response with different temperature
                    text, finish = await retry_api_call(
                        generate_content, client, MODEL_ID, prompt,
                        api_config(temp, SCHEMAS["extract"], enable_thinking=enable_thinking),
                        disable_streaming, pbar=pbar, chunk_idx=idx
                    )
            
            items = [obj for obj in data if isinstance(obj, dict) and "acronym" in obj and "definition" in obj]
            
            if pbar: pbar.write(f"✅ [{idx:3d}] {str(finish).replace('FinishReason.', '')[:12]} | {len(items):2d} defs")
            return idx, items, prompt_tokens, tlen(text)
        except Exception as e:
            if pbar: pbar.write(f"❌ [{idx:3d}] {str(e)[:50]}")
            return idx, [], prompt_tokens, 0

async def two_pass_call(client, pro_prompt: str, flash_prompt_template: str, schema: dict, name: str, enable_thinking: bool, disable_streaming: bool, batch_info: str = "", pbar=None) -> Any:
        """Generic two-pass Pro→Flash call with retry logic"""
        pro_tokens = tlen(pro_prompt)
        batch_prefix = f"{batch_info} " if batch_info else ""
        logging.info(f"🧠 {batch_prefix}[{name}] Pass 1: Pro reasoning ({pro_tokens:,} tokens)...")
        
        if pro_tokens > 100000:
            logging.warning(f"⚠️  {batch_prefix}[{name}] Very large prompt ({pro_tokens:,} tokens)")
        
        # Pass 1: Pro
        for temp in [0.3, 0.7]:  # Try normal then higher temp if needed
            try:
                pro_text, _ = await retry_api_call(
                    generate_content, client, MODEL_PRO, pro_prompt, api_config(temp, enable_thinking=enable_thinking), disable_streaming, pbar=pbar
                )
                if pro_text.strip(): break
                logging.info(f"🔄 {batch_prefix}[{name}] Retrying Pro with higher temperature...")
            except Exception as e:
                logging.warning(f"❌ {batch_prefix}[{name}] Pro failed: {str(e)[:100]}")
                return {} if "disamb" in name.lower() else []
        
        if not pro_text.strip():
            logging.error(f"❌ {batch_prefix}[{name}] Pro failed to generate text")
            return {} if "disamb" in name.lower() else []
            
        logging.info(f"✅ {batch_prefix}[{name}] Pass 1 complete ({tlen(pro_text):,} tokens)")
        
        # Debug for disambiguation
        if "disamb" in name.lower() and pro_text:
            final_choices = pro_text.count("Final choice:")
            section_headers = pro_text.count("###")
            if final_choices > 0:
                logging.info(f"   {batch_prefix}[{name}] Found {final_choices} decisions for {section_headers} sections")
            else:
                logging.warning(f"   {batch_prefix}[{name}] No 'Final choice' format found")
        
        # Pass 2: Flash with retry strategies
        logging.info(f"⚡ {batch_prefix}[{name}] Pass 2: Flash structuring...")
        flash_prompt = flash_prompt_template.format(pro_analysis=pro_text)
        
        for temp in [0.1, 0.0, 0.3]:  # Multiple retry strategies
            try:
                flash_text, _ = await retry_api_call(
                    generate_content, client, MODEL_FLASH, flash_prompt, api_config(temp, schema, enable_thinking=enable_thinking), disable_streaming, pbar=pbar
                )
                data = json.loads(flash_text or "[]")
                
                if "disamb" in name.lower():
                    result = {}
                    for item in data:
                        if isinstance(item, dict) and "acronym" in item and "definition" in item:
                            result[item["acronym"].upper()] = item["definition"]
                    logging.info(f"✅ {batch_prefix}[{name}] Pass 2 complete: {len(result)} acronyms extracted")
                    return result
                else:
                    valid_items = [d for d in data if isinstance(d, dict) and "acronym" in d and "action" in d]
                    logging.info(f"✅ {batch_prefix}[{name}] Pass 2 complete: {len(valid_items)} decisions")
                    return valid_items
                    
            except (json.JSONDecodeError, Exception) as e:
                if temp == 0.3:  # Last attempt
                    logging.warning(f"❌ {batch_prefix}[{name}] Flash failed: {str(e)[:100]}")
                    return {} if "disamb" in name.lower() else []
                logging.warning(f"⚠️  {batch_prefix}[{name}] Flash retry needed: {str(e)[:50]}")
        
        return {} if "disamb" in name.lower() else []

async def batch_process(items: dict, process_func, name: str, client, enable_thinking: bool, disable_streaming: bool, run_dir: _pl.Path, parallel: int, max_tokens: int = 80000) -> dict:
        """Process items in batches with parallel processing"""
        if not items: return {}
        
        total_tokens = sum(tlen(f"{k}: {v}") for k, v in items.items())
        base_tokens = 5000
        max_items = 150
        
        if total_tokens + base_tokens <= max_tokens and len(items) <= max_items:
            logging.info(f"📦 [{name}] Processing {len(items)} items in single batch")
            return await process_func(items)
        
        # Split into batches
        logging.info(f"📦 [{name}] Splitting {len(items)} items into batches")
        batches = []
        current_batch = {}
        current_tokens = 0
        
        for k, v in items.items():
            item_tokens = tlen(f"{k}: {v}")
            if (current_tokens + item_tokens + base_tokens > max_tokens) or (len(current_batch) >= max_items):
                if current_batch:
                    batches.append(current_batch)
                current_batch = {k: v}
                current_tokens = item_tokens
            else:
                current_batch[k] = v
                current_tokens += item_tokens
        
        if current_batch:
            batches.append(current_batch)
        
        logging.info(f"   📦 Created {len(batches)} batches")
        
        # Process batches in parallel with progress bar
        sem = asyncio.Semaphore(parallel)
        pbar = tqdm.tqdm(total=len(batches), desc=f"🔄 {name}", file=sys.stdout)
        
        async def process_batch_with_semaphore(batch_idx, batch):
            async with sem:
                batch_tokens = sum(tlen(f'{k}: {v}') for k, v in batch.items())
                pbar.write(f"   🔄 Processing batch {batch_idx + 1}/{len(batches)} ({len(batch)} items, ~{batch_tokens:,} tokens)")
                try:
                    # Pass batch info and pbar to process_func if it supports them
                    batch_info = f"Batch {batch_idx + 1}/{len(batches)}"
                    func_args = process_func.__code__.co_varnames if hasattr(process_func, '__code__') else []
                    
                    if 'pbar' in func_args and 'batch_info' in func_args:
                        result = await process_func(batch, batch_info, pbar)
                    elif 'batch_info' in func_args:
                        result = await process_func(batch, batch_info)
                    elif 'pbar' in func_args:
                        result = await process_func(batch, pbar)
                    else:
                        result = await process_func(batch)
                    pbar.write(f"   ✅ Batch {batch_idx + 1} complete: {len(result)} results")
                    pbar.update(1)
                    return result
                except Exception as e:
                    pbar.write(f"   ❌ Batch {batch_idx + 1} failed: {str(e)[:100]}")
                    pbar.update(1)
                    return {}
        
        batch_results = await asyncio.gather(*[
            process_batch_with_semaphore(i, batch) for i, batch in enumerate(batches)
        ])
        
        pbar.close()
        
        # Combine results
        results = {}
        for batch_result in batch_results:
            if batch_result:
                results.update(batch_result)
        
        logging.info(f"✅ [{name}] Total: {len(results)}/{len(items)} processed across {len(batches)} batches")
        
        if len(results) < len(items):
            lost_items = set(items.keys()) - set(results.keys())
            logging.warning(f"⚠️  [{name}] Lost {len(lost_items)} items during processing!")
        
        return results

def write_report(filepath: _pl.Path, title: str, data: dict):
    """Write decision report"""
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(f"{title}\n{'=' * 80}\n\n")
        for section, items in data.items():
            if items:
                f.write(f"{section.upper()} ({len(items)}):\n{'-' * 40}\n")
                for item in sorted(items):
                    f.write(f"{item}\n")
                f.write("\n")

async def disambiguate_batch(candidates: Dict[str, List[Tuple[str, str]]], client, enable_thinking: bool, disable_streaming: bool, run_dir: _pl.Path, parallel: int) -> Dict[str, str]:
    """Disambiguate acronyms with parallel batch processing"""
    all_decisions = []
    
    async def process(batch, batch_info="", pbar=None):
        sections = "\n".join([f"### {ac}\n" + "\n".join([f"{i}. \"{d}\" — context: {c[:MAX_CTX_LEN]}" 
                             for i, (d, c) in enumerate(defs, 1)]) for ac, defs in batch.items()])
        
        single_def_count = sum(1 for defs in batch.values() if len(defs) == 1)
        if single_def_count > 0:
            logging.warning(f"   ⚠️  {single_def_count} acronyms with only 1 definition in ambiguous set!")
        
        result = await two_pass_call(
            client, PROMPTS["disamb_pro"].format(acronym_sections=sections),
            PROMPTS["disamb_flash"],
            SCHEMAS["disamb"],
            "Disambiguation", enable_thinking, disable_streaming, batch_info, pbar
        )
        
        # Track decisions and handle missing
        missing_count = 0
        for ac in batch:
            if ac not in result:
                result[ac] = batch[ac][0][0]  # Use first as fallback
                missing_count += 1
            
            if len(batch[ac]) > 1:  # Track multi-option decisions
                decision = {
                    "acronym": ac,
                    "selected": result[ac],
                    "discarded": [d for d, _ in batch[ac] if d != result[ac]],
                    "note": "Fallback: first definition" if missing_count > 0 and ac not in result else None
                }
                all_decisions.append(decision)
        
        if missing_count > 0:
            logging.warning(f"   ⚠️  {missing_count} acronyms used fallback")
        
        return result
    
    final_result = await batch_process(candidates, process, "Disambiguation", client, enable_thinking, disable_streaming, run_dir, parallel)
    
    # Ensure all candidates get results
    for ac, defs in candidates.items():
        if ac not in final_result:
            final_result[ac] = defs[0][0]
            logging.warning(f"   ⚠️  {ac} missing, using first definition")
    
    # Write report
    if all_decisions:
        report_data = {}
        for decision in all_decisions:
            ac = decision["acronym"]
            entry = f"{ac}: ✅ {decision['selected']}"
            if decision["discarded"]:
                entry += f" | ❌ Discarded: {', '.join(decision['discarded'])}"
            if decision.get("note"):
                entry += f" | ℹ️ {decision['note']}"
            report_data.setdefault("decisions", []).append(entry)
        
        write_report(run_dir / "disambiguation_decisions.txt", "ACRONYM DISAMBIGUATION DECISIONS REPORT", report_data)
        logging.info(f"📊 Disambiguation: {len(final_result)}/{len(candidates)} resolved")
        logging.info(f"📄 Report: {run_dir}/disambiguation_decisions.txt")
    
    return final_result

async def filter_acronyms_batch(acronyms: Dict[str, str], client, enable_thinking: bool, disable_streaming: bool, run_dir: _pl.Path, parallel: int) -> Dict[str, str]:
    """Filter acronyms with parallel batch processing"""
    all_decisions = {"kept": [], "discarded": [], "merged": [], "missing": []}
    
    async def process(batch, batch_info="", pbar=None):
        text = "\n".join([f'"{k}": "{v}"' for k, v in batch.items()])
        if not text.strip():
            return {}
            
        data = await two_pass_call(
            client, PROMPTS["filter_pro"].format(acronym_text=text),
            PROMPTS["filter_flash"],
            SCHEMAS["filter"],
            "Filtering", enable_thinking, disable_streaming, batch_info, pbar
        )
        
        filtered = {}
        processed_acronyms = set()
        
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict):
                    acronym = item.get("acronym", "").upper()
                    action = item.get("action", "").lower()
                    reason = item.get("reason", "")
                    processed_acronyms.add(acronym)
                    
                    if action == "keep":
                        for k, v in batch.items():
                            if k.upper() == acronym:
                                filtered[acronym] = v
                                all_decisions["kept"].append(f"{acronym}: {v[:80]}{'...' if len(v) > 80 else ''}")
                                break
                    elif action == "discard":
                        all_decisions["discarded"].append(f"{acronym}: {reason}")
                    elif action == "merge":
                        merge_with = item.get("merge_with", "").upper()
                        all_decisions["merged"].append(f"{acronym} → {merge_with}: {reason}")
        
        # Track missing items
        for k, v in batch.items():
            if k.upper() not in processed_acronyms:
                all_decisions["missing"].append(f"{k.upper()}: '{v[:60]}{'...' if len(v) > 60 else ''}' (not in response)")
        
        return filtered
    
    result = await batch_process(acronyms, process, "Filtering", client, enable_thinking, disable_streaming, run_dir, parallel, 60000)
    
    # Write report
    write_report(run_dir / "filtering_decisions.txt", "ACRONYM FILTERING DECISIONS REPORT", all_decisions)
    
    total_removed = len(all_decisions["discarded"]) + len(all_decisions["merged"]) + len(all_decisions["missing"])
    logging.info(f"📊 FILTERING: {len(result)} kept, {total_removed} removed")
    logging.info(f"📄 Report: {run_dir}/filtering_decisions.txt")
    
    return result


async def run_full_pipeline(input_path: str, acronyms_path: str, run_dir: _pl.Path, json_chunks: bool, client,
                           window: int, parallel: int, enable_thinking: bool, disable_streaming: bool,
                           skip_extraction: bool, skip_disambiguation: bool, skip_filtering: bool):
    """Run the full pipeline: extraction → disambiguation → filtering"""
    
    ontology = {}
    
    if not skip_extraction:
        print("🚀 Starting extraction...")
        # Use the existing run_extract logic but simplified for pipeline
        acronyms = load_acronyms(acronyms_path)
        if not acronyms: 
            print("No acronyms.")
            return
        
        acr_list = ", ".join(acronyms)
        max_chunk_tokens = window - PROMPT_TOK - ANSWER_TOK - tlen(acr_list)
        
        # Load content (preserve original chunks if json_chunks=True)
        content = load_content(input_path, json_chunks)
        
        # Group chunks intelligently to stay within token limits
        if json_chunks and len(content) > 1:
            # Preserve content-aware chunks, group them by token size
            print(f"📦 Loaded {len(content)} content-aware chunks")
            chunks = []
            current_group = ""
            current_tokens = 0
            
            for i, chunk_text in enumerate(content):
                chunk_tokens = tlen(chunk_text)
                
                # If single chunk exceeds limit, split it (unavoidable)
                if chunk_tokens > max_chunk_tokens:
                    if current_group:
                        chunks.append(current_group)
                        current_group = ""
                        current_tokens = 0
                    
                    # Split oversized chunk
                    splitter = TokenTextSplitter(encoding_name="cl100k_base", chunk_size=max_chunk_tokens, chunk_overlap=OVERLAP_TOK)
                    split_chunks = splitter.split_text(chunk_text)
                    chunks.extend(split_chunks)
                    print(f"   ⚠️  Chunk {i+1} ({chunk_tokens:,} tokens) split into {len(split_chunks)} parts")
                
                # If adding this chunk would exceed limit, start new group
                elif current_tokens + chunk_tokens > max_chunk_tokens:
                    if current_group:
                        chunks.append(current_group)
                    current_group = chunk_text
                    current_tokens = chunk_tokens
                
                # Add chunk to current group
                else:
                    if current_group:
                        current_group += "\n\n" + chunk_text
                        current_tokens += chunk_tokens + 2  # +2 for "\n\n"
                    else:
                        current_group = chunk_text
                        current_tokens = chunk_tokens
            
            # Add final group
            if current_group:
                chunks.append(current_group)
            
            print(f"📊 Grouped into {len(chunks)} processing chunks (preserving content-aware boundaries)")
        else:
            # Fallback to original splitting for non-chunk content
            splitter = TokenTextSplitter(encoding_name="cl100k_base", chunk_size=max_chunk_tokens, chunk_overlap=OVERLAP_TOK)
            chunks = []
            for text in content:
                chunks.extend(splitter.split_text(text))
            print(f"📊 Split into {len(chunks)} chunks with token-based splitting")
        
        print(f"🔄 Processing {len(chunks)} chunks with {parallel} workers")

        # Process chunks
        sem = asyncio.Semaphore(parallel)
        
        # Process chunks with proper progress tracking
        pbar = tqdm.tqdm(total=len(chunks), desc="🔄 Extracting", file=sys.stdout)
        tasks = [call_chunk(client, i, ck, acr_list, sem, enable_thinking, disable_streaming, pbar) for i, ck in enumerate(chunks)]
        results = []
        
        for task in asyncio.as_completed(tasks):
            result = await task
            results.append(result)
            pbar.update(1)
        
        pbar.close()
        
        # Aggregate results
        total_prompt_tokens = total_response_tokens = successful = failed = 0
        for idx, items, prompt_tok, response_tok in results:
            total_prompt_tokens += prompt_tok
            total_response_tokens += response_tok
            
            if items:
                successful += 1
            else:
                failed += 1
            
            for item in items:
                ac = item["acronym"].upper()
                rec = ontology.setdefault(ac, {"definitions": [], "contexts": [], "confidence_levels": []})
                if item["definition"] not in rec["definitions"]:
                    rec["definitions"].append(item["definition"])
                    rec["contexts"].append(item.get("context", ""))
                    rec["confidence_levels"].append(item["confidence"])
        
        print(f"✅ Extraction complete: {successful} successful, {failed} failed")
        print(f"📊 Tokens: {total_prompt_tokens:,} → {total_response_tokens:,}")
        print(f"💰 Cost: ~${(total_prompt_tokens * 0.075 + total_response_tokens * 0.3) / 1_000_000:.3f}")
    else:
        print("⏭️  Skipping extraction, loading existing ontology...")
        ontology = json.loads(_pl.Path(input_path).read_text(encoding="utf-8"))
        print(f"📖 Loaded {len(ontology)} acronyms from ontology")
    
    # Create simple mapping with sophisticated disambiguation
    simple, ambiguous = {}, {}
    
    # Identify clear vs ambiguous cases
    for ac, d in ontology.items():
        hi_defs = [(df, ctx) for df, conf, ctx in zip(d["definitions"], d["confidence_levels"], d["contexts"]) if conf == "high"]
        if len(hi_defs) == 1:
            simple[ac] = hi_defs[0][0]
        elif len(hi_defs) > 1:
            ambiguous[ac] = hi_defs
        else:
            simple[ac] = d["definitions"][0] if d["definitions"] else ""
    
    print(f"\n📋 Analysis: {len(simple)} clear, {len(ambiguous)} ambiguous")
    
    # Sophisticated disambiguation for ambiguous cases
    if not skip_disambiguation and ambiguous:
        print(f"🎯 Starting disambiguation for {len(ambiguous)} acronyms...")
        picks = await disambiguate_batch(ambiguous, client, enable_thinking, disable_streaming, run_dir, parallel)
        simple.update(picks)
        print(f"✅ Disambiguation complete: {len(picks)}/{len(ambiguous)} resolved")
    elif skip_disambiguation:
        print("⏭️  Skipping disambiguation, using first definitions...")
        for ac, defs in ambiguous.items():
            simple[ac] = defs[0][0]  # Use first definition
    
    # Sophisticated filtering
    final_simple = simple
    if not skip_filtering:
        print(f"\n🧹 Starting filtering for {len(simple)} acronyms...")
        final_simple = await filter_acronyms_batch(simple, client, enable_thinking, disable_streaming, run_dir, parallel)
        removed_count = len(simple) - len(final_simple)
        print(f"✅ Filtering complete: {len(final_simple)}/{len(simple)} kept ({removed_count} removed)")
    else:
        print("⏭️  Skipping filtering...")
    
    # Save results
    full_path = run_dir / "ontology.json"
    simple_path = run_dir / "ontology_simple.json"
    
    full_path.write_text(json.dumps(ontology, indent=2, ensure_ascii=False))
    simple_path.write_text(json.dumps(final_simple, indent=2, ensure_ascii=False))
    
    print(f"\n🏁 Pipeline complete!")
    print(f"📄 Full ontology: {full_path}")
    print(f"📄 Filtered mapping: {simple_path}")
    print(f"📊 Final result: {len(ontology)} extracted → {len(final_simple)} filtered")

if __name__ == "__main__":
    main()