cAIc/routers/chat.py

"""JarvisChat routers - /api/chat streaming endpoint."""
import asyncio
import json
import logging
import uuid
from datetime import datetime, timezone

import httpx
from fastapi import APIRouter, HTTPException, Request
from fastapi.responses import StreamingResponse

from config import DEFAULT_MODEL, LLAMA_SERVER_BASE
from crypto import encrypt_text, decrypt_text
from db import get_db, get_upload_context
from memory import process_remember_command, auto_detect_facts, check_fact_conflicts
from rag import build_system_prompt, ingest_auto_fact
from search import (calculate_perplexity, is_uncertain, is_refusal,
                    clean_hedging, format_search_results, format_direct_answer,
                    extract_search_query, query_searxng)
from security import read_json_body, log_incident, BODY_LIMIT_CHAT_BYTES
from config import MAX_CHAT_MESSAGE_CHARS, MODEL_CONTEXT_LENGTH

log = logging.getLogger("caic")
router = APIRouter()


def parse_llama_stream_chunk(line: str) -> tuple:
    if line.startswith("data: "):
        line = line[6:]
    if line.strip() == "[DONE]":
        return None, True, {}, []
    try:
        chunk = json.loads(line)
        choices = chunk.get("choices", [])
        if choices:
            delta = choices[0].get("delta", {})
            token = delta.get("content")
            finish = choices[0].get("finish_reason")
            stats = {}
            logprobs_list = []
            logprobs_info = choices[0].get("logprobs")
            if logprobs_info:
                content_logprobs = logprobs_info.get("content", [])
                for entry in content_logprobs:
                    if "logprob" in entry:
                        logprobs_list.append({"logprob": entry["logprob"]})
            if finish == "stop":
                usage = chunk.get("usage", {})
                stats["tokens_per_sec"] = usage.get("tokens_per_second", 0.0)
                stats["completion_tokens"] = usage.get("completion_tokens", 0)
                stats["prompt_tokens"] = usage.get("prompt_tokens", 0)
            return token, finish == "stop", stats, logprobs_list
        if "message" in chunk and "content" in chunk["message"]:
            token = chunk["message"]["content"]
            done = chunk.get("done", False)
            stats = {}
            if done:
                eval_count = chunk.get("eval_count", 0)
                eval_duration = chunk.get("eval_duration", 0)
                stats["tokens_per_sec"] = (eval_count / (eval_duration / 1e9)) if eval_duration > 0 else 0
                stats["completion_tokens"] = eval_count
            return token, done, stats, []
    except json.JSONDecodeError:
        pass
    return None, False, {}, []


@router.post("/api/chat")
async def chat(request: Request):
    body = await read_json_body(request, BODY_LIMIT_CHAT_BYTES)
    conv_id = body.get("conversation_id")
    user_message = body.get("message", "").strip()
    if len(user_message) > MAX_CHAT_MESSAGE_CHARS:
        raise HTTPException(status_code=413, detail="Chat message is too long")
    model = body.get("model", DEFAULT_MODEL)
    preset_prompt = body.get("system_prompt", "")
    upload_context_id = body.get("upload_context_id")
    private_chat = body.get("private_chat", False)

    if not user_message:
        raise HTTPException(status_code=400, detail="Empty message")

    db = get_db()
    now = datetime.now(timezone.utc).isoformat()
    settings = {row["key"]: row["value"] for row in db.execute("SELECT key, value FROM settings").fetchall()}
    search_enabled = settings.get("search_enabled", "true") == "true" and not private_chat

    upload_doc = None
    if upload_context_id and not private_chat:
        ctx = get_upload_context(db, upload_context_id)
        if ctx:
            upload_doc = f"[ATTACHED DOCUMENT: {ctx['filename']}]\n{ctx['content']}\n[END DOCUMENT]"
        else:
            log.warning(f"upload_context_id {upload_context_id} not found or expired, continuing without it")

    remember_response = None if private_chat else process_remember_command(user_message)

    if private_chat:
        if not conv_id:
            conv_id = str(uuid.uuid4())
        system_prompt = ""
        messages = []
        if preset_prompt:
            messages.append({"role": "system", "content": preset_prompt})
        messages.append({"role": "user", "content": user_message})
        history_rows = []
        db.close()
    else:
        if not conv_id:
            conv_id = str(uuid.uuid4())
            title = user_message[:80] + ("..." if len(user_message) > 80 else "")
            db.execute("INSERT INTO conversations (id, title, model, created_at, updated_at) VALUES (?, ?, ?, ?, ?)",
                       (conv_id, encrypt_text(title), model, now, now))
        else:
            db.execute("UPDATE conversations SET updated_at = ? WHERE id = ?", (now, conv_id))

        db.execute("INSERT INTO messages (conversation_id, role, content, created_at) VALUES (?, ?, ?, ?)",
                   (conv_id, "user", encrypt_text(user_message), now))
        db.commit()

        raw_rows = db.execute(
            "SELECT role, content FROM messages WHERE conversation_id = ? ORDER BY id ASC", (conv_id,)
        ).fetchall()
        history_rows = []
        for row in raw_rows:
            history_rows.append({"role": row["role"], "content": decrypt_text(row["content"])})
        extra_prompt = preset_prompt
        if upload_doc:
            extra_prompt = (extra_prompt + "\n\n" + upload_doc) if extra_prompt else upload_doc
        system_prompt = await build_system_prompt(db, extra_prompt, user_message)
        db.close()

        messages = []
        if system_prompt:
            messages.append({"role": "system", "content": system_prompt})
        for row in history_rows:
            messages.append({"role": row["role"], "content": row["content"]})

    upstream_payload = {"model": model, "messages": messages, "stream": True, "logprobs": True}

    async def stream_response():
        full_response = []
        all_logprobs = []
        tokens_per_sec = 0.0
        completion_tokens = 0
        prompt_tokens = 0
        rag_update = None

        if remember_response:
            yield f"data: {json.dumps({'token': remember_response + chr(10) + chr(10), 'conversation_id': conv_id})}\n\n"

        async with httpx.AsyncClient() as client:
            try:
                async with client.stream(
                    "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions",
                    json=upstream_payload,
                    timeout=httpx.Timeout(300.0, connect=10.0),
                ) as resp:
                    async for line in resp.aiter_lines():
                        if line.strip():
                            token, done, stats, chunk_logprobs = parse_llama_stream_chunk(line)
                            if chunk_logprobs:
                                all_logprobs.extend(chunk_logprobs)
                            if token:
                                full_response.append(token)
                                yield f"data: {json.dumps({'token': token, 'conversation_id': conv_id})}\n\n"
                            if done:
                                tokens_per_sec = stats.get("tokens_per_sec", 0.0)
                                completion_tokens = stats.get("completion_tokens", 0)
                                prompt_tokens = stats.get("prompt_tokens", 0)

                assistant_msg = "".join(full_response)
                perplexity = calculate_perplexity(all_logprobs) if all_logprobs else 0.0
                should_search = is_uncertain(all_logprobs) or is_refusal(assistant_msg)

                if search_enabled and should_search:
                    yield f"data: {json.dumps({'searching': True, 'conversation_id': conv_id})}\n\n"
                    search_query = extract_search_query(user_message)
                    search_results = await query_searxng(search_query)

                    if search_results:
                        search_context = format_search_results(search_results)
                        augmented_messages = []
                        if system_prompt:
                            augmented_messages.append({"role": "system", "content": system_prompt + "\n\n" + search_context})
                        else:
                            augmented_messages.append({"role": "system", "content": search_context})
                        for row in history_rows[:-1]:
                            augmented_messages.append({"role": row["role"], "content": row["content"]})
                        augmented_messages.append({"role": "user", "content": user_message})

                        yield f"data: {json.dumps({'search_results': len(search_results), 'conversation_id': conv_id})}\n\n"

                        augmented_response = []
                        async with client.stream(
                            "POST", f"{LLAMA_SERVER_BASE}/v1/chat/completions",
                            json={"model": model, "messages": augmented_messages, "stream": True},
                            timeout=httpx.Timeout(300.0, connect=10.0),
                        ) as resp2:
                            async for line in resp2.aiter_lines():
                                if line.strip():
                                    token2, done2, _, _ = parse_llama_stream_chunk(line)
                                    if token2:
                                        augmented_response.append(token2)
                                    if done2:
                                        break

                        raw_response = "".join(augmented_response) or assistant_msg
                        cleaned_response = clean_hedging(raw_response)
                        if is_refusal(cleaned_response) or len(cleaned_response) < 20:
                            cleaned_response = format_direct_answer(user_message, search_results)

                        yield f"data: {json.dumps({'token': cleaned_response, 'conversation_id': conv_id, 'augmented': True})}\n\n"

                        if not private_chat:
                            saved_msg = cleaned_response + "\n\n---\n*🔍 Enhanced with web search results*"
                            if remember_response:
                                saved_msg = remember_response + "\n\n" + saved_msg

                            db2 = get_db()
                            db2.execute("INSERT INTO messages (conversation_id, role, content, created_at) VALUES (?, ?, ?, ?)",
                                        (conv_id, "assistant", encrypt_text(saved_msg), datetime.now(timezone.utc).isoformat()))
                            db2.commit()
                            db2.close()

                            facts = auto_detect_facts(user_message, cleaned_response)
                            if facts:
                                conflicts = check_fact_conflicts(facts)
                                if conflicts:
                                    rag_update = {"conflicts": conflicts}
                                else:
                                    asyncio.ensure_future(ingest_auto_fact(facts, user_message, cleaned_response))

                        yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'searched': True, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH, **(rag_update and {'rag_update_suggestion': rag_update} or {})})}\n\n"
                        return

                if not private_chat:
                    saved_msg = assistant_msg
                    if remember_response:
                        saved_msg = remember_response + "\n\n" + saved_msg

                    db2 = get_db()
                    db2.execute("INSERT INTO messages (conversation_id, role, content, created_at) VALUES (?, ?, ?, ?)",
                                (conv_id, "assistant", encrypt_text(saved_msg), datetime.now(timezone.utc).isoformat()))
                    db2.commit()
                    db2.close()

                    facts = auto_detect_facts(user_message, assistant_msg)
                    if facts:
                        conflicts = check_fact_conflicts(facts)
                        if conflicts:
                            rag_update = {"conflicts": conflicts}
                        else:
                            asyncio.ensure_future(ingest_auto_fact(facts, user_message, assistant_msg))

                yield f"data: {json.dumps({'done': True, 'conversation_id': conv_id, 'perplexity': round(perplexity, 2), 'tokens_per_sec': round(tokens_per_sec, 1), 'prompt_tokens': prompt_tokens, 'completion_tokens': completion_tokens, 'context_length': MODEL_CONTEXT_LENGTH, **(rag_update and {'rag_update_suggestion': rag_update} or {})})}\n\n"

            except httpx.RemoteProtocolError:
                pass
            except httpx.ConnectError:
                yield f"data: {json.dumps({'error': 'Cannot connect to inference server. Is it running?'})}\n\n"
            except Exception as e:
                incident_key = log_incident("chat_stream", message="Inference stream failure during chat response",
                                            request=request, exc=e)
                yield f"data: {json.dumps({'error': 'Chat response generation failed before completion. Use the incident key for support lookup.', 'error_key': incident_key})}\n\n"

    return StreamingResponse(stream_response(), media_type="text/event-stream")