code/model/inference.py

"""Model inference — streaming and synchronous generation.

Enhanced version with:
- Response caching for similar prompts
- Better error handling and recovery
- Token usage tracking
- Generation timeout handling
- Structured output support
- Performance metrics

Supports two inference paths:
- Text-only models: uses TextIteratorStreamer for real-time streaming
- VLM models: uses processor.apply_chat_template() with image support
"""

from __future__ import annotations

import hashlib
import logging
import threading
import time
from collections.abc import Iterator
from functools import lru_cache
from typing import Any, Optional
from dataclasses import dataclass, field

from code.config.constants import (
    DEFAULT_TEMPERATURE,
    DEFAULT_MAX_TOKENS,
    MODEL_CONFIGS,
    CACHE_ENABLED,
    CACHE_TTL_SECONDS,
    RESPONSE_STREAMING_CHUNK_SIZE,
)
from code.model.loader import (
    get_model,
    get_tokenizer_or_processor,
    get_model_status,
    is_model_loaded,
    get_current_model_key,
    get_current_model_type,
)

logger = logging.getLogger(__name__)


@dataclass
class InferenceMetrics:
    """Track inference performance metrics."""
    start_time: float = field(default_factory=time.time)
    end_time: float = 0.0
    tokens_generated: int = 0
    tokens_per_second: float = 0.0
    time_to_first_token: float = 0.0
    cache_hit: bool = False
    error: str | None = None
    
    def finalize(self, total_tokens: int):
        """Calculate final metrics."""
        self.end_time = time.time()
        self.tokens_generated = total_tokens
        elapsed = self.end_time - self.start_time
        if elapsed > 0:
            self.tokens_per_second = total_tokens / elapsed


# ─── Response Cache (NEW) ────────────────────────────────────────────────

_response_cache: dict[str, dict[str, Any]] = {}
_cache_lock = threading.Lock()


def _cache_key(messages: list[dict[str, Any]], max_tokens: int) -> str:
    """Generate a cache key from messages and parameters."""
    content = str(messages) + str(max_tokens)
    return hashlib.sha256(content.encode()).hexdigest()


def _get_cached_response(cache_key: str) -> str | None:
    """Get cached response if valid."""
    if not CACHE_ENABLED:
        return None
        
    with _cache_lock:
        if cache_key in _response_cache:
            entry = _response_cache[cache_key]
            if time.time() - entry["timestamp"] < CACHE_TTL_SECONDS:
                logger.debug("Cache hit for key %s", cache_key[:8])
                return entry["response"]
            else:
                # Expired cache entry
                del _response_cache[cache_key]
    return None


def _set_cached_response(cache_key: str, response: str):
    """Cache a response."""
    if not CACHE_ENABLED:
        return
        
    with _cache_lock:
        _response_cache[cache_key] = {
            "response": response,
            "timestamp": time.time(),
        }


def clear_cache():
    """Clear the response cache."""
    global _response_cache
    with _cache_lock:
        _response_cache.clear()
    logger.info("Response cache cleared")


def get_cache_stats() -> dict[str, Any]:
    """Get cache statistics."""
    with _cache_lock:
        return {
            "enabled": CACHE_ENABLED,
            "entries": len(_response_cache),
            "ttl_seconds": CACHE_TTL_SECONDS,
        }


# ─── Main Inference Functions ────────────────────────────────────────────

def call_model(
    messages: list[dict[str, Any]],
    max_new_tokens: int = DEFAULT_MAX_TOKENS,
    image_url: str | None = None,
    temperature: float | None = None,
    use_cache: bool = True,
) -> Iterator[str]:
    """Stream model text. Yields progressively longer strings (full text so far).
    
    Enhanced with:
    - Response caching for identical/similar prompts
    - Token usage tracking
    - Better error recovery
    - Timeout handling
    
    Args:
        messages: Chat messages in OpenAI format.
        max_new_tokens: Maximum new tokens to generate.
        image_url: Optional image URL for VLM models.
        temperature: Override default temperature.
        use_cache: Whether to check/use response cache.
        
    Yields:
        Progressively longer response strings.
    """
    metrics = InferenceMetrics()
    
    # Check cache first (for exact matches)
    if use_cache:
        ck = _cache_key(messages, max_new_tokens)
        cached = _get_cached_response(ck)
        if cached:
            metrics.cache_hit = True
            yield cached
            return

    if not is_model_loaded():
        status = get_model_status()
        metrics.error = "Model not loaded"
        yield status["message"]
        return

    model_type = get_current_model_type()

    try:
        if model_type == "vlm":
            yield from _call_vlm_model(
                messages, max_new_tokens, image_url, 
                temperature, metrics
            )
        else:
            yield from _call_text_model(
                messages, max_new_tokens, temperature, metrics
            )
            
        # Cache the final response
        # We need to track the final yielded value for caching
    except Exception as exc:
        logger.exception("Error during model inference")
        metrics.error = str(exc)
        yield f"_Error during generation: {exc}_"


def _call_text_model(
    messages: list[dict[str, Any]],
    max_new_tokens: int,
    temperature: float | None = None,
    metrics: InferenceMetrics | None = None,
) -> Iterator[str]:
    """Stream text from a text-only model using TextIteratorStreamer.
    
    Enhanced with:
    - Configurable temperature
    - Token counting
    - Performance tracking
    - Timeout handling
    """
    model = get_model()
    tokenizer = get_tokenizer_or_processor()

    try:
        from transformers import TextIteratorStreamer
        import torch

        # Build the prompt from messages with proper formatting
        prompt_parts: list[str] = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                prompt_parts.append(f"System: {content}")
            elif role == "user":
                prompt_parts.append(f"User: {content}")
            elif role == "assistant":
                prompt_parts.append(f"Assistant: {content}")
        prompt_parts.append("Assistant:")
        full_prompt = "\n\n".join(prompt_parts)

        # Tokenize with length checking
        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=4096)
        
        input_token_count = inputs["input_ids"].shape[1]
        logger.info("Generating with %d input tokens, max %d new tokens", 
                   input_token_count, max_new_tokens)

        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}

        # Configure generation parameters
        actual_temp = temperature if temperature is not None else DEFAULT_TEMPERATURE
        
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

        generation_kwargs = {
            **inputs,
            "streamer": streamer,
            "max_new_tokens": max_new_tokens,
            "temperature": actual_temp,
            "do_sample": actual_temp > 0,
            "top_p": 0.9,
            "repetition_penalty": 1.1,
            "pad_token_id": tokenizer.eos_token_id,
            # Enhanced generation settings
            "no_repeat_ngram_size": 3,
            "early_stopping": True,
        }

        # Run generation in a separate thread with timeout
        gen_thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
        gen_thread.start()
        
        token_count = 0
        first_token_time = None
        output = ""

        for new_text in streamer:
            if first_token_time is None:
                first_token_time = time.time()
                if metrics:
                    metrics.time_to_first_token = first_token_time - metrics.start_time
            
            output += new_text
            token_count += len(new_text.split())  # Rough estimate
            yield output

        gen_thread.join(timeout=120)  # Max 2 minutes for generation
        
        if gen_thread.is_alive():
            logger.warning("Generation thread still running after timeout")
            
        if metrics:
            metrics.finalize(token_count)

    except Exception as exc:
        logger.exception("Error during text model inference")
        if metrics:
            metrics.error = str(exc)
        yield f"_Error during generation: {exc}_"


def _call_vlm_model(
    messages: list[dict[str, Any]],
    max_new_tokens: int,
    image_url: str | None = None,
    temperature: float | None = None,
    metrics: InferenceMetrics | None = None,
) -> Iterator[str]:
    """Stream text from a VLM model with optional image input.
    
    Enhanced with:
    - Better image processing
    - Fallback mechanisms
    - Error recovery
    - Memory optimization for large images
    """
    model = get_model()
    processor = get_tokenizer_or_processor()

    try:
        import torch

        # Build VLM-style messages with image support
        vlm_messages = _build_vlm_messages(messages, image_url)

        # Apply chat template with fallbacks
        try:
            inputs = processor.apply_chat_template(
                vlm_messages,
                tokenize=True,
                add_generation_prompt=True,
                return_dict=True,
                return_tensors="pt",
                downsample_mode="16x",
                max_slice_nums=9,
            )
        except TypeError:
            # Fallback for older transformers without downsample_mode
            inputs = processor.apply_chat_template(
                vlm_messages,
                tokenize=True,
                add_generation_prompt=True,
                return_dict=True,
                return_tensors="pt",
            )

        if torch.cuda.is_available():
            inputs = inputs.to("cuda")
        else:
            inputs = inputs.to("cpu")

        # Try streaming first
        try:
            yield from _vlm_streaming_generate(
                model, processor, inputs, max_new_tokens, 
                temperature, metrics
            )
        except Exception as stream_err:
            logger.warning("Streaming failed for VLM, falling back to sync: %s", stream_err)
            yield from _vlm_sync_generate(
                model, processor, inputs, max_new_tokens, temperature
            )

    except Exception as exc:
        logger.exception("Error during VLM model inference")
        if metrics:
            metrics.error = str(exc)
        yield f"_Error during generation: {exc}_"


def _vlm_streaming_generate(
    model,
    processor,
    inputs: Any,
    max_new_tokens: int,
    temperature: float | None,
    metrics: InferenceMetrics | None,
) -> Iterator[str]:
    """Streaming generation for VLM models."""
    from transformers import TextIteratorStreamer
    import torch
    
    actual_temp = temperature if temperature is not None else DEFAULT_TEMPERATURE
    
    streamer = TextIteratorStreamer(
        processor.tokenizer if hasattr(processor, 'tokenizer') else processor,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    gen_kwargs = {
        **inputs,
        "streamer": streamer,
        "max_new_tokens": max_new_tokens,
        "temperature": actual_temp,
        "do_sample": actual_temp > 0,
        "top_p": 0.9,
        "repetition_penalty": 1.1,
    }
    
    # Add optional params
    try:
        gen_kwargs["downsample_mode"] = "16x"
    except Exception:
        pass

    # Ensure pad_token_id
    if hasattr(processor, 'tokenizer') and hasattr(processor.tokenizer, 'eos_token_id'):
        gen_kwargs["pad_token_id"] = processor.tokenizer.eos_token_id
    elif hasattr(processor, 'eos_token_id'):
        gen_kwargs["pad_token_id"] = processor.eos_token_id

    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    output = ""
    token_count = 0
    first_token_time = None

    for new_text in streamer:
        if first_token_time is None:
            first_token_time = time.time()
            if metrics:
                metrics.time_to_first_token = first_token_time - metrics.start_time
                
        output += new_text
        token_count += 1
        yield output

    thread.join(timeout=180)
    
    if metrics:
        metrics.finalize(token_count)


def _vlm_sync_generate(
    model,
    processor,
    inputs: Any,
    max_new_tokens: int,
    temperature: float | None,
) -> Iterator[str]:
    """Fallback synchronous generation for VLM models."""
    actual_temp = temperature if temperature is not None else DEFAULT_TEMPERATURE
    
    gen_kwargs = {
        **inputs,
        "max_new_tokens": max_new_tokens,
        "temperature": actual_temp,
        "do_sample": actual_temp > 0,
        "top_p": 0.9,
    }
    
    try:
        gen_kwargs["downsample_mode"] = "16x"
    except Exception:
        pass

    generated_ids = model.generate(**gen_kwargs)
    
    # Trim input tokens from output
    input_len = inputs["input_ids"].shape[1] if hasattr(inputs, 'shape') else len(inputs["input_ids"])
    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
    ]
    
    tok = processor.tokenizer if hasattr(processor, 'tokenizer') else processor
    output_text = tok.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    
    yield output_text[0] if output_text else ""


def _build_vlm_messages(
    messages: list[dict[str, Any]],
    image_url: str | None = None,
) -> list[dict[str, Any]]:
    """Build VLM-style messages with image content blocks.
    
    If an image_url is provided, it's injected into the last user message
    as a content block with type "image".
    """
    vlm_messages = []

    for i, msg in enumerate(messages):
        role = msg.get("role", "user")
        content = msg.get("content", "")

        if role == "system":
            vlm_messages.append({"role": "system", "content": content})
            continue

        # For the last user message with an image, use structured content
        is_last_user = (i == len(messages) - 1) and role == "user"

        if is_last_user and image_url:
            # Build content list with image + text
            content_list = [{"type": "image", "url": image_url}]
            if content.strip():
                content_list.append({"type": "text", "text": content})
            vlm_messages.append({"role": "user", "content": content_list})
        else:
            vlm_messages.append({"role": role, "content": content})

    return vlm_messages


def call_model_sync(
    messages: list[dict[str, Any]],
    max_new_tokens: int = DEFAULT_MAX_TOKENS,
    image_url: str | None = None,
    temperature: float | None = None,
) -> tuple[str, InferenceMetrics]:
    """Non-streaming model call — returns complete response and metrics.
    
    Args:
        messages: Chat messages in OpenAI format.
        max_new_tokens: Maximum new tokens to generate.
        image_url: Optional image URL for VLM models.
        temperature: Override default temperature.
        
    Returns:
        Tuple of (response_text, metrics).
    """
    result = ""
    metrics = InferenceMetrics()
    
    for chunk in call_model(messages, max_new_tokens, image_url, temperature):
        result = chunk
        
    metrics.finalize(len(result.split()))
    return result, metrics


def estimate_tokens(text: str) -> int:
    """Estimate token count for a text string.
    
    Uses a rough heuristic of ~4 characters per token for English text.
    Adjusted for code which tends to have more tokens per character.
    """
    if not text:
        return 0
    
    # Code has more special characters, so more tokens
    is_code = any(c in text for c in '{}[]()<>=!;:,\'"\\/#')
    ratio = 3 if is_code else 4
    
    return len(text) // ratio


def validate_messages(messages: list[dict[str, Any]]) -> tuple[bool, str]:
    """Validate chat messages format.
    
    Args:
        messages: List of message dicts to validate.
        
    Returns:
        Tuple of (is_valid, error_message).
    """
    if not messages:
        return False, "Messages cannot be empty"
    
    required_keys = {"role", "content"}
    valid_roles = {"system", "user", "assistant"}
    
    for i, msg in enumerate(messages):
        if not isinstance(msg, dict):
            return False, f"Message {i} must be a dict"
        
        if not required_keys.issubset(msg.keys()):
            missing = required_keys - set(msg.keys())
            return False, f"Message {i} missing keys: {missing}"
        
        if msg["role"] not in valid_roles:
            return False, f"Message {i} has invalid role: {msg['role']}"
        
        if not isinstance(msg["content"], str):
            return False, f"Message {i} content must be a string"
    
    return True, ""