ai-video-fullstack/backend/services/openai_responses/state.py

"""Persistence and context reconstruction for stateful Responses calls."""

from __future__ import annotations

import json
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any

from db.models import ConversationMessage, ConversationSession, OpenAIResponse
from db.session import SessionLocal
from models import AssistantConfig
from services.openai_realtime.events import assistant_id_from_model
from services.openai_responses.protocol import (
    ResponseRequest,
    ResponsesAPIError,
    new_id,
)
from services.realtime.launcher import resolve_assistant_config
from sqlalchemy import func, select, update
from sqlalchemy.exc import IntegrityError


CONTINUABLE_STATUSES = {"completed"}
IDEMPOTENCY_CONSTRAINT = "uq_openai_response_idempotency"


@dataclass(frozen=True)
class PreparedResponse:
    row: OpenAIResponse
    config: AssistantConfig
    history_messages: list[dict[str, Any]]
    current_messages: list[dict[str, Any]]
    runtime_state: dict[str, Any]
    client_tool_results: list[dict[str, Any]]


def _integrity_constraint_name(error: IntegrityError) -> str | None:
    current = getattr(error, "orig", None)
    seen: set[int] = set()
    while current is not None and id(current) not in seen:
        seen.add(id(current))
        direct = getattr(current, "constraint_name", None)
        if direct:
            return str(direct)
        diagnostic = getattr(current, "diag", None)
        value = getattr(diagnostic, "constraint_name", None)
        if value:
            return str(value)
        current = getattr(current, "__cause__", None) or getattr(
            current,
            "__context__",
            None,
        )
    return None


def _content_text(item: dict[str, Any]) -> str:
    return "\n".join(
        str(part.get("text") or "")
        for part in item.get("content") or []
        if isinstance(part, dict)
    ).strip()


def items_to_messages(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Convert the supported Responses item subset to OpenAI LLM messages."""
    messages: list[dict[str, Any]] = []
    pending_calls: list[dict[str, Any]] = []

    def flush_calls() -> None:
        if not pending_calls:
            return
        messages.append(
            {"role": "assistant", "content": "", "tool_calls": list(pending_calls)}
        )
        pending_calls.clear()

    for item in items:
        item_type = item.get("type")
        if item_type == "function_call":
            pending_calls.append(
                {
                    "id": str(item.get("call_id") or ""),
                    "type": "function",
                    "function": {
                        "name": str(item.get("name") or ""),
                        "arguments": str(item.get("arguments") or "{}"),
                    },
                }
            )
            continue
        flush_calls()
        if item_type == "message":
            messages.append(
                {
                    "role": str(item.get("role") or "user"),
                    "content": _content_text(item),
                }
            )
        elif item_type == "function_call_output":
            messages.append(
                {
                    "role": "tool",
                    "tool_call_id": str(item.get("call_id") or ""),
                    "content": str(item.get("output") or ""),
                }
            )
    flush_calls()
    return messages


async def _load_chain(
    response_id: str,
    *,
    api_key_id: str,
) -> list[OpenAIResponse]:
    chain: list[OpenAIResponse] = []
    seen: set[str] = set()
    async with SessionLocal() as db:
        current_id: str | None = response_id
        is_requested_response = True
        while current_id:
            if current_id in seen or len(chain) >= 1000:
                raise ResponsesAPIError(
                    "Response chain is invalid", code="invalid_response_chain"
                )
            seen.add(current_id)
            row = await db.get(OpenAIResponse, current_id)
            if (
                row is None
                or row.api_key_id != api_key_id
                or (is_requested_response and row.deleted_at is not None)
            ):
                raise ResponsesAPIError(
                    "Response not found", code="response_not_found", status_code=404
                )
            chain.append(row)
            current_id = row.previous_response_id
            is_requested_response = False
    chain.reverse()
    return chain


def _validate_tool_outputs(
    chain: list[OpenAIResponse],
    current_items: list[dict[str, Any]],
) -> None:
    unresolved: dict[str, dict[str, Any]] = {}
    for row in chain:
        for item in row.output_items or []:
            if item.get("type") == "function_call":
                unresolved[str(item.get("call_id") or "")] = item
        for item in row.input_items or []:
            if item.get("type") == "function_call_output":
                unresolved.pop(str(item.get("call_id") or ""), None)

    supplied = [
        str(item.get("call_id") or "")
        for item in current_items
        if item.get("type") == "function_call_output"
    ]
    if len(set(supplied)) != len(supplied):
        raise ResponsesAPIError(
            "A function call output can only be supplied once",
            code="duplicate_tool_output",
            param="input",
        )
    unknown = sorted(set(supplied) - set(unresolved))
    if unknown:
        raise ResponsesAPIError(
            f"Unknown or already resolved call_id: {unknown[0]}",
            code="invalid_tool_output",
            param="input",
        )
    if unresolved and set(supplied) != set(unresolved):
        missing = sorted(set(unresolved) - set(supplied))
        raise ResponsesAPIError(
            f"Missing function_call_output for call_id: {missing[0]}",
            code="missing_tool_output",
            param="input",
        )


def _client_tool_results(
    chain: list[OpenAIResponse],
    current_items: list[dict[str, Any]],
) -> list[dict[str, Any]]:
    calls = [
        item
        for row in chain
        for item in row.output_items or []
        if item.get("type") == "function_call"
    ]
    outputs = {
        str(item.get("call_id") or ""): item
        for item in current_items
        if item.get("type") == "function_call_output"
    }
    results: list[dict[str, Any]] = []
    for call in calls:
        call_id = str(call.get("call_id") or "")
        item = outputs.get(call_id)
        if item is None:
            continue
        output = str(item.get("output") or "")
        try:
            value = json.loads(output)
        except json.JSONDecodeError:
            value = output
        if not isinstance(value, dict) or "status" not in value:
            value = {"status": "ok", "data": value}
        results.append(
            {
                "call_id": call_id,
                "name": call.get("name"),
                "arguments": call.get("arguments"),
                "result": value,
            }
        )
    return results


async def prepare_response(
    request: ResponseRequest,
    *,
    api_key_id: str,
    idempotency_key: str | None = None,
    request_hash: str | None = None,
) -> PreparedResponse:
    chain: list[OpenAIResponse] = []
    parent: OpenAIResponse | None = None
    if request.previous_response_id:
        chain = await _load_chain(
            request.previous_response_id,
            api_key_id=api_key_id,
        )
        parent = chain[-1]
        if parent.status not in CONTINUABLE_STATUSES:
            raise ResponsesAPIError(
                f"Previous response is {parent.status}",
                code="response_not_ready",
                param="previous_response_id",
                status_code=409,
            )
        if (parent.runtime_state or {}).get("status") == "ended":
            raise ResponsesAPIError(
                "The workflow conversation has ended",
                code="conversation_ended",
                param="previous_response_id",
                status_code=409,
            )
        if not parent.assistant_id:
            raise ResponsesAPIError(
                "The assistant used by the previous response no longer exists",
                code="invalid_model",
                param="previous_response_id",
            )
        assistant_id = parent.assistant_id
        if request.model:
            try:
                requested_id = assistant_id_from_model(request.model)
            except ValueError as exc:
                raise ResponsesAPIError(
                    str(exc),
                    code="invalid_model",
                    param="model",
                ) from exc
            if requested_id != assistant_id:
                raise ResponsesAPIError(
                    "model must match the previous response",
                    code="invalid_model",
                    param="model",
                )
        conversation_id = parent.conversation_id
        variables = dict(parent.dynamic_variables or {})
        variables.update(request.dynamic_variables)
        _validate_tool_outputs(chain, request.input_items)
        client_tool_results = _client_tool_results(chain, request.input_items)
        runtime_state = dict(parent.runtime_state or {})
    else:
        try:
            assistant_id = assistant_id_from_model(request.model)
        except ValueError as exc:
            raise ResponsesAPIError(
                str(exc), code="invalid_model", param="model"
            ) from exc
        conversation_id = new_id("conv")[:40]
        variables = dict(request.dynamic_variables)
        client_tool_results = []
        runtime_state = {}
        if any(
            item.get("type") == "function_call_output"
            for item in request.input_items
        ):
            raise ResponsesAPIError(
                "function_call_output requires previous_response_id",
                code="invalid_tool_output",
                param="input",
            )

    try:
        config = await resolve_assistant_config(
            assistant_id,
            dynamic_variables=variables,
            conversation_id=conversation_id,
        )
    except ValueError as exc:
        raise ResponsesAPIError(
            str(exc), code="invalid_model", param="model"
        ) from exc
    if config.runtimeMode != "pipeline":
        raise ResponsesAPIError(
            "Responses API only supports pipeline assistants",
            code="unsupported_runtime_mode",
            param="model",
        )
    if config.type not in {"prompt", "workflow", "dify", "fastgpt"}:
        raise ResponsesAPIError(
            f"Assistant type {config.type} is not supported by Responses API",
            code="unsupported_assistant_type",
            param="model",
        )
    if config.conversation_id != conversation_id:
        config = config.model_copy(update={"conversation_id": conversation_id})
    effective_variables = {
        name: value
        for name, value in (config.dynamic_variables or {}).items()
        if not name.startswith(("system__", "secret__"))
    }

    history_items: list[dict[str, Any]] = []
    for item in chain:
        history_items.extend(item.input_items or [])
        history_items.extend(item.output_items or [])
    history_messages = (
        list(parent.context_messages or [])
        if parent is not None and parent.context_messages
        else items_to_messages(history_items)
    )

    row = OpenAIResponse(
        id=new_id("resp")[:40],
        api_key_id=api_key_id,
        assistant_id=assistant_id,
        conversation_id=conversation_id,
        previous_response_id=request.previous_response_id,
        status="in_progress",
        input_items=request.input_items,
        output_items=[],
        context_messages=[],
        runtime_state={},
        dynamic_variables=effective_variables,
        idempotency_key=idempotency_key,
        request_hash=request_hash,
    )
    async with SessionLocal() as db:
        if parent is None:
            db.add(
                ConversationSession(
                    id=conversation_id,
                    assistant_id=assistant_id,
                    assistant_name=config.name,
                    channel="responses",
                    runtime_mode="pipeline",
                    status="active",
                    message_count=0,
                    analysis_status="none",
                    analysis_data=(
                        {"plan": config.analysis_config}
                        if config.analysis_config.get("enabled")
                        else {}
                    ),
                    extra={},
                )
            )
        else:
            conversation = await db.get(ConversationSession, conversation_id)
            if conversation:
                conversation.status = "active"
                conversation.ended_at = None
        db.add(row)
        try:
            await db.commit()
        except IntegrityError as exc:
            await db.rollback()
            if (
                idempotency_key is not None
                and _integrity_constraint_name(exc) == IDEMPOTENCY_CONSTRAINT
            ):
                raise ResponsesAPIError(
                    "Idempotency-Key is already in use",
                    code="idempotency_key_in_use",
                    param="Idempotency-Key",
                    status_code=409,
                ) from exc
            raise
        await db.refresh(row)

    return PreparedResponse(
        row=row,
        config=config,
        history_messages=history_messages,
        current_messages=items_to_messages(request.input_items),
        runtime_state=runtime_state,
        client_tool_results=client_tool_results,
    )


async def get_response(
    response_id: str,
    *,
    api_key_id: str,
    include_deleted: bool = False,
) -> OpenAIResponse:
    async with SessionLocal() as db:
        row = await db.get(OpenAIResponse, response_id)
        if (
            row is None
            or row.api_key_id != api_key_id
            or (row.deleted_at is not None and not include_deleted)
        ):
            raise ResponsesAPIError(
                "Response not found", code="response_not_found", status_code=404
            )
        return row


async def get_idempotent_response(
    *,
    api_key_id: str,
    idempotency_key: str | None,
    request_hash: str,
) -> OpenAIResponse | None:
    """Return a terminal replay or reject conflicting reuse of a request key."""

    if idempotency_key is None:
        return None
    async with SessionLocal() as db:
        row = (
            await db.execute(
                select(OpenAIResponse).where(
                    OpenAIResponse.api_key_id == api_key_id,
                    OpenAIResponse.idempotency_key == idempotency_key,
                )
            )
        ).scalar_one_or_none()
        if row is None:
            return None
        if row.request_hash != request_hash:
            raise ResponsesAPIError(
                "Idempotency-Key was already used with a different request",
                code="idempotency_key_conflict",
                param="Idempotency-Key",
                status_code=409,
            )
        if row.deleted_at is not None:
            raise ResponsesAPIError(
                "The response created by this Idempotency-Key was deleted",
                code="idempotency_response_deleted",
                param="Idempotency-Key",
                status_code=409,
            )
        if row.status == "in_progress":
            raise ResponsesAPIError(
                "A request with this Idempotency-Key is still in progress",
                code="idempotency_key_in_use",
                param="Idempotency-Key",
                status_code=409,
            )
        return row


async def delete_response(
    response_id: str,
    *,
    api_key_id: str,
) -> OpenAIResponse:
    """Create a tombstone while retaining state needed by existing branches."""

    async with SessionLocal() as db:
        row = await db.get(
            OpenAIResponse,
            response_id,
            with_for_update=True,
        )
        if (
            row is None
            or row.api_key_id != api_key_id
            or row.deleted_at is not None
        ):
            raise ResponsesAPIError(
                "Response not found", code="response_not_found", status_code=404
            )
        if row.status == "in_progress":
            raise ResponsesAPIError(
                "An in-progress response cannot be deleted",
                code="response_not_ready",
                status_code=409,
            )
        row.deleted_at = datetime.now(UTC)
        await db.commit()
        await db.refresh(row)
        return row


async def finish_response(
    response_id: str,
    *,
    status: str,
    output_items: list[dict[str, Any]],
    context_messages: list[dict[str, Any]] | None = None,
    dynamic_variables: dict[str, Any] | None = None,
    runtime_state: dict[str, Any] | None = None,
    error: dict[str, Any] | None = None,
    usage: dict[str, Any] | None = None,
) -> OpenAIResponse:
    now = datetime.now(UTC)
    async with SessionLocal() as db:
        row = await db.get(OpenAIResponse, response_id)
        if row is None:
            raise RuntimeError(f"Response disappeared while running: {response_id}")
        row.status = status
        row.output_items = output_items
        row.context_messages = context_messages or []
        if dynamic_variables is not None:
            row.dynamic_variables = dynamic_variables
        if runtime_state is not None:
            row.runtime_state = runtime_state
        row.error = error
        row.usage = usage
        row.completed_at = now

        conversation = (
            await db.execute(
                select(ConversationSession)
                .where(ConversationSession.id == row.conversation_id)
                .with_for_update()
            )
        ).scalar_one_or_none()
        if conversation:
            max_sequence = await db.scalar(
                select(func.max(ConversationMessage.sequence)).where(
                    ConversationMessage.session_id == row.conversation_id
                )
            )
            sequence = int(max_sequence or 0)
            for item in row.input_items or []:
                if item.get("type") != "message" or item.get("role") != "user":
                    continue
                text = _content_text(item)
                if not text:
                    continue
                sequence += 1
                db.add(
                    ConversationMessage(
                        id=new_id("msg")[:40],
                        session_id=row.conversation_id,
                        sequence=sequence,
                        role="user",
                        content_type="text",
                        content=text,
                        extra={"response_id": row.id},
                    )
                )
            for item in output_items:
                if item.get("type") == "message":
                    text = _content_text(item)
                elif item.get("type") == "function_call":
                    text = f"[function_call] {item.get('name') or ''}"
                else:
                    continue
                if not text:
                    continue
                sequence += 1
                db.add(
                    ConversationMessage(
                        id=new_id("msg")[:40],
                        session_id=row.conversation_id,
                        sequence=sequence,
                        role="assistant",
                        content_type="text",
                        content=text,
                        extra={"response_id": row.id, "item": item},
                    )
                )
            conversation.message_count = sequence
            conversation.status = "completed" if status == "completed" else status
            conversation.ended_at = now
        await db.commit()
        await db.refresh(row)
        return row


async def recover_interrupted_responses() -> None:
    """Mark generations interrupted by a process restart as cancelled."""
    now = datetime.now(UTC)
    async with SessionLocal() as db:
        conversation_ids = select(OpenAIResponse.conversation_id).where(
            OpenAIResponse.status == "in_progress"
        )
        await db.execute(
            update(ConversationSession)
            .where(ConversationSession.id.in_(conversation_ids))
            .values(status="cancelled", ended_at=now)
        )
        await db.execute(
            update(OpenAIResponse)
            .where(OpenAIResponse.status == "in_progress")
            .values(
                status="cancelled",
                error={
                    "code": "server_restarted",
                    "message": "Response was interrupted by a server restart",
                },
                completed_at=now,
            )
        )
        await db.commit()