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

debase/enzyme_lineage_extractor.py

@@ -465,7 +465,7 @@ def get_model():
         "temperature": 0.0,  # Deterministic: always pick the most likely token
         "top_p": 1.0,      # Consider all tokens (but temperature=0 will pick the best)
         "top_k": 1,        # Only consider the single most likely token
-        "max_output_tokens": 32768,  # Increased from 8192 to handle larger sequence extractions
+        "max_output_tokens": 65536,  # Increased to 2x for handling larger lineage tables and sequences
     }
     
     # For Gemini 2.5 Flash, disable thinking tokens to save costs
@@ -758,13 +758,24 @@ mutations were introduced){campaign_specific}. Pay attention to the provided con
 ensure the location you return are actually lineage location with variants and mutations.
 
 Respond with a JSON array of objects, each containing:
-- "location": the figure/table identifier (e.g. "Table S1", "Figure 2B", "Table 1", "Figure 3")
+- "location": the figure/table identifier EXACTLY as it appears in the caption (e.g. "Table S1", "Figure 2B", "Table 1", "Figure 3")
 - "type": one of "table", "figure", "section"  
 - "confidence": your confidence score (0-100) that this location contains lineage data
 - "reason": brief explanation of why this location likely contains lineage
+- "source": one of "manuscript" or "si" - indicate whether this is in the main manuscript or supplementary information
+- "caption": the FULL caption text (include at least the first 200-300 characters of the caption to enable fuzzy matching)
 {campaign_field}
-IMPORTANT: Return ONLY figure/table identifiers like "Figure 2" or "Table S1", 
-NOT page numbers. Focus on the actual figure/table titles and numbers.
+CRITICAL INSTRUCTIONS:
+1. 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
+2. Include the FULL caption text in the "caption" field to enable fuzzy matching when extracting
+   - This should be the complete caption as it appears in the document
+   - Include at least 200-300 characters to ensure unique matching
+3. For each location, specify whether it's in the main manuscript or supplementary information (SI):
+   - Items like "Table S1", "Figure S2", etc. are typically in the SI
+   - Items like "Table 1", "Figure 2", etc. are typically in the main manuscript
+   - If uncertain, use context clues from the text
 
 Order by confidence score (highest first). Tables showing complete variant lineages or 
 mutation lists should be ranked higher than figures showing complete variant lineages.
@@ -774,9 +785,9 @@ Don't include oligonucleotide results or result from only one round.
 
 Example output:
 [
-  {{"location": "Table S1", "type": "table", "confidence": 95, "reason": "Variant lineage table"{campaign_example}}},
-  {{"location": "Figure 2B", "type": "figure", "confidence": 70, "reason": "Phylogenetic tree diagram"{campaign_example}}},
-  {{"location": "Section 3.2", "type": "section", "confidence": 60, "reason": "Evolution description"{campaign_example}}}
+  {{"location": "Table S1.", "type": "table", "confidence": 95, "reason": "Variant lineage table", "source": "si", "caption": "Table S1. Summary of mutations introduced during directed evolution of PA-G8. The table shows all variants tested in each round of SSM with their corresponding mutations and activities..."{campaign_example}}},
+  {{"location": "Figure 2B", "type": "figure", "confidence": 70, "reason": "Phylogenetic tree diagram", "source": "manuscript", "caption": "Figure 2B Phylogenetic tree showing the evolutionary relationships between enzyme variants. Each node represents a variant with mutations indicated on branches..."{campaign_example}}},
+  {{"location": "Section 3.2", "type": "section", "confidence": 60, "reason": "Evolution description", "source": "manuscript", "caption": "Section 3.2 Directed Evolution Campaign. We performed eight rounds of site-saturation mutagenesis..."{campaign_example}}}
 ]
 """.strip()
 
@@ -1461,10 +1472,114 @@ def _extract_text_at_locations(text: str, locations: List[Union[str, dict]], con
 
 # ---- 6.4  Public API -------------------------------------------------------
 
-def _extract_location_text(full_text: str, location: str, location_type: str) -> Optional[str]:
-    """Extract text from a specific location (table, section, etc.) in the full text."""
+def _extract_location_text(full_text: str, location: str, location_type: str, caption_hint: Optional[str] = None) -> Optional[str]:
+    """Extract text from a specific location (table, section, etc.) in the full text.
+    
+    Args:
+        full_text: The full text to search in
+        location: The location identifier (e.g., "Table S1")
+        location_type: Type of location ("table", "figure", "section")
+        caption_hint: Optional full caption text for fuzzy matching
+    """
     import re
     
+    # If caption hint is provided, try fuzzy matching first
+    if caption_hint and len(caption_hint) > 20:
+        log.info(f"Using caption hint for fuzzy matching: {caption_hint[:100]}...")
+        
+        # Normalize texts for better matching (similar to reaction_info_extractor)
+        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[:150])  # Use first 150 chars
+        normalized_text = normalize_for_matching(full_text)
+        
+        # Try to find ALL caption matches using character-based fuzzy matching
+        all_matches = []
+        
+        # Slide through the text looking for all matches above threshold
+        hint_len = len(normalized_hint)
+        for i in range(len(normalized_text) - hint_len + 1):
+            snippet = normalized_text[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 > 0.7:  # 70% similarity threshold
+                all_matches.append({
+                    'norm_pos': i,
+                    'score': score
+                })
+        
+        # If we found matches, extract from all of them
+        if all_matches:
+            log.info(f"Found {len(all_matches)} caption matches with fuzzy matching")
+            
+            # Collect all occurrences from fuzzy matches
+            all_occurrences = []
+            seen_positions = set()
+            
+            for match_info in all_matches:
+                # Get the matched text from normalized version
+                matched_normalized = normalized_text[match_info['norm_pos']:match_info['norm_pos'] + hint_len]
+                
+                # Find where this appears in the original text
+                best_original_pos = -1
+                
+                # Search in the original text for this specific match
+                for i in range(len(full_text) - len(caption_hint) + 1):
+                    if i in seen_positions:
+                        continue
+                        
+                    original_snippet = full_text[i:i + len(caption_hint)]
+                    # Normalize and compare
+                    normalized_snippet = normalize_for_matching(original_snippet)
+                    if normalized_snippet[:hint_len] == matched_normalized:
+                        # Found exact match after normalization
+                        best_original_pos = i
+                        seen_positions.add(i)
+                        break
+                
+                if best_original_pos >= 0:
+                    # Extract generous context from this match position
+                    start = max(0, best_original_pos - 1000)
+                    end = min(len(full_text), best_original_pos + 10000)
+                    context = full_text[start:end]
+                    
+                    all_occurrences.append({
+                        'position': best_original_pos,
+                        'context': context,
+                        'score': match_info['score']
+                    })
+                    log.info(f"Fuzzy match at position {best_original_pos} with {match_info['score']*100:.1f}% similarity")
+            
+            if all_occurrences:
+                # Sort by position to maintain document order
+                all_occurrences.sort(key=lambda x: x['position'])
+                
+                # Combine all occurrences
+                combined_text = f"=== All occurrences of {location} (fuzzy matched) ===\n\n"
+                
+                for i, occurrence in enumerate(all_occurrences, 1):
+                    combined_text += f"--- Occurrence {i} at position {occurrence['position']} (similarity: {occurrence['score']*100:.1f}%) ---\n"
+                    combined_text += occurrence['context']
+                    combined_text += "\n\n"
+                
+                # Apply same limit as table extraction
+                if len(combined_text) > 150000:
+                    combined_text = combined_text[:150000] + "\n\n[Truncated due to length...]"
+                
+                log.info(f"Extracted {len(combined_text)} chars using fuzzy caption matching from {len(all_occurrences)} locations")
+                return combined_text
+            else:
+                log.warning(f"Could not map any fuzzy matches back to original text")
+        else:
+            log.warning(f"No fuzzy matches found for caption above 70% threshold")
+    
     if location_type == 'table':
         # Find ALL mentions of this table and combine them
         location_clean = location.strip()
@@ -1506,6 +1621,9 @@ def _extract_location_text(full_text: str, location: str, location_type: str) ->
         
         log.info(f"Found {len(all_occurrences)} occurrences of table '{location_clean}'")
         
+        # Sort occurrences by position to maintain document order
+        all_occurrences.sort(key=lambda x: x['position'])
+        
         # Combine all occurrences into one text for Gemini to analyze
         combined_text = f"=== All occurrences of {location_clean} ===\n\n"
         
@@ -1515,8 +1633,9 @@ def _extract_location_text(full_text: str, location: str, location_type: str) ->
             combined_text += "\n\n"
         
         # Limit total length to avoid overwhelming the model
-        if len(combined_text) > 50000:
-            combined_text = combined_text[:50000] + "\n\n[Truncated due to length...]"
+        # Increased limit to ensure actual table content is included
+        if len(combined_text) > 150000:
+            combined_text = combined_text[:150000] + "\n\n[Truncated due to length...]"
         
         return combined_text
     
@@ -1600,6 +1719,8 @@ def get_lineage(
     *,
     pdf_paths: Optional[List[Path]] = None,
     debug_dir: str | Path | None = None,
+    manuscript_text: Optional[str] = None,
+    si_text: Optional[str] = None,
 ) -> Tuple[List[Variant], List[Campaign]]:
     """
     High-level wrapper used by the pipeline.
@@ -1713,8 +1834,21 @@ def get_lineage(
                 if location_type in ['table', 'text', 'section'] and not extracted_variants:
                     log.info(f"Attempting text extraction for {location_type}: {location_str}")
                     
-                    # Extract the specific section/table from full text
-                    section_text = _extract_location_text(full_text, location_str, location_type)
+                    # Determine which text to use based on source
+                    location_source = location.get('source', 'manuscript')
+                    if location_source == 'si' and si_text:
+                        text_to_search = si_text
+                        log.info(f"Using SI text for location {location_str}")
+                    elif location_source == 'manuscript' and manuscript_text:
+                        text_to_search = manuscript_text
+                        log.info(f"Using manuscript text for location {location_str}")
+                    else:
+                        text_to_search = full_text
+                        log.info(f"Using combined text for location {location_str} (fallback)")
+                    
+                    # Extract the specific section/table from appropriate text
+                    caption_hint = location.get('caption', '')
+                    section_text = _extract_location_text(text_to_search, location_str, location_type, caption_hint)
                     if section_text:
                         log.info(f"Extracted {len(section_text)} chars from {location_type}: {location_str}")
                         # Save extracted section if debug enabled
@@ -2028,17 +2162,24 @@ PRIORITY: Protein/amino acid sequences are preferred over DNA sequences.
 
 Look for table of contents entries or section listings that mention sequences.
 Return a JSON array where each element has:
-- "section": the section heading or description
+- "section": the section heading or description EXACTLY as it appears
 - "page": the page number (IMPORTANT: Return ONLY the number, e.g., "53" not "p. 53" or "page 53")
+- "source": one of "manuscript" or "si" - indicate whether this is in the main manuscript or supplementary information
+- "caption": the FULL section heading or table of contents entry (at least 100-200 characters for fuzzy matching)
 
 Focus on:
 - Table of contents or entries about "Sequence Information" or "Nucleotide and amino acid sequences"
 - For supplementary pages, use "S" prefix (e.g., "S53" not "p. S53")
 - Prioritize sections that mention "protein" or "amino acid" sequences
 
-CRITICAL: Page numbers must be returned as plain numbers or S-prefixed numbers only:
-- Correct: "53", "S12", "147"
-- Wrong: "p. 53", "P. 53", "page 53", "pg 53"
+CRITICAL: 
+1. Page numbers must be returned as plain numbers or S-prefixed numbers only:
+   - Correct: "53", "S12", "147"
+   - Wrong: "p. 53", "P. 53", "page 53", "pg 53"
+2. For each location, specify whether it's in the main manuscript or supplementary information (SI):
+   - Pages with "S" prefix (e.g., "S53") are typically in the SI
+   - Regular page numbers (e.g., "53") are typically in the main manuscript
+   - Use context clues from the document structure
 
 Return [] if no sequence sections are found.
 Absolutely don't include nucleotides or primer sequences, it is better to return nothing then incomplete sequence, use your best judgement.
@@ -2278,44 +2419,34 @@ def validate_sequence_locations(text: str, locations: list, model, *, pdf_paths:
 
 # --- 7.3  Main extraction prompt ---------------------------------------------
 _SEQ_EXTRACTION_PROMPT = """
-Extract EVERY distinct enzyme-variant sequence you can find in the text.
-
-IMPORTANT: Prioritize amino acid (protein) sequences over DNA sequences:
-- If an amino acid sequence exists for a variant, extract ONLY the aa_seq (set dna_seq to null)
-- Only extract dna_seq if NO amino acid sequence is available for that variant
-- This reduces redundancy since protein sequences are usually more relevant
-
-CRITICAL: Use the EXACT variant identifier as it appears with each sequence:
-- Papers often use different naming conventions in different sections
-- DO NOT normalize or simplify variant IDs
-- Extract the variant_id exactly as written where the sequence appears
-- Common patterns include numeric IDs, generation labels, full descriptive names, or combinations
-
-SEQUENCE EXTRACTION RULES:
-- Copy sequences EXACTLY as they appear in the text
-- Pay careful attention to repeated amino acids, and nucleotides (e.g., "AAA" should remain "AAA", not become "A")
-- Do NOT add, remove, or modify any amino acids, or nucleotides
-- Preserve the exact length and character sequence
-- If a sequence has line breaks or spacing, remove only formatting (spaces, newlines) but keep all amino acids
-- Double-check that consecutive identical amino acids or nucleotides  are copied correctly
-
-For each variant return:
-  * variant_id  - the EXACT label as it appears with the sequence (preserve all formatting)
-  * aa_seq      - amino-acid sequence (uppercase), or null - COPY EXACTLY FROM TEXT
-  * dna_seq     - DNA sequence (A/C/G/T), or null (ONLY if no aa_seq exists) - COPY EXACTLY FROM TEXT
-
-Respond ONLY with **minified JSON** that matches the schema below.
-NO markdown, no code fences, no commentary.
-
-Schema:
-```json
-{schema}
-```
+Extract ALL enzyme variant sequences from the text.
+
+Rules:
+1. Use EXACT variant IDs as they appear with each sequence
+2. Copy sequences EXACTLY - preserve all amino acids/nucleotides including repeats
+3. For each variant:
+   - If amino acid sequence exists: set aa_seq to the sequence, set dna_seq to null
+   - If ONLY DNA sequence exists: set dna_seq to the sequence, set aa_seq to null
+   - NEVER include both aa_seq and dna_seq for the same variant
+   - IMPORTANT: Always prefer amino acid sequences over DNA sequences when both are available
+4. Return ONLY minified JSON, no markdown or commentary
+
+CRITICAL SEQUENCE PRIORITY RULE:
+- If you find BOTH amino acid sequence AND DNA sequence for the same variant, ONLY return the amino acid sequence
+- Set dna_seq to null when aa_seq is available, even if DNA sequence is present in the text
+- Only return dna_seq when NO amino acid sequence exists for that variant
+
+CRITICAL ACCURACY REQUIREMENTS:
+- Extract ONLY sequences that are explicitly present in the provided text
+- DO NOT generate, infer, or hallucinate any sequences
+- Every character in the sequence must be directly copied from the text
+- If a sequence appears truncated or incomplete in the text, extract only what is shown
+- Be extremely careful and accurate - sequence accuracy is critical for scientific validity
+
+Schema: {schema}
 
-TEXT (may be truncated):
-```
+TEXT:
 {text}
-```
 """.strip()
 
 def _check_sequence_responses_match(resp1: Union[list, dict], resp2: Union[list, dict]) -> bool:
@@ -2390,7 +2521,7 @@ def _check_sequence_responses_match(resp1: Union[list, dict], resp2: Union[list,
 
 
 def _extract_sequences_with_triple_validation(model, prompt: str, schema_hint: str, *, debug_dir: str | Path | None = None) -> Optional[Any]:
-    """Extract sequence JSON using Gemini with up to 6 attempts, returning most common result.
+    """Extract sequence JSON using Gemini with up to 3 attempts, returning most common result.
     
     Can exit early after 2 attempts if the responses match exactly.
     
@@ -2404,9 +2535,9 @@ def _extract_sequences_with_triple_validation(model, prompt: str, schema_hint: s
         The most common sequence JSON data or None if all attempts failed
     """
     responses = []
-    max_attempts = 6
+    max_attempts = 3  # Reduced from 6 to 3 for performance
     
-    # Try 6 times with early match detection
+    # Try 3 times with early match detection
     for attempt in range(max_attempts):
         try:
             log.info(f"Sequence extraction attempt {attempt + 1}/{max_attempts}")
@@ -2432,8 +2563,13 @@ def _extract_sequences_with_triple_validation(model, prompt: str, schema_hint: s
             
             # Try to parse as JSON
             try:
-                parsed = json.loads(raw)
-            except json.JSONDecodeError:
+                # First clean the response - remove any BOM or invisible characters
+                raw_clean = raw.strip()
+                if raw_clean.startswith('\ufeff'):  # Remove BOM if present
+                    raw_clean = raw_clean[1:]
+                parsed = json.loads(raw_clean)
+            except json.JSONDecodeError as e:
+                log.debug(f"Initial JSON parsing failed: {e}. Response starts with: {repr(raw[:100])}")
                 # Look for JSON array or object in the response
                 json_start = -1
                 json_end = -1
@@ -2482,17 +2618,22 @@ def _extract_sequences_with_triple_validation(model, prompt: str, schema_hint: s
             responses.append(parsed)
             log.info(f"Sequence extraction attempt {attempt + 1}: {len(parsed) if isinstance(parsed, list) else 'object'} sequences")
             
-            # Early match detection after 2 attempts
-            if attempt >= 1:  # After 2nd attempt (0-indexed)
-                valid_responses_so_far = [r for r in responses if r is not None]
-                if len(valid_responses_so_far) >= 2:
-                    # Check if the last two valid responses match
-                    if _check_sequence_responses_match(valid_responses_so_far[-2], valid_responses_so_far[-1]):
-                        log.info(f"Early match detected after {attempt + 1} attempts - sequences are consistent")
-                        # Add the matching response 4 more times to simulate consensus
-                        for _ in range(max_attempts - attempt - 1):
-                            responses.append(valid_responses_so_far[-1])
-                        break
+            # If we got a good response with sequences, we can check for early termination
+            if isinstance(parsed, list) and len(parsed) > 0:
+                # Early match detection after 2 attempts
+                if attempt >= 1:  # After 2nd attempt (0-indexed)
+                    valid_responses_so_far = [r for r in responses if r is not None and isinstance(r, list) and len(r) > 0]
+                    if len(valid_responses_so_far) >= 2:
+                        # Check if the last two valid responses match
+                        if _check_sequence_responses_match(valid_responses_so_far[-2], valid_responses_so_far[-1]):
+                            log.info(f"Early match detected after {attempt + 1} attempts - sequences are consistent")
+                            # Add the matching response to fill remaining attempts
+                            for _ in range(max_attempts - attempt - 1):
+                                responses.append(valid_responses_so_far[-1])
+                            break
+                # If this is the first attempt and we got sequences, continue to validate with at least one more
+                elif attempt == 0 and len(parsed) > 5:  # Got substantial sequences on first try
+                    log.info("Got substantial sequences on first attempt, will validate with one more")
             
         except Exception as e:
             log.warning(f"Sequence extraction attempt {attempt + 1} failed: {e}")
@@ -2852,9 +2993,9 @@ def get_sequences(text: str, model, *, pdf_paths: List[Path] = None, debug_dir:
             focused_text = ""
             if pdf_paths and isinstance(best_location, dict) and 'page' in best_location:
                 page_num = best_location['page']
-                # Extract current page plus next 15 pages
+                # Extract current page plus next 5 pages (6 total) to prevent hallucination
                 all_pages = []
-                for i in range(16):  # Current + next 15
+                for i in range(6):  # Current + next 5 (6 pages total)
                     if isinstance(page_num, str) and page_num.upper().startswith('S'):
                         next_page = f"S{int(page_num[1:]) + i}"
                     else:
@@ -2866,7 +3007,7 @@ def get_sequences(text: str, model, *, pdf_paths: List[Path] = None, debug_dir:
                         break
                 if all_pages:
                     focused_text = "\n".join(all_pages)
-                    log.info("Extracted %d chars from pages %s through %d more pages", 
+                    log.info("Extracted %d chars from pages %s through %d more pages (limited to 6 pages total)", 
                              len(focused_text), page_num, len(all_pages) - 1)
             
             # Fallback to text search if page extraction didn't work
@@ -3152,6 +3293,83 @@ def fetch_sequence_by_name(enzyme_name: str) -> Dict[str, str]:
         return {}
 
 
+def _match_variant_ids_with_gemini(
+    lineage_variant_ids: List[str], 
+    pdb_variant_ids: List[str], 
+    model
+) -> Dict[str, str]:
+    """Use Gemini to match variant IDs that may have slight formatting differences.
+    
+    Args:
+        lineage_variant_ids: List of variant IDs from the lineage
+        pdb_variant_ids: List of variant IDs from PDB matching
+        model: Gemini model for matching
+        
+    Returns:
+        Dictionary mapping lineage_variant_id -> pdb_variant_id
+    """
+    if not lineage_variant_ids or not pdb_variant_ids or not model:
+        return {}
+    
+    # If the lists are identical, return direct mapping
+    if set(lineage_variant_ids) == set(pdb_variant_ids):
+        return {vid: vid for vid in lineage_variant_ids if vid in pdb_variant_ids}
+    
+    # Use Gemini to match variant IDs that may have formatting differences
+    prompt = f"""Match variant IDs between two lists that may have slight formatting differences (whitespace, encoding, etc.).
+These represent the same enzyme variants but may be formatted differently.
+
+Lineage variant IDs:
+{json.dumps(lineage_variant_ids, indent=2)}
+
+PDB variant IDs:
+{json.dumps(pdb_variant_ids, indent=2)}
+
+Match variants that represent the SAME enzyme variant, accounting for:
+- Whitespace differences (extra spaces, tabs)
+- Character encoding differences
+- Minor formatting variations
+
+Return ONLY a JSON object mapping lineage IDs to PDB IDs.
+Format: {{"lineage_id": "pdb_id", ...}}
+Only include matches you are confident represent the same variant.
+Return an empty object {{}} if no matches can be confidently made.
+"""
+    
+    try:
+        response = model.generate_content(prompt)
+        text = _extract_text(response).strip()
+        
+        # Parse JSON response
+        if text.startswith("```"):
+            text = text.split("```")[1].strip()
+            if text.startswith("json"):
+                text = text[4:].strip()
+        
+        # Clean up the text
+        text = text.strip()
+        if not text or text == "{}":
+            return {}
+        
+        matches = json.loads(text)
+        log.info(f"Gemini matched {len(matches)} variant IDs for PDB assignment")
+        
+        # Validate matches
+        valid_matches = {}
+        for lineage_id, pdb_id in matches.items():
+            if lineage_id in lineage_variant_ids and pdb_id in pdb_variant_ids:
+                valid_matches[lineage_id] = pdb_id
+                log.info(f"Variant ID match: {lineage_id} -> {pdb_id}")
+            else:
+                log.warning(f"Invalid match ignored: {lineage_id} -> {pdb_id}")
+        
+        return valid_matches
+        
+    except Exception as e:
+        log.warning(f"Failed to match variant IDs with Gemini: {e}")
+        return {}
+
+
 def match_pdb_to_variants(
     pdb_sequences: Dict[str, str],
     variants: List[Variant],
@@ -3235,24 +3453,76 @@ Return ONLY the variant_id as a JSON string, e.g.: "ApePgb GLVRSQL"
         text = _extract_text(response).strip()
         
         # Parse JSON response (expecting a single string)
-        if text.startswith("```"):
+        # Look for JSON code blocks first
+        if "```json" in text:
+            # Extract content between ```json and ```
+            import re
+            json_match = re.search(r'```json\s*\n?(.*?)\n?```', text, re.DOTALL)
+            if json_match:
+                json_content = json_match.group(1).strip()
+                try:
+                    # Parse as JSON and extract the string value
+                    parsed = json.loads(json_content)
+                    matched_variant = str(parsed).strip('"\'')
+                except:
+                    # If JSON parsing fails, try to extract the quoted string
+                    quoted_match = re.search(r'"([^"]+)"', json_content)
+                    if quoted_match:
+                        matched_variant = quoted_match.group(1)
+                    else:
+                        matched_variant = json_content.strip('"\'')
+            else:
+                matched_variant = text.strip('"\'')
+        elif text.startswith("```"):
+            # Handle other code blocks
             text = text.split("```")[1].strip()
             if text.startswith("json"):
                 text = text[4:].strip()
+            matched_variant = text.strip('"\'')
+        else:
+            # Look for quoted strings in the response
+            import re
+            quoted_match = re.search(r'"([^"]+)"', text)
+            if quoted_match:
+                matched_variant = quoted_match.group(1)
+            else:
+                # Remove quotes if present
+                matched_variant = text.strip('"\'')
         
-        # Remove quotes if present
-        text = text.strip('"\'')
-        
-        matched_variant = text
+        log.info(f"Extracted variant name: '{matched_variant}' from response")
         log.info(f"PDB {pdb_id} matched to variant: {matched_variant}")
         
         # Return mapping with all chains pointing to the same variant
         mapping = {}
-        if matched_variant and any(v.variant_id == matched_variant for v in variants):
-            for chain_id in pdb_sequences:
-                mapping[matched_variant] = chain_id
-                break  # Only use the first chain
+        if matched_variant:
+            # Debug logging
+            variant_ids = [v.variant_id for v in variants]
+            log.info(f"Looking for variant '{matched_variant}' in lineage variants: {variant_ids}")
+            
+            # Check if the matched variant exists in the lineage
+            found_variant = any(v.variant_id == matched_variant for v in variants)
+            log.info(f"Variant '{matched_variant}' found in lineage: {found_variant}")
+            
+            if found_variant:
+                for chain_id in pdb_sequences:
+                    mapping[matched_variant] = chain_id
+                    log.info(f"Created mapping: {matched_variant} -> {chain_id}")
+                    break  # Only use the first chain
+            else:
+                log.warning(f"Variant '{matched_variant}' not found in lineage variants")
+                # Try fuzzy matching
+                for variant in variants:
+                    if variant.variant_id.strip() == matched_variant.strip():
+                        log.info(f"Found fuzzy match: '{variant.variant_id}' == '{matched_variant}'")
+                        for chain_id in pdb_sequences:
+                            mapping[variant.variant_id] = chain_id
+                            log.info(f"Created fuzzy mapping: {variant.variant_id} -> {chain_id}")
+                            break
+                        break
+        else:
+            log.warning("No matched variant extracted from response")
         
+        log.info(f"Final mapping result: {mapping}")
         return mapping
         
     except Exception as e:
@@ -3634,14 +3904,28 @@ def run_pipeline(
     caption_text = limited_caption_concat(*pdf_paths)
     full_text = limited_concat(*pdf_paths)
     
+    # Also load separate texts for manuscript and SI
+    manuscript_text = limited_concat(manuscript) if manuscript else None
+    si_text = limited_concat(si_path) if si_path else None
+    
     log.info("Loaded %d chars of captions for identification and %d chars of full text for extraction", 
              len(caption_text), len(full_text))
+    if manuscript_text:
+        log.info("Loaded %d chars from manuscript", len(manuscript_text))
+    if si_text:
+        log.info("Loaded %d chars from SI", len(si_text))
 
     # 2. Connect to Gemini -----------------------------------------------------
     model = get_model()
 
     # 3. Extract lineage (Section 6) ------------------------------------------
-    lineage, campaigns = get_lineage(caption_text, full_text, model, pdf_paths=pdf_paths, debug_dir=debug_dir)
+    lineage, campaigns = get_lineage(
+        caption_text, full_text, model, 
+        pdf_paths=pdf_paths, 
+        debug_dir=debug_dir,
+        manuscript_text=manuscript_text,
+        si_text=si_text
+    )
 
     if not lineage:
         raise RuntimeError("Pipeline aborted: failed to extract any lineage data")
@@ -3721,12 +4005,40 @@ def run_pipeline(
                         pdb_sequences, lineage, full_text, model, pdb_id
                     )
                     
+                    log.info(f"PDB matching result: {variant_to_chain}")
+                    log.info(f"Available PDB sequences: {list(pdb_sequences.keys())}")
+                    log.info(f"Lineage variants: {[v.variant_id for v in lineage]}")
+                    
                     # Convert to SequenceBlock objects
                     pdb_seq_blocks = []
-                    for variant in lineage:
-                        if variant.variant_id in variant_to_chain:
-                            chain_id = variant_to_chain[variant.variant_id]
-                            if chain_id in pdb_sequences:
+                    
+                    # Use Gemini-based matching for robust variant ID comparison
+                    if variant_to_chain and model:
+                        # Create a mapping using Gemini for robust string matching
+                        gemini_mapping = _match_variant_ids_with_gemini(
+                            lineage_variant_ids=[v.variant_id for v in lineage],
+                            pdb_variant_ids=list(variant_to_chain.keys()),
+                            model=model
+                        )
+                        
+                        for variant in lineage:
+                            log.info(f"Processing variant: {variant.variant_id}")
+                            
+                            # Try direct match first
+                            chain_id = variant_to_chain.get(variant.variant_id)
+                            log.info(f"Direct match for {variant.variant_id}: {chain_id}")
+                            
+                            # If no direct match, try Gemini-based matching
+                            if not chain_id:
+                                matched_pdb_variant = gemini_mapping.get(variant.variant_id)
+                                log.info(f"Gemini match for {variant.variant_id}: {matched_pdb_variant}")
+                                if matched_pdb_variant:
+                                    chain_id = variant_to_chain.get(matched_pdb_variant)
+                                    log.info(f"Chain ID from Gemini match: {chain_id}")
+                            
+                            if chain_id and chain_id in pdb_sequences:
+                                seq_length = len(pdb_sequences[chain_id])
+                                log.info(f"Creating sequence block for {variant.variant_id} with {seq_length} residues from chain {chain_id}")
                                 seq_block = SequenceBlock(
                                     variant_id=variant.variant_id,
                                     aa_seq=pdb_sequences[chain_id],
@@ -3737,6 +4049,26 @@ def run_pipeline(
                                 )
                                 pdb_seq_blocks.append(seq_block)
                                 log.info(f"Added PDB sequence for {variant.variant_id} from {pdb_id}:{chain_id}")
+                            else:
+                                log.warning(f"No chain_id found for variant {variant.variant_id} or chain not in sequences")
+                    else:
+                        # Fallback to direct matching if no model or no matches
+                        for variant in lineage:
+                            if variant.variant_id in variant_to_chain:
+                                chain_id = variant_to_chain[variant.variant_id]
+                                if chain_id in pdb_sequences:
+                                    seq_block = SequenceBlock(
+                                        variant_id=variant.variant_id,
+                                        aa_seq=pdb_sequences[chain_id],
+                                        dna_seq=None,
+                                        confidence=1.0,  # High confidence for PDB sequences
+                                        truncated=False,
+                                        metadata={"source": "PDB", "pdb_id": pdb_id, "chain": chain_id}
+                                    )
+                                    pdb_seq_blocks.append(seq_block)
+                                    log.info(f"Added PDB sequence for {variant.variant_id} from {pdb_id}:{chain_id}")
+                    
+                    log.info(f"PDB sequence blocks created: {len(pdb_seq_blocks)}")
                     
                     if pdb_seq_blocks:
                         # Update the dataframe with PDB sequences
@@ -3746,8 +4078,13 @@ def run_pipeline(
                                 df_merged.loc[mask, 'aa_seq'] = seq_block.aa_seq
                                 df_merged.loc[mask, 'seq_confidence'] = seq_block.confidence
                                 df_merged.loc[mask, 'seq_source'] = seq_block.metadata.get('source', 'PDB')
+                                log.info(f"Updated dataframe with sequence for {seq_block.variant_id}")
+                            else:
+                                log.warning(f"No matching row in dataframe for variant {seq_block.variant_id}")
                         log.info(f"Successfully extracted {len(pdb_seq_blocks)} sequences from PDB {pdb_id}")
                         break
+                    else:
+                        log.warning(f"No PDB sequence blocks were created for {pdb_id}")
                 else:
                     log.warning(f"No sequences found in PDB {pdb_id}")
         else: