uxoxo/eb2ab

1#!/usr/bin/env python32"""3LLM Processor - Process text chunks with LLM providers4 5Supports:6- OpenRouter (unified API for multiple models)7- OpenAI (direct)8- Anthropic (direct)9- Ollama (local)10 11Author: Claude Code12Date: 2025-10-1713"""14 15import os16import time17import traceback18from dataclasses import dataclass19from typing import List, Dict, Any, Callable, Optional20 21try:22    import openai23    HAS_OPENAI = True24except ImportError:25    HAS_OPENAI = False26 27try:28    import anthropic29    HAS_ANTHROPIC = True30except ImportError:31    HAS_ANTHROPIC = False32 33 34@dataclass35class TextChunk:36    """Text chunk for processing"""37    content: str38    index: int39    tokens: int40    metadata: Dict[str, Any]41 42 43@dataclass44class ProcessingProgress:45    """Progress update during processing"""46    current_chunk: int47    total_chunks: int48    progress_percent: float49    elapsed_time: float50    estimated_remaining: float51    current_cost: float52 53 54@dataclass55class ProcessingResult:56    """Result from LLM processing"""57    success: bool58    processed_chunks: List[str]59    total_cost: float60    total_time_seconds: float61    total_input_tokens: int62    total_output_tokens: int63    error: Optional[str] = None64 65 66# OpenRouter model pricing (per 1M tokens)67# Updated to match OpenRouter's current model IDs (without version dates)68OPENROUTER_PRICING = {69    "anthropic/claude-3.5-sonnet": {"input": 3.00, "output": 15.00},70    "anthropic/claude-3-haiku": {"input": 0.25, "output": 1.25},71    "openai/gpt-4o": {"input": 2.50, "output": 10.00},72    "openai/gpt-4o-mini": {"input": 0.15, "output": 0.60},73    "openai/gpt-3.5-turbo": {"input": 0.50, "output": 1.50},74    "meta-llama/llama-3.1-8b-instruct": {"input": 0.05, "output": 0.05},75    "meta-llama/llama-3.1-70b-instruct": {"input": 0.35, "output": 0.40},76}77 78 79def process_with_llm(80    chunks: List[TextChunk],81    instruction: str,82    provider: str,83    progress_callback: Optional[Callable] = None,84    max_retries: int = 385) -> ProcessingResult:86    """87    Process text chunks with LLM88 89    Args:90        chunks: List of text chunks to process91        instruction: Processing instruction (e.g., "Convert to modern English")92        provider: Provider name (e.g., "Claude 3 Haiku (Fastest, Cheapest)")93        progress_callback: Optional callback for progress updates94        max_retries: Max retry attempts per chunk95 96    Returns:97        ProcessingResult with success status and processed text98    """99 100    # Map UI provider names to actual implementations101    provider_map = {102        "Claude 3 Haiku (Fastest, Cheapest)": "openrouter/claude-haiku",103        "Claude 3.5 Sonnet (Balanced)": "openrouter/claude-sonnet",104        "GPT-4o (Best Quality)": "openrouter/gpt-4o",105        "GPT-3.5 Turbo (Fast & Cheap)": "openrouter/gpt-3.5-turbo",106        "Ollama (Free, Local - Slower)": "ollama/local",107    }108 109    provider_key = provider_map.get(provider, "openrouter/claude-haiku")110 111    # Check for API keys112    openrouter_key = os.environ.get("OPENROUTER_API_KEY")113 114    if not openrouter_key and provider_key.startswith("openrouter"):115        return ProcessingResult(116            success=False,117            processed_chunks=[],118            total_cost=0,119            total_time_seconds=0,120            total_input_tokens=0,121            total_output_tokens=0,122            error="OPENROUTER_API_KEY not found in environment variables"123        )124 125    # Process chunks126    processed_chunks = []127    total_input_tokens = 0128    total_output_tokens = 0129    total_cost = 0.0130    start_time = time.time()131 132    for i, chunk in enumerate(chunks):133        chunk_start = time.time()134 135        # Progress update136        if progress_callback:137            elapsed = time.time() - start_time138            avg_time_per_chunk = elapsed / max(i, 1)139            remaining = avg_time_per_chunk * (len(chunks) - i)140 141            progress = ProcessingProgress(142                current_chunk=i + 1,143                total_chunks=len(chunks),144                progress_percent=(i / len(chunks)) * 100,145                elapsed_time=elapsed,146                estimated_remaining=remaining,147                current_cost=total_cost148            )149            progress_callback(progress)150 151        # Process chunk with retries152        for attempt in range(max_retries):153            try:154                if provider_key.startswith("openrouter"):155                    result = _process_with_openrouter(156                        chunk.content,157                        instruction,158                        provider_key,159                        openrouter_key160                    )161                elif provider_key.startswith("ollama"):162                    result = _process_with_ollama(chunk.content, instruction)163                else:164                    raise ValueError(f"Unsupported provider: {provider_key}")165 166                processed_chunks.append(result["processed_text"])167                total_input_tokens += result.get("input_tokens", chunk.tokens)168                total_output_tokens += result.get("output_tokens", chunk.tokens)169                total_cost += result.get("cost", 0)170                break171 172            except Exception as e:173                if attempt == max_retries - 1:174                    # Final attempt failed175                    return ProcessingResult(176                        success=False,177                        processed_chunks=processed_chunks,178                        total_cost=total_cost,179                        total_time_seconds=time.time() - start_time,180                        total_input_tokens=total_input_tokens,181                        total_output_tokens=total_output_tokens,182                        error=f"Failed on chunk {i+1}/{len(chunks)}: {str(e)}"183                    )184 185                # Exponential backoff186                time.sleep(2 ** attempt)187 188    # Final progress update189    if progress_callback:190        progress = ProcessingProgress(191            current_chunk=len(chunks),192            total_chunks=len(chunks),193            progress_percent=100,194            elapsed_time=time.time() - start_time,195            estimated_remaining=0,196            current_cost=total_cost197        )198        progress_callback(progress)199 200    return ProcessingResult(201        success=True,202        processed_chunks=processed_chunks,203        total_cost=total_cost,204        total_time_seconds=time.time() - start_time,205        total_input_tokens=total_input_tokens,206        total_output_tokens=total_output_tokens207    )208 209 210def _process_with_openrouter(211    text: str,212    instruction: str,213    provider_key: str,214    api_key: str215) -> Dict[str, Any]:216    """217    Process text with OpenRouter API218 219    Args:220        text: Text to process221        instruction: Processing instruction222        provider_key: Provider key (e.g., "openrouter/claude-haiku")223        api_key: OpenRouter API key224 225    Returns:226        Dict with processed_text, input_tokens, output_tokens, cost227    """228 229    # Map provider_key to OpenRouter model name230    # OpenRouter uses format: provider/model (without version dates)231    model_map = {232        "openrouter/claude-haiku": "anthropic/claude-3-haiku",233        "openrouter/claude-sonnet": "anthropic/claude-3.5-sonnet",234        "openrouter/gpt-4o": "openai/gpt-4o",235        "openrouter/gpt-4o-mini": "openai/gpt-4o-mini",236        "openrouter/gpt-3.5-turbo": "openai/gpt-3.5-turbo",237    }238 239    model = model_map.get(provider_key, "anthropic/claude-3-haiku")240 241    # Use OpenAI-compatible client for OpenRouter242    if not HAS_OPENAI:243        raise ImportError("openai package required for OpenRouter. Install with: pip install openai")244 245    client = openai.OpenAI(246        base_url="https://openrouter.ai/api/v1",247        api_key=api_key,248    )249 250    # Create messages251    messages = [252        {253            "role": "user",254            "content": f"{instruction}\n\n{text}"255        }256    ]257 258    # Call API259    response = client.chat.completions.create(260        model=model,261        messages=messages,262        max_tokens=4000,263        extra_headers={264            "HTTP-Referer": "https://github.com/your-repo",  # Required by OpenRouter265            "X-Title": "LLM Ebook Processor"266        }267    )268 269    # Extract results270    processed_text = response.choices[0].message.content271    input_tokens = response.usage.prompt_tokens272    output_tokens = response.usage.completion_tokens273 274    # Calculate cost275    pricing = OPENROUTER_PRICING.get(model, {"input": 0, "output": 0})276    cost = (277        (input_tokens / 1_000_000) * pricing["input"] +278        (output_tokens / 1_000_000) * pricing["output"]279    )280 281    return {282        "processed_text": processed_text,283        "input_tokens": input_tokens,284        "output_tokens": output_tokens,285        "cost": cost286    }287 288 289def _process_with_ollama(text: str, instruction: str) -> Dict[str, Any]:290    """291    Process text with local Ollama292 293    Args:294        text: Text to process295        instruction: Processing instruction296 297    Returns:298        Dict with processed_text, input_tokens, output_tokens, cost299    """300 301    try:302        import requests303    except ImportError:304        raise ImportError("requests package required for Ollama")305 306    # Check if Ollama is running307    try:308        response = requests.get("http://localhost:11434/api/tags", timeout=5)309        response.raise_for_status()310    except Exception as e:311        raise ConnectionError(f"Ollama not running on localhost:11434. Please start Ollama first. Error: {e}")312 313    # Call Ollama API314    response = requests.post(315        "http://localhost:11434/api/generate",316        json={317            "model": "llama3.1",318            "prompt": f"{instruction}\n\n{text}",319            "stream": False320        }321    )322 323    response.raise_for_status()324    result = response.json()325 326    return {327        "processed_text": result["response"],328        "input_tokens": result.get("prompt_eval_count", 0),329        "output_tokens": result.get("eval_count", 0),330        "cost": 0  # Ollama is free331    }332 333 334# Export main functions335__all__ = [336    'process_with_llm',337    'ProcessingResult',338    'ProcessingProgress',339    'TextChunk',340]341