llm-in-text/backend/pro_completions.py

import asyncio
import contextlib
import json
import logging
import os
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

from fastapi import FastAPI, HTTPException, Request, Security
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel

from geoip import get_ip_location_text
from llm import stream_ollama_events
from models import UserPreferences
from prompt import build_pro_completion_prompts

logger = logging.getLogger("api.pro")

PRO_COMPLETION_TIMEOUT = float(os.getenv("PRO_COMPLETION_TIMEOUT", "3600"))
PRO_QUEUE_TIMEOUT = float(os.getenv("PRO_QUEUE_TIMEOUT", "600"))
PRO_MAX_CONCURRENCY = max(1, int(os.getenv("PRO_MAX_CONCURRENCY", "1")))
PRO_QUEUE_MAX_SIZE = max(0, int(os.getenv("PRO_QUEUE_MAX_SIZE", "5")))
PRO_STATUS_RETENTION_SECONDS = float(os.getenv("PRO_STATUS_RETENTION_SECONDS", "600"))
PRO_CANCEL_ACK_TIMEOUT = 5.0
STREAM_HEARTBEAT_SECONDS = float(os.getenv("STREAM_HEARTBEAT_SECONDS", "2"))
PUBLIC_PRO_ERROR = "PRO generation failed. Please retry or adjust the instruction."


class ProCompletionRequest(BaseModel):
    prefix: str
    suffix: str
    languageId: str = "markdown"
    instruction: str = ""
    pro_thinking: str = "medium"
    privacy_mode: bool = False
    user_preferences: Optional[UserPreferences] = None


class ProCancelRequest(BaseModel):
    request_id: str
    reason: str = "abort"


@dataclass
class ProRequestState:
    request_id: str
    status: str = "queued"
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    error: str = ""
    task: asyncio.Task | None = None
    cancel_requested: bool = False
    done_event: asyncio.Event = field(default_factory=asyncio.Event)

    def touch(self, status: str | None = None, error: str = "") -> None:
        if status:
            self.status = status
        if error:
            self.error = error
        self.updated_at = time.time()

    def request_cancel(self) -> None:
        self.cancel_requested = True
        self.touch("cancelled")


PRO_STATES: dict[str, ProRequestState] = {}
PRO_STATES_LOCK = asyncio.Lock()
PRO_SEMAPHORE = asyncio.Semaphore(PRO_MAX_CONCURRENCY)


def _iso_timestamp(value: float) -> str:
    return datetime.fromtimestamp(value, tz=timezone.utc).isoformat()


def _clamp_thinking(value: str | None) -> str | None:
    normalized = (value or "medium").strip().lower()
    if normalized in {"none", "off", "false"}:
        return None
    if normalized in {"low", "medium", "high"}:
        return normalized
    return "medium"


def _queued_states() -> list[ProRequestState]:
    return [state for state in PRO_STATES.values() if state.status == "queued"]


def _queue_position(request_id: str) -> int | None:
    queued = sorted(_queued_states(), key=lambda item: item.created_at)
    for index, state in enumerate(queued, start=1):
        if state.request_id == request_id:
            return index
    return None


async def _cleanup_states() -> None:
    now = time.time()
    expired = [
        request_id
        for request_id, state in PRO_STATES.items()
        if state.status in {"done", "error", "cancelled"}
        and now - state.updated_at > PRO_STATUS_RETENTION_SECONDS
    ]
    for request_id in expired:
        PRO_STATES.pop(request_id, None)


def _state_payload(state: ProRequestState) -> dict:
    return {
        "request_id": state.request_id,
        "status": state.status,
        "queue_position": _queue_position(state.request_id),
        "created_at": _iso_timestamp(state.created_at),
        "updated_at": _iso_timestamp(state.updated_at),
        "error": state.error,
    }


def _build_pro_prompts(
    *,
    prefix: str,
    suffix: str,
    language_id: str,
    instruction: str,
    pro_thinking: str = "medium",
    location: str = "",
    preferences: UserPreferences | None = None,
) -> tuple[str, str]:
    return build_pro_completion_prompts(
        prefix=prefix,
        suffix=suffix,
        instruction=instruction,
        language_id=language_id,
        location=location,
        pro_thinking_level=pro_thinking,
        preferences=preferences,
    )


def _get_client_ip(request: Request) -> str:
    if request.client:
        return request.headers.get("X-Client-IP") or request.client.host
    return request.headers.get("X-Client-IP") or "unknown"


async def _send_sse_event(queue: asyncio.Queue, event_name: str, data: dict) -> None:
    await queue.put((event_name, json.dumps(data, ensure_ascii=False)))


async def _wait_for_cancel_cleanup(state: ProRequestState, request_tag: str, reason: str) -> None:
    if state.done_event.is_set():
        return
    try:
        await asyncio.wait_for(state.done_event.wait(), timeout=PRO_CANCEL_ACK_TIMEOUT)
    except asyncio.TimeoutError:
        logger.warning(
            "[%s] /v1/pro/completions cancel cleanup not confirmed request_id=%s reason=%s",
            request_tag,
            state.request_id,
            reason,
        )


def register_pro_completion_routes(app: FastAPI, get_api_key):
    @app.post("/v1/pro/completions")
    async def create_pro_completion(
        request: Request,
        req: ProCompletionRequest,
        api_key: str = Security(get_api_key),
    ):
        request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4())
        request_tag = request_id[:8]
        event_queue: asyncio.Queue[tuple[str, str] | None] = asyncio.Queue()
        previous_state: ProRequestState | None = None

        async with PRO_STATES_LOCK:
            await _cleanup_states()
            queued_count = len(_queued_states())
            if queued_count >= PRO_QUEUE_MAX_SIZE:
                logger.info("[%s] /v1/pro/completions rejected queue_full request_id=%s", request_tag, request_id)
                return JSONResponse(
                    content={"error": "PRO queue is full", "request_id": request_id},
                    status_code=429,
                )

            existing = PRO_STATES.get(request_id)
            if existing and existing.task and not existing.task.done():
                existing.request_cancel()
                existing.task.cancel()
                previous_state = existing

            state = ProRequestState(request_id=request_id)
            PRO_STATES[request_id] = state

        if previous_state:
            await _wait_for_cancel_cleanup(previous_state, request_tag, "replace")

        client_ip = "hidden"
        location = ""
        if not req.privacy_mode:  # pragma: no cover
            client_ip = _get_client_ip(request)
            location = get_ip_location_text(client_ip)

        prefix = req.prefix or ""
        suffix = req.suffix or ""

        system_prompt, user_prompt = _build_pro_prompts(
            prefix=prefix,
            suffix=suffix,
            language_id=req.languageId,
            instruction=req.instruction,
            pro_thinking=req.pro_thinking,
            location=location,
            preferences=req.user_preferences,
        )

        logger.info(
            "[%s] /v1/pro/completions request_id=%s client_ip=%s prefix_chars=%d suffix_chars=%d instruction_chars=%d lang=%s thinking=%s",
            request_tag,
            request_id,
            client_ip,
            len(prefix),
            len(suffix),
            len(req.instruction or ""),
            req.languageId,
            req.pro_thinking,
        )

        async def producer() -> None:
            acquired = False
            chunks: list[str] = []
            try:
                async with PRO_STATES_LOCK:
                    if state.cancel_requested:
                        raise asyncio.CancelledError()
                    state.touch("queued")
                    queue_position = _queue_position(request_id)
                await _send_sse_event(event_queue, "queued", {"request_id": request_id, "queue_position": queue_position})

                await asyncio.wait_for(PRO_SEMAPHORE.acquire(), timeout=PRO_QUEUE_TIMEOUT)
                acquired = True

                async with PRO_STATES_LOCK:
                    if state.cancel_requested:
                        raise asyncio.CancelledError()
                    state.touch("started")
                await _send_sse_event(event_queue, "started", {"request_id": request_id})

                async for event_type, payload in stream_ollama_events(
                    user_prompt,
                    system_prompt=system_prompt,
                    tag=f"{request_tag}-pro",
                    temperature=0.7,
                    thinking=_clamp_thinking(req.pro_thinking),
                    use_pro_model=True,
                    enable_thinking=True,
                    timeout=PRO_COMPLETION_TIMEOUT,
                ):
                    # Handle 'thinking' event - just update state, don't accumulate
                    if event_type == "thinking":
                        await _send_sse_event(event_queue, "thinking", {"request_id": request_id})
                        continue

                    # Handle 'chunk' event - accumulate content
                    if not payload:
                        continue
                    chunks.append(payload)
                    await _send_sse_event(event_queue, "chunk", {"delta": payload, "request_id": request_id})

                # Handle 'done' event - return full content
                content = "".join(chunks)
                async with PRO_STATES_LOCK:
                    if state.cancel_requested:
                        raise asyncio.CancelledError()
                if not content:
                    raise ValueError("PRO returned empty content")

                async with PRO_STATES_LOCK:
                    state.touch("done")
                logger.info("[%s] /v1/pro/completions done request_id=%s content_chars=%d", request_tag, request_id, len(content))
                await _send_sse_event(event_queue, "done", {"content": content, "request_id": request_id})
            except asyncio.CancelledError:
                async with PRO_STATES_LOCK:
                    state.request_cancel()
                logger.info("[%s] /v1/pro/completions cancelled request_id=%s", request_tag, request_id)
                await _send_sse_event(event_queue, "cancelled", {"cancelled": True, "request_id": request_id})
                raise
            except Exception as exc:
                async with PRO_STATES_LOCK:
                    state.touch("error", PUBLIC_PRO_ERROR)
                logger.exception("[%s] /v1/pro/completions failed request_id=%s", request_tag, request_id)
                await _send_sse_event(event_queue, "error", {"error": PUBLIC_PRO_ERROR, "request_id": request_id})
            finally:
                if acquired:
                    PRO_SEMAPHORE.release()
                state.done_event.set()
                await event_queue.put(None)

        producer_task = asyncio.create_task(producer())
        async with PRO_STATES_LOCK:
            state.task = producer_task

        async def event_stream():
            try:
                while True:
                    try:
                        item = await asyncio.wait_for(event_queue.get(), timeout=STREAM_HEARTBEAT_SECONDS)
                    except asyncio.TimeoutError:
                        yield ": keepalive\n\n"
                        continue
                    if item is None:
                        break
                    event_name, data = item
                    yield f"event: {event_name}\ndata: {data}\n\n"
                    if event_name in {"done", "error", "cancelled"}:
                        break
            except asyncio.CancelledError:
                async with PRO_STATES_LOCK:
                    state.request_cancel()
                producer_task.cancel()
                raise
            finally:
                if not producer_task.done() and not state.done_event.is_set():
                    async with PRO_STATES_LOCK:
                        state.request_cancel()
                    producer_task.cancel()
                    with contextlib.suppress(asyncio.TimeoutError):
                        await asyncio.wait_for(state.done_event.wait(), timeout=PRO_CANCEL_ACK_TIMEOUT)

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream; charset=utf-8",
            headers={
                "Cache-Control": "no-cache, no-transform",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    @app.post("/v1/pro/completions/cancel")
    async def cancel_pro_completion(req: ProCancelRequest, api_key: str = Security(get_api_key)):
        request_id = req.request_id or ""
        request_tag = request_id[:8]
        state_to_wait: ProRequestState | None = None
        async with PRO_STATES_LOCK:
            await _cleanup_states()
            state = PRO_STATES.get(request_id)
            if not state:
                return {"cancelled": False, "status": "not_found"}
            if state.task and not state.task.done():
                state.request_cancel()
                state.task.cancel()
                state_to_wait = state
            if state.status in {"done", "error", "cancelled"}:
                if not state_to_wait:
                    return {"cancelled": False, "status": state.status}
            else:
                state.request_cancel()

        if state_to_wait:
            await _wait_for_cancel_cleanup(state_to_wait, request_tag, req.reason)

        return {"cancelled": True, "status": "ok"}

    @app.get("/v1/pro/completions/status/{request_id}")
    async def get_pro_completion_status(request_id: str, api_key: str = Security(get_api_key)):
        async with PRO_STATES_LOCK:
            await _cleanup_states()
            state = PRO_STATES.get(request_id)
            if not state:
                raise HTTPException(status_code=404, detail="PRO request not found")
            return _state_payload(state)