app.py

import html
import json
import logging
import os
import shutil
import tempfile
import time
from datetime import datetime, timezone

import gradio as gr
import spaces
from dotenv import load_dotenv
from langchain.agents import create_agent
from langchain_core.messages import AIMessage, ToolMessage
from langchain_openai.chat_models.base import OpenAIContextOverflowError
from nemotron_llm import NemotronChatOpenAI
from nemotron_llm import message_content_to_string
from nemotron_llm import strip_thinking_for_display

from context_management import (
    MAX_TOKENS,
    build_context_middleware,
    context_settings,
)
from diagnostic_planning import build_todo_middleware
from llama_server import N_CTX, SERVER_CONFIG, start_llama_server
from elm_server import ELM_CONFIG, ensure_elm_emulator
from obd_connection import init_obd_session, obd_session_lock
from obd_connection import VITALS_POLL_INTERVAL
from obd_vitals import refresh_vitals_hud, start_vitals_poller
from vitals_store import DB_PATH, PURGE_INTERVAL_SECONDS, RETENTION_HOURS
from fault_simulation import (
    apply_fault,
    clear_faults,
    describe_fault,
    fault_investigation_prompt,
    get_active_fault,
)
from tools import DIAGNOSTIC_TOOLS
from ui.layout import build_ui
from ui.theme import build_cockpit_theme, load_cockpit_css
from trace_store import (
    append_trace,
    build_turn_trace,
    trace_file_exists,
    trace_file_path,
)

load_dotenv()
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger("car-diagnostic-agent")

GPU_DURATION = int(os.environ.get("SPACES_GPU_DURATION", "180"))
DEFAULT_THINKING_ENABLED = os.environ.get("THINKING_ENABLED", "false").lower() in (
    "1",
    "true",
    "yes",
)

SYSTEM_PROMPT = (
    "You are a helpful car diagnostic assistant for a Toyota Auris Hybrid (2012). "
    "Help users understand symptoms, possible causes, and suggested next steps. "
    "Be clear and practical.\n\n"
    "OBD vitals and DTCs are recorded automatically every few seconds. DTCs may "
    "not appear immediately when a problem starts—the ECU can set them after "
    "conditions persist. During an active diagnosis, keep checking live and recent "
    "data across multiple turns (wait a few seconds between checks when watching "
    "for new codes) rather than relying on a single reading.\n\n"
    "You MUST call tools for vehicle data (never invent readings or codes). "
    "Tool names and parameters are available via the tool schemas.\n\n"
    "**Before investigating any symptoms**, call `obd_vitals_reference` to learn "
    "stored labels, command names, and how to use live vs historical tools. "
    "Then call `diagnosis_methodology` for DTC and workflow rules.\n\n"
    "Always call obd_vitals_reference and diagnosis_methodology before reading vitals or DTCs. "
    "If live vitals fail or before the first live read, call check_obd_vitals_link. "
    "Stored DTCs may remain after a fault clears — "
    "confirm any code with live vitals and get_stored_vital_history trends before concluding. "
    "When you suspect a specific failure, call get_diagnosis_guide with the matching fault_id.\n\n"
    "When investigating symptoms, use get_stored_vital_history to trend one sensor over time, "
    "get_recent_vitals and read_dtc_codes across the conversation to catch developing faults. "
    "Compare snapshots over time for trends (fuel trims, temps, RPM). "
    "Actually invoke tools, then summarize results for the user.\n\n"
    "For investigations that need many tool calls in one turn, use `write_todos` "
    "to track steps (reference, vitals, DTCs, trends, guides) before you start."
)

CONTEXT_OVERFLOW_REPLY = (
    "The conversation context is full (model limit ~{n_ctx} tokens). "
    "Please start a new chat, or ask for fewer vitals at a time "
    "(e.g. one PID or a smaller recent count)."
)

_agent = None
_llm: NemotronChatOpenAI | None = None


def _truncate(text: str, limit: int = 1200) -> str:
    text = text.strip()
    if len(text) <= limit:
        return text
    return text[:limit] + "…"


def _pretty_payload(raw: str | dict | list) -> str:
    if isinstance(raw, (dict, list)):
        return json.dumps(raw, indent=2, ensure_ascii=False)
    text = str(raw).strip()
    try:
        return json.dumps(json.loads(text), indent=2, ensure_ascii=False)
    except (json.JSONDecodeError, TypeError):
        return text


def _message_key(msg) -> int:
    return id(msg)


def _append_tool_trace(lines: list[str], msg, seen: set[int]) -> None:
    key = _message_key(msg)
    if key in seen:
        return
    seen.add(key)

    if isinstance(msg, AIMessage):
        for tc in msg.tool_calls or []:
            name = tc.get("name", "?")
            args = tc.get("args", {})
            lines.append(f"**Call `{name}`**")
            lines.append(f"```json\n{_pretty_payload(args)}\n```")
            logger.info("Tool call: %s args=%s", name, args)
    elif isinstance(msg, ToolMessage):
        tool_name = getattr(msg, "name", None) or msg.tool_call_id or "tool"
        lines.append(f"**Result `{tool_name}`**")
        lines.append(f"```json\n{_truncate(_pretty_payload(msg.content), 4000)}\n```")
        logger.info("Tool result %s: %s", tool_name, msg.content)


def _format_trace(lines: list[str], *, in_progress: bool = False) -> str:
    if not lines:
        status = "*Waiting for tool activity…*" if not in_progress else "*Running…*"
        return f"### Tool calls & results\n\n{status}"
    body = "\n\n".join(lines)
    suffix = "\n\n---\n*In progress…*" if in_progress else ""
    return f"### Tool calls & results\n\n{body}{suffix}"


def _normalize_chat_history(history: list | None) -> list[dict[str, str]]:
    """Gradio 6 messages format: list of {role, content} dicts with string content."""
    normalized: list[dict[str, str]] = []
    for entry in history or []:
        if isinstance(entry, dict):
            role = entry.get("role")
            content = entry.get("content")
            if role and content is not None:
                normalized.append(
                    {"role": str(role), "content": message_content_to_string(content)}
                )
            continue
        if isinstance(entry, (list, tuple)) and len(entry) == 2:
            user_msg, assistant_msg = entry
            if user_msg:
                normalized.append(
                    {"role": "user", "content": message_content_to_string(user_msg)}
                )
            if assistant_msg:
                normalized.append(
                    {
                        "role": "assistant",
                        "content": message_content_to_string(assistant_msg),
                    }
                )
    return normalized


def _final_reply(result_messages: list, input_count: int) -> str:
    final = strip_thinking_for_display(
        message_content_to_string(result_messages[-1].content)
    )
    return final or "(no response)"


def _model_label() -> str:
    return f"{SERVER_CONFIG.hf_model_repo_id}:{SERVER_CONFIG.model}"


def _trace_params(enable_thinking: bool) -> dict:
    return {
        "temperature": 0.6,
        "max_tokens": MAX_TOKENS,
        "n_ctx": N_CTX,
        "thinking_enabled": enable_thinking,
    }


def _persist_turn_trace(
    *,
    user_message: str,
    prior_turns: int,
    result_messages: list,
    start_index: int,
    assistant_message: str,
    enable_thinking: bool,
    latency_ms: int,
    status: str,
    error: str | None = None,
    turn_started: datetime,
) -> None:
    try:
        append_trace(
            build_turn_trace(
                user_message=user_message,
                prior_turns=prior_turns,
                result_messages=result_messages,
                start_index=start_index,
                assistant_message=assistant_message,
                model=_model_label(),
                params=_trace_params(enable_thinking),
                latency_ms=latency_ms,
                status=status,
                error=error,
                ts=turn_started,
            )
        )
    except Exception:
        logger.exception("Failed to persist agent trace")


def _download_traces_file() -> str | None:
    path = trace_file_path()
    if not trace_file_exists():
        return None
    dest = os.path.join(tempfile.gettempdir(), "agent_traces.jsonl")
    shutil.copy2(path, dest)
    return dest


def refresh_vitals_panel() -> str:
    """Poll ECU now, store snapshot, and update the vitals HUD."""
    try:
        return refresh_vitals_hud()
    except Exception as exc:
        logger.exception("Vitals panel refresh failed")
        return (
            '<div class="vitals-hud"><div class="vitals-hud__error">'
            f"Could not load vitals: {html.escape(str(exc))}"
            "</div></div>"
        )


def on_fault_select(fault_id: str | None) -> str:
    return describe_fault(fault_id)


def on_apply_fault(
    fault_id: str | None,
    history: list | None,
    enable_thinking: bool,
):
    """Apply fault, refresh vitals, and ask the agent to investigate (visible in chat)."""
    try:
        if not fault_id or fault_id == "none":
            status = clear_faults()
            vitals = refresh_vitals_panel()
            yield status, vitals, history, gr.update()
            return

        status = apply_fault(fault_id)
        vitals = refresh_vitals_panel()
        prompt = fault_investigation_prompt()
        yield status, vitals, history, gr.update()
        for chat_history, trace in chat(prompt, history, enable_thinking):
            yield status, vitals, chat_history, trace
    except Exception as exc:
        logger.exception("Apply fault failed")
        yield (
            f"*Could not apply fault: {exc}*",
            refresh_vitals_panel(),
            history,
            gr.update(),
        )


def on_clear_fault() -> tuple[str, str, None]:
    try:
        status = clear_faults()
        return status, refresh_vitals_panel(), gr.update(value="none")
    except Exception as exc:
        logger.exception("Clear fault failed")
        return f"*Could not clear faults: {exc}*", refresh_vitals_panel(), gr.update()


def active_fault_status() -> str:
    fault = get_active_fault()
    if fault is None:
        return "**Active fault:** none (healthy `car` scenario)"
    return (
        f"**Active fault:** {fault.name} (`{fault.id}`)  \n"
        f"**Typical DTC:** `{fault.typical_dtc}`"
    )


@spaces.GPU(duration=GPU_DURATION)
def load_agent() -> None:
    """Start llama.cpp server (model load) and build the agent inside GPU context."""
    global _agent, _llm
    if _agent is not None:
        return

    print("Loading model via llama.cpp server...")
    logger.info(
        "Starting agent (model=%s, max_tokens=%d, thinking=%s, tools=%s)",
        SERVER_CONFIG.model_alias,
        MAX_TOKENS,
        DEFAULT_THINKING_ENABLED,
        [t.name for t in DIAGNOSTIC_TOOLS],
    )
    start_llama_server(SERVER_CONFIG)
    ctx = context_settings()
    logger.info(
        "Context budget: N_CTX=%d, MAX_TOKENS=%d, input_budget=%d, message_budget=%d, "
        "summary_trigger=%d, summary_keep_tokens=%d, tool_clear_trigger=%d",
        ctx.n_ctx,
        ctx.max_tokens,
        ctx.input_budget,
        ctx.message_budget,
        ctx.summary_trigger_tokens,
        ctx.summary_keep_tokens,
        ctx.tool_clear_trigger_tokens,
    )
    _llm = NemotronChatOpenAI(
        base_url=SERVER_CONFIG.base_url,
        api_key=os.environ.get("OPENAI_API_KEY", "local"),
        model=SERVER_CONFIG.model_alias,
        temperature=0.6,
        top_p=0.95,
        max_tokens=MAX_TOKENS,
        enable_thinking=DEFAULT_THINKING_ENABLED,
        profile={"max_input_tokens": N_CTX},
    )
    summary_llm = NemotronChatOpenAI(
        base_url=SERVER_CONFIG.base_url,
        api_key=os.environ.get("OPENAI_API_KEY", "local"),
        model=SERVER_CONFIG.model_alias,
        temperature=0.3,
        max_tokens=512,
        enable_thinking=False,
        profile={"max_input_tokens": N_CTX},
    )
    _agent = create_agent(
        model=_llm,
        tools=DIAGNOSTIC_TOOLS,
        system_prompt=SYSTEM_PROMPT,
        middleware=[build_todo_middleware(), *build_context_middleware(summary_llm)],
    )
    print("Agent ready.")
    logger.info("Agent ready.")


@spaces.GPU(duration=GPU_DURATION)
def chat(message: str, history: list | None, enable_thinking: bool):
    if _agent is None:
        load_agent()

    if _llm is not None:
        _llm.enable_thinking = enable_thinking

    history = _normalize_chat_history(history)
    agent_messages = [
        {"role": entry["role"], "content": entry["content"]} for entry in history
    ]
    agent_messages.append({"role": "user", "content": message})

    trace_lines: list[str] = []
    seen: set[int] = set()
    pending = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": "…"},
    ]
    yield pending, _format_trace(trace_lines, in_progress=True)

    logger.info("User message: %s (thinking=%s)", message, enable_thinking)
    result_messages = list(agent_messages)
    turn_started = datetime.now(timezone.utc)
    turn_start_mono = time.monotonic()
    prior_turns = sum(1 for entry in history if entry.get("role") == "user")
    start_index = len(agent_messages)
    assistant_reply = ""

    try:
        for step in _agent.stream(
            {"messages": agent_messages},
            stream_mode="updates",
        ):
            for _node, update in step.items():
                if not isinstance(update, dict):
                    continue
                new_msgs = update.get("messages", [])
                result_messages.extend(new_msgs)
                for msg in new_msgs:
                    _append_tool_trace(trace_lines, msg, seen)
                yield pending, _format_trace(trace_lines, in_progress=True)
    except OpenAIContextOverflowError as exc:
        logger.warning("Agent context overflow during stream")
        trace_lines.append("*Context limit reached*")
        assistant_reply = CONTEXT_OVERFLOW_REPLY.format(n_ctx=N_CTX)
        latency_ms = int((time.monotonic() - turn_start_mono) * 1000)
        _persist_turn_trace(
            user_message=message,
            prior_turns=prior_turns,
            result_messages=result_messages,
            start_index=start_index,
            assistant_message=assistant_reply,
            enable_thinking=enable_thinking,
            latency_ms=latency_ms,
            status="context_overflow",
            error=str(exc),
            turn_started=turn_started,
        )
        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": assistant_reply},
        ]
        yield history, _format_trace(trace_lines, in_progress=False)
        return
    except Exception:
        logger.exception("Agent stream failed; falling back to invoke")
        trace_lines.append("*Non-streaming fallback*")
        yield pending, _format_trace(trace_lines, in_progress=True)
        try:
            final_state = _agent.invoke({"messages": agent_messages})
        except OpenAIContextOverflowError as exc:
            logger.warning("Agent context overflow during invoke")
            trace_lines.append("*Context limit reached*")
            assistant_reply = CONTEXT_OVERFLOW_REPLY.format(n_ctx=N_CTX)
            latency_ms = int((time.monotonic() - turn_start_mono) * 1000)
            _persist_turn_trace(
                user_message=message,
                prior_turns=prior_turns,
                result_messages=result_messages,
                start_index=start_index,
                assistant_message=assistant_reply,
                enable_thinking=enable_thinking,
                latency_ms=latency_ms,
                status="context_overflow",
                error=str(exc),
                turn_started=turn_started,
            )
            history = history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": assistant_reply},
            ]
            yield history, _format_trace(trace_lines, in_progress=False)
            return
        except Exception as exc:
            logger.exception("Agent invoke failed")
            trace_lines.append(f"*Error:* `{exc}`")
            assistant_reply = "Sorry, the diagnostic agent failed. Please try again."
            latency_ms = int((time.monotonic() - turn_start_mono) * 1000)
            _persist_turn_trace(
                user_message=message,
                prior_turns=prior_turns,
                result_messages=result_messages,
                start_index=start_index,
                assistant_message=assistant_reply,
                enable_thinking=enable_thinking,
                latency_ms=latency_ms,
                status="error",
                error=str(exc),
                turn_started=turn_started,
            )
            history = history + [
                {"role": "user", "content": message},
                {"role": "assistant", "content": assistant_reply},
            ]
            yield history, _format_trace(trace_lines, in_progress=False)
            return
        result_messages = final_state["messages"]
        for msg in result_messages[len(agent_messages) :]:
            _append_tool_trace(trace_lines, msg, seen)

    reply = _final_reply(result_messages, len(agent_messages))
    assistant_reply = reply
    latency_ms = int((time.monotonic() - turn_start_mono) * 1000)
    _persist_turn_trace(
        user_message=message,
        prior_turns=prior_turns,
        result_messages=result_messages,
        start_index=start_index,
        assistant_message=assistant_reply,
        enable_thinking=enable_thinking,
        latency_ms=latency_ms,
        status="ok",
        turn_started=turn_started,
    )
    history = history + [
        {"role": "user", "content": message},
        {"role": "assistant", "content": reply},
    ]
    if not trace_lines:
        trace_lines.append("*No tools were called.*")
    logger.info("Reply length: %d chars", len(reply))
    yield history, _format_trace(trace_lines, in_progress=False)


def _submit(message, history, enable_thinking):
    if not message or not str(message).strip():
        yield history, gr.update()
        return
    yield from chat(message.strip(), history, enable_thinking)


try:
    ensure_elm_emulator()
    with obd_session_lock:
        init_obd_session()
except Exception:
    logger.exception("ELM/OBD singleton failed to start at boot (poller will retry)")

demo = build_ui(
    default_thinking_enabled=DEFAULT_THINKING_ENABLED,
    active_fault_status_fn=active_fault_status,
    on_fault_select_fn=on_fault_select,
    on_apply_fault_fn=on_apply_fault,
    on_clear_fault_fn=on_clear_fault,
    chat_submit_fn=_submit,
    download_traces_fn=_download_traces_file,
)

load_agent()
start_vitals_poller()
logger.info(
    "ELM327-emulator (scenario=%s) at %s; vitals poll every ~%ss; db=%s; "
    "purge every %ss (retention %.1fh)",
    ELM_CONFIG.scenario,
    ELM_CONFIG.connect_uri,
    VITALS_POLL_INTERVAL,
    DB_PATH,
    PURGE_INTERVAL_SECONDS,
    RETENTION_HOURS,
)

if __name__ == "__main__":
    allowed = ["/data", tempfile.gettempdir()]
    data_dir = os.environ.get("VITALS_DATA_DIR", "/data")
    if os.path.isdir(data_dir):
        allowed.append(data_dir)
    demo.launch(
        allowed_paths=allowed,
        theme=build_cockpit_theme(),
        css=load_cockpit_css(),
    )