llm_fill.py

#!/usr/bin/env python3

import os
import re
import json
import argparse
import requests
import time
import logging
import concurrent.futures
from typing import List, Tuple, Dict, Any
import sys
from dotenv import load_dotenv

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

DEFAULT_MODEL = "claude-3-7-sonnet-20250219"
CLAUDE_API_URL = "https://api.anthropic.com/v1/messages"

# For rate limiting
MIN_REQUESTS_REMAINING = 5
RATE_LIMIT_PAUSE = 1.0  # seconds to pause when rate limit is approaching

def validate_api_access(api_key: str, model: str) -> bool:
    """Validate API access by making a simple test request.
    Returns True if successful, False otherwise."""
    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json"
    }

    data = {
        "model": model,
        "max_tokens": 1,
        "messages": [{
            "role": "user",
            "content": "test"
        }]
    }

    try:
        response = requests.post(CLAUDE_API_URL, headers=headers, json=data)
        if response.status_code == 200:
            logger.info("API validation successful")
            return True
        else:
            logger.error(f"API validation failed with status code: {response.status_code}")
            logger.error(f"Response: {response.text}")
            return False
    except requests.exceptions.RequestException as e:
        logger.error(f"API validation failed with error: {str(e)}")
        return False

def get_llm_response_with_rate_limit(prompt: str, model: str, request_id: int) -> Tuple[int, str, int]:
    """Get a single response from Claude API with rate limit handling.
    Returns a tuple of (request_id, response_text, requests_remaining)
    """
    api_key = os.getenv("LLM_API_KEY")
    if not api_key:
        raise ValueError("LLM_API_KEY environment variable is not set")

    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json"
    }

    data = {
        "model": model,
        "max_tokens": 1024,
        "messages": [{
            "role": "user",
            "content": prompt
        }]
    }

    logger.debug(f"Sending request {request_id} to Claude API")
    response = requests.post(CLAUDE_API_URL, headers=headers, json=data)
    response.raise_for_status()
    
    # Extract rate limit information from headers
    requests_remaining = int(response.headers.get('anthropic-ratelimit-requests-remaining', '999'))
    logger.debug(f"Request {request_id} completed. Requests remaining: {requests_remaining}")
    
    result = response.json()
    return request_id, result["content"][0]["text"], requests_remaining

def get_llm_responses_concurrent(llm_requests: List[Dict[str, Any]], model: str) -> Dict[int, str]:
    """Get responses from Claude API using concurrent execution with rate limiting."""
    results = {}
    requests_remaining = 999  # Initial high value
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        # Submit initial batch of tasks
        future_to_request = {}
        
        for i, request in enumerate(llm_requests):
            # If we're getting close to the rate limit, pause before continuing
            if requests_remaining < MIN_REQUESTS_REMAINING:
                logger.info(f"Approaching rate limit ({requests_remaining} remaining). Pausing for {RATE_LIMIT_PAUSE} seconds...")
                time.sleep(RATE_LIMIT_PAUSE)
            
            future = executor.submit(get_llm_response_with_rate_limit, request["prompt"], model, i)
            future_to_request[future] = i
        
        # Process results as they complete
        for future in concurrent.futures.as_completed(future_to_request):
            request_idx = future_to_request[future]
            try:
                request_id, response_text, new_requests_remaining = future.result()
                results[request_id] = response_text
                # Update the requests remaining counter
                requests_remaining = new_requests_remaining
                logger.debug(f"Successfully processed request {request_id}, {requests_remaining} requests remaining")
            except Exception as e:
                logger.warning(f"Request {request_idx} failed with error: {str(e)}")
    
    logger.info(f"Received {len(results)} successful responses")
    return results

def find_llm_fill_comments(content: str) -> List[Tuple[int, int, str, str]]:
    """Find all LLM_PLEASE_FILL comments in the content.
    Returns list of tuples (start_index, end_index, prompt, indentation)
    """
    pattern = r'^(\s*)<!--\s*LLM_PLEASE_FILL:\s*(.*?)\s*-->'
    matches = []
    for match in re.finditer(pattern, content, re.MULTILINE):
        start = match.start()
        end = match.end()
        indentation = match.group(1)
        prompt = match.group(2).strip()
        matches.append((start, end, prompt, indentation))
    logger.debug(f"Found {len(matches)} LLM_PLEASE_FILL comments")
    return matches

def find_llm_replace_comments(content: str) -> List[Tuple[int, int, str, str, str]]:
    """Find all LLM_PLEASE_REPLACE_BELOW/ABOVE comment pairs in the content.
    Returns list of tuples (start_index, end_index, prompt, existing_content, indentation)
    """
    pattern = r'^(\s*)<!--\s*LLM_PLEASE_REPLACE_BELOW:\s*(.*?)\s*-->(.*?)<!--\s*LLM_PLEASE_REPLACE_ABOVE\s*-->'
    matches = []
    for match in re.finditer(pattern, content, re.MULTILINE | re.DOTALL):
        start = match.start()
        end = match.end()
        indentation = match.group(1)
        prompt = match.group(2).strip()
        existing_content = match.group(3).strip()
        matches.append((start, end, prompt, existing_content, indentation))
    logger.debug(f"Found {len(matches)} LLM_PLEASE_REPLACE comment pairs")
    return matches

def process_file(input_file: str, output_file: str) -> None:
    """Process the input file and generate the output file with LLM-filled content."""
    logger.info(f"Processing file: {input_file}")
    
    # Check environment variables
    api_key = os.getenv("LLM_API_KEY")
    if not api_key:
        raise ValueError("LLM_API_KEY environment variable is not set")
    
    model = os.getenv("LLM_MODEL", DEFAULT_MODEL)
    logger.info(f"Using model: {model}")
    
    # Validate API access before proceeding
    if not validate_api_access(api_key, model):
        raise ValueError("Failed to validate API access. Please check your API key and model configuration.")
    
    with open(input_file, 'r') as f:
        content = f.read()

    fill_matches = find_llm_fill_comments(content)
    replace_matches = find_llm_replace_comments(content)
    # strip all the <!-- --> comments but not the content inside them
    content_without_comments = re.sub(r'<!--[\s\S]*?-->', '', content)
    if not fill_matches and not replace_matches:
        logger.info("No LLM comments found in the file.")
        return

    # Collect all requests for concurrent processing
    llm_requests = []
    replacements = []
    
    # Process fill comments based on the full report
    for start, end, prompt, indentation in fill_matches:
        # Create a prompt with explicit instructions
        full_prompt = f"""You are a helpful AI assistant. Please try to understand
the following request and fill in the content based on the request.

Request: {prompt}

Content to use for filling:

-- Begin of content --
{content_without_comments}
-- End of content --

Please provide only the content that should be filled in, without any 
additional explanation or formatting. The content should be indented
with the same level as the original content. When asked to summarize,
please make sure to use all and entire content to make the summary.
You also need to make judgement call on what is the most important
part of the content and what is not.
"""
        llm_requests.append({
            "type": "fill",
            "start": start,
            "end": end,
            "indentation": indentation,
            "prompt": full_prompt
        })
    
    content_without_llm_comments = re.sub(r'(\s*)<!--\s*LLM_PLEASE_FILL:\s*(.*?)\s*-->', '', content, flags=re.MULTILINE)
    content_without_llm_comments = re.sub(
        r'(\s*)<!--\s*LLM_PLEASE_REPLACE_BELOW:\s*(.*?)\s*-->[\s\S]*?<!--\s*LLM_PLEASE_REPLACE_ABOVE\s*-->', 
        '', 
        content_without_llm_comments, 
        flags=re.MULTILINE
    )
    
    # Process replace comments
    for start, end, prompt, existing_content, indentation in replace_matches:
        # Create a prompt with explicit instructions
        full_prompt = f"""You are a helpful AI assistant. Please replace the following content based on the prompt provided.

Existing Content to Replace:
{existing_content}

Request:
{prompt}

Please provide only the new content that should replace the existing content, without any additional explanation or formatting. The content should be indented with the same level as the original content."""
        
        llm_requests.append({
            "type": "replace",
            "start": start,
            "end": end,
            "indentation": indentation,
            "prompt": full_prompt,
            "existing_content": existing_content,  # Store the existing content for logging
            "original_prompt": prompt  # Store the original prompt for debugging
        })

    logger.info(f"Found {len(llm_requests)} LLM comments to process")

    # Get responses using concurrent execution
    responses = get_llm_responses_concurrent(llm_requests, model)
    
    # Process responses and create replacements
    for i, request in enumerate(llm_requests):
        if i in responses:
            response = responses[i]
            # Apply indentation to each line of the response
            indented_response = '\n'.join(request["indentation"] + line for line in response.split('\n'))
            replacements.append((request["start"], request["end"], indented_response))
            
            # Log original and replacement content for LLM_PLEASE_REPLACE operations
            if request["type"] == "replace":
                logger.info(f"\nReplacing LLM_PLEASE_REPLACE content:")
                logger.info("Original: ")
                logger.info(f"<content>\n{request['existing_content']}\n</content>\n")
                logger.info("LLM replacement:")
                logger.info(f"<replaced content>\n{response}\n</replaced content>")
                
                # Debug: Show the prompt that generated this replacement
                logger.info(f"\nOriginal prompt from LLM_PLEASE_REPLACE_BELOW comment:")
                logger.info(f"<prompt>\n{request['original_prompt']}\n</prompt>")
                
                # Check if the response matches what's expected based on the content
                if len(set(response.lower().split()).intersection(set(request['existing_content'].lower().split()))) < 3:
                    logger.warning("WARNING: Replacement content appears to be unrelated to original content!")
                    logger.warning("This might indicate the prompt is instructing the LLM to generate new content rather than modifying existing content.")
            
            logger.debug(f"Processed response for request {i}")
        else:
            logger.warning(f"No response received for request {i}")

    # Apply replacements in reverse order to maintain correct indices
    for start, end, replacement in reversed(replacements):
        content = content[:start] + replacement + content[end:]

    # Write the processed content to the output file
    with open(output_file, 'w') as f:
        f.write(content)
    logger.info(f"Successfully wrote output to {output_file}")

def main():
    parser = argparse.ArgumentParser(description='Process files and fill LLM_PLEASE_FILL comments using Claude API')
    parser.add_argument('--input', required=True, help='Input file path')
    parser.add_argument('--output', required=True, help='Output file path')
    parser.add_argument('--debug', action='store_true', help='Enable debug logging')
    
    args = parser.parse_args()
    
    if args.debug:
        logger.setLevel(logging.DEBUG)
        logger.debug("Debug logging enabled")
    
    try:
        load_dotenv()
        process_file(args.input, args.output)
        logger.info(f"Successfully processed file. Output written to {args.output}")
    except ValueError as e:
        logger.error(f"Configuration error: {str(e)}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Error processing file: {str(e)}")
        raise e

if __name__ == "__main__":
    load_dotenv()
    main()