from __future__ import annotations import asyncio from collections.abc import Sequence from datetime import UTC, datetime from time import perf_counter from uuid import uuid4 from fastapi import HTTPException, status from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.core.enums import ( PlaygroundMessageKind, PlaygroundMessageRole, PlaygroundMessageStatus, ) from app.core.logging import get_logger from app.db.session import async_session_factory from app.models.playground_message import PlaygroundMessage from app.models.playground_session import PlaygroundSession from app.schemas.ai import ( PlaygroundMessageEditRequest, PlaygroundMessageActionResponse, PlaygroundMessageCreateRequest, PlaygroundMessageRecord, PlaygroundMessageResendRequest, PlaygroundMessageStopRequest, PlaygroundSessionState, PlaygroundSessionUpsertRequest, PlaygroundThreadResponse, SituationalAnalysisRequest, ) from app.services.ai_client import AIProviderClient from app.services.business_logs import emit_business_log, exception_context from app.services.playground_session_store import _to_response as session_to_response from app.services.playground_session_store import upsert_playground_session logger = get_logger(__name__, service="ai") STREAM_CHUNK_SIZE = 24 STREAM_INTERVAL_SECONDS = 0.08 THINKING_PREVIEW_SECONDS = 2.6 ORPHANED_RUN_MESSAGE = "后台生成任务已中断，请点击上一条用户消息的重试按钮重新生成。" ACTIVE_MESSAGE_STATUSES = frozenset( { PlaygroundMessageStatus.PENDING.value, PlaygroundMessageStatus.THINKING.value, PlaygroundMessageStatus.ANSWERING.value, } ) class _ActiveRun: def __init__(self, task: asyncio.Task[None]) -> None: self.task = task self.stop_requested = asyncio.Event() _ACTIVE_RUNS: dict[str, _ActiveRun] = {} async def _get_session_by_key( db: AsyncSession, *, user_id: int, session_key: str, ) -> PlaygroundSession | None: result = await db.execute( select(PlaygroundSession).where( PlaygroundSession.user_id == user_id, PlaygroundSession.session_key == session_key, ) ) return result.scalar_one_or_none() async def _require_session( db: AsyncSession, *, user_id: int, session_key: str, ) -> PlaygroundSession: session = await _get_session_by_key(db, user_id=user_id, session_key=session_key) if session is None: raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Playground session not found") return session async def _require_visible_message( db: AsyncSession, *, user_id: int, public_id: str, role: str | None = None, ) -> PlaygroundMessage: conditions = [ PlaygroundMessage.user_id == user_id, PlaygroundMessage.public_id == public_id, PlaygroundMessage.is_visible.is_(True), ] if role is not None: conditions.append(PlaygroundMessage.role == role) result = await db.execute(select(PlaygroundMessage).where(*conditions)) message = result.scalar_one_or_none() if message is None: detail = ( "User message not found" if role == PlaygroundMessageRole.USER.value else "Playground message not found" ) raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=detail) return message def _message_to_record(message: PlaygroundMessage, parent_public_id: str | None = None) -> PlaygroundMessageRecord: return PlaygroundMessageRecord( id=message.public_id, role=message.role, kind=message.kind, status=message.status, title=message.title, content=message.content or "", thinking_content=message.thinking_content or "", meta=list(message.meta or []), markdown=message.role != PlaygroundMessageRole.SYSTEM.value, provider=message.provider, model=message.model, request_id=message.request_id, raw_response=dict(message.raw_response or {}), content_blocks=list(message.content_blocks or []), text_blocks=list(message.text_blocks or []), thinking_blocks=list(message.thinking_blocks or []), parent_message_id=parent_public_id, created_at=message.created_at.isoformat(), updated_at=message.updated_at.isoformat(), ) async def _ensure_session( db: AsyncSession, *, user_id: int, session_key: str, title: str, state: PlaygroundSessionState | None = None, ) -> PlaygroundSession: result = await db.execute( select(PlaygroundSession).where( PlaygroundSession.user_id == user_id, PlaygroundSession.session_key == session_key, ) ) session = result.scalar_one_or_none() if session is not None: if title: session.title = title[:200] if state is not None: session.state = state.model_dump(mode="json") await db.flush() await db.refresh(session) return session payload = PlaygroundSessionUpsertRequest( session_key=session_key, title=title[:200], state=state or PlaygroundSessionState(title=title[:200]), ) await upsert_playground_session(db, user_id=user_id, payload=payload) result = await db.execute( select(PlaygroundSession).where( PlaygroundSession.user_id == user_id, PlaygroundSession.session_key == session_key, ) ) return result.scalar_one() async def _list_visible_messages( db: AsyncSession, *, session_id: int, ) -> list[PlaygroundMessage]: result = await db.execute( select(PlaygroundMessage) .where( PlaygroundMessage.session_id == session_id, PlaygroundMessage.is_visible.is_(True), ) .order_by(PlaygroundMessage.sort_order.asc(), PlaygroundMessage.id.asc()) ) return list(result.scalars().all()) async def _build_thread_response( db: AsyncSession, *, session: PlaygroundSession, ) -> PlaygroundThreadResponse: messages = await _list_visible_messages(db, session_id=session.id) messages = await _reconcile_orphaned_active_messages(db, messages) id_map = {item.id: item.public_id for item in messages} return PlaygroundThreadResponse( session=session_to_response(session), messages=[_message_to_record(item, id_map.get(item.parent_message_id)) for item in messages], ) async def _reconcile_orphaned_active_messages( db: AsyncSession, messages: list[PlaygroundMessage], ) -> list[PlaygroundMessage]: changed = False for item in messages: if item.status not in ACTIVE_MESSAGE_STATUSES: continue if item.public_id in _ACTIVE_RUNS: continue item.status = PlaygroundMessageStatus.ERROR.value item.content = item.content or ORPHANED_RUN_MESSAGE orphan_meta = "错误: 后台任务已中断" if orphan_meta not in (item.meta or []): item.meta = [*(item.meta or []), orphan_meta] changed = True if changed: await db.flush() await db.commit() for item in messages: await db.refresh(item) return messages async def get_thread( db: AsyncSession, *, user_id: int, session_key: str, ) -> PlaygroundThreadResponse | None: session = await _get_session_by_key(db, user_id=user_id, session_key=session_key) if session is None: return None return await _build_thread_response(db, session=session) async def _build_action_response( db: AsyncSession, *, session: PlaygroundSession, active_message_id: str | None = None, ) -> PlaygroundMessageActionResponse: thread = await _build_thread_response(db, session=session) return PlaygroundMessageActionResponse( session=thread.session, messages=thread.messages, active_message_id=active_message_id, ) def _collect_constraints(raw_constraints: str) -> list[str]: return [item.strip() for item in raw_constraints.split("\n") if item.strip()] async def _next_sort_order(db: AsyncSession, session_id: int) -> int: result = await db.execute( select(func.max(PlaygroundMessage.sort_order)).where(PlaygroundMessage.session_id == session_id) ) current = result.scalar_one_or_none() return int(current or 0) async def _set_session_state( db: AsyncSession, *, session: PlaygroundSession, payload: PlaygroundMessageCreateRequest, ) -> PlaygroundSession: session.state = PlaygroundSessionState( messages=[], selectedPresetKey=payload.selected_preset_key, title=payload.title, objective=payload.objective, constraints=payload.constraints, inputValue="", analysis=None, latestAnalysisMessageId=None, analysisMeta={}, helpExpanded=payload.help_expanded, ).model_dump(mode="json") session.title = payload.title[:200] await db.flush() await db.refresh(session) return session def _spawn_assistant_run( *, user_id: int, session_id: int, session_key: str, user_message_id: int, assistant_message_id: int, assistant_public_id: str, payload: PlaygroundMessageCreateRequest, provider_client: AIProviderClient, ) -> None: task = asyncio.create_task( _run_assistant_message( user_id=user_id, session_id=session_id, session_key=session_key, user_message_id=user_message_id, assistant_message_id=assistant_message_id, payload=payload, provider_client=provider_client, ) ) _ACTIVE_RUNS[assistant_public_id] = _ActiveRun(task) async def create_turn( db: AsyncSession, *, user_id: int, payload: PlaygroundMessageCreateRequest, provider_client: AIProviderClient, ) -> PlaygroundMessageActionResponse: session = await _ensure_session( db, user_id=user_id, session_key=payload.session_key, title=payload.title, state=PlaygroundSessionState( selectedPresetKey=payload.selected_preset_key, title=payload.title, objective=payload.objective, constraints=payload.constraints, inputValue="", helpExpanded=payload.help_expanded, ), ) session = await _set_session_state(db, session=session, payload=payload) base_order = await _next_sort_order(db, session.id) user_message = PlaygroundMessage( public_id=uuid4().hex, session_id=session.id, user_id=user_id, role=PlaygroundMessageRole.USER.value, kind=PlaygroundMessageKind.MESSAGE.value, status=PlaygroundMessageStatus.DONE.value, title=payload.selected_preset_key, content=payload.input, meta=[payload.title], sort_order=base_order + 10, ) assistant_message = PlaygroundMessage( public_id=uuid4().hex, session_id=session.id, user_id=user_id, parent_message_id=None, role=PlaygroundMessageRole.ASSISTANT.value, kind=PlaygroundMessageKind.THINKING.value, status=PlaygroundMessageStatus.PENDING.value, title="AI 回应", content="", thinking_content="", meta=[], sort_order=base_order + 20, ) db.add(user_message) await db.flush() assistant_message.parent_message_id = user_message.id db.add(assistant_message) await db.flush() await db.refresh(user_message) await db.refresh(assistant_message) await db.commit() await db.refresh(session) await db.refresh(user_message) await db.refresh(assistant_message) _spawn_assistant_run( user_id=user_id, session_id=session.id, session_key=payload.session_key, user_message_id=user_message.id, assistant_message_id=assistant_message.id, assistant_public_id=assistant_message.public_id, payload=payload, provider_client=provider_client, ) return await _build_action_response( db, session=session, active_message_id=assistant_message.public_id, ) async def _create_assistant_retry_turn( db: AsyncSession, *, user_id: int, session: PlaygroundSession, user_message: PlaygroundMessage, payload: PlaygroundMessageCreateRequest, provider_client: AIProviderClient, ) -> PlaygroundMessageActionResponse: base_order = await _next_sort_order(db, session.id) assistant_message = PlaygroundMessage( public_id=uuid4().hex, session_id=session.id, user_id=user_id, parent_message_id=user_message.id, role=PlaygroundMessageRole.ASSISTANT.value, kind=PlaygroundMessageKind.THINKING.value, status=PlaygroundMessageStatus.PENDING.value, title="AI 回应", content="", thinking_content="", meta=[], sort_order=base_order + 10, ) db.add(assistant_message) await db.flush() await db.commit() await db.refresh(session) await db.refresh(assistant_message) _spawn_assistant_run( user_id=user_id, session_id=session.id, session_key=payload.session_key, user_message_id=user_message.id, assistant_message_id=assistant_message.id, assistant_public_id=assistant_message.public_id, payload=payload, provider_client=provider_client, ) return await _build_action_response( db, session=session, active_message_id=assistant_message.public_id, ) async def stop_message( db: AsyncSession, *, user_id: int, payload: PlaygroundMessageStopRequest, ) -> PlaygroundMessageActionResponse: session = await _require_session(db, user_id=user_id, session_key=payload.session_key) message = await _require_visible_message(db, user_id=user_id, public_id=payload.message_id) if message.status not in ACTIVE_MESSAGE_STATUSES: return await _build_action_response(db, session=session) active_run = _ACTIVE_RUNS.get(message.public_id) if active_run is not None: active_run.stop_requested.set() active_run.task.cancel() message.status = PlaygroundMessageStatus.STOPPED.value if "已手动停止生成" not in (message.meta or []): message.meta = [*(message.meta or []), "已手动停止生成"] await db.flush() await db.commit() await db.refresh(message) return await _build_action_response(db, session=session) async def resend_turn( db: AsyncSession, *, user_id: int, payload: PlaygroundMessageResendRequest, provider_client: AIProviderClient, ) -> PlaygroundMessageActionResponse: session = await _require_session(db, user_id=user_id, session_key=payload.session_key) user_message = await _require_visible_message( db, user_id=user_id, public_id=payload.user_message_id, role=PlaygroundMessageRole.USER.value, ) later_messages = await db.execute( select(PlaygroundMessage).where( PlaygroundMessage.session_id == session.id, PlaygroundMessage.sort_order > user_message.sort_order, PlaygroundMessage.is_visible.is_(True), ) ) for item in later_messages.scalars().all(): item.is_visible = False if item.status in ACTIVE_MESSAGE_STATUSES: active_run = _ACTIVE_RUNS.get(item.public_id) if active_run is not None: active_run.stop_requested.set() active_run.task.cancel() await db.flush() await db.commit() session_state = PlaygroundSessionState.model_validate(session.state or {}) create_payload = PlaygroundMessageCreateRequest( session_key=payload.session_key, title=session_state.title or session.title, objective=session_state.objective or "继续当前对话", constraints=session_state.constraints or "", input=user_message.content, selected_preset_key=session_state.selectedPresetKey or "bgp-brief", help_expanded=session_state.helpExpanded, ) return await _create_assistant_retry_turn( db, user_id=user_id, session=session, user_message=user_message, payload=create_payload, provider_client=provider_client, ) async def edit_user_message( db: AsyncSession, *, user_id: int, payload: PlaygroundMessageEditRequest, ) -> PlaygroundMessageActionResponse: session = await _require_session(db, user_id=user_id, session_key=payload.session_key) user_message = await _require_visible_message( db, user_id=user_id, public_id=payload.user_message_id, role=PlaygroundMessageRole.USER.value, ) user_message.content = payload.content.strip() await db.flush() await db.commit() await db.refresh(user_message) return await _build_action_response(db, session=session) async def _append_meta_if_missing(db: AsyncSession, message_id: int, meta_line: str) -> None: result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == message_id)) message = result.scalar_one_or_none() if message is None: return if meta_line not in (message.meta or []): message.meta = [*(message.meta or []), meta_line] await db.flush() async def _should_stop(message_public_id: str) -> bool: active_run = _ACTIVE_RUNS.get(message_public_id) return active_run.stop_requested.is_set() if active_run is not None else False async def _mark_message_state( db: AsyncSession, *, message_id: int, **updates, ) -> PlaygroundMessage: result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == message_id)) message = result.scalar_one() for key, value in updates.items(): setattr(message, key, value) await db.flush() await db.refresh(message) return message def _build_conversation_history(messages: Sequence[PlaygroundMessage], current_user_message_id: int) -> list[dict]: history: list[dict] = [] for item in messages: if item.id >= current_user_message_id: break if item.role == PlaygroundMessageRole.SYSTEM.value: continue history.append( { "role": item.role, "kind": item.kind or PlaygroundMessageKind.MESSAGE.value, "title": item.title, "content": item.content or "", } ) return history[-8:] def _format_run_exception(exc: Exception) -> str: if isinstance(exc, HTTPException): detail = exc.detail if isinstance(detail, str): return detail return str(detail) return str(exc) or type(exc).__name__ async def _run_assistant_message( *, user_id: int, session_id: int, session_key: str, user_message_id: int, assistant_message_id: int, payload: PlaygroundMessageCreateRequest, provider_client: AIProviderClient, ) -> None: request_id = str(uuid4()) started_at = perf_counter() assistant_public_id: str | None = None try: async with async_session_factory() as db: session = await db.get(PlaygroundSession, session_id) user_message = await db.get(PlaygroundMessage, user_message_id) assistant_message = await db.get(PlaygroundMessage, assistant_message_id) if session is None or user_message is None or assistant_message is None: return assistant_public_id = assistant_message.public_id visible_messages = await _list_visible_messages(db, session_id=session_id) conversation_history = _build_conversation_history(visible_messages, user_message_id) request_payload = SituationalAnalysisRequest( title=payload.title, objective=payload.objective, observations=[item.strip() for item in payload.input.split("\n") if item.strip()], constraints=_collect_constraints(payload.constraints), context={ "source": "playground", "preset": payload.selected_preset_key, "conversation_history": conversation_history, "history_size": len(conversation_history), }, thinking={"type": "enabled"}, ) await emit_business_log( logger, event="ai.playground.run.start", message="Playground AI run started", category="ai", service="ai", module=__name__, request_id=request_id, user_id=user_id, context={ "session_id": session_id, "session_key": session_key, "user_message_id": user_message_id, "assistant_message_id": assistant_message_id, "preset": payload.selected_preset_key, }, ) analysis = await provider_client.analyze(request_payload, request_id=request_id) await emit_business_log( logger, event="ai.playground.run.success", message="Playground AI run completed", category="ai", service="ai", module=__name__, request_id=request_id, user_id=user_id, context={ "session_id": session_id, "session_key": session_key, "provider": analysis.provider, "model": analysis.model, "content_block_count": len(analysis.content_blocks or []), "thinking_block_count": len(analysis.thinking_blocks or []), }, ) async with async_session_factory() as db: assistant_message = await _mark_message_state( db, message_id=assistant_message_id, status=( PlaygroundMessageStatus.THINKING.value if analysis.thinking_blocks else PlaygroundMessageStatus.ANSWERING.value ), title=f"{analysis.provider} / {analysis.model}", provider=analysis.provider, model=analysis.model, request_id=request_id, raw_response=analysis.raw_response, content_blocks=[item.model_dump(mode="json") for item in analysis.content_blocks], text_blocks=analysis.text_blocks, thinking_blocks=analysis.thinking_blocks, thinking_content="\n\n".join(analysis.thinking_blocks).strip(), ) session = await db.get(PlaygroundSession, session_id) if session is not None: session_state = PlaygroundSessionState.model_validate(session.state or {}) session.state = session_state.model_copy( update={ "latestAnalysisMessageId": assistant_message.public_id, "analysis": analysis.model_dump(mode="json"), } ).model_dump(mode="json") await db.flush() await db.commit() if assistant_public_id and analysis.thinking_blocks: await asyncio.sleep(THINKING_PREVIEW_SECONDS) if await _should_stop(assistant_public_id): return content = analysis.content or "" cursor = 0 while cursor < len(content): if assistant_public_id and await _should_stop(assistant_public_id): return cursor = min(len(content), cursor + STREAM_CHUNK_SIZE) async with async_session_factory() as db: await _mark_message_state( db, message_id=assistant_message_id, status=PlaygroundMessageStatus.ANSWERING.value, content=content[:cursor], ) await db.commit() await asyncio.sleep(STREAM_INTERVAL_SECONDS) duration_ms = round((perf_counter() - started_at) * 1000) async with async_session_factory() as db: assistant_message = await _mark_message_state( db, message_id=assistant_message_id, status=PlaygroundMessageStatus.DONE.value, content=content, meta=[ f"Request ID: {request_id}", f"耗时: {duration_ms} ms", f"完成时间: {datetime.now(UTC).astimezone().isoformat()}", ], ) session = await db.get(PlaygroundSession, session_id) if session is not None: session_state = PlaygroundSessionState.model_validate(session.state or {}) session.state = session_state.model_copy( update={ "latestAnalysisMessageId": assistant_message.public_id, "analysis": analysis.model_dump(mode="json"), "analysisMeta": { "requestId": request_id, "durationMs": duration_ms, "completedAt": datetime.now().isoformat(), }, } ).model_dump(mode="json") await db.flush() await db.commit() except asyncio.CancelledError: await emit_business_log( logger, event="ai.playground.run.cancelled", message="Playground AI run cancelled", category="ai", level="warning", service="ai", module=__name__, request_id=request_id, user_id=user_id, context={ "session_id": session_id, "session_key": session_key, "assistant_message_id": assistant_message_id, "duration_ms": round((perf_counter() - started_at) * 1000), }, ) async with async_session_factory() as db: result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == assistant_message_id)) message = result.scalar_one_or_none() if message is not None and message.status in ACTIVE_MESSAGE_STATUSES: message.status = PlaygroundMessageStatus.STOPPED.value if "已手动停止生成" not in (message.meta or []): message.meta = [*(message.meta or []), "已手动停止生成"] await db.flush() await db.commit() raise except Exception as exc: await emit_business_log( logger, event="ai.playground.run.failed", message="Playground AI run failed", category="ai", level="error", service="ai", module=__name__, request_id=request_id, user_id=user_id, context=exception_context( exc, { "session_id": session_id, "session_key": session_key, "assistant_message_id": assistant_message_id, "duration_ms": round((perf_counter() - started_at) * 1000), }, ), ) error_message = _format_run_exception(exc) async with async_session_factory() as db: result = await db.execute(select(PlaygroundMessage).where(PlaygroundMessage.id == assistant_message_id)) message = result.scalar_one_or_none() if message is not None: message.status = PlaygroundMessageStatus.ERROR.value message.content = message.content or f"分析失败：{error_message}" message.meta = [ *(message.meta or []), f"Request ID: {request_id}", f"错误: {error_message}", ] await db.flush() await db.commit() finally: if assistant_public_id: _ACTIVE_RUNS.pop(assistant_public_id, None)