PMFlow/backend/app/schemas/rag.py

"""RAG / 知识库 / QA 记录请求模型。"""

from __future__ import annotations

from datetime import datetime
from typing import ClassVar, Optional

from pydantic import BaseModel, Field, field_validator


class RAGAskRequest(BaseModel):
    _ALLOWED_ASKED_VIA: ClassVar[set[str]] = {
        "pmflow",
        "feishu",
        "lark",
        "feishu_group",
        "feishu_dm",
        "pmflow_manual",
        "/ai-pmo/process-text",
        "ai-pmo-process-text",
        "ai_pmo",
        "ai_pmo_ui",
    }
    _ASKED_VIA_ALIAS: ClassVar[dict[str, str]] = {
        "ai-pmo-process-text": "/ai-pmo/process-text",
    }
    _ALLOWED_CONSULTANT_MODES: ClassVar[set[str]] = {
        "clarify",
        "advise",
        "challenge",
        "brainstorm",
        "execute",
    }
    question: str = Field(..., min_length=1, max_length=4000)
    project_id: Optional[int] = Field(default=None)
    asked_via: str = Field(default="pmflow", min_length=1, max_length=32)
    source_chat_id: Optional[str] = Field(default=None, max_length=128)
    session_id: Optional[int] = Field(default=None, ge=1)
    new_session: bool = Field(default=False)
    thinking_mode: str = Field(default="deep", pattern="^(normal|deep)$")
    action_mode: str = Field(default="auto", pattern="^(auto|read|write|read_write)$")
    sync_requirement: bool = Field(default=True)
    message_intent: Optional[str] = Field(default=None, max_length=32)
    mixed_signal_part: Optional[str] = Field(default=None, max_length=2000)
    consultant_mode_preference: Optional[str] = Field(default=None, max_length=16)

    @field_validator("asked_via", mode="before")
    @classmethod
    def _normalize_asked_via(cls, value: object) -> str:
        text = str(value or "pmflow").strip().lower()
        text = cls._ASKED_VIA_ALIAS.get(text, text)
        if text not in cls._ALLOWED_ASKED_VIA:
            allowed = ", ".join(sorted(cls._ALLOWED_ASKED_VIA))
            raise ValueError(f"asked_via must be one of: {allowed}")
        return text

    @field_validator("message_intent", mode="before")
    @classmethod
    def _normalize_message_intent(cls, value: object) -> Optional[str]:
        if value is None:
            return None
        text = str(value).strip().lower()
        if not text:
            return None
        return text

    @field_validator("mixed_signal_part", mode="before")
    @classmethod
    def _normalize_mixed_signal_part(cls, value: object) -> Optional[str]:
        if value is None:
            return None
        text = str(value).strip()
        return text or None

    @field_validator("consultant_mode_preference", mode="before")
    @classmethod
    def _normalize_consultant_mode_preference(cls, value: object) -> Optional[str]:
        if value is None:
            return None
        text = str(value).strip().lower()
        if not text:
            return None
        if text not in cls._ALLOWED_CONSULTANT_MODES:
            allowed = ", ".join(sorted(cls._ALLOWED_CONSULTANT_MODES))
            raise ValueError(f"consultant_mode_preference must be one of: {allowed}")
        return text


class KnowledgeDocCreateRequest(BaseModel):
    title: str = Field(..., min_length=1, max_length=255)
    file_type: str = Field(default="txt", max_length=32)
    content: str = Field(..., min_length=1, max_length=200000)
    source_type: str = Field(default="file", max_length=32)
    file_url: Optional[str] = Field(default=None, max_length=1024)
    project_id: Optional[int] = Field(default=None)


class KnowledgeURLCreateRequest(BaseModel):
    url: str = Field(..., min_length=5, max_length=1024)
    title: Optional[str] = Field(default=None, max_length=255)
    project_id: Optional[int] = Field(default=None)


class QAFeedbackCreateRequest(BaseModel):
    rating: str = Field(..., pattern="^(positive|negative)$")
    helpfulness: Optional[str] = Field(default=None, pattern="^(fully|partially|could_be_better)$")
    error_types: list[str] = Field(default_factory=list)
    correct_answer: Optional[str] = Field(default=None, max_length=4000)


class QAListQuery(BaseModel):
    project_id: Optional[int] = None
    date_from: Optional[datetime] = None
    date_to: Optional[datetime] = None
    rating: Optional[str] = Field(default=None, pattern="^(positive|negative)$")
    page: int = Field(default=1, ge=1)
    page_size: int = Field(default=20, ge=1, le=100)


class QAExportQuery(BaseModel):
    include_unrated: bool = False
    project_id: Optional[int] = None
    rating: Optional[str] = Field(default=None, pattern="^(positive|negative|unrated)$")
    domain: Optional[str] = None
    date_from: Optional[datetime] = None
    date_to: Optional[datetime] = None


class QASessionListQuery(BaseModel):
    date_from: Optional[datetime] = None
    date_to: Optional[datetime] = None
    page: int = Field(default=1, ge=1)
    page_size: int = Field(default=20, ge=1, le=100)


class QASessionRenameRequest(BaseModel):
    title: str = Field(..., min_length=1, max_length=120)


class EntityResolutionRequest(BaseModel):
    question: str = Field(..., min_length=1, max_length=4000)
    project_id: Optional[int] = Field(default=None)
    limit: int = Field(default=10, ge=1, le=50)


class MultiStepSearchRequest(BaseModel):
    question: str = Field(..., min_length=1, max_length=4000)
    project_id: Optional[int] = Field(default=None)
    source_chat_id: Optional[str] = Field(default=None, max_length=128)


class SymbolSyncRequest(BaseModel):
    base_dir: Optional[str] = Field(default=None, max_length=500)
    max_files: int = Field(default=500, ge=1, le=3000)
    force_reindex: bool = Field(default=False)


class SymbolSearchRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=255)
    language: Optional[str] = Field(default=None, max_length=20)
    symbol_type: Optional[str] = Field(default=None, max_length=32)
    limit: int = Field(default=30, ge=1, le=200)