simbarag/blueprints/conversation/__init__.py

import asyncio
import json
import logging
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 utils.image_process import image_content_block
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 get_memories_for_user
from .models import (
    Conversation,
    PydConversation,
    PydListConversation,
    Speaker,
)
from .prompts import SIMBA_SYSTEM_PROMPT

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


_SYSTEM_PROMPT = SIMBA_SYSTEM_PROMPT


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)


async def _build_system_prompt_with_memories(user_id: str) -> str:
    """Append user memories to the base system prompt."""
    memories = await get_memories_for_user(user_id)
    if not memories:
        return _SYSTEM_PROMPT
    memory_block = "\n".join(f"- {m}" for m in memories)
    return f"{_SYSTEM_PROMPT}\n\nUSER MEMORIES (facts the user has asked you to remember):\n{memory_block}"


def _build_messages_payload(
    conversation,
    query_text: str,
    image_block: dict | None = None,
    system_prompt: str | None = None,
) -> list:
    recent_messages = (
        conversation.messages[-20:]
        if len(conversation.messages) > 20
        else conversation.messages
    )
    messages_payload = [{"role": "system", "content": system_prompt or _SYSTEM_PROMPT}]
    for msg in recent_messages[:-1]:  # Exclude the message we just added
        role = "user" if msg.speaker == Speaker.USER else "assistant"
        text = msg.text
        if msg.image_key and role == "user":
            text = f"[User sent an image]\n{text}"
        messages_payload.append({"role": role, "content": text})

    # Attach the image directly to the current message so the (multimodal)
    # agent model sees it in one call, instead of a separate describe pass.
    if image_block:
        content = [
            {"type": "text", "text": query_text or "The user sent this image."},
            image_block,
        ]
    else:
        content = query_text
    messages_payload.append({"role": "user", "content": content})
    return messages_payload


@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 conversation.fetch_related("messages")
    await add_message_to_conversation(
        conversation=conversation,
        message=query,
        speaker="user",
        user=user,
    )

    system_prompt = await _build_system_prompt_with_memories(str(user.id))
    messages_payload = _build_messages_payload(
        conversation, query, system_prompt=system_prompt
    )
    payload = {"messages": messages_payload}
    agent_config = {
        "configurable": {"user_id": str(user.id), "is_admin": user.is_admin()}
    }

    response = await main_agent.ainvoke(payload, config=agent_config)
    message = 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 conversation.fetch_related("messages")
    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))
    messages_payload = _build_messages_payload(
        conversation, query_text or "", image_block, system_prompt=system_prompt
    )
    payload = {"messages": messages_payload}
    agent_config = {
        "configurable": {"user_id": str(user.id), "is_admin": user.is_admin()}
    }

    async def event_generator():
        final_message = None
        streamed = ""
        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()
            total_ms = round((now - t0) * 1000)
            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": total_ms,
                        "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:
            async for event in main_agent.astream_events(
                payload, version="v2", config=agent_config
            ):
                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")
                    content = getattr(chunk, "content", None)
                    if isinstance(content, str) and content:
                        if first_content_at is None:
                            first_content_at = time.monotonic()
                        turn_has_content = True
                        n_tokens += 1
                        streamed += content
                        yield f"data: {json.dumps({'type': 'content', 'delta': content})}\n\n"
                    elif not turn_has_content:
                        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 = content

        except asyncio.CancelledError:
            cancelled = True
            _log_summary()
            text = final_message or streamed
            if text:
                _persist_in_background(persist(text))
            raise
        except Exception as e:
            _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:
            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"

        _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})