code/server/chat_helpers.py

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

from __future__ import annotations

from typing import Any

from code.config.constants import SYSTEM_PROMPT
from code.execution.code_extractor import strip_thinking_blocks


def chat_history_to_messages(history: list[dict[str, str]]) -> list[dict[str, Any]]:
    """Convert chat history list to messages format for the model.

    Prepends the system prompt and strips thinking blocks from
    assistant messages.
    """
    messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
    for item in 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})
    return messages


def clip_context(text: str, limit: int = 4_000) -> str:
    """Truncate text to a character limit with a note."""
    if len(text) <= limit:
        return text
    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.
    """
    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."""
    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"
        )

    return (
        f"Target: {target_language}{framework_hint}. Generate a complete, runnable application. "
        "If the user asks for a web app, include all HTML/CSS/JS. "
        "If they ask for a backend, include the server code and any API definitions. "
        "For single-file apps, use a single code block. For multi-file projects, use the @@FILE: format. "
        "Make the code complete, working, and well-structured."
        f"{gradio_hint}"
        f"{search_block}"
        f"{context_block}\n\n"
        f"User request:\n{prompt}"
    )