Add context length error handling and buffer management for prompt construction in LLM inference process

app.py

@@ -45,7 +45,10 @@ st.set_page_config(
 
 # Constants
 MAX_CONTEXT_SIZE = 131072  # Default max context
-MAX_TOKENS_BUFFER = 0.9  # Use 90% of max tokens
+MAX_TOKENS_BUFFER = 0.9  # Initial buffer (existing)
+MIN_TOKENS_BUFFER = 0.5  # Minimum buffer before giving up
+BUFFER_REDUCTION_STEP = 0.05  # How much to reduce each retry (5%)
+MAX_BUFFER_RETRIES = 10  # Maximum retry attempts
 
 
 def _redact_secrets(text: str, secrets: list[str | None]) -> str:
@@ -68,6 +71,107 @@ def _safe_model_config_for_session(model_config: dict | None) -> dict | None:
     return safe
 
 
+def _is_context_length_error(error_msg: str) -> bool:
+    """
+    Check if an error message indicates a context length error.
+    
+    Args:
+        error_msg: The error message string
+        
+    Returns:
+        True if it's a context length error, False otherwise
+    """
+    error_lower = error_msg.lower()
+    return (
+        "maximum context length" in error_lower
+        or "requested about" in error_lower
+        or ("context length is" in error_lower and "you requested" in error_lower)
+    )
+
+
+def _build_prompt_with_buffer(
+    buffer_factor: float,
+    paper_text: str,
+    code_loader: CodeLoader | None,
+    code_text: str | None,
+    model_config: dict,
+    token_counter: TokenCounter,
+) -> tuple[str, str, int, int]:
+    """
+    Build prompt with a specific buffer factor.
+    
+    Args:
+        buffer_factor: Buffer factor to use (e.g., 0.9 for 90%)
+        paper_text: The paper text
+        code_loader: CodeLoader instance (if using GitHub repo)
+        code_text: Raw code text (if using uploaded file)
+        model_config: Model configuration dictionary
+        token_counter: TokenCounter instance
+        
+    Returns:
+        Tuple of (final_prompt, code_prompt, final_tokens, max_tokens_for_completion)
+    """
+    max_context = model_config["max_context"]
+    
+    # Calculate tokens for paper + prompt template
+    prompt_template = Prompt("discrepancy_generation")
+    intermediate_prompt = prompt_template(paper=paper_text, code="")
+    tokens_intermediate_prompt = token_counter(intermediate_prompt)
+    
+    # Calculate remaining tokens for code using the provided buffer factor
+    max_total_tokens = int(max_context * buffer_factor)
+    remaining_code_tokens = max_total_tokens - tokens_intermediate_prompt
+    
+    logger.info(f"Tokens in intermediate prompt: {tokens_intermediate_prompt}")
+    logger.info(f"Remaining tokens for code (buffer {buffer_factor:.1%}): {remaining_code_tokens}")
+    
+    # Get code prompt with token limit
+    if code_loader:
+        # Use CodeLoader for GitHub repos
+        code_prompt = code_loader.get_code_prompt(
+            token_counter=token_counter,
+            max_tokens=remaining_code_tokens,
+        )
+    else:
+        # Truncate code text to fit within token limit
+        code_prompt = ""
+        code_tokens = 0
+        code_lines = code_text.split('\n')
+        
+        for line in code_lines:
+            line_with_newline = line + '\n'
+            line_tokens = token_counter(line_with_newline)
+            if code_tokens + line_tokens > remaining_code_tokens:
+                logger.warning(f"Truncating code at {code_tokens} tokens (limit: {remaining_code_tokens})")
+                break
+            code_prompt += line_with_newline
+            code_tokens += line_tokens
+    
+    # Construct final prompt
+    final_prompt = prompt_template(paper=paper_text, code=code_prompt)
+    final_tokens = token_counter(final_prompt)
+    logger.info(f"Total tokens in final prompt: {final_tokens}")
+    
+    # Calculate max_tokens for completion (respecting model's context limit)
+    # Leave some buffer for safety (use 95% of remaining context)
+    remaining_for_completion = max_context - final_tokens
+    
+    if remaining_for_completion <= 0:
+        raise ValueError(
+            f"Prompt too long: {final_tokens} tokens exceeds model's context limit of {max_context} tokens"
+        )
+    
+    # Use 95% of remaining to be safe, but ensure at least some tokens
+    max_tokens_for_completion = max(1, int(remaining_for_completion * 0.95))
+    
+    logger.info(
+        f"Max context: {max_context}, Input tokens: {final_tokens}, "
+        f"Remaining: {remaining_for_completion}, Max completion tokens: {max_tokens_for_completion}"
+    )
+    
+    return final_prompt, code_prompt, final_tokens, max_tokens_for_completion
+
+
 def validate_urls(arxiv_url: str, github_url: str) -> tuple[bool, str]:
     """Validate input URLs."""
     if not arxiv_url:
@@ -210,76 +314,27 @@ def process_discrepancy_detection(
                         state="running",
                     )
 
-                # Step 5: Calculate tokens and prepare prompt
+                # Step 5: Calculate tokens and prepare prompt (initial build)
                 step_start = time.time()
                 status.update(label="📝 Preparing prompt...", state="running")
+                
+                # Create token counter (needed for both Step 5 and Step 6 retry loop)
+                tokenizer_name = model_config["tokenizer"]
+                token_counter = TokenCounter(model=tokenizer_name)
+                
                 try:
-                    # Use provided model config
-                    tokenizer_name = model_config["tokenizer"]
-                    max_context = model_config["max_context"]
-                    
-                    token_counter = TokenCounter(model=tokenizer_name)
-
-                    # Calculate tokens for paper + prompt template
-                    prompt_template = Prompt("discrepancy_generation")
-                    intermediate_prompt = prompt_template(paper=paper_text, code="")
-                    tokens_intermediate_prompt = token_counter(intermediate_prompt)
-
-                    # Calculate remaining tokens for code
-                    max_total_tokens = int(max_context * MAX_TOKENS_BUFFER)
-                    remaining_code_tokens = max_total_tokens - tokens_intermediate_prompt
-
-                    logger.info(f"Tokens in intermediate prompt: {tokens_intermediate_prompt}")
-                    logger.info(f"Remaining tokens for code: {remaining_code_tokens}")
-
-                    # Get code prompt with token limit
-                    if code_loader:
-                        # Use CodeLoader for GitHub repos
-                        code_prompt = code_loader.get_code_prompt(
-                            token_counter=token_counter,
-                            max_tokens=remaining_code_tokens,
-                        )
-                    else:
-                        # Truncate code text to fit within token limit
-                        # Simple approach: count tokens as we add content
-                        code_prompt = ""
-                        code_tokens = 0
-                        code_lines = code_text.split('\n')
-                        
-                        for line in code_lines:
-                            line_with_newline = line + '\n'
-                            line_tokens = token_counter(line_with_newline)
-                            if code_tokens + line_tokens > remaining_code_tokens:
-                                logger.warning(f"Truncating code at {code_tokens} tokens (limit: {remaining_code_tokens})")
-                                break
-                            code_prompt += line_with_newline
-                            code_tokens += line_tokens
+                    # Build initial prompt with default buffer
+                    final_prompt, code_prompt, final_tokens, max_tokens_for_completion = _build_prompt_with_buffer(
+                        buffer_factor=MAX_TOKENS_BUFFER,
+                        paper_text=paper_text,
+                        code_loader=code_loader,
+                        code_text=code_text,
+                        model_config=model_config,
+                        token_counter=token_counter,
+                    )
                     
                     results["code_prompt"] = code_prompt
-
-                    # Construct final prompt
-                    final_prompt = prompt_template(paper=paper_text, code=code_prompt)
                     results["prompt"] = final_prompt
-
-                    final_tokens = token_counter(final_prompt)
-                    logger.info(f"Total tokens in final prompt: {final_tokens}")
-                    
-                    # Calculate max_tokens for completion (respecting model's context limit)
-                    # Leave some buffer for safety (use 95% of remaining context)
-                    max_context = model_config["max_context"]
-                    remaining_for_completion = max_context - final_tokens
-                    
-                    if remaining_for_completion <= 0:
-                        error_msg = f"Prompt too long: {final_tokens} tokens exceeds model's context limit of {max_context} tokens"
-                        logger.error(error_msg)
-                        results["error"] = error_msg
-                        status.update(label="❌ Prompt too long", state="error")
-                        return results
-                    
-                    # Use 95% of remaining to be safe, but ensure at least some tokens
-                    max_tokens_for_completion = max(1, int(remaining_for_completion * 0.95))
-                    
-                    logger.info(f"Max context: {max_context}, Input tokens: {final_tokens}, Remaining: {remaining_for_completion}, Max completion tokens: {max_tokens_for_completion}")
                     
                     step_time = time.time() - step_start
                     step_timings["Prompt Preparation"] = step_time
@@ -295,75 +350,135 @@ def process_discrepancy_detection(
                     status.update(label="❌ Error preparing prompt", state="error")
                     return results
 
-                # Step 6: Detect discrepancies with LLM
+                # Step 6: Detect discrepancies with LLM (with retry on context length errors)
                 step_start = time.time()
                 status.update(label="🤖\uFE0F Detecting discrepancies (this may take a while)...", state="running")
-                try:
-                    # Extract model configuration
-                    model = model_config["model"]
-                    api_key = model_config.get("api_key")
-                    api_base = model_config.get("api_base")
-                    max_context = model_config.get("max_context")
-                    
-                    llm = LLM(
-                        model=model,
-                        api_key=api_key,
-                        api_base=api_base,
-                        temperature=1.0,
-                        top_p=1.0,
-                        reasoning_effort="high",
-                        max_context=max_context,
-                        max_tokens=max_tokens_for_completion,  # Respect model's context limit
-                    )
-
-                    response = llm(final_prompt)
-                    results["llm_response"] = response
-
-                    # Extract content from response
-                    choices = response.get("choices", [])
-                    if not choices:
-                        raise ValueError("No choices in LLM response")
-                    
-                    content = (
-                        choices[0]
-                        .get("message", {})
-                        .get("content", "")
-                    )
-                    
-                    if not content:
-                        raise ValueError("Empty content in LLM response")
-
-                    # Parse discrepancies
-                    discrepancies = parse_discrepancies(content)
-                    results["discrepancies"] = discrepancies
-
-                    step_time = time.time() - step_start
-                    step_timings["LLM Inference"] = step_time
-                    total_time = sum(step_timings.values())
-                    
-                    st.write(f"✅ LLM inference: {step_time:.1f}s")
-                    st.write("---")
-                    st.write(f"**Total time: {total_time:.1f}s**")
-                    
-                    if discrepancies:
-                        count = len(discrepancies)
-                        discrepancy_text = "discrepancy" if count == 1 else "discrepancies"
-                        status.update(
-                            label=f"✅ Complete! Found {count} {discrepancy_text} ({total_time:.1f}s total)",
-                            state="complete",
+                
+                # Retry configuration
+                initial_buffer = MAX_TOKENS_BUFFER  # 0.9
+                min_buffer = MIN_TOKENS_BUFFER  # 0.5
+                buffer_reduction_step = BUFFER_REDUCTION_STEP  # 0.05
+                max_retries = MAX_BUFFER_RETRIES  # 5
+                
+                current_buffer = initial_buffer
+                retry_count = 0
+                success = False
+                current_final_prompt = final_prompt
+                current_max_tokens_for_completion = max_tokens_for_completion
+                
+                while not success and current_buffer >= min_buffer and retry_count < max_retries:
+                    try:
+                        # Rebuild prompt with current buffer (if retry, otherwise use existing)
+                        if retry_count > 0:
+                            status.update(
+                                label=f"🔄 Retrying with reduced buffer ({current_buffer:.1%})...",
+                                state="running"
+                            )
+                            st.write(f"🔄 Retrying with reduced buffer ({current_buffer:.1%})...")
+                            
+                            # Rebuild prompt with reduced buffer
+                            current_final_prompt, code_prompt, final_tokens, current_max_tokens_for_completion = _build_prompt_with_buffer(
+                                buffer_factor=current_buffer,
+                                paper_text=paper_text,
+                                code_loader=code_loader,
+                                code_text=code_text,
+                                model_config=model_config,
+                                token_counter=token_counter,
+                            )
+                            results["code_prompt"] = code_prompt
+                            results["prompt"] = current_final_prompt
+                        
+                        # Extract model configuration
+                        model = model_config["model"]
+                        api_key = model_config.get("api_key")
+                        api_base = model_config.get("api_base")
+                        max_context = model_config.get("max_context")
+                        
+                        llm = LLM(
+                            model=model,
+                            api_key=api_key,
+                            api_base=api_base,
+                            temperature=1.0,
+                            top_p=1.0,
+                            reasoning_effort="high",
+                            max_context=max_context,
+                            max_tokens=current_max_tokens_for_completion,  # Respect model's context limit
                         )
-                    else:
-                        status.update(
-                            label=f"✅ Complete! No discrepancies found ({total_time:.1f}s total)",
-                            state="complete",
+                        
+                        response = llm(current_final_prompt)
+                        results["llm_response"] = response
+                        
+                        # Extract content from response
+                        choices = response.get("choices", [])
+                        if not choices:
+                            raise ValueError("No choices in LLM response")
+                        
+                        content = (
+                            choices[0]
+                            .get("message", {})
+                            .get("content", "")
                         )
-
-                except Exception as e:
-                    api_key = model_config.get("api_key") if isinstance(model_config, dict) else None
-                    error_msg = f"Error during LLM inference: {_redact_secrets(str(e), [api_key])}"
+                        
+                        if not content:
+                            raise ValueError("Empty content in LLM response")
+                        
+                        # Parse discrepancies
+                        discrepancies = parse_discrepancies(content)
+                        results["discrepancies"] = discrepancies
+                        
+                        step_time = time.time() - step_start
+                        step_timings["LLM Inference"] = step_time
+                        total_time = sum(step_timings.values())
+                        
+                        st.write(f"✅ LLM inference: {step_time:.1f}s")
+                        st.write("---")
+                        st.write(f"**Total time: {total_time:.1f}s**")
+                        
+                        if discrepancies:
+                            count = len(discrepancies)
+                            discrepancy_text = "discrepancy" if count == 1 else "discrepancies"
+                            status.update(
+                                label=f"✅ Complete! Found {count} {discrepancy_text} ({total_time:.1f}s total)",
+                                state="complete",
+                            )
+                        else:
+                            status.update(
+                                label=f"✅ Complete! No discrepancies found ({total_time:.1f}s total)",
+                                state="complete",
+                            )
+                        
+                        success = True
+                        
+                    except Exception as e:
+                        error_msg = str(e)
+                        api_key = model_config.get("api_key") if isinstance(model_config, dict) else None
+                        redacted_error = _redact_secrets(error_msg, [api_key])
+                        
+                        # Check if it's a context length error
+                        if _is_context_length_error(error_msg) and current_buffer > min_buffer:
+                            retry_count += 1
+                            current_buffer -= buffer_reduction_step
+                            logger.warning(
+                                f"Context length error detected. Retrying with buffer {current_buffer:.1%} "
+                                f"(attempt {retry_count}/{max_retries})"
+                            )
+                            continue  # Retry with smaller buffer
+                        else:
+                            # Not a context length error, or we've exhausted retries
+                            logger.error(f"Error during LLM inference: {redacted_error}")
+                            results["error"] = f"Error during LLM inference: {redacted_error}"
+                            status.update(label="❌ Error during inference", state="error")
+                            return results
+                
+                # If we exhausted retries or hit minimum buffer
+                if not success:
+                    error_msg = (
+                        f"Could not fit prompt within context limits after {retry_count} retries. "
+                        f"Minimum buffer ({min_buffer:.1%}) reached."
+                    )
                     logger.error(error_msg)
                     results["error"] = error_msg
-                    status.update(label="❌ Error during inference", state="error")
+                    status.update(label="❌ Prompt too large for model", state="error")
                     return results
 
             except Exception as e: