X-Financial/server/src/app/services/user_agent.py

from __future__ import annotations

import json
import re
from datetime import UTC, datetime, timedelta
from decimal import Decimal, InvalidOperation
from typing import Any

from sqlalchemy import or_, select
from sqlalchemy.orm import Session, selectinload

from app.api.deps import CurrentUserContext
from app.core.agent_enums import AgentAssetStatus, AgentAssetType
from app.models.employee import Employee
from app.models.financial_record import ExpenseClaim
from app.schemas.agent_asset import AgentAssetListItem
from app.schemas.reimbursement import TravelReimbursementCalculatorRequest
from app.schemas.user_agent import (
    UserAgentCitation,
    UserAgentDraftPayload,
    UserAgentExpenseQueryRecord,
    UserAgentQueryPayload,
    UserAgentQueryStatusGroup,
    UserAgentReviewAction,
    UserAgentReviewEditField,
    UserAgentReviewClaimGroup,
    UserAgentReviewDocumentCard,
    UserAgentReviewDocumentField,
    UserAgentReviewPayload,
    UserAgentReviewRiskBrief,
    UserAgentReviewSlotCard,
    UserAgentRequest,
    UserAgentResponse,
    UserAgentSuggestedAction,
)
from app.services.agent_assets import AgentAssetService
from app.services.agent_foundation import AgentFoundationService
from app.services.expense_claims import ExpenseClaimService
from app.services.expense_rule_runtime import ExpenseRuleRuntimeService, RuntimeTravelPolicy, resolve_document_type_label
from app.services.risk_ontology_bridge import resolve_rule_codes_for_risk_check
from app.services.runtime_chat import RuntimeChatService
from app.services.travel_reimbursement_calculator import TravelReimbursementCalculatorService
from app.services.user_agent_documents import UserAgentDocumentService
from app.services.user_agent_knowledge import UserAgentKnowledgeMixin

from app.services.user_agent_constants import *
from app.services.user_agent_response import UserAgentResponseMixin
from app.services.user_agent_review_core import UserAgentReviewCoreMixin
from app.services.user_agent_review_messages import UserAgentReviewMessageMixin
from app.services.user_agent_review_profile import UserAgentReviewProfileMixin
from app.services.user_agent_review_slots import UserAgentReviewSlotMixin
from app.services.user_agent_review_travel_policy import UserAgentReviewTravelPolicyMixin
from app.services.user_agent_review_travel_receipts import UserAgentReviewTravelReceiptMixin


class UserAgentService(
    UserAgentResponseMixin,
    UserAgentKnowledgeMixin,
    UserAgentReviewCoreMixin,
    UserAgentReviewTravelPolicyMixin,
    UserAgentReviewTravelReceiptMixin,
    UserAgentReviewMessageMixin,
    UserAgentReviewProfileMixin,
    UserAgentReviewSlotMixin,
):
    def __init__(self, db: Session) -> None:
        self.db = db
        self.asset_service = AgentAssetService(db)
        self.runtime_chat_service = RuntimeChatService(db)
        self._document_service = UserAgentDocumentService(group_scene_labels=GROUP_SCENE_LABELS)

    def respond(self, payload: UserAgentRequest) -> UserAgentResponse:
        AgentFoundationService(self.db).ensure_foundation_ready()
        citations = self._build_citations(payload)
        suggested_actions = self._build_suggested_actions(payload)
        if self._should_prompt_expense_scene_selection(payload):
            return UserAgentResponse(
                answer=self._build_expense_scene_selection_answer(payload),
                citations=citations,
                suggested_actions=suggested_actions,
                query_payload=None,
                draft_payload=None,
                review_payload=None,
                risk_flags=[],
                requires_confirmation=False,
            )
        risk_flags = self._resolve_risk_flags(payload)
        query_payload = self._build_query_payload(payload)
        draft_payload = (
            self._build_draft_payload(payload)
            if self._should_build_draft_payload(payload)
            else None
        )
        review_payload = self._build_review_payload(
            payload,
            citations=citations,
            draft_payload=draft_payload,
        )
        review_answer = self._build_review_body_answer(
            payload,
            review_payload=review_payload,
            draft_payload=draft_payload,
        )

        if payload.degraded and payload.tool_payload.get("message"):
            return UserAgentResponse(
                answer=review_answer or str(payload.tool_payload["message"]),
                citations=citations,
                suggested_actions=suggested_actions,
                query_payload=query_payload,
                draft_payload=draft_payload,
                review_payload=review_payload,
                risk_flags=risk_flags,
                requires_confirmation=payload.requires_confirmation,
            )

        if review_answer:
            return UserAgentResponse(
                answer=review_answer,
                citations=citations,
                suggested_actions=suggested_actions,
                query_payload=query_payload,
                draft_payload=draft_payload,
                review_payload=review_payload,
                risk_flags=risk_flags,
                requires_confirmation=payload.requires_confirmation,
            )

        guided_answer = None
        if draft_payload is None or draft_payload.claim_id is None:
            guided_answer = self._build_guided_answer(payload)
        if guided_answer:
            return UserAgentResponse(
                answer=guided_answer,
                citations=citations,
                suggested_actions=suggested_actions,
                query_payload=query_payload,
                draft_payload=draft_payload,
                review_payload=review_payload,
                risk_flags=risk_flags,
                requires_confirmation=payload.requires_confirmation,
            )

        fast_knowledge_answer = self._build_fast_knowledge_answer(
            payload,
            citations=citations,
        )
        if fast_knowledge_answer:
            return UserAgentResponse(
                answer=fast_knowledge_answer,
                citations=citations,
                suggested_actions=suggested_actions,
                query_payload=query_payload,
                draft_payload=draft_payload,
                review_payload=review_payload,
                risk_flags=risk_flags,
                requires_confirmation=payload.requires_confirmation,
            )

        fallback_answer = self._build_fallback_answer(
            payload,
            citations=citations,
            draft_payload=draft_payload,
        )
        answer = None
        if not self._should_skip_model_answer(payload, review_payload):
            answer = self._generate_answer_with_model(
                payload,
                citations=citations,
                suggested_actions=suggested_actions,
                risk_flags=risk_flags,
                draft_payload=draft_payload,
                fallback_answer=fallback_answer,
            )

        return UserAgentResponse(
            answer=answer or fallback_answer,
            citations=citations,
            suggested_actions=suggested_actions,
            query_payload=query_payload,
            draft_payload=draft_payload,
            review_payload=review_payload,
            risk_flags=risk_flags,
            requires_confirmation=payload.requires_confirmation,
        )

    def _classify_document(
        self,
        item: dict[str, object],
        payload: UserAgentRequest,
    ) -> dict[str, str]:
        entity_values = self._collect_entity_values(payload)
        return self._document_service.classify_document(
            item,
            expense_type_code=entity_values.get("expense_type_code", ""),
            has_customer=bool(entity_values.get("customer")),
        )

    @staticmethod
    def _normalize_group_code(expense_type_code: str) -> str:
        return UserAgentDocumentService.normalize_group_code(expense_type_code)

    def _extract_document_fields(self, item: dict[str, object]) -> dict[str, str]:
        return self._document_service.extract_document_fields(item)

    @staticmethod
    def _resolve_document_time_display_label(
        *,
        document_type: str,
        key: str,
        label: str,
        normalized_label: str,
    ) -> str:
        return UserAgentDocumentService.resolve_document_time_display_label(
            document_type=document_type,
            key=key,
            label=label,
            normalized_label=normalized_label,
        )

    @staticmethod
    def _normalize_document_field_label(*, key: str, label: str) -> str:
        return UserAgentDocumentService.normalize_document_field_label(key=key, label=label)

    def _normalize_document_field_value(self, *, label: str, value: str) -> str:
        return self._document_service.normalize_document_field_value(label=label, value=value)

    def _extract_amount_text_from_value(self, value: str) -> str:
        return self._document_service.extract_amount_text_from_value(value)

    def _extract_document_merchant_name(self, item: dict[str, object]) -> str:
        return self._document_service.extract_document_merchant_name(item)

    @staticmethod
    def _is_hotel_document_item(item: dict[str, object]) -> bool:
        return UserAgentDocumentService.is_hotel_document_item(item)

    @staticmethod
    def _extract_document_merchant_name_from_text(text: str) -> str:
        return UserAgentDocumentService.extract_document_merchant_name_from_text(text)

    @staticmethod
    def _extract_amount_from_card(card: UserAgentReviewDocumentCard) -> float:
        return UserAgentDocumentService.extract_amount_from_card(card)

    def _resolve_amount_value(self, payload: UserAgentRequest) -> float:
        return self._document_service.resolve_amount_value(payload)

    def _sum_ocr_amounts(self, ocr_documents: list[dict[str, object]]) -> float:
        return self._document_service.sum_ocr_amounts(ocr_documents)

    def _infer_expense_type_from_documents(
        self,
        payload: UserAgentRequest,
        ocr_documents: list[dict[str, object]],
    ) -> str:
        entity_values = self._collect_entity_values(payload)
        return self._document_service.infer_expense_type_from_documents(
            ocr_documents,
            expense_type_code=entity_values.get("expense_type_code", ""),
            has_customer=bool(entity_values.get("customer")),
        )