code/server/chat_helpers.py

"""Chat helper functions — history conversion, prompt building, iteration context.

Enhanced version with:
- Better context management
- Conversation summarization for long histories
- Prompt optimization
- Token counting
- Session awareness
"""

from __future__ import annotations

import logging
from typing import Any, Optional
from dataclasses import dataclass

from code.config.constants import (
    SYSTEM_PROMPT,
    MAX_SESSION_HISTORY,
)
from code.execution.code_extractor import strip_thinking_blocks
from code.model.inference import estimate_tokens

logger = logging.getLogger(__name__)


@dataclass
class ContextInfo:
    """Information about the current conversation context."""
    message_count: int = 0
    estimated_tokens: int = 0
    has_images: bool = False
    has_code_blocks: bool = False
    is_truncated: bool = False
    summary: str | None = None


def chat_history_to_messages(
    history: list[dict[str, str]],
    max_messages: int | None = None,
) -> list[dict[str, Any]]:
    """Convert chat history list to messages format for the model.
    
    Enhanced with:
    - Message limit handling
    - Context info tracking
    - Automatic truncation for long conversations
    
    Args:
        history: Chat history from frontend.
        max_messages: Maximum messages to include (None for default).
        
    Returns:
        Messages in OpenAI format with system prompt prepended.
    """
    if max_messages is None:
        max_messages = MAX_SESSION_HISTORY
        
    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    # Take only the most recent messages if history is too long
    working_history = history[-max_messages:] if len(history) > max_messages else history
    
    context_info = ContextInfo()
    
    for item in working_history:
        role = item.get("role")
        content = str(item.get("content") or "").strip()
        if role not in {"user", "assistant"} or not content:
            continue
            
        if role == "assistant":
            content = strip_thinking_blocks(content)
            
        messages.append({"role": role, "content": content})
        
        # Track context info
        context_info.message_count += 1
        context_info.estimated_tokens += estimate_tokens(content)
        
        if not context_info.has_images:
            context_info.has_images = "image" in content.lower() or "![" in content
        if not context_info.has_code_blocks:
            context_info.has_code_blocks = "```" in content
    
    context_info.is_truncated = len(history) > max_messages
    
    if context_info.is_truncated:
        logger.info(
            "History truncated from %d to %d messages", 
            len(history), len(working_history)
        )
    
    return messages


def clip_context(text: str, limit: int = 4_000) -> str:
    """Truncate text to a character limit with a note.
    
    Preserves code blocks when possible by truncating outside of them.
    
    Args:
        text: Text to truncate.
        limit: Maximum character count.
        
    Returns:
        Truncated text with ellipsis note if truncated.
    """
    if len(text) <= limit:
        return text
    
    # Simple truncation (could be enhanced to preserve code blocks)
    return text[:limit] + f"\n... truncated {len(text) - limit} characters ..."


def iteration_context(execution_context: dict[str, Any] | None) -> str:
    """Build a context string from previous execution results.
    
    This allows the model to reference prior code, stdout, and stderr
    when the user asks to iterate or debug.
    
    Args:
        execution_context: Dict with previous execution details.
        
    Returns:
        Formatted context string for inclusion in prompts.
    """
    if not execution_context or not execution_context.get("code"):
        return ""

    code = clip_context(str(execution_context.get("code") or ""), 6_000)
    target = str(execution_context.get("target") or "code")
    fence_lang = str(execution_context.get("fence_lang") or target)
    status = str(execution_context.get("status") or "")
    stdout = clip_context(str(execution_context.get("stdout") or ""), 2_000)
    stderr = clip_context(str(execution_context.get("stderr") or ""), 2_000)

    parts = [
        "Previous generated code and run result are available for iteration.",
        f"Previous target: {target}",
        f"Previous status: {status}",
        f"Previous code:\n```{fence_lang}\n{code}\n```",
    ]
    if stdout:
        parts.append(f"Previous stdout:\n{stdout}")
    if stderr:
        parts.append(f"Previous stderr / traceback:\n{stderr}")
    parts.append(
        "If the user asks to revise, debug, extend, or explain the prior code, use this context."
    )
    return "\n\n".join(parts)


def targeted_prompt(
    prompt: str,
    target_language: str,
    target_framework: str = "",
    execution_context: dict[str, Any] | None = None,
    search_context: str = "",
) -> str:
    """Build the full user prompt with language, framework, search, and iteration context.
    
    Enhanced with:
    - Better framework-specific hints
    - Code quality instructions
    - Security reminders
    
    Args:
        prompt: User's original prompt.
        target_language: Target programming language.
        target_framework: Optional target framework.
        execution_context: Previous execution context for iterations.
        search_context: Web search results to incorporate.
        
    Returns:
        Complete formatted prompt for the model.
    """
    iter_ctx = iteration_context(execution_context)
    context_block = f"\n\n{iter_ctx}" if iter_ctx else ""

    search_block = ""
    if search_context:
        search_block = (
            f"\n\n{search_context}\n\n"
            "Use the above search results to inform your code generation if relevant."
        )

    framework_hint = f" using {target_framework}" if target_framework else ""

    gradio_hint = ""
    if target_framework == "Gradio":
        gradio_hint = (
            "\n\nIMPORTANT: This is a Gradio app. Create a complete Python script that:\n"
            "- Imports gradio as gr\n"
            "- Defines the UI using gr.Interface() or gr.Blocks()\n"
            "- Includes all processing logic inline\n"
            "- Calls .launch(server_name='0.0.0.0', server_port=7860) at the end\n"
            "- Uses only standard library + gradio + common packages (PIL, matplotlib, numpy)\n"
            "- Make the UI clean, modern, and functional\n"
            "- Include proper error handling and loading states"
        )
    
    # Framework-specific hints
    react_hint = ""
    if target_framework == "React":
        react_hint = (
            "\n\nREACT SPECIFIC:\n"
            "- Use functional components and hooks\n"
            "- Include proper TypeScript types if applicable\n"
            "- Use CSS modules or styled-components for styling\n"
            "- Make components reusable and composable"
        )
    
    flask_hint = ""
    if target_framework == "Flask":
        flask_hint = (
            "\n\nFLASK SPECIFIC:\n"
            "- Use Flask blueprints for larger apps\n"
            "- Include proper error handlers\n"
            "- Add input validation using marshmallow or similar\n"
            "- Structure with separate routes, models, and templates directories"
        )
    
    security_reminder = (
        "\n\nSECURITY REMINDERS:\n"
        "- Validate all user inputs\n"
        "- Use parameterized queries for database operations\n"
        "- Never hardcode secrets or API keys\n"
        "- Implement proper authentication/authorization where needed\n"
        "- Sanitize outputs to prevent XSS"
    )

    return (
        f"Target: {target_language}{framework_hint}. Generate a complete, runnable application. "
        "Use the `write_file` tool to save each file to the workspace. "
        "Do NOT paste code in markdown blocks — always use `write_file`. "
        "For multi-file projects, call `write_file` once per file. "
        "After writing files, give a short summary of what you created. "
        "Include proper error handling, comments, and follow best practices. "
        f"{gradio_hint}"
        f"{react_hint}"
        f"{flask_hint}"
        f"{security_reminder}"
        f"{search_block}"
        f"{context_block}\n\n"
        f"User request:\n{prompt}"
    )


def summarize_conversation(
    history: list[dict[str, str]],
    max_summary_length: int = 500,
) -> str | None:
    """Create a summary of the conversation for context preservation.
    
    This is useful when the conversation gets too long and needs to be
    compressed while preserving important context.
    
    Args:
        history: Full conversation history.
        max_summary_length: Maximum length of the summary.
        
    Returns:
        Summary string or None if summarization not needed.
    """
    if len(history) < 10:  # Only summarize longer conversations
        return None
    
    # Extract key information
    user_requests = []
    files_created = []
    
    for msg in history:
        role = msg.get("role")
        content = msg.get("content", "")
        
        if role == "user":
            # Get first sentence as summary of request
            first_sentence = content.split('.')[0].split('\n')[0]
            if first_sentence:
                user_requests.append(first_sentence)
        
        elif role == "assistant":
            # Look for file creation mentions
            if "created" in content.lower() or "wrote" in content.lower():
                # Simple extraction - could be enhanced with NLP
                pass
    
    if not user_requests:
        return None
    
    summary_parts = [
        "Conversation summary:",
        f"User made {len(user_requests)} requests.",
        "Key requests: " + "; ".join(user_requests[-5:]),  # Last 5 requests
    ]
    
    summary = "\n".join(summary_parts)
    
    if len(summary) > max_summary_length:
        summary = summary[:max_summary_length] + "..."
    
    return summary


def build_system_prompt_with_context(
    custom_instructions: str | None = None,
    active_skills: list[str] | None = None,
    active_agent: str | None = None,
) -> str:
    """Build system prompt with additional context.
    
    Args:
        custom_instructions: Additional custom instructions.
        active_skills: List of currently active skill names.
        active_agent: Name of the active agent if any.
        
    Returns:
        Complete system prompt string.
    """
    base_prompt = SYSTEM_PROMPT
    
    additions = []
    
    if active_agent:
        additions.append(f"You are currently operating as the '{active_agent}' agent.")
    
    if active_skills:
        skills_str = ", ".join(active_skills)
        additions.append(f"Active skills: {skills_str}. Apply these skill guidelines.")
    
    if custom_instructions:
        additions.append(f"ADDITIONAL INSTRUCTIONS:\n{custom_instructions}")
    
    if additions:
        base_prompt += "\n\n" + "\n\n".join(additions)
    
    return base_prompt