diff --git "a/page_files/categorized/Backend/PDF_DataExtraction.py" "b/page_files/categorized/Backend/PDF_DataExtraction.py" new file mode 100644--- /dev/null +++ "b/page_files/categorized/Backend/PDF_DataExtraction.py" @@ -0,0 +1,2410 @@ +""" +DocToDB — Materials Science PDF Extractor (Dual LLM Consensus Pipeline) +======================================================================= +PRIMARY: Gemini 2.0 Flash (first preference, no 1.5) +SECONDARY: GPT-4o + +FLOW: + 1. Extract all chunks (tables + text) + 2. Index into ChromaDB for semantic ranking + 3. Rank ALL chunks via single schema-derived retrieval query + 4. Build batches + 5. Run Gemini + GPT-4o in PARALLEL (ThreadPoolExecutor) + - Each LLM respects its own internal rate-limit delays independently + - No shared state → no interference between the two APIs + 6. Consensus filter — only rows matching in BOTH results are kept + (fuzzy: same property_name+section+material_name, value within 5%) + 7. DOI auto-extracted from PDF by both LLMs; merged with override support + 8. DOI rendered as clickable link in Streamlit UI + +Rate limit handling: + - Gemini: exponential backoff 60s→120s→180s, 65s between batches + - GPT-4o: exponential backoff 30s→60s→90s, 10s between batches + - ChromaDB persistence (no re-embedding for same PDF) + - JSON cache (no re-calling LLMs for same PDF) +""" + +from __future__ import annotations + +import hashlib +import io +import json +import logging +import os +import re +import tempfile +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Tuple + +import fitz # pymupdf +import numpy as np +import pandas as pd +import pdfplumber +import requests +from dotenv import load_dotenv +import threading + +load_dotenv() # loads .env from the working directory + + +# ── optional deps ───────────────────────────────────────────────────────────── +try: + from docling.document_converter import DocumentConverter + DOCLING_AVAILABLE = True +except ImportError: + DOCLING_AVAILABLE = False + +try: + import chromadb + CHROMA_AVAILABLE = True +except ImportError: + CHROMA_AVAILABLE = False + logging.warning("chromadb not installed — pip install chromadb") + +try: + from sentence_transformers import SentenceTransformer + ST_AVAILABLE = True +except ImportError: + ST_AVAILABLE = False + logging.warning("sentence-transformers not installed — pip install sentence-transformers") + +try: + import camelot + CAMELOT_AVAILABLE = True +except ImportError: + CAMELOT_AVAILABLE = False + +try: + import pytesseract + from PIL import Image + OCR_AVAILABLE = True +except ImportError: + OCR_AVAILABLE = False + +# ───────────────────────────────────────────────────────────────────────────── +# CONFIG +# ───────────────────────────────────────────────────────────────────────────── + +logging.basicConfig(level=logging.INFO, format="%(levelname)s │ %(message)s") +log = logging.getLogger(__name__) + +if OCR_AVAILABLE: + pytesseract.pytesseract.tesseract_cmd = r"C:\Program Files\Tesseract-OCR\tesseract.exe" + +# ── API Keys ────────────────────────────────────────────────────────────────── +GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "") +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "") + +# ── Gemini model selection (no 1.5 models) ──────────────────────────────────── +_cached_gemini_model: Optional[str] = None + +GEMINI_PREFERRED_MODELS = [ + "gemini-3.5-flash", # GA, 1M context, 65k output — first choice + "gemini-3.1-flash-lite", # fast/cheap fallback + "gemini-3.1-pro-preview", # pro fallback + "gemini-2.5-flash", # last resort only + # gemini-2.0-*, gemini-1.5-* excluded — deprecated/gone +] + +def get_gemini_model(api_key: str) -> str: + global _cached_gemini_model + if _cached_gemini_model: + return _cached_gemini_model + try: + resp = requests.get( + f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}", + timeout=10, + ) + if not resp.ok: + _cached_gemini_model = "gemini-2.0-flash" + _run_startup_probes(_cached_gemini_model) + return _cached_gemini_model + models = resp.json().get("models", []) + available = [ + m["name"].replace("models/", "") + for m in models + if "generateContent" in m.get("supportedGenerationMethods", []) + ] + for pref in GEMINI_PREFERRED_MODELS: + for name in available: + if name.startswith(pref): + log.info(f"Gemini auto-selected: {name}") + _cached_gemini_model = name + return _cached_gemini_model + flash_models = [ + n for n in available + if "flash" in n + and "1.5" not in n + and "2.0" not in n + and "tts" not in n # exclude audio models + and "audio" not in n # exclude audio models + and "image" not in n # exclude image-gen models + ] + _cached_gemini_model = flash_models[0] if flash_models else "gemini-3.2-flash" + _run_startup_probes(_cached_gemini_model) + return _cached_gemini_model + except Exception as e: + log.warning(f"Gemini model auto-detect failed: {e}") + _cached_gemini_model = "gemini-2.0-flash" + _run_startup_probes(_cached_gemini_model) + return _cached_gemini_model + + +GEMINI_MODEL = get_gemini_model(GEMINI_API_KEY) +GEMINI_API_URL = ( + f"https://generativelanguage.googleapis.com/v1beta/" + f"models/{GEMINI_MODEL}:generateContent" + f"?key={GEMINI_API_KEY}" +) + +# ── OpenAI / GPT-4o ─────────────────────────────────────────────────────────── +GPT_MODEL = "gpt-4o" +GPT_API_URL = "https://api.openai.com/v1/chat/completions" +#GPT_MAX_TOKENS = 16384 + +# ── ChromaDB ───────────────────────────────────────────────────────────────── +CHROMA_PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "./chroma_store") +CHROMA_COLLECTION = "doctodb_chunks" +EMBED_MODEL_NAME = "all-MiniLM-L6-v2" +NEIGHBOR_OVERLAP = 0 +# Pre-warm the embedding model at import time so the first PDF doesn't pay +# the cold-load cost. Runs in a daemon thread — doesn't block startup. +def _prewarm_embed_model() -> None: + try: + _get_embed_model() + log.info("Embedding model pre-warmed.") + except Exception as e: + log.warning(f"Embed model pre-warm failed: {e}") + +threading.Thread(target=_prewarm_embed_model, daemon=True).start() +# ── Batching ───────────────────────────────────────────────────────────────── + +# ── Adaptive config (probed at startup, falls back to safe defaults) ────── +_GEMINI_DEFAULTS = {"max_input_chars": 60_000, "batch_delay": 5, "max_output_tokens": 32_768} +_GPT_DEFAULTS = {"max_input_chars": 50_000, "batch_delay": 2, "max_tokens": 4096} +_config_lock = threading.Lock() +_GEMINI_CONFIG = dict(_GEMINI_DEFAULTS) +_GPT_CONFIG = dict(_GPT_DEFAULTS) + + +def _probe_gemini_limits(model_name: str) -> None: + """Fetch real token limits from Gemini model metadata and update _GEMINI_CONFIG.""" + try: + url = ( + f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}" + f"?key={GEMINI_API_KEY}" + ) + r = requests.get(url, timeout=10) + if r.status_code != 200: + return + info = r.json() + input_tok = info.get("inputTokenLimit", 200_000) + output_tok = info.get("outputTokenLimit", 32_768) + rpm = info.get("rpm", 15) # not always present; default free-tier + + # Use 70% of input limit, leaving headroom for prompt overhead + max_chars = min(int(input_tok * 3 * 0.70), 500_000) + with _config_lock: + _GEMINI_CONFIG["max_input_chars"] = max_chars + _GEMINI_CONFIG["max_output_tokens"] = min(output_tok, 32768) + _GEMINI_CONFIG["batch_delay"] = max(2, round(60 / rpm)) + log.info("Gemini limits probed: input=%d tokens → %d chars, delay=%ds", + input_tok, max_chars, _GEMINI_CONFIG["batch_delay"]) + except Exception as exc: + log.warning("Gemini limit probe failed (%s) — using defaults", exc) + + +def _probe_gpt_limits() -> None: + """Make a 1-token dummy call to GPT-4o and read rate-limit headers.""" + try: + headers = { + "Authorization": f"Bearer {OPENAI_API_KEY}", + "Content-Type": "application/json", + } + body = { + "model": "gpt-4o", + "max_tokens": 1, + "messages": [{"role": "user", "content": "hi"}], + } + r = requests.post( + "https://api.openai.com/v1/chat/completions", + headers=headers, json=body, timeout=15 + ) + # Headers are present even on successful responses + tpm = int(r.headers.get("x-ratelimit-limit-tokens", 30_000)) + rpm = int(r.headers.get("x-ratelimit-limit-requests", 500)) + + # 50% of TPM for input (rest is output budget), 15% of TPM for max_tokens + max_chars = int(tpm * 0.50 * 3) # tokens → chars + max_tok = int(tpm * 0.15) + + with _config_lock: + _GPT_CONFIG["max_input_chars"] = max_chars + _GPT_CONFIG["max_tokens"] = max_tok + _GPT_CONFIG["batch_delay"] = max(1, round(60 / rpm)) + log.info("GPT limits probed: tpm=%d → %d chars, rpm=%d, delay=%ds", + tpm, max_chars, rpm, _GPT_CONFIG["batch_delay"]) + except Exception as exc: + log.warning("GPT limit probe failed (%s) — using defaults", exc) + + +def _run_startup_probes(model_name: str) -> None: + """Run both probes in parallel — called once after model selection.""" + with ThreadPoolExecutor(max_workers=2) as ex: + ex.submit(_probe_gemini_limits, model_name) + ex.submit(_probe_gpt_limits) + + +# MAX_BATCH_CHARS is the shared floor — set after probes complete +def _get_max_batch_chars() -> int: + return min(_GEMINI_CONFIG["max_input_chars"], _GPT_CONFIG["max_input_chars"]) +GEMINI_RETRY_DELAYS = [60, 120, 180] +GPT_RETRY_DELAYS = [30, 60, 90] + +# ── Extraction ──────────────────────────────────────────────────────────────── +MIN_TABLE_ROWS = 2 +MIN_TABLE_COLS = 2 +MIN_CELL_CHARS = 2 +MIN_TEXT_CHARS = 40 +OCR_THRESHOLD = 50 + +# ── Consensus ──────────────────────────────────────────────────────────────── +CONSENSUS_VALUE_TOL = 0.05 # 5 % numeric tolerance for value matching + +CACHE_FILE = "pdf_extraction_cache.json" + +# ───────────────────────────────────────────────────────────────────────────── +# SINGLE RETRIEVAL QUERY +# ───────────────────────────────────────────────────────────────────────────── + +RETRIEVAL_QUERY = ( + "Material property data including section category, property name, " + "measured value, SI unit, imperial unit, test condition standard " + "such as ASTM ISO DIN, and comments. Properties include mechanical " + "thermal electrical physical rheological optical categories. " + "Values in MPa GPa percent density conductivity temperature modulus " + "strength elongation hardness viscosity flammability. " + "Material name, manufacturer, trade grade, abbreviation. " + "DOI digital object identifier paper reference link." +) + +# ───────────────────────────────────────────────────────────────────────────── +# GEMINI SCHEMA + PROMPT +# ───────────────────────────────────────────────────────────────────────────── + +SCHEMA = { + "type": "OBJECT", + "properties": { + "material_name": {"type": "STRING"}, + "material_abbreviation": {"type": "STRING"}, + #"trade_grade": {"type": "STRING", "maxLength": 100}, + "manufacturer": {"type": "STRING", "maxLength": 100}, + "doi": {"type": "STRING"}, + "mechanical_properties": { + "type": "ARRAY", + "items": { + "type": "OBJECT", + "properties": { + "section": {"type": "STRING"}, + "property_name": {"type": "STRING"}, + "value": {"type": "STRING"}, + "unit": {"type": "STRING"}, + "english": {"type": "STRING"}, + "test_condition": {"type": "STRING"}, + "comments": {"type": "STRING"}, + "material_name": {"type": "STRING"}, + "source_page": {"type": "STRING"}, + "chunk_type": {"type": "STRING"}, + "source_text": {"type": "STRING"}, + }, + "required": [ + "section", "property_name", "value", "unit", + "english", "test_condition", "comments", + "material_name", "source_page", "chunk_type", + "source_text", + ], + }, + }, + }, +} + +# EXTRACTION_PROMPT = ( +# "You are an expert materials scientist. " +# "The content below is extracted from a materials datasheet or research paper. " +# "Each block is tagged with its page number and type (TABLE or TEXT).\n\n" +# "Extract every material property. For each property record:\n" +# " - section : category (Mechanical, Thermal, Electrical, Physical, Rheological, etc.)\n" +# " - property_name : exact name as written\n" +# " - value : exact value or range\n" +# " - unit : SI unit\n" +# " - english : imperial equivalent if shown, else ''\n" +# " - test_condition : ASTM/ISO/DIN standard or conditions, else ''\n" +# " - comments : footnotes or qualifications, else ''\n" +# " - material_name : exact material name (never leave blank)\n" +# " - source_page : page number from the block header (digits only, e.g. '3')\n" +# " - chunk_type : 'table' if from a TABLE block, 'text' if from a TEXT block\n\n" +# "Also extract at the top level:\n" +# " - doi : the DOI of the paper/datasheet (e.g. '10.1016/j.polymer.2023.01.001'). " +# " If a clickable 'click here' link or https://doi.org/... URL appears, " +# " extract just the DOI identifier (without the https://doi.org/ prefix). " +# " Leave empty string '' if not found.\n\n" +# "RULES:\n" +# " - Extract ONLY measured/specified properties, NOT equations or simulation params.\n" +# " - If multiple materials appear, create separate entries for each.\n" +# " - Preserve source_page exactly from the block header.\n" +# " - Respond ONLY with valid JSON matching the schema.\n" +# "\n\nCONTENT:\n" +# ) +EXTRACTION_PROMPT = ( + "You are an expert materials scientist and information extraction engine.\n" + "The content below is extracted from a materials datasheet or research paper.\n" + "Each block is tagged with its page number and type (TABLE or TEXT).\n\n" + "Extract EVERY material property with a numeric value. This includes properties " + "found in tables, prose sentences, captions, and comparative statements.\n" + "Example: 'The tensile strength of PLA/WF composites was 45.3 MPa' is a valid extraction.\n\n" + "Extract the following TOP-LEVEL fields (one value for the whole document):\n" + " - material_name : material name only, not a full description paragraph.\n" + " - material_abbreviation : short abbreviation (e.g. 'PLA', 'WF/PLA').\n" + " - trade_grade : ONLY a short grade code under 20 characters (e.g. 'PLA002', 'Grade A'). " + " If no specific product grade code exists in the text, leave exactly ''. " + " DO NOT put descriptions, study context, sentences, or multiple grades here. " + " If multiple grades exist, pick the primary one only.\n" + " - manufacturer : company name only (e.g. 'Kuraray Company'). " + "Do not include descriptions or test conditions here.\n" + " - doi : DOI identifier only (e.g. '10.1016/j.polymer.2023.01.001'). " + "Strip any https://doi.org/ prefix. Leave '' if not found.\n\n" + "For each property record extract:\n" + " - section : must be one of: Mechanical, Thermal, Electrical, " + "Physical, Rheological, Optical, Chemical, Barrier, Morphological, Surface, Other\n" + " - property_name : exact name as written in the source\n" + " - value : exact value or midpoint of range\n" + " - unit : unit exactly as written\n" + " - english : imperial equivalent if shown, else ''\n" + " - test_condition : ASTM/ISO/DIN standard or test conditions, else ''\n" + " - comments : footnotes, qualifications, or ± uncertainty, else ''\n" + " - material_name : exact material name for this property (never leave blank)\n" + " - source_page : page number from the block header (digits only, e.g. '3')\n" + " - chunk_type : 'table' if from a TABLE block, 'text' if from a TEXT block\n\n" + " - source_text : copy the exact sentence or table row this property " + "was extracted from. Never leave blank if the property came from text.\n" + "INCLUDE ALL PROPERTY TYPES — not just mechanical:\n" + "tensile strength, modulus, elongation, impact strength, hardness, density, " + "glass transition temperature, melting temperature, decomposition temperature, " + "thermal conductivity, crystallinity, molecular weight, melt flow index, " + "degradation rate, water absorption, electrical conductivity, viscosity, " + "contact angle, and any other numeric material characteristic.\n\n" + "RULES:\n" + " - Extract from BOTH tables and prose — never skip text blocks.\n" + " - Extract values explicitly stated in the text — do not invent values not present.\n" + " - If multiple materials appear, create separate entries for each.\n" + " - NEVER return an empty properties array if any numeric value exists.\n" + " - Preserve source_page exactly from the block header.\n" + " - Respond ONLY with valid JSON matching the schema.\n" + "\n\nCONTENT:\n" +) +# GPT-4o uses plain-text prompt (no structured output schema enforcement, +# but we instruct it to return the same JSON shape) +# GPT_SYSTEM_PROMPT = ( +# "You are an expert materials scientist. " +# "Extract material properties from the provided PDF content blocks. " +# "Respond ONLY with valid JSON in exactly this shape — no markdown, no extra text:\n" +# "{\n" +# ' "material_name": "...",\n' +# ' "material_abbreviation": "...",\n' +# ' "trade_grade": "...",\n' +# ' "manufacturer": "...",\n' +# ' "doi": "...",\n' +# ' "mechanical_properties": [\n' +# ' {\n' +# ' "section": "Mechanical",\n' +# ' "property_name": "Tensile Strength",\n' +# ' "value": "85",\n' +# ' "unit": "MPa",\n' +# ' "english": "",\n' +# ' "test_condition": "ASTM D638",\n' +# ' "comments": "",\n' +# ' "material_name": "ABS",\n' +# ' "source_page": "3",\n' +# ' "chunk_type": "table"\n' +# ' }\n' +# ' ]\n' +# "}\n\n" +# "Rules:\n" +# "- section: Mechanical / Thermal / Electrical / Physical / Rheological / Optical\n" +# "- Extract ONLY measured/specified properties, NOT equations or simulation params.\n" +# "- doi: extract DOI identifier only (e.g. '10.1016/j.mat.2023.01.001'). " +# " If a 'click here' link or https://doi.org/... URL appears in the text, strip the prefix.\n" +# "- source_page: digits only from block header.\n" +# "- chunk_type: 'table' or 'text' from block header.\n" +# "- Leave fields as '' if unknown." +# ) +GPT_SYSTEM_PROMPT = ( + "You are an expert materials-science information extraction engine.\n" + "Your task is to exhaustively extract ALL material-property data from scientific PDF text.\n\n" + "The PDF may contain narrative prose, tables, captions, figure descriptions, " + "comparative statements, and ranges. Extract properties from ALL of them.\n\n" + "OUTPUT RULES:\n" + "- Return ONLY valid JSON. No markdown, no explanations, no comments, no trailing commas.\n" + "- Output MUST exactly match the schema below.\n" + "- Never omit required fields. Use \"\" for unknown or missing values.\n\n" + "- value: ALWAYS fill this. Put the exact number here (e.g. '51.8'). " + "Never leave value blank if a number exists for this property. " + "For ranges like '40–52': put '40–52' in value AND '40' in min_value AND '52' in max_value.\n" + "JSON SCHEMA:\n" + "{\n" + ' "material_name": "",\n' + ' "material_abbreviation": "",\n' + ' "trade_grade": "",\n' + ' "manufacturer": "",\n' + ' "doi": "",\n' + ' "properties": [\n' + ' {\n' + ' "section": "",\n' + ' "property_name": "",\n' + ' "value": "",\n' + ' "min_value": "",\n' + ' "max_value": "",\n' + ' "unit": "",\n' + ' "test_condition": "",\n' + ' "comments": "",\n' + ' "source_text": "",\n' + ' "material_name": "",\n' + ' "source_page": "",\n' + ' "chunk_type": ""\n' + ' }\n' + ' ]\n' + "}\n\n" + "CRITICAL EXTRACTION RULES:\n" + "- Extract EVERY material property with a numeric value and unit.\n" + "- Extract from BOTH prose and tables — sentences like 'The tensile strength was 45.3 MPa' " + "are valid extractions.\n" + "- Extract ranges and ± values: put the midpoint or reported value in 'value', " + "lower bound in 'min_value', upper bound in 'max_value'.\n" + "- Extract repeated measurements only once unless conditions differ.\n" + "- NEVER ignore a property because it is outside mechanical testing.\n" + "- NEVER return an empty properties array if any numeric material property exists.\n\n" + "INCLUDE ALL PROPERTY TYPES:\n" + "Mechanical, Thermal, Electrical, Physical, Rheological, Optical, " + "Chemical, Barrier, Morphological, Surface, Crystallinity, Biodegradation, Other\n\n" + "EXAMPLES OF VALID PROPERTIES:\n" + "tensile strength, elastic modulus, Young's modulus, flexural modulus, " + "elongation at break, impact strength, hardness, density, melt flow index, " + "viscosity, glass transition temperature, melting temperature, decomposition temperature, " + "thermal conductivity, electrical conductivity, crystallinity, molecular weight, " + "degradation rate, water absorption, contact angle\n\n" + "FIELD RULES:\n" + "- section must be one of: Mechanical, Thermal, Electrical, Physical, Rheological, " + "Optical, Chemical, Barrier, Morphological, Surface, Other\n" + "- source_page: digits only from the block header.\n" + "- chunk_type: 'table' or 'text' from the block header.\n" + "- source_text: copy the exact sentence or table row this property was extracted from.\n" + "- material_name: exact material name, never leave blank.\n" + "- Leave all other fields as \"\" if unknown — never omit them.\n\n" + "NORMALIZATION RULES:\n" + "- Preserve original units and property names exactly as written.\n" + "- Do NOT convert units, infer missing values, or hallucinate standards.\n" + "- If multiple materials appear, extract the primary material being characterized.\n\n" + "ANTI-HALLUCINATION:\n" + "- Only extract values explicitly stated in the provided text.\n" + "- If a value is not present, do not invent it.\n" + "- Missing a weakly-labeled property is acceptable. Inventing one is not.\n" +) +# ───────────────────────────────────────────────────────────────────────────── +# BOILERPLATE DETECTION +# ───────────────────────────────────────────────────────────────────────────── + +_SKIP_HEADING_RE = re.compile( + r"^(references|bibliography|acknowledgements?|table\s+of\s+contents|" + r"copyright|legal\s+notice|disclaimer|index|appendix\s+[a-z]$)", + re.IGNORECASE, +) + +def _is_boilerplate(text: str) -> bool: + return bool(_SKIP_HEADING_RE.match(text.strip().split("\n")[0].strip())) + +# ───────────────────────────────────────────────────────────────────────────── +# DATA STRUCTURES +# ───────────────────────────────────────────────────────────────────────────── + +@dataclass +class Chunk: + page_num: int + chunk_type: str + source: str + raw_rows: Optional[List[List[str]]] = None + raw_text: Optional[str] = None + text: str = field(init=False) + score: float = 0.0 + relevant: bool = False + + def __post_init__(self): + if self.chunk_type == "table" and self.raw_rows: + self.text = _rows_to_text(self.raw_rows) + elif self.raw_text: + self.text = self.raw_text.strip() + else: + self.text = "" + + @property + def block_header(self) -> str: + return ( + f"\n\n{'─'*60}\n" + f"[{self.chunk_type.upper()} | Page {self.page_num} | " + f"score={self.score:.3f}]\n" + f"{'─'*60}\n" + ) + +# ───────────────────────────────────────────────────────────────────────────── +# HELPERS +# ───────────────────────────────────────────────────────────────────────────── + +def _rows_to_text(rows: List[List[Any]]) -> str: + lines = [] + for row in rows: + cells = [str(c).strip() if c is not None else "" for c in row] + if any(len(c) >= MIN_CELL_CHARS for c in cells): + lines.append(" | ".join(cells)) + return _fix_spaced_text("\n".join(lines)) + + +def _is_valid_table(rows: List[List[Any]]) -> bool: + if not rows or len(rows) < MIN_TABLE_ROWS: + return False + if max((len(r) for r in rows), default=0) < MIN_TABLE_COLS: + return False + for row in rows[1:]: + if any(re.search(r"\d", str(c)) for c in row): + return True + return False + + +def _split_paragraphs(raw_text: str) -> List[str]: + paragraphs: List[str] = [] + current: List[str] = [] + for line in raw_text.splitlines(): + stripped = line.strip() + if not stripped: + if current: + paragraphs.append(" ".join(current)) + current = [] + else: + current.append(stripped) + if current: + paragraphs.append(" ".join(current)) + return [_fix_spaced_text(p) for p in paragraphs if len(p) >= MIN_TEXT_CHARS] + + +def _normalise_doi(raw: str) -> str: + """Strip URL prefix and whitespace, return bare DOI or empty string.""" + if not raw: + return "" + doi = raw.strip() + for prefix in ("https://doi.org/", "http://doi.org/", "doi.org/", "DOI:", "doi:"): + if doi.lower().startswith(prefix.lower()): + doi = doi[len(prefix):] + return doi.strip() + + +def _doi_url(doi: str) -> str: + """Return full clickable URL for a bare DOI.""" + doi = _normalise_doi(doi) + if not doi: + return "" + return f"https://doi.org/{doi}" + +# ───────────────────────────────────────────────────────────────────────────── +# CACHE +# ───────────────────────────────────────────────────────────────────────────── + +def _pdf_hash(pdf_bytes: bytes) -> str: + return hashlib.sha256(pdf_bytes).hexdigest()[:16] + +def _load_cache() -> Dict: + try: + if os.path.exists(CACHE_FILE): + with open(CACHE_FILE) as f: + return json.load(f) + except Exception: + pass + return {} + +def _save_cache(cache: Dict): + try: + with open(CACHE_FILE, "w") as f: + json.dump(cache, f) + except Exception as e: + log.warning(f"Cache save failed: {e}") + +def cache_get(pdf_bytes: bytes) -> Optional[Dict]: + return _load_cache().get(_pdf_hash(pdf_bytes)) + +def cache_set(pdf_bytes: bytes, result: Dict): + cache = _load_cache() + cache[_pdf_hash(pdf_bytes)] = result + _save_cache(cache) + +# ───────────────────────────────────────────────────────────────────────────── +# CHROMADB CLIENT +# ───────────────────────────────────────────────────────────────────────────── + +_chroma_client: Optional[Any] = None +_chroma_collection: Optional[Any] = None +_embed_model: Optional[Any] = None + + +def _get_embed_model() -> Any: + global _embed_model + if _embed_model is None: + if not ST_AVAILABLE: + raise ImportError("sentence-transformers not installed.\nRun: pip install sentence-transformers") + log.info(f"Loading embedding model '{EMBED_MODEL_NAME}' …") + _embed_model = SentenceTransformer(EMBED_MODEL_NAME) + log.info("Embedding model ready.") + return _embed_model + + +def _get_chroma_collection() -> Any: + global _chroma_client, _chroma_collection + if _chroma_collection is not None: + return _chroma_collection + if not CHROMA_AVAILABLE: + raise ImportError("chromadb not installed. Run: pip install chromadb") + os.makedirs(CHROMA_PERSIST_DIR, exist_ok=True) + _chroma_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR) + _chroma_collection = _chroma_client.get_or_create_collection( + name=CHROMA_COLLECTION, + metadata={"hnsw:space": "cosine"}, + ) + log.info(f"ChromaDB ready — '{CHROMA_COLLECTION}' ({_chroma_collection.count()} existing vectors)") + return _chroma_collection + + +def _chroma_pdf_exists(pdf_hash: str) -> bool: + try: + col = _get_chroma_collection() + res = col.get(where={"pdf_hash": pdf_hash}, limit=1) + return len(res["ids"]) > 0 + except Exception: + return False + + +def _chroma_store_chunks(chunks: List[Chunk], pdf_hash: str) -> None: + col = _get_chroma_collection() + model = _get_embed_model() + texts = [c.text for c in chunks] + metadatas = [ + {"pdf_hash": pdf_hash, "page_num": c.page_num, + "chunk_type": c.chunk_type, "source": c.source} + for c in chunks + ] + ids = [f"{pdf_hash}_{i}" for i in range(len(chunks))] + # all_embeddings: List[List[float]] = [] + # for start in range(0, len(texts), 64): + # batch = texts[start : start + 64] + # vecs = model.encode(batch, normalize_embeddings=True) + # all_embeddings.extend(vecs.tolist()) + vecs = model.encode( + texts, + normalize_embeddings=True, + batch_size=256, # internal mini-batch size for GPU memory + show_progress_bar=False, + ) + all_embeddings = vecs.tolist() + + col.upsert(ids=ids, documents=texts, embeddings=all_embeddings, metadatas=metadatas) + log.info(f"ChromaDB: stored {len(chunks)} chunks for pdf_hash={pdf_hash}") + + +def _chroma_rank_all(pdf_hash: str) -> List[Chunk]: + col = _get_chroma_collection() + model = _get_embed_model() + all_ids = col.get(where={"pdf_hash": pdf_hash}, include=[]) + n_total = len(all_ids["ids"]) + if n_total == 0: + log.warning(f"No chunks in ChromaDB for pdf_hash={pdf_hash}") + return [] + query_vec = model.encode([RETRIEVAL_QUERY], normalize_embeddings=True)[0].tolist() + results = col.query( + query_embeddings=[query_vec], + n_results=n_total, + where={"pdf_hash": pdf_hash}, + include=["documents", "metadatas", "distances"], + ) + ranked_chunks: List[Chunk] = [] + for doc, meta, dist in zip( + results.get("documents", [[]])[0], + results.get("metadatas", [[]])[0], + results.get("distances", [[]])[0], + ): + similarity = 1.0 - float(dist) + chunk = Chunk( + page_num=int(meta.get("page_num", 1)), + chunk_type=meta.get("chunk_type", "text"), + source=meta.get("source", ""), + raw_text=doc, + score=similarity, + relevant=True, + ) + ranked_chunks.append(chunk) + table_n = sum(1 for c in ranked_chunks if c.chunk_type == "table") + text_n = sum(1 for c in ranked_chunks if c.chunk_type == "text") + log.info(f"ChromaDB ranked {len(ranked_chunks)} chunks ({table_n} tables + {text_n} text)") + return ranked_chunks + +# ───────────────────────────────────────────────────────────────────────────── +# EXTRACTION — TABLES +# ───────────────────────────────────────────────────────────────────────────── + +def _extract_tables_docling(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + if not DOCLING_AVAILABLE: + return chunks + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes); tmp_path = tmp.name + try: + converter = DocumentConverter() + result = converter.convert(tmp_path) + doc = result.document + for table in doc.tables: + rows = [[cell.text for cell in row] for row in table.data.grid] + page_num = table.prov[0].page_no if table.prov else 1 + if _is_valid_table(rows): + chunks.append(Chunk(page_num=page_num, chunk_type="table", source="docling", raw_rows=rows)) + log.info(f"Docling tables: {len(chunks)}") + except Exception as e: + log.error(f"Docling table extraction failed: {e}") + finally: + os.unlink(tmp_path) + return chunks + + +def _extract_tables_pdfplumber(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + try: + with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf: + for page_num, page in enumerate(pdf.pages, start=1): + for strategy in ( + {"vertical_strategy": "lines_strict", "horizontal_strategy": "lines_strict", + "snap_tolerance": 3, "join_tolerance": 3}, + {"vertical_strategy": "text", "horizontal_strategy": "text"}, + ): + tables = page.extract_tables(table_settings=strategy) or [] + for rows in tables: + cleaned = [[str(c).strip() if c else "" for c in row] for row in rows] + if _is_valid_table(cleaned): + chunks.append(Chunk(page_num=page_num, chunk_type="table", + source="pdfplumber", raw_rows=cleaned)) + if tables: + break + log.info(f"pdfplumber tables: {len(chunks)}") + except Exception as e: + log.error(f"pdfplumber table extraction failed: {e}") + return chunks + + +def _extract_tables_camelot(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + if not CAMELOT_AVAILABLE: + return chunks + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes); tmp_path = tmp.name + try: + for flavor in ("lattice", "stream"): + try: + tables = camelot.read_pdf(tmp_path, pages="all", flavor=flavor) + if flavor == "stream": + tables = [t for t in tables if t.parsing_report.get("accuracy", 0) > 65] + for table in tables: + rows = [[str(c).strip() for c in row] for row in table.df.values.tolist()] + if _is_valid_table(rows): + chunks.append(Chunk(page_num=table.page, chunk_type="table", + source=f"camelot-{flavor}", raw_rows=rows)) + if chunks: + break + except Exception as e: + log.warning(f"camelot {flavor}: {e}") + finally: + os.unlink(tmp_path) + log.info(f"camelot tables: {len(chunks)}") + return chunks + +# ───────────────────────────────────────────────────────────────────────────── +# EXTRACTION — TEXT +# ───────────────────────────────────────────────────────────────────────────── + +def _extract_text_docling(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + if not DOCLING_AVAILABLE: + return chunks + with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: + tmp.write(pdf_bytes); tmp_path = tmp.name + try: + converter = DocumentConverter() + result = converter.convert(tmp_path) + doc = result.document + for item, _ in doc.iterate_items(): + text = getattr(item, "text", "").strip() + page_num = item.prov[0].page_no if item.prov else 1 + if not text or len(text) < MIN_TEXT_CHARS or _is_boilerplate(text): + continue + chunks.append(Chunk(page_num=page_num, chunk_type="text", source="docling", raw_text=text)) + log.info(f"Docling text chunks: {len(chunks)}") + except Exception as e: + log.error(f"Docling text extraction failed: {e}") + finally: + os.unlink(tmp_path) + return chunks + + +def _extract_text_pymupdf(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + try: + with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: + for page_idx, page in enumerate(doc): + page_num = page_idx + 1 + raw = page.get_text("text", sort=True) or "" + for para in _split_paragraphs(raw): + if not _is_boilerplate(para): + chunks.append(Chunk(page_num=page_num, chunk_type="text", + source="pymupdf", raw_text=para)) + log.info(f"pymupdf text chunks: {len(chunks)}") + except Exception as e: + log.error(f"pymupdf text extraction failed: {e}") + return chunks + + +def _extract_text_pdfplumber(pdf_bytes: bytes) -> List[Chunk]: + chunks: List[Chunk] = [] + try: + with pdfplumber.open(io.BytesIO(pdf_bytes)) as pdf: + for page_num, page in enumerate(pdf.pages, start=1): + raw = page.extract_text(x_tolerance=3, y_tolerance=3) or "" + for para in _split_paragraphs(raw): + if not _is_boilerplate(para): + chunks.append(Chunk(page_num=page_num, chunk_type="text", + source="pdfplumber-text", raw_text=para)) + log.info(f"pdfplumber text chunks: {len(chunks)}") + except Exception as e: + log.error(f"pdfplumber text extraction failed: {e}") + return chunks + + +def _ocr_page(pdf_bytes: bytes, page_num: int) -> str: + if not OCR_AVAILABLE: + return "" + try: + with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: + page = doc[page_num - 1] + mat = fitz.Matrix(300 / 72, 300 / 72) + pix = page.get_pixmap(matrix=mat) + img = Image.open(io.BytesIO(pix.tobytes("png"))) + return pytesseract.image_to_string(img, lang="eng") or "" + except Exception as e: + log.warning(f"OCR page {page_num}: {e}") + return "" + + +def _verify_page_coverage(pdf_bytes: bytes, chunks: List[Chunk]) -> List[Chunk]: + with fitz.open(stream=pdf_bytes, filetype="pdf") as doc: + total_pages = set(range(1, len(doc) + 1)) + covered = {c.page_num for c in chunks} + missing = total_pages - covered + if not missing: + return chunks + log.warning(f"Page coverage: {len(missing)} pages missing — OCRing: {sorted(missing)}") + for page_num in sorted(missing): + ocr_text = _ocr_page(pdf_bytes, page_num) + for para in _split_paragraphs(ocr_text): + if not _is_boilerplate(para): + chunks.append(Chunk(page_num=page_num, chunk_type="text", + source="ocr-fallback", raw_text=para)) + return chunks + + +def _dedup(chunks: List[Chunk]) -> List[Chunk]: + seen: set = set() + unique: List[Chunk] = [] + for c in chunks: + key = (c.chunk_type, re.sub(r"\s+", " ", c.text.strip())[:200]) + if key not in seen: + seen.add(key) + unique.append(c) + return unique + + +# def extract_all_chunks(pdf_bytes: bytes) -> List[Chunk]: +# table_chunks: List[Chunk] = [] +# text_chunks: List[Chunk] = [] +# docling_tables = _extract_tables_docling(pdf_bytes) +# table_chunks.extend(docling_tables) +# table_chunks.extend(_extract_tables_pdfplumber(pdf_bytes)) +# table_chunks.extend(_extract_tables_camelot(pdf_bytes)) +# table_chunks = _dedup(table_chunks) +# if DOCLING_AVAILABLE and docling_tables: +# text_chunks.extend(_extract_text_docling(pdf_bytes)) +# pymupdf_text = _extract_text_pymupdf(pdf_bytes) +# text_chunks.extend(pymupdf_text) +# if sum(len(c.text) for c in pymupdf_text) < 500: +# log.warning("pymupdf sparse — supplementing with pdfplumber") +# text_chunks.extend(_extract_text_pdfplumber(pdf_bytes)) +# text_chunks = _dedup(text_chunks) +# all_chunks = table_chunks + text_chunks +# all_chunks = _verify_page_coverage(pdf_bytes, all_chunks) +# log.info(f"Total chunks: {len(all_chunks)} ({len(table_chunks)} tables + {len(text_chunks)} text)") +# return all_chunks + +def extract_all_chunks(pdf_bytes: bytes) -> List[Chunk]: + with ThreadPoolExecutor(max_workers=5) as ex: + futures = { + ex.submit(_extract_tables_docling, pdf_bytes): "docling_tables", + ex.submit(_extract_tables_pdfplumber, pdf_bytes): "pdfplumber_tables", + ex.submit(_extract_tables_camelot, pdf_bytes): "camelot_tables", + ex.submit(_extract_text_pymupdf, pdf_bytes): "pymupdf_text", + ex.submit(_extract_text_pdfplumber, pdf_bytes): "pdfplumber_text", + } + results = {} + for future in as_completed(futures): + name = futures[future] + try: + results[name] = future.result() + except Exception as e: + log.error(f"{name} failed: {e}") + results[name] = [] + + docling_tables = results.get("docling_tables", []) + pdfplumber_tables= results.get("pdfplumber_tables", []) + camelot_tables = results.get("camelot_tables", []) + pymupdf_text = results.get("pymupdf_text", []) + pdfplumber_text = results.get("pdfplumber_text", []) + + table_chunks = _dedup(docling_tables + pdfplumber_tables + camelot_tables) + + text_chunks: List[Chunk] = [] + if DOCLING_AVAILABLE and docling_tables: + # docling text runs after we know docling tables succeeded + text_chunks.extend(_extract_text_docling(pdf_bytes)) + text_chunks.extend(pymupdf_text) + if sum(len(c.text) for c in pymupdf_text) < 500: + log.warning("pymupdf sparse — supplementing with pdfplumber text") + text_chunks.extend(pdfplumber_text) + text_chunks = _dedup(text_chunks) + + all_chunks = table_chunks + text_chunks + all_chunks = _verify_page_coverage(pdf_bytes, all_chunks) + log.info( + f"Total chunks: {len(all_chunks)} " + f"({len(table_chunks)} tables + {len(text_chunks)} text)" + ) + return all_chunks + +def _fix_spaced_text(text: str) -> str: + """Fix PDF kerning artifacts like 'C O M P O S I T E' → 'COMPOSITE'.""" + # Pattern: single chars separated by spaces forming a word + return re.sub( + r'\b([A-Z]) (?=[A-Z] |[A-Z]\b)', + r'\1', + text + ) +# ───────────────────────────────────────────────────────────────────────────── +# CHROMADB INDEX +# ───────────────────────────────────────────────────────────────────────────── + +def _build_overlapping_chunks(chunks: List[Chunk], neighbor_window: int = NEIGHBOR_OVERLAP) -> List[Chunk]: + if neighbor_window <= 0: + return chunks + eligible = [(i, c) for i, c in enumerate(chunks) if not _is_boilerplate(c.text)] + overlap_chunks: List[Chunk] = [] + seen_content: set = set() + for pos, (_, centre) in enumerate(eligible): + if centre.chunk_type == "table": + continue + parts: List[str] = [] + for offset in range(-neighbor_window, neighbor_window + 1): + nb_pos = pos + offset + if nb_pos < 0 or nb_pos >= len(eligible): + continue + _, nb_chunk = eligible[nb_pos] + if nb_chunk.chunk_type == "table": + continue + parts.append(nb_chunk.text.strip()) + merged = " \n\n ".join(p for p in parts if p) + if not merged or merged in seen_content: + continue + seen_content.add(merged) + overlap_chunks.append(Chunk(page_num=centre.page_num, chunk_type="text", + source=f"{centre.source}+overlap", raw_text=merged)) + return chunks + overlap_chunks + + +def index_chunks_in_chroma(chunks: List[Chunk], pdf_hash: str) -> None: + if _chroma_pdf_exists(pdf_hash): + log.info(f"ChromaDB: pdf_hash={pdf_hash} already indexed — skipping.") + return + all_chunks = _build_overlapping_chunks(chunks, neighbor_window=NEIGHBOR_OVERLAP) + storable = [c for c in all_chunks if not _is_boilerplate(c.text) and c.text.strip()] + _chroma_store_chunks(storable, pdf_hash) + +# ───────────────────────────────────────────────────────────────────────────── +# BATCH BUILDER +# ────────────────────────────────────────��──────────────────────────────────── + +def build_batches(chunks: List[Chunk]) -> List[str]: + if not chunks: + return [] + batches: List[str] = [] + current_batch: List[Chunk] = [] + current_chars: int = 0 + for chunk in chunks: + entry_len = len(chunk.block_header) + len(chunk.text) + _max = _get_max_batch_chars() + if entry_len > _max: + chunk.text = chunk.text[: _max - len(chunk.block_header) - 20] + entry_len = len(chunk.block_header) + len(chunk.text) + + if current_chars + entry_len > _max and current_batch: + current_batch.sort(key=lambda c: c.page_num) + batches.append("\n".join(c.block_header + c.text for c in current_batch)) + current_batch = [] + current_chars = 0 + current_batch.append(chunk) + current_chars += entry_len + if current_batch: + current_batch.sort(key=lambda c: c.page_num) + batches.append("\n".join(c.block_header + c.text for c in current_batch)) + log.info(f"Built {len(batches)} batch(es) from {len(chunks)} chunks") + return batches + +# ───────────────────────────────────────────────────────────────────────────── +# GEMINI EXTRACTION +# ───────────────────────────────────────────────────────────────────────────── +def _call_gemini(text_payload: str) -> Tuple[Optional[Dict], str]: + log.info(f"Gemini calling with maxOutputTokens={_GEMINI_CONFIG['max_output_tokens']}") + log.info(f"Gemini input (first 300 chars): {repr(text_payload[:300])}") # add this + + payload = { + "contents": [{"parts": [{"text": EXTRACTION_PROMPT + text_payload}]}], + "generationConfig": { + "temperature": 0, + "responseMimeType": "application/json", + "responseSchema": SCHEMA, + "maxOutputTokens": _GEMINI_CONFIG["max_output_tokens"], + }, + } + for attempt, wait in enumerate([0] + GEMINI_RETRY_DELAYS): + if wait: + log.warning(f"Gemini 429 — waiting {wait}s (retry {attempt})") + time.sleep(wait) + try: + resp = requests.post(GEMINI_API_URL, json=payload, timeout=600) + if resp.status_code == 429: + continue + if not resp.ok: + return None, f"Gemini HTTP {resp.status_code}: {resp.text[:2000]}" + data = resp.json() + candidates = data.get("candidates", []) + if not candidates: + return None, f"Gemini no candidates: {json.dumps(data)[:400]}" + parts = candidates[0].get("content", {}).get("parts", []) + for part in parts: + text = part.get("text", "").strip() + if text.startswith("{"): + try: + # Layer 2 — regex truncation of abused top-level string fields + text = re.sub( + r'("(?:trade_grade|manufacturer)"\s*:\s*")' + r'([^"]{150})[^"]*"', + r'\1\2"', + text + ) + + parsed = json.loads(text) + + # Layer 3 — post-parse semantic validation + tg = parsed.get("trade_grade", "") + if tg and ( + len(tg) > 30 + or tg.count(" ") > 3 + or tg.count("_") > 2 + or any(w in tg.lower() for w in + ["study", "used", "modified", "provided", "analysis", + "document", "without", "unless", "specified", "otherwise", + "percent", "weight", "fiber", "composite", "matrix"]) + ): + log.warning(f"Clearing abusive trade_grade: {repr(tg[:80])}") + parsed["trade_grade"] = "" + + mfr = parsed.get("manufacturer", "") + if mfr and (len(mfr) > 100 or mfr.count("_") > 2): + log.warning(f"Truncating abusive manufacturer: {repr(mfr[:80])}") + parsed["manufacturer"] = mfr[:100].rsplit(" ", 1)[0] + + return parsed, "" + + except json.JSONDecodeError as e: + log.error(f"Gemini raw response (first 500 chars): {repr(text[:500])}") + return None, f"Gemini JSON parse error: {e}" + + return None, "No JSON in Gemini response." + except requests.Timeout: + return None, "Gemini timeout after 600s." + except Exception as e: + return None, f"Gemini unexpected: {e}" + return None, "Gemini max retries exceeded." + + +def run_gemini_batches(batches: List[str], progress_callback=None) -> Tuple[List[Dict], List[str]]: + results: List[Dict] = [] + errors: List[str] = [] + total = len(batches) + for idx, batch_text in enumerate(batches): + pct = 0.60 + 0.17 * (idx / max(total, 1)) + msg = f"[Gemini] batch {idx+1}/{total}…" + log.info(msg) + if progress_callback: + progress_callback(msg, pct) + if idx > 0: + time.sleep(_GEMINI_CONFIG["batch_delay"]) + + result, err = _call_gemini(batch_text) + if result: + results.append(result) + if err: + errors.append(f"Gemini batch {idx+1}: {err}") + log.error(f"Gemini batch {idx+1} error: {err}") + return results, errors + +# ───────────────────────────────────────────────────────────────────────────── +# GPT-4o EXTRACTION +# ───────────────────────────────────────────────────────────────────────────── + +def _call_gpt(text_payload: str) -> Tuple[Optional[Dict], str]: + if not OPENAI_API_KEY: + return None, "OPENAI_API_KEY not set." + headers = { + "Authorization": f"Bearer {OPENAI_API_KEY}", + "Content-Type": "application/json", + } + payload = { + "model": GPT_MODEL, + "temperature": 0, + "response_format": {"type": "json_object"}, + #"max_tokens": GPT_MAX_TOKENS, + "messages": [ + {"role": "system", "content": GPT_SYSTEM_PROMPT}, + {"role": "user", "content": EXTRACTION_PROMPT + text_payload}, + ], + } + for attempt, wait in enumerate([0] + GPT_RETRY_DELAYS): + if wait: + log.warning(f"GPT-4o 429 — waiting {wait}s (retry {attempt})") + time.sleep(wait) + try: + resp = requests.post(GPT_API_URL, headers=headers, json=payload, timeout=300) + if resp.status_code == 429: + continue + if not resp.ok: + return None, f"GPT HTTP {resp.status_code}: {resp.text[:2000]}" + data = resp.json() + content = data.get("choices", [{}])[0].get("message", {}).get("content", "").strip() + # Strip markdown code fences if present + content = re.sub(r"^```(?:json)?\s*", "", content) + content = re.sub(r"\s*```$", "", content) + if content.startswith("{"): + try: + return json.loads(content), "" + except json.JSONDecodeError as e: + return None, f"GPT JSON parse error: {e} | raw: {content[:300]}" + return None, f"GPT response not JSON: {content[:300]}" + except requests.Timeout: + return None, "GPT timeout after 300s." + except Exception as e: + return None, f"GPT unexpected: {e}" + return None, "GPT max retries exceeded." + + +def run_gpt_batches(batches: List[str], progress_callback=None) -> Tuple[List[Dict], List[str]]: + results: List[Dict] = [] + errors: List[str] = [] + total = len(batches) + for idx, batch_text in enumerate(batches): + pct = 0.60 + 0.17 * (idx / max(total, 1)) + msg = f"[GPT-4o] batch {idx+1}/{total}…" + log.info(msg) + if progress_callback: + progress_callback(msg, pct) + if idx > 0: + time.sleep(_GPT_CONFIG["batch_delay"]) + result, err = _call_gpt(batch_text) + if result: + results.append(result) + if err: + errors.append(f"GPT batch {idx+1}: {err}") + log.error(f"GPT batch {idx+1} error: {err}") + return results, errors + +# ───────────────────────────────────────────────────────────────────────────── +# PARALLEL DUAL-LLM RUNNER +# Both LLMs get the same batches, run in parallel threads. +# Each respects its own rate-limit delays — they never share state. +# ───────────────────────────────────────────────────────────────────────────── + +def run_dual_llm_parallel( + batches: List[str], + progress_callback=None, +) -> Tuple[List[Dict], List[Dict], List[str]]: + """ + Returns (gemini_results, gpt_results, all_errors). + Both LLMs run concurrently in separate threads. + """ + gemini_results: List[Dict] = [] + gpt_results: List[Dict] = [] + all_errors: List[str] = [] + + def _gemini_task(): + return run_gemini_batches(batches, progress_callback=None) + + def _gpt_task(): + return run_gpt_batches(batches, progress_callback=None) + + if progress_callback: + progress_callback("Running Gemini + GPT-4o in parallel…", 0.62) + + with ThreadPoolExecutor(max_workers=2) as executor: + future_gemini = executor.submit(_gemini_task) + future_gpt = executor.submit(_gpt_task) + + for future in as_completed([future_gemini, future_gpt]): + if future is future_gemini: + try: + res, errs = future.result() + gemini_results.extend(res) + all_errors.extend(errs) + log.info(f"Gemini done — {len(res)} batch result(s)") + except Exception as e: + all_errors.append(f"Gemini thread error: {e}") + else: + try: + res, errs = future.result() + gpt_results.extend(res) + all_errors.extend(errs) + log.info(f"GPT-4o done — {len(res)} batch result(s)") + except Exception as e: + all_errors.append(f"GPT thread error: {e}") + + return gemini_results, gpt_results, all_errors + +# ───────────────────────────────────────────────────────────────────────────── +# MERGE RESULTS INTO DATAFRAME (per-LLM) +# ───────────────────────────────────────────────────────────────────────────── + +FIXED_COLS = ["material_name", "material_abbreviation", "manufacturer", "doi"] + + +def _make_abbreviation(name: str) -> str: + abbr = "".join(c for c in name if c.isupper()) + return abbr or name[:4].upper() + + +def _norm(s: str) -> str: + return re.sub(r"[^a-z0-9]", "", str(s).lower().strip()) + + +def _richness(row: Dict) -> int: + return sum(1 for v in row.values() if str(v).strip() not in ("", "N/A")) + + +def _fuzzy_dedup(rows: List[Dict]) -> List[Dict]: + seen_fuzzy: Dict[tuple, int] = {} + kept: List[Dict] = [] + for row in rows: + key = ( + _norm(row.get("material_name", "")), + _norm(row.get("property_name", "")), + _norm(row.get("section", "")), + ) + if key in seen_fuzzy: + existing_idx = seen_fuzzy[key] + existing = kept[existing_idx] + if _richness(row) > _richness(existing): + merged = dict(row) + for k, v in existing.items(): + if str(merged.get(k, "")).strip() in ("", "N/A") and str(v).strip() not in ("", "N/A"): + merged[k] = v + kept[existing_idx] = merged + else: + seen_fuzzy[key] = len(kept) + kept.append(row) + return kept + + +def _extract_doi_from_results(results: List[Dict]) -> str: + """Find the first non-empty DOI across all batch results.""" + for r in results: + doi = _normalise_doi(r.get("doi", "")) + if doi: + return doi + return "" + + +def merge_to_dataframe(results: List[Dict]) -> Tuple[pd.DataFrame, str]: + """ + Returns (DataFrame, doi_string). + """ + all_rows: List[Dict] = [] + seen_exact: set = set() + + fallback = {k: "" for k in FIXED_COLS} + for r in results: + for k in FIXED_COLS: + if not fallback[k] and r.get(k): + fallback[k] = r[k] + if not fallback["material_abbreviation"] and fallback["material_name"]: + fallback["material_abbreviation"] = _make_abbreviation(fallback["material_name"]) + + doi = _normalise_doi(fallback.get("doi", "")) + + for r in results: + r_doi = _normalise_doi(r.get("doi", "")) + if r_doi and not doi: + doi = r_doi + + r_identity = { + "material_name": r.get("material_name", "") or fallback["material_name"], + "material_abbreviation": r.get("material_abbreviation", "") or fallback["material_abbreviation"], + "manufacturer": r.get("manufacturer", "") or fallback["manufacturer"], + } + if not r_identity["material_abbreviation"] and r_identity["material_name"]: + r_identity["material_abbreviation"] = _make_abbreviation(r_identity["material_name"]) + + for item in r.get("properties", r.get("mechanical_properties", [])): + prop_mat = item.get("material_name", "").strip() + identity = dict(r_identity) + if prop_mat: + identity["material_name"] = prop_mat + identity["material_abbreviation"] = _make_abbreviation(prop_mat) + + raw_page = str(item.get("source_page", "")).strip() + page_label = ( + f"Page {raw_page}" if raw_page.isdigit() + else raw_page if raw_page + else "Unknown" + ) + chunk_type = item.get("chunk_type", "").strip() or "unknown" + + key = ( + _norm(identity["material_name"]), + _norm(item.get("section", "")), + _norm(item.get("property_name", "")), + _norm(item.get("value", "")), + ) + if key in seen_exact: + continue + seen_exact.add(key) + + # all_rows.append({ + # **identity, + # "section": item.get("section", "") or "General", + # "property_name": item.get("property_name", "") or "Unknown", + # "value": item.get("value", "") or "N/A", + # "unit": item.get("unit", "") or "", + # "english": item.get("english", "") or "", + # "test_condition": item.get("test_condition", "") or "", + # "comments": item.get("comments", "") or "", + # "source_page": page_label, + # "chunk_type": chunk_type, + # }) + value = item.get("value", "").strip() + min_v = item.get("min_value", "").strip() + max_v = item.get("max_value", "").strip() + + # if GPT left value blank but filled min/max, reconstruct it + if not value and min_v and max_v: + value = f"{min_v}–{max_v}" + elif not value and min_v: + value = min_v + elif not value and max_v: + value = max_v + + all_rows.append({ + **identity, + "section": item.get("section", "") or "General", + "property_name": item.get("property_name", "") or "Unknown", + "value": value or "N/A", + "min_value": min_v, + "max_value": max_v, + "unit": item.get("unit", "") or "", + "english": item.get("english", "") or "", + "test_condition": item.get("test_condition", "") or "", + "comments": item.get("comments", "") or "", + "source_text": item.get("source_text", "") or "", + "source_page": page_label, + "chunk_type": chunk_type, + "source_text": item.get("source_text", "") or "", + }) + + all_rows = _fuzzy_dedup(all_rows) + df = pd.DataFrame(all_rows) + if not df.empty: + base_cols = [c for c in df.columns if c not in ("source_page", "chunk_type")] + df = df[base_cols + ["source_page", "chunk_type"]] + + return df, doi + +# ───────────────────────────────────────────────────────────────────────────── +# CONSENSUS FILTER +# Only rows present in BOTH Gemini AND GPT outputs are kept. +# Match key: normalised (section, property_name, material_name) +# Value match: exact OR within 5% numeric tolerance. +# ───────────────────────────────────────────────────────────────────────────── + +def _first_num(s: str) -> Optional[float]: + m = re.search(r"[\d.]+", str(s)) + return float(m.group()) if m else None + + +def _value_match(val_a: str, val_b: str, tol: float = CONSENSUS_VALUE_TOL) -> bool: + na, nb = _norm(val_a), _norm(val_b) + if na == nb: + return True + fa, fb = _first_num(na), _first_num(nb) + if fa is not None and fb is not None and fb != 0: + if abs(fa - fb) / abs(fb) <= tol: + return True + # Range check: if one is a range and the other a point inside it + range_m = re.match(r"([\d.]+)[^\d.]+([\d.]+)", nb) + if range_m and fa is not None: + lo, hi = float(range_m.group(1)), float(range_m.group(2)) + if lo <= fa <= hi: + return True + return False + + +def _prop_key(row: Dict) -> Tuple[str, str, str]: + return ( + _norm(row.get("section", "")), + _norm(row.get("property_name", "")), + _norm(row.get("material_name", "")), + ) + + +def _token_set(s: str) -> set: + """Split normalised string into tokens of 3+ chars.""" + return {t for t in re.split(r"[^a-z0-9]+", _norm(s)) if len(t) >= 3} + + +def _token_overlap(a: str, b: str) -> float: + """Jaccard similarity between token sets of two strings.""" + sa, sb = _token_set(a), _token_set(b) + if not sa and not sb: + return 1.0 + if not sa or not sb: + return 0.0 + return len(sa & sb) / len(sa | sb) + + +def _embed_similarity(a: str, b: str) -> float: + """Cosine similarity between two strings using the loaded SentenceTransformer.""" + try: + model = _get_embed_model() + vecs = model.encode([a, b], normalize_embeddings=True) + return float(np.dot(vecs[0], vecs[1])) + except Exception: + return 0.0 + + +def _llm_adjudicate(pairs: List[Tuple[str, str]]) -> List[bool]: + """ + Ask GPT-4o whether each (gemini_key, gpt_key) pair refers to the same + material property. Returns a list of bools, one per pair. + Falls back to False on any error. + """ + if not pairs: + return [] + lines = "\n".join( + f"{i+1}. Gemini: \"{a}\" | GPT: \"{b}\"" + for i, (a, b) in enumerate(pairs) + ) + prompt = ( + "You are a materials science expert.\n" + "For each numbered pair below, decide if both sides refer to the " + "SAME material property (accounting for abbreviations, synonyms, " + "different units appended, or different word order).\n" + "Reply ONLY with a JSON array of booleans, one per pair, in order.\n" + "Example for 3 pairs: [true, false, true]\n\n" + f"{lines}" + ) + try: + headers = { + "Authorization": f"Bearer {OPENAI_API_KEY}", + "Content-Type": "application/json", + } + body = { + "model": GPT_MODEL, + "temperature": 0, + "max_tokens": max(len(pairs) * 15 + 100, 500), + + "messages": [{"role": "user", "content": prompt}], + } + resp = requests.post(GPT_API_URL, headers=headers, json=body, timeout=30) + if not resp.ok: + return [False] * len(pairs) + content = resp.json()["choices"][0]["message"]["content"].strip() + content = re.sub(r"^```(?:json)?\s*", "", content) + content = re.sub(r"\s*```$", "", content) + result = json.loads(content) + if isinstance(result, list) and len(result) == len(pairs): + return [bool(v) for v in result] + except Exception as e: + log.warning(f"LLM adjudication failed: {e}") + return [False] * len(pairs) + + +def _key_str(row: Dict) -> str: + """Human-readable key string for embedding/LLM — not normalised.""" + return f"{row.get('section','')} {row.get('property_name','')} {row.get('material_name','')}".strip() + + +def _cascade_match(gem_row: Dict, gpt_rows: List[Dict]) -> Optional[Dict]: + """ + Try to find a matching GPT row for a Gemini row using 4-layer cascade: + 1. Exact normalised key match + 2. Token overlap (Jaccard ≥ 0.5) + 3. Embedding cosine similarity ≥ 0.85 (clear match) + or flag as ambiguous if 0.65–0.85 (needs LLM) + 4. LLM adjudication for ambiguous pairs (called in batch outside) + Returns the best matching GPT row or None. + Also tags the row with which layer matched it. + """ + gem_key = _prop_key(gem_row) + gem_str = _key_str(gem_row) + + # Layer 1 — exact + for gpt_row in gpt_rows: + if _prop_key(gpt_row) == gem_key: + return {**gpt_row, "_match_layer": "exact"} + + # Layer 2 — token overlap + best_tok, best_tok_row = 0.0, None + for gpt_row in gpt_rows: + score = _token_overlap(gem_str, _key_str(gpt_row)) + if score > best_tok: + best_tok, best_tok_row = score, gpt_row + if best_tok >= 0.50: + return {**best_tok_row, "_match_layer": "token"} + + # Layer 3 — embedding + best_emb, best_emb_row = 0.0, None + for gpt_row in gpt_rows: + score = _embed_similarity(gem_str, _key_str(gpt_row)) + if score > best_emb: + best_emb, best_emb_row = score, gpt_row + if best_emb >= 0.85: + return {**best_emb_row, "_match_layer": "embedding"} + if best_emb >= 0.65: + # Mark as ambiguous for LLM batch — caller handles this + return {**best_emb_row, "_match_layer": "ambiguous", "_emb_score": best_emb} + + return None + +def _text_cascade_match(term: str, source_text: str, embed_threshold: float = 0.75) -> bool: + """ + Check if a term is present in source_text using cascade: + 1. Exact substring match + 2. Token overlap — any meaningful token from term appears in source text + 3. Embedding similarity between term and source text + """ + if not term or not source_text: + return False + + term_lower = term.lower().strip() + src_lower = source_text.lower().strip() + + # Layer 1 — exact substring + if term_lower in src_lower: + return True + + # Layer 2 — token overlap + tokens = [t for t in re.split(r"[^a-z0-9]+", _norm(term)) if len(t) >= 3] + if tokens and any(t in src_lower for t in tokens): + return True + + # Layer 3 — embedding similarity + try: + score = _embed_similarity(term, source_text) + if score >= embed_threshold: + return True + except Exception: + pass + + return False + +def _source_text_passes(row: Dict) -> bool: + """Both property name and value must be verifiable in source_text.""" + src = str(row.get("source_text", "")).strip() + val = str(row.get("value", "")).strip() + prop = str(row.get("property_name", "")).strip() + if not src: + return False + # value check — exact numeric match only + val_ok = False + m = re.search(r"[\d.]+", val) + if m and m.group() in src: + val_ok = True + if not val_ok: + return False + # property name check — cascade + return _text_cascade_match(prop, src) + +def _resolve_value_disagreement( + gem_row: Dict, + gpt_row: Dict, + layer: str, + tol: float = CONSENSUS_VALUE_TOL, +) -> List[Dict]: + """ + Resolves a matched pair where values may differ. + Returns a list of 0, 1, or 2 rows to add to confirmed. + 0 rows = rejected + 1 row = single consensus row + 2 rows = both kept (condition variant or both unverified) + """ + gem_val = str(gem_row.get("value", "")).strip() + gpt_val = str(gpt_row.get("value", "")).strip() + gem_src = str(gem_row.get("source_text", "")).strip() + gpt_src = str(gpt_row.get("source_text", "")).strip() + gem_unit = str(gem_row.get("unit", "")).strip() + gpt_unit = str(gpt_row.get("unit", "")).strip() + gem_cond = _norm(gem_row.get("test_condition", "")) + gpt_cond = _norm(gpt_row.get("test_condition", "")) + + gem_null = not gem_val or gem_val == "N/A" + gpt_null = not gpt_val or gpt_val == "N/A" + + def _clean_gpt(row: Dict) -> Dict: + return {k: v for k, v in row.items() if not k.startswith("_")} + + def _make_row(base: Dict, other: Dict, needs_review: bool = False) -> Dict: + row = dict(base) + for k, v in other.items(): + if k.startswith("_"): + continue + if str(row.get(k, "")).strip() in ("", "N/A") and str(v).strip() not in ("", "N/A"): + row[k] = v + row["confirmed_by_gpt"] = True + row["match_layer"] = layer + row["needs_review"] = needs_review + return row + + # ── Step 1: Null fallback ───────────────────────────────────────────────── + if gem_null and not gpt_null: + if _source_text_passes(gpt_row): + row = _make_row(gem_row, gpt_row) + row["value"] = gpt_val + return [row] + else: + row = _make_row(gem_row, gpt_row, needs_review=True) + row["value"] = gpt_val + log.info(f"Null fallback unverified — '{gem_row.get('property_name','')}': GPT={gpt_val}") + return [row] + + if gpt_null and not gem_null: + if _source_text_passes(gem_row): + return [_make_row(gem_row, gpt_row)] + else: + row = _make_row(gem_row, gpt_row, needs_review=True) + log.info(f"Null fallback unverified — '{gem_row.get('property_name','')}': Gemini={gem_val}") + return [row] + + if gem_null and gpt_null: + return [] + + # ── Step 2: Both non-null, values agree ─────────────────────────────────── + if _value_match(gem_val, gpt_val, tol): + return [_make_row(gem_row, gpt_row)] + + # ── Step 3: Values disagree — check conditions ──────────────────────────── + conditions_differ = gem_cond and gpt_cond and gem_cond != gpt_cond + if conditions_differ: + gem_out = dict(gem_row) + gem_out["confirmed_by_gpt"] = True + gem_out["match_layer"] = layer + gem_out["needs_review"] = False + + gpt_out = _clean_gpt(gpt_row) + gpt_out["confirmed_by_gpt"] = True + gpt_out["match_layer"] = layer + gpt_out["needs_review"] = False + + log.info( + f"Condition variant — '{gem_row.get('property_name','')}': " + f"Gemini={gem_val} ({gem_cond}), GPT={gpt_val} ({gpt_cond})" + ) + return [gem_out, gpt_out] + + # ── Step 4: Check units ─────────────────────────────────────────────────── + units_differ = gem_unit and gpt_unit and _norm(gem_unit) != _norm(gpt_unit) + if units_differ: + row = dict(gem_row) + row["confirmed_by_gpt"] = True + row["match_layer"] = layer + row["needs_review"] = False + log.info( + f"Unit mismatch — keeping Gemini '{gem_unit}' over GPT '{gpt_unit}' " + f"for '{gem_row.get('property_name','')}'" + ) + return [row] + + # ── Step 5: Source text check ───────────────────────────────────────────── + gem_verified = _source_text_passes(gem_row) + gpt_verified = _source_text_passes(gpt_row) + + if gem_verified and not gpt_verified: + log.info(f"Source text: Gemini verified — '{gem_row.get('property_name','')}': {gem_val}") + return [_make_row(gem_row, gpt_row)] + + if gpt_verified and not gem_verified: + row = _make_row(gem_row, gpt_row) + row["value"] = gpt_val + log.info(f"Source text: GPT verified — '{gem_row.get('property_name','')}': {gpt_val}") + return [row] + + # Neither verified — flag both for review + gem_out = dict(gem_row) + gem_out["confirmed_by_gpt"] = True + gem_out["match_layer"] = layer + gem_out["needs_review"] = True + + gpt_out = _clean_gpt(gpt_row) + gpt_out["confirmed_by_gpt"] = True + gpt_out["match_layer"] = layer + gpt_out["needs_review"] = True + + log.info( + f"Both unverified — '{gem_row.get('property_name','')}': " + f"Gemini={gem_val}, GPT={gpt_val} — flagged for review" + ) + return [gem_out, gpt_out] +def _post_consensus_unmatched( + df_gemini: pd.DataFrame, + df_gpt: pd.DataFrame, + df_consensus: pd.DataFrame, + tol: float = CONSENSUS_VALUE_TOL, +) -> pd.DataFrame: + """ + After consensus, collect unmatched rows from both LLMs. + Run source text check on each — keep if passes, drop if fails. + Mix kept rows silently into the consensus output. + """ + + + + def _already_in_consensus(row: Dict, consensus_rows: List[Dict]) -> bool: + """Check if this row is already represented in consensus output.""" + key = ( + _norm(row.get("material_name", "")), + _norm(row.get("property_name", "")), + _norm(row.get("value", "")), + ) + for cr in consensus_rows: + ck = ( + _norm(cr.get("material_name", "")), + _norm(cr.get("property_name", "")), + _norm(cr.get("value", "")), + ) + if key == ck: + return True + return False + + consensus_rows = df_consensus.to_dict(orient="records") if not df_consensus.empty else [] + gem_rows = df_gemini.to_dict(orient="records") if not df_gemini.empty else [] + gpt_rows = df_gpt.to_dict(orient="records") if not df_gpt.empty else [] + + kept: List[Dict] = [] + + # ── Unmatched Gemini rows ───────────────────────────────────────────────── + for row in gem_rows: + if _already_in_consensus(row, consensus_rows): + continue + if _source_text_passes(row): + clean = dict(row) + clean["confirmed_by_gpt"] = False + clean["match_layer"] = "single_gemini" + clean["needs_review"] = False + kept.append(clean) + else: + log.info( + f"Unmatched Gemini row dropped — source text failed: " + f"'{row.get('property_name','')}' = {row.get('value','')} " + f"[{row.get('material_name','')}]" + ) + + # ── Unmatched GPT rows ──────────────────────────────────────────────────── + for row in gpt_rows: + if _already_in_consensus(row, consensus_rows): + continue + if _source_text_passes(row): + clean = dict(row) + clean["confirmed_by_gpt"] = False + clean["match_layer"] = "single_gpt" + clean["needs_review"] = False + kept.append(clean) + else: + log.info( + f"Unmatched GPT row dropped — source text failed: " + f"'{row.get('property_name','')}' = {row.get('value','')} " + f"[{row.get('material_name','')}]" + ) + + if not kept: + return df_consensus + + df_extra = pd.DataFrame(kept) + + # Align columns — add missing columns from consensus with empty values + if not df_consensus.empty: + for col in df_consensus.columns: + if col not in df_extra.columns: + df_extra[col] = "" + df_extra = df_extra[df_consensus.columns] + + result = pd.concat([df_consensus, df_extra], ignore_index=True) + log.info( + f"Post-consensus: {len(kept)} unmatched rows passed source text check " + f"({sum(1 for r in kept if r['match_layer']=='single_gemini')} Gemini, " + f"{sum(1 for r in kept if r['match_layer']=='single_gpt')} GPT)" + ) + return result + +def consensus_filter( + df_gemini: pd.DataFrame, + df_gpt: pd.DataFrame, + tol: float = CONSENSUS_VALUE_TOL, +) -> pd.DataFrame: + if df_gemini.empty and df_gpt.empty: + return pd.DataFrame() + if df_gemini.empty: + return df_gpt.assign(confirmed_by_gpt=False, match_layer="none", needs_review=False) + if df_gpt.empty: + return df_gemini.assign(confirmed_by_gpt=False, match_layer="none", needs_review=False) + + from collections import defaultdict + + gem_rows = df_gemini.to_dict(orient="records") + gpt_rows = df_gpt.to_dict(orient="records") + + def _group_by_material(rows: List[Dict]) -> Dict[str, List[Dict]]: + groups: Dict[str, List[Dict]] = defaultdict(list) + for row in rows: + groups[_norm(row.get("material_name", ""))].append(row) + return groups + + gem_by_mat = _group_by_material(gem_rows) + gpt_by_mat = _group_by_material(gpt_rows) + + def _best_material_match(mat_key: str, other_groups: Dict[str, List[Dict]]) -> Optional[str]: + if mat_key in other_groups: + return mat_key + best_key, best_score = None, 0.0 + for other_key in other_groups: + score = _token_overlap(mat_key, other_key) + if score > best_score: + best_score, best_key = score, other_key + if best_score >= 0.3: + return best_key + for other_key in other_groups: + score = _embed_similarity(mat_key, other_key) + if score > best_score: + best_score, best_key = score, other_key + return best_key if best_score >= 0.70 else None + + confirmed: List[Dict] = [] + ambiguous: List[Tuple[int, Dict, Dict]] = [] + main_resolved: List[Dict] = [] + + for gem_mat_key, gem_mat_rows in gem_by_mat.items(): + matched_gpt_mat_key = _best_material_match(gem_mat_key, gpt_by_mat) + if matched_gpt_mat_key is None: + log.info(f"Material '{gem_mat_key}' found only in Gemini — skipping") + continue + + gpt_candidate_pool = gpt_by_mat[matched_gpt_mat_key] + log.info( + f"Material match: Gemini '{gem_mat_key}' → GPT '{matched_gpt_mat_key}' " + f"({len(gem_mat_rows)} Gemini rows, {len(gpt_candidate_pool)} GPT rows)" + ) + + for gem_row in gem_mat_rows: + candidate = _cascade_match(gem_row, gpt_candidate_pool) + if candidate is None: + continue + + layer = candidate.get("_match_layer", "") + if layer == "ambiguous": + ambiguous.append((len(confirmed), gem_row, candidate)) + confirmed.append(None) # slot placeholder — index must stay stable + continue + + resolved = _resolve_value_disagreement(gem_row, candidate, layer, tol) + main_resolved.extend(resolved) + + # ── LLM adjudication ───────────────────────────────────────────────────── + if ambiguous: + pairs = [(_key_str(g), _key_str(c)) for _, g, c in ambiguous] + verdicts = _llm_adjudicate(pairs) + for (slot_idx, gem_row, candidate), verdict in zip(ambiguous, verdicts): + if verdict: + resolved = _resolve_value_disagreement(gem_row, candidate, "llm", tol) + if resolved: + confirmed[slot_idx] = resolved[0] + if len(resolved) == 2: + main_resolved.append(resolved[1]) + else: + confirmed[slot_idx] = None + else: + confirmed[slot_idx] = None + + matched = [r for r in confirmed if r is not None] + main_resolved + if not matched: + log.warning("Consensus cascade: no rows agreed between Gemini and GPT-4o.") + return pd.DataFrame() + + df = pd.DataFrame(matched) + if "needs_review" not in df.columns: + df["needs_review"] = False + + layer_counts = df["match_layer"].value_counts().to_dict() + review_count = int(df["needs_review"].sum()) + log.info( + f"Consensus cascade: {len(df_gemini)} Gemini + {len(df_gpt)} GPT " + f"→ {len(df)} agreed {layer_counts} | {review_count} need review" + ) + return _post_consensus_unmatched(df_gemini, df_gpt, df, tol) + + +# ───────────────────────────────────────────────────────────────────────────── +# DOI RESOLUTION (merge from both LLMs + manual override) +# ───────────────────────────────────────────────────────────────────────────── + +def resolve_doi(doi_gemini: str, doi_gpt: str, doi_manual: str = "") -> str: + """Return the best DOI: manual override > Gemini > GPT.""" + for candidate in (doi_manual, doi_gemini, doi_gpt): + d = _normalise_doi(candidate) + if d: + return d + return "" + +# ───────────────────────────────────────────────────────────────────────────── +# CACHE HELPERS +# ───────────────────────────────────────────────────────────────────────────── + +def _df_to_cache_dict(df: pd.DataFrame, doi: str = "") -> Dict: + if df.empty: + return {} + row0 = df.iloc[0] + return { + "material_name": str(row0.get("material_name", "")), + "material_abbreviation": str(row0.get("material_abbreviation", "")), + #"trade_grade": str(row0.get("trade_grade", "")), + "manufacturer": str(row0.get("manufacturer", "")), + "doi": doi, + "mechanical_properties": df.drop( + columns=[c for c in FIXED_COLS if c in df.columns], errors="ignore" + ).to_dict(orient="records"), + } + +# ───────────────────────────────────────────────────────────────────────────── +# TOP-LEVEL PIPELINE ORCHESTRATOR +# ───────────────────────────────────────────────────────────────────────────── + +def run_pipeline( + pdf_bytes: bytes, + doi_override: str = "", + progress_callback: Any = None, +) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, List[Chunk], List[str], Dict]: + """ + Full dual-LLM consensus pipeline. + + Returns + ------- + df_consensus : rows agreed by both LLMs ← primary output + df_gemini : Gemini-only output + df_gpt : GPT-only output + all_chunks : all Chunk objects with scores + api_errors : list of error strings + meta : pipeline stats dict + """ + def _prog(msg: str, pct: float): + log.info(f"[{pct*100:.0f}%] {msg}") + if progress_callback: + progress_callback(msg, pct) + + meta: Dict[str, Any] = {} + pdf_hash = _pdf_hash(pdf_bytes) + + # ── Cache check ─────────────────────────────────────────────────────────── + _prog("Checking cache…", 0.0) + + + cached = cache_get(pdf_bytes) + if cached: + _prog("Cache hit.", 1.0) + meta["path"] = "cache" + df, doi = merge_to_dataframe([cached]) + doi = resolve_doi(doi, "", doi_override) + df["doi_url"] = _doi_url(doi) + # Return early — no embedding, no ChromaDB, no LLM calls + return df, df.copy(), df.copy(), [], [], meta + + # ── Stage 1: Extract ────────────────────────────────────────────────────── + _prog("Stage 1 — extracting tables + text…", 0.05) + all_chunks = extract_all_chunks(pdf_bytes) + meta["chunks_total"] = len(all_chunks) + meta["chunks_tables"] = sum(1 for c in all_chunks if c.chunk_type == "table") + meta["chunks_text"] = sum(1 for c in all_chunks if c.chunk_type == "text") + + if not all_chunks: + _prog("No content extracted.", 1.0) + meta["path"] = "failed" + return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), [], ["No content extracted."], meta + + _prog(f"Stage 1 done — {meta['chunks_tables']} tables, {meta['chunks_text']} text.", 0.20) + + # ── Stage 2: Index into ChromaDB ───────────────────────────────────────── + _prog("Stage 2 — indexing into ChromaDB…", 0.25) + try: + index_chunks_in_chroma(all_chunks, pdf_hash) + except Exception as e: + log.error(f"ChromaDB indexing failed: {e}") + + # ── Stage 3: Rank ALL chunks ────────────────────────────────────────────── + _prog("Stage 3 — ranking all chunks…", 0.40) + ranked_chunks = _chroma_rank_all(pdf_hash) + if not ranked_chunks: + # Fallback: use all chunks unranked + ranked_chunks = all_chunks + for c in ranked_chunks: + c.relevant = True + meta["chunks_ranked"] = len(ranked_chunks) + _prog(f"Stage 3 done — {len(ranked_chunks)} chunks ranked.", 0.50) + + # ── Stage 4: Build batches ──────────────────────────────────────────────── + _prog("Stage 4 — building batches…", 0.55) + batches = build_batches(ranked_chunks) + meta["batches"] = len(batches) + _prog(f"Stage 4 done — {len(batches)} batch(es).", 0.60) + + # ── Stage 5: Dual LLM in parallel ──────────────────────────────────────── + _prog("Stage 5 — Gemini + GPT-4o running in parallel…", 0.60) + gemini_raw, gpt_raw, api_errors = run_dual_llm_parallel(batches, progress_callback) + + if not gemini_raw and not gpt_raw: + meta["path"] = "failed" + return pd.DataFrame(), pd.DataFrame(), pd.DataFrame(), all_chunks, api_errors, meta + + # ── Stage 6: Merge each LLM's results ──────────────────────────────────── + _prog("Stage 6 — merging results…", 0.90) + df_gemini, doi_gemini = merge_to_dataframe(gemini_raw) + df_gpt, doi_gpt = merge_to_dataframe(gpt_raw) + + doi = resolve_doi(doi_gemini, doi_gpt, doi_override) + doi_url = _doi_url(doi) + + meta["gemini_properties"] = len(df_gemini) + meta["gpt_properties"] = len(df_gpt) + meta["doi"] = doi + log.info(f"GPT raw batch count: {len(gpt_raw)}") + for i, r in enumerate(gpt_raw): + props = r.get("properties", r.get("mechanical_properties", [])) + log.info(f" GPT batch {i+1}: {len(props)} properties") + for p in props[:3]: # show first 3 properties from each batch + log.info(f" → {p.get('property_name','?')} = {p.get('value','?')} {p.get('unit','?')}") + + log.info(f"Gemini raw batch count: {len(gemini_raw)}") + for i, r in enumerate(gemini_raw): + props = r.get("properties", r.get("mechanical_properties", [])) + log.info(f" Gemini batch {i+1}: {len(props)} properties") + for p in props[:3]: # show first 3 properties from each batch + log.info(f" → {p.get('property_name','?')} = {p.get('value','?')} {p.get('unit','?')}") + # ── Stage 7: Consensus filter ───────────────────────────────────────────── + _prog("Stage 7 — consensus filtering…", 0.95) + df_consensus = consensus_filter(df_gemini, df_gpt) + meta["consensus_properties"] = len(df_consensus) + meta["path"] = "dual-llm-consensus" + + # Attach doi_url to all frames + for df in (df_consensus, df_gemini, df_gpt): + if not df.empty: + df["doi_url"] = doi_url + + if not df_consensus.empty: + cache_set(pdf_bytes, _df_to_cache_dict(df_consensus, doi)) + + _prog(f"Done — {len(df_consensus)} consensus properties.", 1.0) + return df_consensus, df_gemini, df_gpt, all_chunks, api_errors, meta + +# ───────────────────────────────────────────────────────────────────────────── +# STREAMLIT UI +# ───────────────────────────────────────────────────────────────────────────── + +def _run_streamlit(): + import streamlit as st + + st.set_page_config( + page_title="DocToDB — Dual LLM Consensus Extractor", + page_icon="🧬", + layout="wide", + ) + st.title("🧬 DocToDB — Dual LLM Consensus Extractor") + st.caption( + f"Extracts material properties using **{GEMINI_MODEL}** + **{GPT_MODEL}** in parallel. " + "Only rows agreed by both LLMs are kept in the consensus output." + ) + + # ── Sidebar ─────────────────────────────────────────────────────────────── + with st.sidebar: + st.header("⚙️ Settings") + st.divider() + st.markdown(f"**Gemini model:** `{GEMINI_MODEL}`") + st.markdown(f"**GPT model:** `{GPT_MODEL}`") + st.markdown(f"**Embedder:** `{EMBED_MODEL_NAME}`") + st.markdown(f"**ChromaDB:** {'✅' if CHROMA_AVAILABLE else '❌ not installed'}") + st.markdown(f"**Docling:** {'✅' if DOCLING_AVAILABLE else '❌'}") + st.markdown(f"**Camelot:** {'✅' if CAMELOT_AVAILABLE else '❌'}") + st.markdown(f"**OCR:** {'✅' if OCR_AVAILABLE else '❌'}") + gemini_ok = bool(GEMINI_API_KEY) + gpt_ok = bool(OPENAI_API_KEY) + st.markdown(f"**Gemini API Key:** {'✅' if gemini_ok else '❌ GEMINI_API_KEY not set'}") + st.markdown(f"**OpenAI API Key:** {'✅' if gpt_ok else '❌ OPENAI_API_KEY not set'}") + st.divider() + st.markdown(f"**Consensus tolerance:** ±{CONSENSUS_VALUE_TOL*100:.0f}%") + st.markdown("**Retrieval:** Schema-derived query, all chunks ranked") + st.divider() + + doi_manual = st.text_input( + "DOI override (optional)", + placeholder="10.1016/j.mat.2023.01.001", + help="Paste a DOI here to override auto-extraction. Leave blank to use the extracted DOI.", + ) + + st.divider() + if st.button("🗑 Clear JSON Cache"): + if os.path.exists(CACHE_FILE): + os.remove(CACHE_FILE) + st.success("JSON cache cleared.") + if st.button("🗑 Clear ChromaDB Collection"): + try: + col = _get_chroma_collection() + col.delete(where={"pdf_hash": {"$ne": ""}}) + st.success("ChromaDB cleared.") + except Exception as e: + st.error(f"ChromaDB clear failed: {e}") + + # ── Upload ──────────────────────────────────────────────────────────────── + uploaded = st.file_uploader("Upload PDF", type=["pdf"]) + if not uploaded: + st.info("Upload a PDF to get started.") + return + + pdf_bytes = uploaded.getvalue() + stem = uploaded.name.rsplit(".", 1)[0] + + if not gemini_ok: + st.error("GEMINI_API_KEY environment variable not set.") + return + if not gpt_ok: + st.warning("OPENAI_API_KEY not set — GPT-4o will be skipped; no consensus possible.") + + if st.button("🚀 Run Extraction", type="primary", use_container_width=True): + bar = st.progress(0.0) + status = st.empty() + + def cb(msg, pct): + bar.progress(min(pct, 1.0)) + status.text(msg) + + with st.spinner("Running dual-LLM pipeline…"): + df_consensus, df_gemini, df_gpt, chunks, errors, meta = run_pipeline( + pdf_bytes, doi_override=doi_manual, progress_callback=cb + ) + + bar.progress(1.0) + status.empty() + + # ── DOI display ─────────────────────────────────────────────────────── + doi = meta.get("doi", "") + doi_url = _doi_url(doi) + if doi_url: + st.markdown( + f"**DOI:** [{doi}]({doi_url})", + help="Extracted from the PDF by both LLMs. Click to open the paper.", + ) + else: + st.info("No DOI found in this PDF.") + + # ── Metrics ─────────────────────────────────────────────────────────── + c1, c2, c3, c4, c5, c6, c7 = st.columns(7) + c1.metric("Tables", meta.get("chunks_tables", "—")) + c2.metric("Text blocks", meta.get("chunks_text", "—")) + c3.metric("Ranked", meta.get("chunks_ranked", "—")) + c4.metric("Batches", meta.get("batches", "—")) + c5.metric("Gemini rows", meta.get("gemini_properties", len(df_gemini))) + c6.metric("GPT-4o rows", meta.get("gpt_properties", len(df_gpt))) + c7.metric("✅ Consensus", meta.get("consensus_properties", len(df_consensus))) + + if errors: + with st.expander(f"⚠️ {len(errors)} API error(s)"): + for e in errors: + st.code(e) + + # ── Helper: render df with clickable DOI ────────────────────────────── + def _show_df_with_doi(df: pd.DataFrame, key: str): + """Display dataframe. DOI URL column rendered as markdown links.""" + if df.empty: + st.info("No data.") + return + # Build a display copy with clickable DOI + display_df = df.copy() + if "doi_url" in display_df.columns: + display_df = display_df.drop(columns=["doi_url"]) + + # Streamlit doesn't natively render HTML in dataframes, so we + # show the DOI link above and the table below. + doi_url_val = df["doi_url"].iloc[0] if "doi_url" in df.columns and not df.empty else "" + doi_val = _normalise_doi(doi_url_val.replace("https://doi.org/", "") if doi_url_val else "") + if doi_url_val: + st.markdown(f"🔗 **Paper DOI:** [{doi_val}]({doi_url_val})") + + st.dataframe(display_df, use_container_width=True, hide_index=True) + + # ── Tabs ────────────────────────────────────────────────────────────── + tab_consensus, tab_gemini, tab_gpt, tab_pages, tab_export = st.tabs([ + "✅ Consensus", + f"🟦 Gemini ({len(df_gemini)})", + f"🟩 GPT-4o ({len(df_gpt)})", + "🗂 By Page", + "📤 Export", + ]) + + with tab_consensus: + st.caption( + f"**{len(df_consensus)} properties** agreed by both Gemini and GPT-4o " + f"(fuzzy match ±{CONSENSUS_VALUE_TOL*100:.0f}%)" + ) + if df_consensus.empty: + st.warning( + "No consensus rows found. " + "Check the Gemini and GPT-4o tabs — both returned results but they don't overlap." + ) + else: + # Sub-tabs by chunk type + t_tbl, t_txt, t_all = st.tabs(["📋 Tables", "📄 Text", "🔢 All"]) + with t_tbl: + _show_df_with_doi(df_consensus[df_consensus.get("chunk_type", "") == "table"] if "chunk_type" in df_consensus.columns else pd.DataFrame(), "ct_tbl") + with t_txt: + _show_df_with_doi(df_consensus[df_consensus.get("chunk_type", "") == "text"] if "chunk_type" in df_consensus.columns else pd.DataFrame(), "ct_txt") + with t_all: + _show_df_with_doi(df_consensus, "ct_all") + + with tab_gemini: + st.caption(f"{len(df_gemini)} properties extracted by Gemini") + _show_df_with_doi(df_gemini, "gem") + + with tab_gpt: + st.caption(f"{len(df_gpt)} properties extracted by GPT-4o") + _show_df_with_doi(df_gpt, "gpt") + + with tab_pages: + df_src = df_consensus if not df_consensus.empty else df_gemini + if "source_page" in df_src.columns: + def _page_sort_key(label: str) -> int: + m = re.search(r"\d+", label) + return int(m.group()) if m else 9999 + for pg in sorted(df_src["source_page"].unique(), key=_page_sort_key): + pg_df = df_src[df_src["source_page"] == pg] + tbl_n = (pg_df.get("chunk_type", pd.Series()) == "table").sum() + txt_n = (pg_df.get("chunk_type", pd.Series()) == "text").sum() + with st.expander( + f"📄 {pg} — {len(pg_df)} " + f"propert{'y' if len(pg_df)==1 else 'ies'} " + f"({tbl_n} table · {txt_n} text)", + expanded=False, + ): + _show_df_with_doi(pg_df, f"page_{pg}") + + with tab_export: + export_df = df_consensus if not df_consensus.empty else df_gemini + if "doi_url" in export_df.columns: + export_df = export_df.copy() + + col1, col2 = st.columns(2) + col1.download_button( + "⬇️ Consensus CSV", + export_df.to_csv(index=False).encode(), + f"{stem}_consensus.csv", + "text/csv", + use_container_width=True, + ) + col2.download_button( + "⬇️ Consensus JSON", + export_df.to_json(orient="records", indent=2).encode(), + f"{stem}_consensus.json", + "application/json", + use_container_width=True, + ) + st.divider() + col3, col4 = st.columns(2) + col3.download_button( + "⬇️ Gemini CSV", + df_gemini.to_csv(index=False).encode() if not df_gemini.empty else b"", + f"{stem}_gemini.csv", + "text/csv", + use_container_width=True, + disabled=df_gemini.empty, + ) + col4.download_button( + "⬇️ GPT-4o CSV", + df_gpt.to_csv(index=False).encode() if not df_gpt.empty else b"", + f"{stem}_gpt4o.csv", + "text/csv", + use_container_width=True, + disabled=df_gpt.empty, + ) + + st.divider() + with st.expander("🔎 Chunk inspector (top-30 by score)"): + show_type = st.radio("Show", ["all", "table", "text"], horizontal=True) + shown = [c for c in chunks if show_type == "all" or c.chunk_type == show_type] + shown = sorted(shown, key=lambda c: c.score, reverse=True)[:30] + for c in shown: + st.code( + f"[{c.chunk_type.upper()}] Page {c.page_num} | score={c.score:.3f}\n\n{c.text[:400]}", + language=None, + ) + +# ───────────────────────────────────────────────────────────────────────────── +# ENTRY POINT +# ───────────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + import sys + + _in_streamlit = False + try: + import streamlit.runtime.scriptrunner as _sr + if _sr.get_script_run_ctx() is not None: + _in_streamlit = True + except Exception: + pass + + if _in_streamlit: + _run_streamlit() + else: + print( + "\nUsage:\n" + " streamlit run doctodb_dual.py\n" + "\nEnvironment variables required:\n" + " GEMINI_API_KEY — Google AI Studio key\n" + " OPENAI_API_KEY — OpenAI key\n" + ) \ No newline at end of file