debase 0.6.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

debase/reaction_info_extractor.py

@@ -37,7 +37,16 @@ from dataclasses import dataclass, field
 from functools import lru_cache
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+# Import universal caption pattern
+try:
+    from .caption_pattern import get_universal_caption_pattern
+    from .campaign_utils import enhance_prompt_with_campaign, get_location_hints_for_campaign
+except ImportError:
+    # Fallback if running as standalone script
+    from caption_pattern import get_universal_caption_pattern
+    from campaign_utils import enhance_prompt_with_campaign, get_location_hints_for_campaign
 
 import fitz  # PyMuPDF - for image extraction
 import google.generativeai as genai  # type: ignore
@@ -451,13 +460,23 @@ Respond with a JSON array where each element contains:
 - "campaign_clues": specific text in the caption that indicates the campaign (enzyme names, substrate types, etc.)
 
 Tables are generally preferred over figures unless you are convinced that only the figure you find have complete lineage reaction matrix information. Some table don't have performance data, check provided context of the specific table.
+
+IMPORTANT FOR TABLES: When evaluating a table, check if the context below the table shows performance values (TTN, yield, ee, etc.). If the table caption mentions enzymes but the table only shows mutations/sequences, look for performance data in the text immediately following the table. If context below the table shows numerical values, use the table location as it likely contains the referenced data.
+
 Do not include too much sources, just return 2 or 3 sources.
 Adjust confidence comparing all locations you will be returning, only rank figure the highest when you are absolutely certain table won't contain complete information.
 When returning confidence scores, be more accurate and avoid scores that are too close together.
+
+CRITICAL: 
+- Return "location" EXACTLY as the first reference identifier appears in the actual caption text
+- Copy the exact characters including all punctuation (periods, colons, pipes, etc.) up to the first space after the identifier
+- Do NOT modify, standardize, or interpret the location - return it verbatim from the document
+- Include "document" field to specify which PDF contains this location: "manuscript" or "supplementary"
+
 Respond ONLY with **minified JSON**. NO markdown fences.
 
-Example:
-[{"location": "Table S1", "type": "table", "confidence": 95, "caption": "Table S1. Detailed information...", "reason": "Complete performance metrics", "lineage_hint": "first enzyme family", "campaign_clues": "PYS lineage, pyrrolidine synthesis"}]
+Format:
+[{"location": "", "type": "", "document": "", "confidence": 0, "caption": "", "reason": "", "lineage_hint": "", "campaign_clues": ""}]
 """)
 
 PROMPT_EXTRACT_METRICS = dedent("""
@@ -734,7 +753,8 @@ class ReactionExtractor:
     _TAB_RE = re.compile(r"(?:supplementary\s+)?tab(?:le)?\s+s?\d+[a-z]?", re.I)
 
     def __init__(self, manuscript: Path, si: Optional[Path], cfg: Config, debug_dir: Optional[Path] = None, 
-                 campaign_filter: Optional[str] = None, all_campaigns: Optional[List[str]] = None):
+                 campaign_filter: Optional[str] = None, all_campaigns: Optional[List[str]] = None,
+                 campaign_info: Optional[Dict[str, Any]] = None):
         self.manuscript = manuscript
         self.si = si
         self.cfg = cfg
@@ -742,6 +762,7 @@ class ReactionExtractor:
         self.debug_dir = debug_dir
         self.campaign_filter = campaign_filter  # Filter for specific campaign
         self.all_campaigns = all_campaigns or []  # List of all campaigns for context
+        self.campaign_info = campaign_info  # Detailed campaign information from campaigns.json
         
         # Cache for extracted figures to avoid redundant extractions (bounded to prevent memory leaks)
         self._figure_cache = LRUCache(maxsize=100)  # Figures are large, so smaller cache
@@ -778,10 +799,8 @@ class ReactionExtractor:
     # ------------------------------------------------------------------
 
     def _collect_captions_and_titles(self) -> str:
-        # Pattern to match Table or Figure with optional leading whitespace and page numbers
-        # This catches all variations including "Supplementary Table", "Table S 2", "Figure S1", etc.
-        # Also handles cases where there's whitespace or page numbers before the caption
-        cap_pattern = re.compile(r"^[\s\d]*\s*(Supplementary\s+Table|Table|Figure).*", re.I | re.M)
+        # Use universal caption pattern that handles all common formats
+        cap_pattern = get_universal_caption_pattern()
         captions: List[str] = []
         
         # Process each page individually to avoid TOC entries
@@ -861,7 +880,27 @@ class ReactionExtractor:
         
         # Add campaign context - always provide context to help model understanding
         campaign_context = ""
-        if self.campaign_filter:
+        
+        # If we have detailed campaign info, use it to provide specific guidance
+        if self.campaign_info:
+            location_hints = get_location_hints_for_campaign(self.campaign_info)
+            campaign_context = f"""
+            IMPORTANT: You are looking for performance data specifically for the {self.campaign_filter} campaign.
+            
+            Campaign Details:
+            - Name: {self.campaign_info.get('campaign_name', '')}
+            - Description: {self.campaign_info.get('description', '')}
+            - Model Substrate: {self.campaign_info.get('model_substrate', '')} (ID: {self.campaign_info.get('substrate_id', '')})
+            - Model Product: {self.campaign_info.get('model_product', '')} (ID: {self.campaign_info.get('product_id', '')})
+            
+            KNOWN DATA LOCATIONS FOR THIS CAMPAIGN: {', '.join(location_hints)}
+            These locations are known to contain relevant data - prioritize them highly.
+            
+            {f"ALL CAMPAIGNS IN THIS PAPER: {chr(10).join([f'- {c}' for c in self.all_campaigns])}" if self.all_campaigns else ""}
+            
+            CRITICAL: Only return locations that contain data for this specific campaign.
+            """
+        elif self.campaign_filter:
             campaigns_warning = ""
             if self.all_campaigns:
                 campaigns_warning = f"""
@@ -1103,11 +1142,114 @@ class ReactionExtractor:
             
         return toc_score >= 2
 
+    def _build_caption_index(self) -> Dict[str, Dict[str, Any]]:
+        """Build an index of all captions for quick lookup."""
+        if hasattr(self, '_caption_index'):
+            return self._caption_index
+            
+        cap_pattern = get_universal_caption_pattern()
+        caption_index = {}
+        
+        for idx, page in enumerate(self.all_pages):
+            source = "manuscript" if idx < len(self.ms_pages) else "supplementary"
+            page_num = idx + 1 if idx < len(self.ms_pages) else idx - len(self.ms_pages) + 1
+            
+            for match in cap_pattern.finditer(page):
+                caption_text = match.group(0).strip()
+                # Extract a normalized key (e.g., "table 5", "figure 3")
+                caption_lower = caption_text.lower()
+                
+                # Store multiple access patterns for the same caption
+                caption_info = {
+                    'full_caption': caption_text,
+                    'page_content': page,
+                    'page_idx': idx,
+                    'source': source,
+                    'page_num': page_num,
+                    'match_start': match.start()
+                }
+                
+                # Create multiple keys for flexible matching
+                # Key 1: Full caption text (first 100 chars)
+                key1 = caption_text[:100].lower().strip()
+                caption_index[key1] = caption_info
+                
+                # Key 2: Simplified reference (e.g., "table 5", "figure s3")
+                ref_match = re.search(r'(table|figure|fig|scheme)\s*s?(\d+[a-z]?)', caption_lower)
+                if ref_match:
+                    key2 = f"{ref_match.group(1)} {ref_match.group(2)}"
+                    caption_index[key2] = caption_info
+                    
+                    # Also store with 's' prefix if in SI
+                    if source == "supplementary" and 's' not in key2:
+                        key3 = f"{ref_match.group(1)} s{ref_match.group(2)}"
+                        caption_index[key3] = caption_info
+        
+        self._caption_index = caption_index
+        return caption_index
+    
     def _page_with_reference(self, ref_id: str) -> Optional[str]:
-        for page in self.all_pages:
-            if ref_id.lower() in page.lower():
-                return page
-        return None
+        """Find page(s) containing a reference using flexible matching."""
+        caption_index = self._build_caption_index()
+        ref_lower = ref_id.lower().strip()
+        
+        # Try multiple matching strategies
+        matches = []
+        
+        # Strategy 1: Direct key lookup
+        if ref_lower in caption_index:
+            matches.append(caption_index[ref_lower])
+        
+        # Strategy 2: Normalized reference lookup (e.g., "table 5", "figure s3")
+        ref_match = re.match(r'(table|figure|fig|scheme)\s*s?(\d+[a-z]?)', ref_lower, re.I)
+        if ref_match:
+            ref_type, ref_num = ref_match.groups()
+            if ref_type == 'fig':
+                ref_type = 'figure'
+            
+            # Try different key formats
+            keys_to_try = [
+                f"{ref_type} {ref_num}",
+                f"{ref_type} s{ref_num}",
+                f"table {ref_num}",  # Sometimes figures are mislabeled
+                f"fig {ref_num}",
+                f"figure {ref_num}"
+            ]
+            
+            for key in keys_to_try:
+                if key in caption_index and caption_index[key] not in matches:
+                    matches.append(caption_index[key])
+        
+        # Strategy 3: Fuzzy matching on caption text
+        if not matches:
+            # Look for any caption containing the reference number
+            for key, info in caption_index.items():
+                if ref_match and ref_num in key and any(t in key for t in ['table', 'figure', 'fig', 'scheme']):
+                    if info not in matches:
+                        matches.append(info)
+        
+        # Return results
+        if not matches:
+            LOGGER.warning(f"No matches found for reference '{ref_id}'")
+            # Last resort: simple text search
+            for page in self.all_pages:
+                if ref_lower in page.lower():
+                    return page
+            return None
+        
+        # If single match, return it
+        if len(matches) == 1:
+            return matches[0]['page_content']
+        
+        # Multiple matches: combine them with source annotations
+        LOGGER.info(f"Found {len(matches)} potential matches for '{ref_id}'")
+        combined_pages = []
+        for match in matches:
+            header = f"\n\n=== {match['source'].upper()} PAGE {match['page_num']} ===\n"
+            header += f"Caption: {match['full_caption'][:200]}...\n"
+            combined_pages.append(header + match['page_content'])
+        
+        return "\n".join(combined_pages)
 
     # ---- Table text helper - now returns full page ----
     def _extract_table_context(self, ref: str) -> str:
@@ -1140,19 +1282,29 @@ class ReactionExtractor:
         return pix
 
     # ---- NEW: Page image helper for both figures and tables ----
-    def _extract_page_png(self, ref: str, extract_figure_only: bool = True) -> Optional[str]:
+    def _extract_page_png(self, ref: str, extract_figure_only: bool = True, caption_hint: str = "", document_hint: str = "") -> Optional[str]:
         """Export the page containing the reference as PNG.
         If extract_figure_only=True, extracts just the figure above the caption.
         If False, extracts the entire page (useful for tables).
-        Returns a base64-encoded PNG or None."""
-        LOGGER.debug("_extract_page_png called with ref='%s', extract_figure_only=%s", ref, extract_figure_only)
+        Returns a base64-encoded PNG or None.
         
-        # Check cache first
-        cache_key = f"{ref}_{extract_figure_only}"
+        Args:
+            ref: The reference string (e.g., "Fig. 3")
+            extract_figure_only: Whether to extract just the figure or the entire page
+            caption_hint: Optional caption text from location data to help find the exact figure
+            document_hint: Optional hint about which document to search ("manuscript" or "supplementary")
+        """
+        LOGGER.info("_extract_page_png called with ref='%s', extract_figure_only=%s, caption_hint='%s', document_hint='%s'", 
+                    ref, extract_figure_only, caption_hint[:50] + "..." if caption_hint else "EMPTY", document_hint)
+        
+        # Check cache first - include document hint in key to avoid cross-document contamination
+        cache_key = f"{ref}_{extract_figure_only}_{document_hint}" if document_hint else f"{ref}_{extract_figure_only}"
         cached_result = self._figure_cache.get(cache_key)
         if cached_result is not None:
-            LOGGER.debug("Using cached figure for %s", ref)
+            LOGGER.info("Using cached figure for %s (cache key: %s)", ref, cache_key)
             return cached_result
+        else:
+            LOGGER.info("Cache miss for %s (cache key: %s)", ref, cache_key)
         
         # For table extraction, use multi-page approach
         if not extract_figure_only:
@@ -1162,13 +1314,28 @@ class ReactionExtractor:
                 return self._extract_multiple_pages_png(pages_with_ref, ref)
             return None
 
-        # For figure extraction, search both documents for actual figure captions
-        docs = list(filter(None, [self.ms_doc, self.si_doc]))
-        LOGGER.debug("Searching for '%s' in %d documents", ref, len(docs))
+        # For figure extraction, prioritize based on document hint
+        if document_hint == "manuscript" and self.ms_doc:
+            # Search manuscript first, then SI as fallback
+            docs = list(filter(None, [self.ms_doc, self.si_doc]))
+            LOGGER.info("Prioritizing manuscript document for '%s' (hint: %s)", ref, document_hint)
+            LOGGER.info("Search order: 1) Manuscript, 2) SI (fallback)")
+        elif document_hint == "supplementary" and self.si_doc:
+            # Search SI first, then manuscript as fallback
+            docs = list(filter(None, [self.si_doc, self.ms_doc]))
+            LOGGER.info("Prioritizing supplementary document for '%s' (hint: %s)", ref, document_hint)
+            LOGGER.info("Search order: 1) SI, 2) Manuscript (fallback)")
+        else:
+            # Default behavior - search both in order
+            docs = list(filter(None, [self.ms_doc, self.si_doc]))
+            LOGGER.info("Searching for '%s' in %d documents (no document hint)", ref, len(docs))
+            LOGGER.info("Search order: 1) Manuscript, 2) SI (default order)")
         
         for doc_idx, doc in enumerate(docs):
-            doc_name = "MS" if doc_idx == 0 else "SI"
-            LOGGER.debug("Searching in %s document with %d pages", doc_name, doc.page_count)
+            # Determine document name based on actual document, not position
+            doc_name = "MS" if doc == self.ms_doc else "SI"
+            LOGGER.info("Searching document %d/%d: %s (has %d pages)", 
+                       doc_idx + 1, len(docs), doc_name, doc.page_count)
             
             for page_number in range(doc.page_count):
                 page = doc.load_page(page_number)
@@ -1181,9 +1348,90 @@ class ReactionExtractor:
                     LOGGER.debug("Skipping page %d - detected as Table of Contents", page_number + 1)
                     continue
                 
+                # If we have a caption hint, try to find it using fuzzy matching
+                if caption_hint:
+                    LOGGER.info("=== CAPTION HINT SEARCH ===")
+                    LOGGER.info("Caption hint provided: %s", caption_hint[:100])
+                    LOGGER.info("Searching in %s document, page %d", doc_name, page_number + 1)
+                    LOGGER.info("Page text length: %d chars", len(page_text))
+                    
+                    # Check if caption exists in raw form
+                    if caption_hint[:50] in page_text:
+                        LOGGER.info("✓ Caption hint found in raw page text!")
+                    else:
+                        LOGGER.info("✗ Caption hint NOT found in raw page text")
+                    
+                    # Normalize texts for better matching
+                    def normalize_for_matching(text):
+                        # Remove extra whitespace, normalize spaces around punctuation
+                        text = ' '.join(text.split())
+                        # Normalize different dash types
+                        text = text.replace('–', '-').replace('—', '-')
+                        return text
+                    
+                    normalized_hint = normalize_for_matching(caption_hint[:100])  # Use first 100 chars
+                    normalized_page = normalize_for_matching(page_text)
+                    
+                    # Try to find the caption using fuzzy matching
+                    best_match_pos = -1
+                    best_match_score = 0
+                    match_found = False
+                    
+                    # Slide through the page text looking for best match
+                    hint_len = len(normalized_hint)
+                    for i in range(len(normalized_page) - hint_len + 1):
+                        snippet = normalized_page[i:i + hint_len]
+                        # Simple character-based similarity
+                        matches = sum(1 for a, b in zip(normalized_hint, snippet) if a == b)
+                        score = matches / hint_len
+                        
+                        if score > best_match_score and score > 0.8:  # 80% similarity threshold
+                            best_match_score = score
+                            best_match_pos = i
+                            match_found = True
+                    
+                    if match_found and best_match_pos >= 0:
+                        LOGGER.info("Found caption match in %s document on page %d with %.1f%% similarity", 
+                                   doc_name, page_number + 1, best_match_score * 100)
+                        
+                        # Instead of complex position mapping, just search for the beginning of the caption
+                        # Use the first 30 chars which should be unique enough
+                        search_text = caption_hint[:30].strip()
+                        LOGGER.info("Searching for caption text: '%s'", search_text)
+                        caption_instances = page.search_for(search_text)
+                        LOGGER.info("Found %d caption instances", len(caption_instances) if caption_instances else 0)
+                        
+                        if caption_instances:
+                            cap_rect = caption_instances[0]
+                            caption_found = True
+                            # Extract figure above this caption
+                            if extract_figure_only:
+                                LOGGER.info("Extracting figure area including caption for %s from %s document", ref, doc_name)
+                                LOGGER.info("Caption found at rect: %s on page %d", cap_rect, page_number + 1)
+                                page_rect = page.rect
+                                
+                                # Include the caption in the extraction
+                                # Add some padding below the caption to ensure we get the full text
+                                caption_padding = 30  # pixels below caption
+                                figure_rect = fitz.Rect(0, 0, page_rect.width, cap_rect.y1 + caption_padding)
+                                LOGGER.info("Page rect: %s, Figure rect including caption: %s", page_rect, figure_rect)
+                                mat = fitz.Matrix(5.0, 5.0)
+                                pix = page.get_pixmap(matrix=mat, clip=figure_rect)
+                                pix = self._ensure_rgb_pixmap(pix)
+                                img_bytes = pix.tobytes("png")
+                                img_b64 = b64encode(img_bytes).decode('utf-8')
+                                self._figure_cache.put(cache_key, img_b64)
+                                LOGGER.info("Successfully extracted figure using caption hint for %s from %s document, page %d", 
+                                           ref, doc_name, page_number + 1)
+                                return img_b64
+                    else:
+                        LOGGER.info("No fuzzy match found for caption hint on page %d (best score: %.1f%%)", 
+                                   page_number + 1, best_match_score * 100)
+                
+                # If caption hint didn't work or wasn't provided, fall back to pattern matching
                 # Look for figure caption pattern more flexibly
                 # Normalize the reference to handle variations
-                figure_num = ref.replace('Figure', '').replace('figure', '').strip()
+                figure_num = ref.replace('Figure', '').replace('figure', '').replace('Fig.', '').replace('Fig', '').strip()
                 
                 # Extract main figure number from subfigure (e.g., "1C" -> "1")
                 main_figure_num = re.match(r'^(\d+)', figure_num)
@@ -1193,10 +1441,10 @@ class ReactionExtractor:
                     main_figure_num = figure_num
                 
                 # Create a flexible pattern that handles various spacing and formatting
-                # This pattern looks for "Figure" (case insensitive) followed by optional spaces
-                # then the figure number, then any of: period, colon, space+capital letter, or end of line
+                # This pattern looks for "Figure" or "Fig" (case insensitive) followed by optional spaces
+                # then the figure number, then any of: period, colon, pipe, space+capital letter, or end of line
                 # Also match at the beginning of a line to catch captions
-                flexible_pattern = rf"(?i)(?:^|\n)\s*figure\s*{re.escape(main_figure_num)}(?:\.|:|(?=\s+[A-Z])|\s*$)"
+                flexible_pattern = rf"(?i)(?:^|\n)\s*(?:figure|fig\.?)\s*{re.escape(main_figure_num)}(?:\.|:|\||\s+\||(?=\s+[A-Z])|\s*$)"
                 
                 LOGGER.debug("Looking for figure caption '%s' with flexible pattern: %s", 
                            main_figure_num, flexible_pattern)
@@ -1252,8 +1500,8 @@ class ReactionExtractor:
                     if text_instances:
                         cap_rect = text_instances[0]
                         caption_found = True
-                        LOGGER.info("Found actual caption for %s: '%s' with following text: '%s...'", 
-                                  ref, caption_text, remaining_text[:50])
+                        LOGGER.info("Found actual caption for %s in %s document on page %d: '%s' with following text: '%s...'", 
+                                  ref, doc_name, page_number + 1, caption_text, remaining_text[:50])
                         break
                 
                 if not caption_found:
@@ -1274,18 +1522,18 @@ class ReactionExtractor:
                     continue
                 
                 if extract_figure_only:
-                    # Extract only the area above the caption (the actual figure)
-                    # This excludes caption text and focuses on visual elements
-                    LOGGER.info("Extracting figure area above caption for %s", ref)
+                    # Extract the figure area including the caption
+                    LOGGER.info("Extracting figure area including caption for %s", ref)
                     
                     # Get the page dimensions
                     page_rect = page.rect
                     
-                    # Extract the area above the caption
+                    # Extract the area including the caption
                     if cap_rect:
-                        # Extract from top of page to top of caption
-                        figure_rect = fitz.Rect(0, 0, page_rect.width, cap_rect.y0)
-                        LOGGER.debug("Extracting figure area: %s (caption at y=%f)", figure_rect, cap_rect.y0)
+                        # Extract from top of page to bottom of caption plus padding
+                        caption_padding = 30  # pixels below caption
+                        figure_rect = fitz.Rect(0, 0, page_rect.width, cap_rect.y1 + caption_padding)
+                        LOGGER.debug("Extracting figure area with caption: %s (caption ends at y=%f)", figure_rect, cap_rect.y1)
                     else:
                         # If no caption found, use top 80% of page
                         figure_rect = fitz.Rect(0, 0, page_rect.width, page_rect.height * 0.8)
@@ -1419,7 +1667,7 @@ class ReactionExtractor:
         # Sort pages by document and page number
         pages.sort(key=lambda x: (id(x[0]), x[1]))
         
-        # Extract the range of pages including one page after
+        # Extract the range of pages including one page after for tables
         all_images = []
         for i, (doc, page_num) in enumerate(pages):
             # Add the current page
@@ -1428,7 +1676,16 @@ class ReactionExtractor:
             pix = self._ensure_rgb_pixmap(pix)
             all_images.append(pix)
             
-            # Only extract the page containing the reference (removed next page logic)
+            # Add the next page as well for tables (in case data continues)
+            next_page_num = page_num + 1
+            if next_page_num < doc.page_count:
+                try:
+                    next_pix = doc.load_page(next_page_num).get_pixmap(matrix=mat)
+                    next_pix = self._ensure_rgb_pixmap(next_pix)
+                    all_images.append(next_pix)
+                    LOGGER.info("Including next page (%d) for table %s", next_page_num + 1, ref)
+                except Exception as e:
+                    LOGGER.warning("Failed to extract next page %d for %s: %s", next_page_num + 1, ref, e)
         
         if not all_images:
             return None
@@ -1533,21 +1790,9 @@ class ReactionExtractor:
     
     def _validate_location_exists(self, ref: str) -> bool:
         """Verify that the referenced location actually exists in the document."""
-        # Search for the actual reference in both manuscript and SI documents
-        docs_to_check = [self.ms_doc]
-        if self.si_doc:
-            docs_to_check.append(self.si_doc)
-            
-        for doc in docs_to_check:
-            for page_num in range(len(doc)):
-                page = doc[page_num]
-                text = page.get_text()
-                
-                # Look for table references like "Table 1", "Table S1", etc.
-                if re.search(rf'\b{re.escape(ref)}\b', text, re.IGNORECASE):
-                    return True
-        
-        return False
+        # Use the caption index to check if location exists
+        result = self._page_with_reference(ref)
+        return result is not None
 
     def _validate_context(self, snippet: str, enzyme_list: List[str], ref: str) -> bool:
         """Validate that the context contains meaningful content for extraction."""
@@ -1596,15 +1841,31 @@ class ReactionExtractor:
         LOGGER.info("Response validated for %s: %d enzymes with data", ref, enzymes_with_data)
         return True
 
-    def extract_metrics_batch(self, enzyme_list: List[str], ref: str) -> List[Dict[str, Any]]:
-        """Extract performance metrics for multiple enzymes from the identified location in batch."""
-        LOGGER.info("extract_metrics_batch called with ref='%s' for %d enzymes", ref, len(enzyme_list))
-        ref_lc = ref.lower()
+    def extract_metrics_batch(self, enzyme_list: List[str], ref: Union[str, Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Extract performance metrics for multiple enzymes from the identified location in batch.
+        
+        Args:
+            enzyme_list: List of enzyme names to extract metrics for
+            ref: Either a string reference (e.g., "Fig. 3") or a location dict with 'location' and optionally 'caption'
+        """
+        # Handle both string and dict inputs
+        if isinstance(ref, dict):
+            location_str = ref['location']
+            caption_hint = ref.get('caption', '')
+            document_hint = ref.get('document', '')
+            LOGGER.info("extract_metrics_batch called with location='%s' (with caption hint, document=%s) for %d enzymes", 
+                       location_str, document_hint, len(enzyme_list))
+        else:
+            location_str = ref
+            caption_hint = ''
+            document_hint = ''
+            LOGGER.info("extract_metrics_batch called with ref='%s' for %d enzymes", location_str, len(enzyme_list))
+        ref_lc = location_str.lower()
         image_b64: Optional[str] = None
         
         # First, validate that the location actually exists in the document
-        if not self._validate_location_exists(ref):
-            LOGGER.warning("Location %s not found in document - skipping", ref)
+        if not self._validate_location_exists(location_str):
+            LOGGER.warning("Location %s not found in document - skipping", location_str)
             return []
         
         # Add campaign context if available
@@ -1614,30 +1875,30 @@ class ReactionExtractor:
         
         if self._TAB_RE.search(ref_lc):
             # For tables, try to extract the page as an image first
-            image_b64 = self._extract_page_png(ref, extract_figure_only=False)
+            image_b64 = self._extract_page_png(location_str, extract_figure_only=False, document_hint=document_hint)
             if not image_b64:
-                LOGGER.debug("No page image found for %s - using full page text", ref)
-                snippet = self._extract_table_context(ref)
+                LOGGER.debug("No page image found for %s - using full page text", location_str)
+                snippet = self._extract_table_context(location_str)
         elif self._FIG_RE.search(ref_lc):
             # For figures, extract just the figure image (same logic as compound mapping)
-            LOGGER.debug("Attempting to extract figure image for '%s'", ref)
-            image_b64 = self._extract_page_png(ref, extract_figure_only=True)
+            LOGGER.info("Attempting to extract figure image for '%s'", location_str)
+            image_b64 = self._extract_page_png(location_str, extract_figure_only=True, caption_hint=caption_hint, document_hint=document_hint)
             if not image_b64:
-                LOGGER.warning("Failed to extract figure image for '%s' - falling back to caption text", ref)
-                snippet = self._extract_figure_caption(ref)
+                LOGGER.warning("Failed to extract figure image for '%s' - falling back to caption text", location_str)
+                snippet = self._extract_figure_caption(location_str)
                 LOGGER.debug("Caption extraction result: %s", 
                            f"'{snippet[:100]}...'" if snippet else "empty")
             else:
-                LOGGER.info("Successfully extracted figure image for '%s'", ref)
+                LOGGER.info("Successfully extracted figure image for '%s'", location_str)
                 # If figure is found, ignore text information - use image only
                 snippet = ""
         else:
-            snippet = self._page_with_reference(ref) or ""
+            snippet = self._page_with_reference(location_str) or ""
 
         # For figures with images, skip text validation and proceed with image extraction
         if image_b64 and self._FIG_RE.search(ref_lc):
-            LOGGER.info("Using figure image for %s - ignoring text context", ref)
-        elif not image_b64 and not self._validate_context(snippet, enzyme_list, ref):
+            LOGGER.info("Using figure image for %s - ignoring text context", location_str)
+        elif not image_b64 and not self._validate_context(snippet, enzyme_list, location_str):
             return []
 
         # Create enhanced enzyme descriptions with parent/mutation context
@@ -1934,6 +2195,19 @@ Different campaigns may use different model reactions.
         prompt = PROMPT_COMPOUND_MAPPING
         if campaign_filter:
             prompt += f"\n\nIMPORTANT: Focus on compound information relevant to the {campaign_filter} campaign/reaction system."
+        
+        # Add campaign info as hints
+        if self.campaign_info:
+            substrate_id = self.campaign_info.get('substrate_id', '')
+            product_id = self.campaign_info.get('product_id', '')
+            model_substrate = self.campaign_info.get('model_substrate', '')
+            model_product = self.campaign_info.get('model_product', '')
+            
+            if substrate_id and model_substrate:
+                prompt += f"\n\nHINT: The model substrate for this campaign is likely '{model_substrate}' (ID: {substrate_id})"
+            if product_id and model_product:
+                prompt += f"\nHINT: The model product for this campaign is likely '{model_product}' (ID: {product_id})"
+        
         if compound_ids:
             prompt += "\n\nCOMPOUNDS TO MAP: " + ", ".join(sorted(compound_ids))
         prompt += "\n\nTEXT:\n" + extraction_text
@@ -2011,6 +2285,22 @@ IMPORTANT CAMPAIGN CONTEXT: Focus on compound information relevant to the {campa
 {campaigns_warning}
 Different campaigns may use different numbering systems for compounds.
 Do NOT include compound information from other campaigns."""
+        
+        # Add campaign info as hints
+        if self.campaign_info:
+            substrate_id = self.campaign_info.get('substrate_id', '')
+            product_id = self.campaign_info.get('product_id', '')
+            model_substrate = self.campaign_info.get('model_substrate', '')
+            model_product = self.campaign_info.get('model_product', '')
+            
+            hints = []
+            if substrate_id and model_substrate:
+                hints.append(f"The model substrate for this campaign is likely '{model_substrate}' (ID: {substrate_id})")
+            if product_id and model_product:
+                hints.append(f"The model product for this campaign is likely '{model_product}' (ID: {product_id})")
+            
+            if hints:
+                prompt += "\n\nHINTS FROM CAMPAIGN INFO:\n" + "\n".join(hints)
 
         prompt += """
 
@@ -2476,6 +2766,7 @@ Do NOT include compound information from other campaigns.
         compound_mappings = {}
         if compound_ids:
             LOGGER.info("Using 3-tier compound mapping approach for compounds: %s", compound_ids)
+            
             # Pass the IUPAC location hint if we have it
             iupac_hint = locations.get("iupac_location") if locations else None
             compound_mappings = self._extract_compound_mappings_adaptive(
@@ -2506,6 +2797,22 @@ CRITICAL WARNING: Do NOT confuse campaigns! Each campaign uses completely differ
 - Different campaigns may use similar enzyme names but different substrates
 - Be extremely careful to only extract data for the {self.campaign_filter} campaign
 - Ignore data from other campaigns even if they seem similar
+"""
+            
+            # Add specific campaign info if available
+            campaign_info_context = ""
+            if self.campaign_info:
+                campaign_info_context = f"""
+
+KNOWN CAMPAIGN INFORMATION:
+- Campaign: {self.campaign_info.get('campaign_name', '')}
+- Model Substrate: {self.campaign_info.get('model_substrate', '')} (ID: {self.campaign_info.get('substrate_id', '')})
+- Model Product: {self.campaign_info.get('model_product', '')} (ID: {self.campaign_info.get('product_id', '')})
+- Known Data Locations: {', '.join(self.campaign_info.get('data_locations', []))}
+
+IMPORTANT: Use this information to guide your extraction. The model reaction should involve:
+- Substrate ID: {self.campaign_info.get('substrate_id', '')}
+- Product ID: {self.campaign_info.get('product_id', '')}
 """
             
             campaign_context = f"""
@@ -2515,6 +2822,7 @@ You are extracting the model reaction used specifically for these enzyme variant
 
 These variants belong to campaign: {self.campaign_filter}
 {campaigns_context}
+{campaign_info_context}
 Focus on extracting the model reaction that was used to evaluate THESE specific variants.
 Different campaigns may use different model reactions and substrates.
 
@@ -2790,7 +3098,7 @@ Different campaigns may use different model reactions and substrates.
                        best_location.get('confidence', 0))
             
             # Extract metrics from the most confident source only
-            metrics_rows = self.extract_metrics_batch(all_enzyme_ids, best_location['location'])
+            metrics_rows = self.extract_metrics_batch(all_enzyme_ids, best_location)
             
             # Filter to valid metrics
             valid_metrics = [m for m in metrics_rows if self._has_valid_metrics(m)]
@@ -3034,7 +3342,7 @@ Different campaigns may use different model reactions and substrates.
                        location['location'], location.get('confidence', 0))
             
             # Extract metrics from this location
-            metrics_rows = self.extract_metrics_batch(list(all_variants), location['location'])
+            metrics_rows = self.extract_metrics_batch(list(all_variants), location)
             
             # Filter to valid metrics
             valid_metrics = [m for m in metrics_rows if self._has_valid_metrics(m)]
@@ -3277,8 +3585,20 @@ def main() -> None:
                     campaign_debug_dir.mkdir(parents=True, exist_ok=True)
                     LOGGER.info("Campaign debug directory: %s", campaign_debug_dir)
                 
+                # Load campaign info from campaigns.json if available
+                campaign_info = None
+                if args.debug_dir:
+                    from .campaign_utils import load_campaigns_from_file, find_campaign_by_id
+                    campaigns_file = Path(args.debug_dir) / "campaigns.json"
+                    if campaigns_file.exists():
+                        campaigns = load_campaigns_from_file(campaigns_file)
+                        campaign_info = find_campaign_by_id(campaigns, campaign_filter)
+                        if campaign_info:
+                            LOGGER.info("Loaded campaign info for %s from campaigns.json", campaign_filter)
+                
                 extractor = ReactionExtractor(args.manuscript, args.si, cfg, debug_dir=campaign_debug_dir, 
-                                            campaign_filter=campaign_filter, all_campaigns=all_campaigns)
+                                            campaign_filter=campaign_filter, all_campaigns=all_campaigns,
+                                            campaign_info=campaign_info)
                 df_metrics = extractor.run(enzyme_df)
                 
             elif len(all_campaigns) > 1:
@@ -3304,8 +3624,20 @@ def main() -> None:
                         campaign_debug_dir.mkdir(parents=True, exist_ok=True)
                         LOGGER.info("Campaign debug directory: %s", campaign_debug_dir)
                     
+                    # Load campaign info from campaigns.json if available
+                    campaign_info = None
+                    if args.debug_dir:
+                        from .campaign_utils import load_campaigns_from_file, find_campaign_by_id
+                        campaigns_file = Path(args.debug_dir) / "campaigns.json"
+                        if campaigns_file.exists():
+                            campaigns = load_campaigns_from_file(campaigns_file)
+                            campaign_info = find_campaign_by_id(campaigns, campaign)
+                            if campaign_info:
+                                LOGGER.info("Loaded campaign info for %s from campaigns.json", campaign)
+                    
                     extractor = ReactionExtractor(args.manuscript, args.si, cfg, debug_dir=campaign_debug_dir, 
-                                                campaign_filter=campaign, all_campaigns=all_campaigns)
+                                                campaign_filter=campaign, all_campaigns=all_campaigns,
+                                                campaign_info=campaign_info)
                     
                     # Run extraction for this campaign
                     campaign_metrics = extractor.run(campaign_df)