simbarag/blueprints/conversation/__init__.py

import asyncio
import json
import logging
import os
import re
import time
import uuid

from quart import Blueprint, jsonify, make_response, request
from quart_jwt_extended import (
    get_jwt_identity,
    jwt_refresh_token_required,
)

import blueprints.users.models
from metrics import count_result, observe_response, time_response
from utils.image_process import image_content_block
from utils.strip_markdown import extract_agent_text
from utils.image_upload import ImageValidationError, process_image
from utils.s3_client import generate_presigned_url as s3_presigned_url
from utils.s3_client import get_image as s3_get_image
from utils.s3_client import upload_image as s3_upload_image

from .agents import main_agent
from .logic import (
    add_message_to_conversation,
    persist_agent_tool_messages,
    tool_result_text,
)
from .memory import build_system_prompt_with_memories
from .models import (
    Conversation,
    PydConversation,
    PydListConversation,
    Speaker,
)

TTFT_TIMEOUT = float(os.getenv("TTFT_TIMEOUT_SECONDS", "30"))

_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)


def _route_content(
    buf: str, text: str, in_thinking: bool
) -> tuple[str, str, bool, str]:
    """Split a streaming content chunk into (thinking_text, content_text).

    Handles tags split across chunk boundaries by keeping a small trailing
    buffer. Returns the updated in_thinking state and any leftover buffer.
    """
    combined = buf + text
    thinking_parts: list[str] = []
    content_parts: list[str] = []

    while combined:
        tag, hold = (
            (_THINK_CLOSE, len(_THINK_CLOSE) - 1)
            if in_thinking
            else (_THINK_OPEN, len(_THINK_OPEN) - 1)
        )
        idx = combined.find(tag)
        if idx == -1:
            # Tag not yet complete — hold back enough bytes to catch a split.
            if len(combined) <= hold:
                return (
                    "".join(thinking_parts),
                    "".join(content_parts),
                    in_thinking,
                    combined,
                )
            safe, leftover = combined[:-hold], combined[-hold:]
            (thinking_parts if in_thinking else content_parts).append(safe)
            return (
                "".join(thinking_parts),
                "".join(content_parts),
                in_thinking,
                leftover,
            )
        (thinking_parts if in_thinking else content_parts).append(combined[:idx])
        combined = combined[idx + len(tag) :]
        in_thinking = not in_thinking

    return "".join(thinking_parts), "".join(content_parts), in_thinking, ""


conversation_blueprint = Blueprint(
    "conversation_api", __name__, url_prefix="/api/conversation"
)

# Detached persistence tasks. A strong reference keeps them from being garbage
# collected mid-flight when the SSE stream is cancelled by a client disconnect.
_persist_tasks: set[asyncio.Task] = set()


def _persist_in_background(coro) -> asyncio.Task:
    task = asyncio.create_task(coro)
    _persist_tasks.add(task)
    task.add_done_callback(_persist_tasks.discard)
    return task


async def _get_owned_conversation(conversation_id, user_id) -> Conversation | None:
    """Fetch a conversation only if it belongs to the given user."""
    return await Conversation.get_or_none(id=conversation_id, user_id=user_id)


def _build_user_message(
    query_text: str,
    image_block: dict | None = None,
) -> dict:
    """Build the single new user message for this turn.

    History is restored by the LangGraph checkpointer; only the current
    message is passed to avoid duplicating prior turns in checkpoint state.
    """
    if image_block:
        content = [
            {"type": "text", "text": query_text or "The user sent this image."},
            image_block,
        ]
    else:
        content = query_text
    return {"role": "user", "content": content}


@conversation_blueprint.post("/query")
@jwt_refresh_token_required
async def query():
    current_user_uuid = get_jwt_identity()
    user = await blueprints.users.models.User.get(id=current_user_uuid)
    data = await request.get_json()
    query = data.get("query")
    conversation_id = data.get("conversation_id")
    conversation = await _get_owned_conversation(conversation_id, current_user_uuid)
    if conversation is None:
        return jsonify({"error": "conversation not found"}), 404
    await add_message_to_conversation(
        conversation=conversation,
        message=query,
        speaker="user",
        user=user,
    )

    with time_response(stream=False):
        system_prompt = await build_system_prompt_with_memories(str(user.id))
        payload = {"messages": [_build_user_message(query)]}
        agent_config = {
            "configurable": {
                "user_id": str(user.id),
                "is_admin": user.is_admin(),
                "thread_id": str(conversation.id),
                "system_prompt": system_prompt,
            }
        }

        response = await main_agent.ainvoke(payload, config=agent_config)
        message = extract_agent_text(response.get("messages", [])[-1].content)
        await persist_agent_tool_messages(
            conversation, user, response.get("messages", [])
        )
        await add_message_to_conversation(
            conversation=conversation,
            message=message,
            speaker="simba",
            user=user,
        )
        return jsonify({"response": message})


@conversation_blueprint.post("/upload-image")
@jwt_refresh_token_required
async def upload_image():
    current_user_uuid = get_jwt_identity()
    await blueprints.users.models.User.get(id=current_user_uuid)

    files = await request.files
    form = await request.form
    file = files.get("file")
    conversation_id = form.get("conversation_id")

    if not file or not conversation_id:
        return jsonify({"error": "file and conversation_id are required"}), 400

    file_bytes = file.read()
    content_type = file.content_type or "image/jpeg"

    try:
        processed_bytes, output_content_type = process_image(file_bytes, content_type)
    except ImageValidationError as e:
        return jsonify({"error": str(e)}), 400

    ext = output_content_type.split("/")[-1]
    if ext == "jpeg":
        ext = "jpg"
    key = f"conversations/{conversation_id}/{uuid.uuid4()}.{ext}"

    await s3_upload_image(processed_bytes, key, output_content_type)

    return jsonify({"image_key": key})


@conversation_blueprint.get("/image/<path:image_key>")
@jwt_refresh_token_required
async def serve_image(image_key: str):
    url = await s3_presigned_url(image_key)
    return jsonify({"url": url})


@conversation_blueprint.post("/stream-query")
@jwt_refresh_token_required
async def stream_query():
    current_user_uuid = get_jwt_identity()
    user = await blueprints.users.models.User.get(id=current_user_uuid)
    data = await request.get_json()
    query_text = data.get("query")
    conversation_id = data.get("conversation_id")
    image_key = data.get("image_key")
    conversation = await _get_owned_conversation(conversation_id, current_user_uuid)
    if conversation is None:
        return jsonify({"error": "conversation not found"}), 404
    await add_message_to_conversation(
        conversation=conversation,
        message=query_text or "",
        speaker="user",
        user=user,
        image_key=image_key,
    )

    # If an image was uploaded, attach it to the agent request directly
    image_block = None
    if image_key:
        try:
            image_bytes, _ = await s3_get_image(image_key)
            image_block = image_content_block(image_bytes)
        except Exception as e:
            logging.error(f"Failed to load image {image_key}: {e}")
            query_text = f"{query_text or ''}\n\n[An image was sent but could not be loaded.]".strip()

    system_prompt = await build_system_prompt_with_memories(str(user.id))
    payload = {"messages": [_build_user_message(query_text or "", image_block)]}
    agent_config = {
        "configurable": {
            "user_id": str(user.id),
            "is_admin": user.is_admin(),
            "thread_id": str(conversation.id),
            "system_prompt": system_prompt,
        }
    }

    async def event_generator():
        final_message = None
        streamed = ""
        # Tag-routing state for <think>…</think> blocks.
        in_thinking = False
        tag_buf = ""
        # Fallback heartbeat for models that emit empty chunks before thinking text.
        last_thinking = 0.0
        turn_has_content = False

        # Timing / observability
        t0 = time.monotonic()
        llm_start_at: float | None = None
        first_content_at: float | None = None
        tool_timers: dict[str, float] = {}  # run_id -> start time
        tool_spans: list[dict] = []
        n_tokens = 0
        model_name: str | None = None
        cancelled = False

        async def persist(text: str) -> None:
            await add_message_to_conversation(
                conversation=conversation,
                message=text,
                speaker="simba",
                user=user,
            )

        def _log_summary(*, error: str | None = None) -> None:
            now = time.monotonic()
            thinking_ms = (
                round((first_content_at - llm_start_at) * 1000)
                if llm_start_at and first_content_at
                else None
            )
            ttft_ms = (
                round((first_content_at - t0) * 1000) if first_content_at else None
            )
            logging.info(
                json.dumps(
                    {
                        "event": "stream_request",
                        "conv": str(conversation.id),
                        "user": user.username,
                        "total_ms": round((now - t0) * 1000),
                        "thinking_ms": thinking_ms,
                        "ttft_ms": ttft_ms,
                        "tokens": n_tokens,
                        "tools": tool_spans,
                        "model": model_name,
                        "cancelled": cancelled,
                        **({"error": error} if error else {}),
                    }
                )
            )

        try:
            aiter = main_agent.astream_events(
                payload, version="v2", config=agent_config
            ).__aiter__()
            while True:
                # Apply a short TTFT timeout only during the window between
                # on_chat_model_start and the first content token. Before the
                # LLM call starts (tool phase) and after the first token, let
                # the httpx read timeout handle stalls instead.
                if llm_start_at is not None and first_content_at is None:
                    ttft_wait = TTFT_TIMEOUT - (time.monotonic() - llm_start_at)
                    timeout = max(ttft_wait, 1.0)
                else:
                    timeout = None

                try:
                    if timeout is not None:
                        event = await asyncio.wait_for(
                            aiter.__anext__(), timeout=timeout
                        )
                    else:
                        event = await aiter.__anext__()
                except StopAsyncIteration:
                    break
                except asyncio.TimeoutError:
                    _log_summary(error=f"TTFT timeout after {TTFT_TIMEOUT:.0f}s")
                    yield f"data: {json.dumps({'type': 'error', 'message': 'Simba took too long to respond. Please try again.', 'stalled': True})}\n\n"
                    return

                event_type = event.get("event")

                if event_type == "on_tool_start":
                    tool_timers[event["run_id"]] = time.monotonic()
                    yield f"data: {json.dumps({'type': 'tool_start', 'tool': event['name']})}\n\n"

                elif event_type == "on_tool_end":
                    start = tool_timers.pop(event["run_id"], None)
                    tool_spans.append(
                        {
                            "tool": event["name"],
                            "ms": round((time.monotonic() - start) * 1000)
                            if start
                            else None,
                        }
                    )
                    yield f"data: {json.dumps({'type': 'tool_end', 'tool': event['name']})}\n\n"
                    # Persist tool output so later turns keep the fetched context
                    output = event.get("data", {}).get("output")
                    if output is not None:
                        await add_message_to_conversation(
                            conversation=conversation,
                            message=tool_result_text(event["name"], output),
                            speaker="tool",
                            user=user,
                        )

                elif event_type == "on_chat_model_start":
                    turn_has_content = False
                    if llm_start_at is None:
                        llm_start_at = time.monotonic()
                    if model_name is None:
                        model_name = event.get("metadata", {}).get(
                            "ls_model_name"
                        ) or event.get("name")

                elif event_type == "on_chat_model_stream":
                    chunk = event.get("data", {}).get("chunk")
                    raw = getattr(chunk, "content", None)
                    if isinstance(raw, str) and raw:
                        think_text, content_text, in_thinking, tag_buf = _route_content(
                            tag_buf, raw, in_thinking
                        )
                        if think_text:
                            yield f"data: {json.dumps({'type': 'thinking_delta', 'delta': think_text})}\n\n"
                        if content_text:
                            if first_content_at is None:
                                first_content_at = time.monotonic()
                            turn_has_content = True
                            n_tokens += 1
                            streamed += content_text
                            yield f"data: {json.dumps({'type': 'content', 'delta': content_text})}\n\n"
                    elif not turn_has_content:
                        # Empty chunk before any content — heartbeat for models
                        # that don't emit <think> tags at all.
                        now = time.monotonic()
                        if now - last_thinking >= 1.0:
                            last_thinking = now
                            yield f"data: {json.dumps({'type': 'thinking'})}\n\n"

                elif event_type == "on_chain_end":
                    output = event.get("data", {}).get("output")
                    if isinstance(output, dict):
                        msgs = output.get("messages", [])
                        if msgs:
                            last_msg = msgs[-1]
                            content = getattr(last_msg, "content", None)
                            if isinstance(content, str) and content:
                                final_message = _THINK_RE.sub("", content).strip()

        except asyncio.CancelledError:
            cancelled = True
            _log_summary()
            text = final_message or streamed
            if text:
                _persist_in_background(persist(text))
            raise
        except Exception as e:
            count_result(stream=True, result="error")
            _log_summary(error=str(e))
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"

        final_message = final_message or streamed
        if final_message:
            # Persist via a detached, shielded task so a disconnect at this
            # point (after generation completes) still commits the answer.
            task = _persist_in_background(persist(final_message))
            await asyncio.shield(task)
            yield f"data: {json.dumps({'type': 'response', 'message': final_message})}\n\n"
        else:
            yield f"data: {json.dumps({'type': 'error', 'message': 'No response generated'})}\n\n"

        observe_response(stream=True, seconds=time.monotonic() - t0)
        count_result(stream=True, result="success")
        _log_summary()
        yield "data: [DONE]\n\n"

    return await make_response(
        event_generator(),
        200,
        {
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )


@conversation_blueprint.route("/<conversation_id>")
@jwt_refresh_token_required
async def get_conversation(conversation_id: str):
    user_uuid = get_jwt_identity()
    conversation = await _get_owned_conversation(conversation_id, user_uuid)
    if conversation is None:
        return jsonify({"error": "conversation not found"}), 404
    await conversation.fetch_related("messages")

    # Manually serialize the conversation with messages
    messages = []
    for msg in conversation.messages:
        # Tool messages hold raw fetched context for the LLM; only surface
        # their header line (e.g. "[web_search result]") in the UI.
        text = msg.text.split("\n", 1)[0] if msg.speaker == Speaker.TOOL else msg.text
        messages.append(
            {
                "id": str(msg.id),
                "text": text,
                "speaker": msg.speaker.value,
                "created_at": msg.created_at.isoformat(),
                "image_key": msg.image_key,
            }
        )
    return jsonify(
        {
            "id": str(conversation.id),
            "name": conversation.name,
            "messages": messages,
            "created_at": conversation.created_at.isoformat(),
            "updated_at": conversation.updated_at.isoformat(),
        }
    )


@conversation_blueprint.post("/")
@jwt_refresh_token_required
async def create_conversation():
    user_uuid = get_jwt_identity()
    user = await blueprints.users.models.User.get(id=user_uuid)
    conversation = await Conversation.create(
        name="New Conversation",
        user=user,
    )

    serialized_conversation = await PydConversation.from_tortoise_orm(conversation)
    return jsonify(serialized_conversation.model_dump())


@conversation_blueprint.get("/")
@jwt_refresh_token_required
async def get_all_conversations():
    user_uuid = get_jwt_identity()
    user = await blueprints.users.models.User.get(id=user_uuid)

    query = Conversation.filter(user=user)

    search = request.args.get("search", "").strip()
    if search:
        query = query.filter(name__icontains=search)

    query = query.order_by("-updated_at")

    # `limit` is optional: when omitted we return the full list (backward
    # compatible with any non-frontend callers). `offset` only applies with it.
    limit = request.args.get("limit", type=int)
    if limit is not None:
        offset = request.args.get("offset", default=0, type=int)
        query = query.offset(offset).limit(limit)

    serialized_conversations = await PydListConversation.from_queryset(query)

    return jsonify(serialized_conversations.model_dump())


@conversation_blueprint.patch("/<conversation_id>")
@jwt_refresh_token_required
async def rename_conversation(conversation_id: str):
    user_uuid = get_jwt_identity()
    body = await request.get_json()
    name = (body or {}).get("name", "").strip()
    if not name:
        return jsonify({"error": "name is required"}), 400

    conversation = await _get_owned_conversation(conversation_id, user_uuid)
    if conversation is None:
        return jsonify({"error": "conversation not found"}), 404

    conversation.name = name[:255]
    await conversation.save()

    serialized_conversation = await PydConversation.from_tortoise_orm(conversation)
    return jsonify(serialized_conversation.model_dump())


@conversation_blueprint.delete("/<conversation_id>")
@jwt_refresh_token_required
async def delete_conversation(conversation_id: str):
    user_uuid = get_jwt_identity()
    conversation = await _get_owned_conversation(conversation_id, user_uuid)
    if conversation is None:
        return jsonify({"error": "conversation not found"}), 404

    await conversation.delete()
    return jsonify({"success": True})