Refactor prompt construction and token management in LLM inference process, enhancing context handling and logging for code prompts

app.py

@@ -34,6 +34,10 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Constants
+CONTEXT_BUFFER_FACTOR = 0.9
+MAX_CONTEXT_SIZE = 131072  # Default max context
+
 # Page configuration
 st.set_page_config(
     page_title="SciCoQA Paper- Code Discrepancy Detection",
@@ -43,14 +47,6 @@ st.set_page_config(
 )
 
 
-# Constants
-MAX_CONTEXT_SIZE = 131072  # Default max context
-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:
     """Best-effort redaction for secrets that may appear in exception strings/logs."""
     redacted = text
@@ -71,37 +67,17 @@ 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,
+def _build_prompt(
     paper_text: str,
     code_loader: CodeLoader | None,
     code_text: str | None,
     model_config: dict,
     token_counter: TokenCounter,
-) -> tuple[str, str, int, int]:
+) -> tuple[str, str, int]:
     """
-    Build prompt with a specific buffer factor.
+    Build prompt by counting tokens and truncating code until prompt + paper + code < CONTEXT_BUFFER_FACTOR * model context length.
     
     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)
@@ -109,21 +85,31 @@ def _build_prompt_with_buffer(
         token_counter: TokenCounter instance
         
     Returns:
-        Tuple of (final_prompt, code_prompt, final_tokens, max_tokens_for_completion)
+        Tuple of (final_prompt, code_prompt, final_tokens)
     """
     max_context = model_config["max_context"]
+    max_total_tokens = int(max_context * CONTEXT_BUFFER_FACTOR)
     
-    # Calculate tokens for paper + prompt template
+    # Build 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
+    # Calculate tokens for template + paper
+    template_with_paper = prompt_template(paper=paper_text, code="")
+    tokens_template_and_paper = token_counter(template_with_paper)
+    
+    # Calculate remaining tokens for code
+    remaining_code_tokens = max_total_tokens - tokens_template_and_paper
     
-    logger.info(f"Tokens in intermediate prompt: {tokens_intermediate_prompt}")
-    logger.info(f"Remaining tokens for code (buffer {buffer_factor:.1%}): {remaining_code_tokens}")
+    if remaining_code_tokens <= 0:
+        raise ValueError(
+            f"Paper text too long: {tokens_template_and_paper} tokens exceeds "
+            f"90% of context limit ({max_total_tokens} tokens)"
+        )
+    
+    logger.info(
+        f"Template + paper tokens: {tokens_template_and_paper}, "
+        f"Remaining for code: {remaining_code_tokens}"
+    )
     
     # Get code prompt with token limit
     if code_loader:
@@ -136,40 +122,31 @@ def _build_prompt_with_buffer(
         # 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
+        if code_text and remaining_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
+    # Construct final prompt and verify it's within limit
     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:
+    if final_tokens > max_total_tokens:
         raise ValueError(
-            f"Prompt too long: {final_tokens} tokens exceeds model's context limit of {max_context} tokens"
+            f"Final prompt too long: {final_tokens} tokens exceeds "
+            f"90% of context limit ({max_total_tokens} 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"Final prompt tokens: {final_tokens} (limit: {max_total_tokens})")
     
-    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
+    return final_prompt, code_prompt, final_tokens
 
 
 def validate_urls(arxiv_url: str, github_url: str) -> tuple[bool, str]:
@@ -314,18 +291,17 @@ def process_discrepancy_detection(
                         state="running",
                     )
 
-                # Step 5: Calculate tokens and prepare prompt (initial build)
+                # Step 5: Calculate tokens and prepare prompt
                 step_start = time.time()
                 status.update(label="📝 Preparing prompt...", state="running")
                 
-                # Create token counter (needed for both Step 5 and Step 6 retry loop)
+                # Create token counter
                 tokenizer_name = model_config["tokenizer"]
                 token_counter = TokenCounter(model=tokenizer_name)
                 
                 try:
-                    # 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,
+                    # Build prompt with simple token counting
+                    final_prompt, code_prompt, final_tokens = _build_prompt(
                         paper_text=paper_text,
                         code_loader=code_loader,
                         code_text=code_text,
@@ -338,7 +314,7 @@ def process_discrepancy_detection(
                     
                     step_time = time.time() - step_start
                     step_timings["Prompt Preparation"] = step_time
-                    st.write(f"✅ Prompt prepared: {step_time:.1f}s ({final_tokens:,} tokens, max output: {max_tokens_for_completion:,} tokens)")
+                    st.write(f"✅ Prompt prepared: {step_time:.1f}s ({final_tokens:,} tokens)")
                     status.update(
                         label=f"✅ Prompt prepared ({step_time:.1f}s, {final_tokens:,} tokens)",
                         state="running",
@@ -350,135 +326,76 @@ def process_discrepancy_detection(
                     status.update(label="❌ Error preparing prompt", state="error")
                     return results
 
-                # Step 6: Detect discrepancies with LLM (with retry on context length errors)
+                # Step 6: Detect discrepancies with LLM
                 step_start = time.time()
                 status.update(label="🤖\uFE0F Detecting discrepancies (this may take a while)...", state="running")
                 
-                # 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
+                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,
+                    )
+                    
+                    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",
                         )
-                        
-                        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", "")
+                    else:
+                        status.update(
+                            label=f"✅ Complete! No discrepancies found ({total_time:.1f}s total)",
+                            state="complete",
                         )
                         
-                        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="❌ Prompt too large for model", state="error")
+                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])
+                    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
 
             except Exception as e:

core/code_loader_demo.py

@@ -255,6 +255,12 @@ class CodeLoader:
         """Generate code prompt with repo tree and file contents."""
         code_prompt = "Repo tree:\n" + self.get_repo_tree() + "\n\n"
         tokens = token_counter(code_prompt) if token_counter is not None else 0
+        
+        if token_counter is not None and max_tokens is not None:
+            logger.info(
+                f"Building code prompt: repo tree tokens={tokens}, max_tokens={max_tokens}, "
+                f"remaining for files={max_tokens - tokens}"
+            )
 
         files_to_replace = {}
         if code_changes:
@@ -275,18 +281,29 @@ class CodeLoader:
             if token_counter is not None:
                 logger.debug(f"Adding file: {file_path}")
                 num_tokens = token_counter(code_file)
+                # Check if adding this file would exceed the limit BEFORE adding it
+                if max_tokens and (tokens + num_tokens) > max_tokens:
+                    logger.warning(
+                        f"Truncating. Max tokens reached for {self.github_url}. "
+                        f"Current tokens: {tokens}, File tokens: {num_tokens}, "
+                        f"Max tokens for code is {max_tokens}"
+                    )
+                    break
                 tokens += num_tokens
                 logger.debug(
                     f"Number of tokens in file: {num_tokens}. "
                     f"Total number of tokens in code prompt: {tokens}"
                 )
-            if max_tokens and tokens > max_tokens:
-                logger.warning(
-                    f"Truncating. Max tokens reached for {self.github_url}. "
-                    f"Max tokens for code is {max_tokens}"
-                )
-                break
             code_prompt += code_file
+        
+        # Log final code prompt size
+        if token_counter is not None:
+            final_code_tokens = token_counter(code_prompt)
+            logger.info(
+                f"Code prompt built: {final_code_tokens} tokens "
+                f"(max was {max_tokens if max_tokens else 'unlimited'})"
+            )
+        
         return code_prompt
 
 

core/llm_demo.py

@@ -72,8 +72,11 @@ class LLM:
             kwargs = {
                 "model": self.model,
                 "messages": [{"role": "user", "content": prompt}],
-                "max_tokens": self.max_tokens,
             }
+            
+            # Only set max_tokens if explicitly provided (let API use default if None)
+            if self.max_tokens is not None:
+                kwargs["max_tokens"] = self.max_tokens
 
             # Let LiteLLM drop unsupported params per-provider/model (e.g., GPT-5 rejecting top_p)
             if self.drop_params: