"""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 CONTINUABLE_STATUSES = {"completed"} @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 _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 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: raise ResponsesAPIError( "Response not found", code="response_not_found", status_code=404 ) chain.append(row) current_id = row.previous_response_id 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 = { str(item.get("call_id") or ""): item for row in chain for item in row.output_items or [] if item.get("type") == "function_call" } results: list[dict[str, Any]] = [] for item in current_items: if item.get("type") != "function_call_output": continue call = calls.get(str(item.get("call_id") or "")) or {} 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": item.get("call_id"), "name": call.get("name"), "result": value, } ) return results async def prepare_response( request: ResponseRequest, *, api_key_id: str, ) -> 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: requested_id = assistant_id_from_model(request.model) 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, ) 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) await db.commit() 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, ) -> 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: raise ResponsesAPIError( "Response not found", code="response_not_found", status_code=404 ) 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, ) -> 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.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()