X-Financial/server/src/app/api/v1/endpoints/steward.py

from __future__ import annotations

import asyncio
import json
from collections.abc import AsyncIterator
from typing import Annotated, Any

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from app.api.deps import get_db
from app.schemas.common import ErrorResponse
from app.schemas.steward import (
    StewardPlanRequest,
    StewardPlanResponse,
    StewardRuntimeDecisionRequest,
    StewardRuntimeDecisionResponse,
    StewardSlotDecisionRequest,
    StewardSlotDecisionResponse,
    StewardThinkingEvent,
)
from app.services.agent_conversations import AgentConversationService
from app.services.runtime_chat import RuntimeChatService
from app.services.steward_flow_state import StewardFlowStateService
from app.services.steward_intent_agent import StewardIntentAgent
from app.services.steward_planner import StewardPlannerService
from app.services.steward_runtime_decision_agent import StewardRuntimeDecisionAgent
from app.services.steward_slot_decision_agent import StewardSlotDecisionAgent

router = APIRouter(prefix="/steward")
DbSession = Annotated[Session, Depends(get_db)]


@router.post(
    "/plans",
    response_model=StewardPlanResponse,
    summary="生成小财管家任务计划",
    description="把首页自然语言和附件元信息拆解为可确认、可追踪、可分派的财务任务计划。",
    responses={
        status.HTTP_400_BAD_REQUEST: {
            "model": ErrorResponse,
            "description": "请求缺少任务描述，无法生成小财管家计划。",
        }
    },
)
def create_steward_plan(payload: StewardPlanRequest, db: DbSession) -> StewardPlanResponse:
    try:
        plan = _build_steward_planner(db).build_plan(payload)
        return _attach_conversation_state(db, payload, plan)
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc


@router.post(
    "/slot-decisions",
    response_model=StewardSlotDecisionResponse,
    summary="判断小财管家当前任务字段缺口",
    description="结合当前任务、本体字段和用户上下文，使用 function calling 判断下一步应先追问用户还是展示核对结果。",
)
def create_steward_slot_decision(
    payload: StewardSlotDecisionRequest,
    db: DbSession,
) -> StewardSlotDecisionResponse:
    return StewardSlotDecisionAgent(RuntimeChatService(db)).decide(payload)


@router.post(
    "/runtime-decisions",
    response_model=StewardRuntimeDecisionResponse,
    summary="判断小财管家运行时下一步动作",
    description="结合任务队列、当前结构化结果和用户输入，使用 function calling 判断应提交当前单据、继续下一任务、补字段或重新规划。",
)
def create_steward_runtime_decision(
    payload: StewardRuntimeDecisionRequest,
    db: DbSession,
) -> StewardRuntimeDecisionResponse:
    hydrated_payload = _hydrate_runtime_decision_payload(db, payload)
    decision = StewardRuntimeDecisionAgent(RuntimeChatService(db)).decide(hydrated_payload)
    return _attach_runtime_conversation_state(db, hydrated_payload, decision)


@router.post(
    "/plans/stream",
    summary="流式生成小财管家任务计划",
    description="以 NDJSON 逐条返回小财管家的过程摘要事件，最后返回完整任务计划。",
)
async def stream_steward_plan(payload: StewardPlanRequest, db: DbSession) -> StreamingResponse:
    return StreamingResponse(
        _iter_steward_plan_events(payload, _build_steward_planner(db), db),
        media_type="application/x-ndjson",
    )


async def _iter_steward_plan_events(
    payload: StewardPlanRequest,
    planner: StewardPlannerService,
    db: Session,
) -> AsyncIterator[str]:
    yield _encode_stream_event(
        "thinking",
        StewardThinkingEvent(
            event_id="intent_agent_stream_start",
            stage="stream_start",
            title="读取用户输入",
            content="我先判断这句话里是否同时包含申请、报销或附件归集事项，再决定处理顺序。",
            status="running",
        ).model_dump(mode="json"),
    )
    await asyncio.sleep(0)

    try:
        plan = planner.build_plan(payload)
        plan = _attach_conversation_state(db, payload, plan)
    except ValueError as exc:
        yield _encode_stream_event("error", {"message": str(exc)})
        return

    for event in plan.thinking_events:
        yield _encode_stream_event("thinking", event.model_dump(mode="json"))
        await asyncio.sleep(0.6)

    yield _encode_stream_event("plan", plan.model_dump(mode="json"))


def _encode_stream_event(event: str, data: dict[str, Any]) -> str:
    return json.dumps({"event": event, "data": data}, ensure_ascii=False) + "\n"


def _build_steward_planner(db: Session) -> StewardPlannerService:
    return StewardPlannerService(
        intent_agent=StewardIntentAgent(RuntimeChatService(db)),
    )


def _attach_conversation_state(
    db: Session,
    payload: StewardPlanRequest,
    plan: StewardPlanResponse,
) -> StewardPlanResponse:
    context_json = dict(payload.context_json or {})
    context_json["session_type"] = str(context_json.get("session_type") or "steward").strip() or "steward"
    conversation_service = AgentConversationService(db)
    conversation = conversation_service.get_or_create_conversation(
        conversation_id=_resolve_conversation_id(context_json),
        user_id=payload.user_id,
        source="user_message",
        context_json=context_json,
    )
    current_state = _resolve_current_steward_state(conversation.state_json, context_json)
    steward_state = StewardFlowStateService().merge_plan(current_state, plan)
    conversation = conversation_service.update_state(
        conversation_id=conversation.conversation_id,
        run_id=None,
        scenario="steward",
        intent="plan",
        context_json={
            **context_json,
            "steward_state": steward_state,
        },
    ) or conversation
    conversation_service.append_message(
        conversation_id=conversation.conversation_id,
        role="user",
        content=payload.message,
        message_json={"source": "steward_plan_request"},
    )
    conversation_service.append_message(
        conversation_id=conversation.conversation_id,
        role="assistant",
        content=plan.summary,
        message_json={
            "source": "steward_plan_response",
            "plan_id": plan.plan_id,
            "steward_state": steward_state,
        },
    )
    return plan.model_copy(
        update={
            "conversation_id": conversation.conversation_id,
            "steward_state": steward_state,
        }
    )


def _attach_runtime_conversation_state(
    db: Session,
    payload: StewardRuntimeDecisionRequest,
    decision: StewardRuntimeDecisionResponse,
) -> StewardRuntimeDecisionResponse:
    steward_state = decision.steward_state
    if not isinstance(steward_state, dict) or not steward_state:
        return decision
    context_json = dict(payload.context_json or {})
    conversation_id = _resolve_conversation_id(context_json)
    if not conversation_id:
        return decision

    conversation_service = AgentConversationService(db)
    conversation_service.update_state(
        conversation_id=conversation_id,
        run_id=None,
        scenario="steward",
        intent="runtime_decision",
        context_json={
            **context_json,
            "steward_state": steward_state,
        },
    )
    return decision


def _hydrate_runtime_decision_payload(
    db: Session,
    payload: StewardRuntimeDecisionRequest,
) -> StewardRuntimeDecisionRequest:
    context_json = dict(payload.context_json or {})
    runtime_state = dict(payload.runtime_state or {})
    if isinstance(runtime_state.get("steward_state"), dict) and runtime_state["steward_state"]:
        return payload
    if isinstance(context_json.get("steward_state"), dict) and context_json["steward_state"]:
        return payload

    conversation_id = _resolve_conversation_id(context_json)
    if not conversation_id:
        return payload
    conversation = AgentConversationService(db).get_conversation(conversation_id)
    stored_state = conversation.state_json.get("steward_state") if conversation and isinstance(conversation.state_json, dict) else None
    if not isinstance(stored_state, dict) or not stored_state:
        return payload

    runtime_state["steward_state"] = stored_state
    conversation_state = dict(context_json.get("conversation_state") or {})
    conversation_state["steward_state"] = stored_state
    context_json["conversation_state"] = conversation_state
    return payload.model_copy(
        update={
            "runtime_state": runtime_state,
            "context_json": context_json,
        }
    )


def _resolve_conversation_id(context_json: dict[str, Any]) -> str | None:
    return str(
        context_json.get("conversation_id")
        or context_json.get("conversationId")
        or ""
    ).strip() or None


def _resolve_current_steward_state(
    conversation_state: dict[str, Any] | None,
    context_json: dict[str, Any],
) -> dict[str, Any]:
    state_json = conversation_state if isinstance(conversation_state, dict) else {}
    stored_state = state_json.get("steward_state")
    if isinstance(stored_state, dict) and stored_state:
        return stored_state
    incoming_state = context_json.get("steward_state") or context_json.get("stewardState")
    return incoming_state if isinstance(incoming_state, dict) else {}