dify/api/fields/conversation_fields.py

from __future__ import annotations

from collections.abc import Sequence
from datetime import datetime
from typing import Annotated, Any

from pydantic import Field, WithJsonSchema, field_validator, model_validator

from core.rag.entities.citation_metadata import RetrievalSourceMetadata
from fields.base import ResponseModel, SessionResponseSource
from graphon.file import File
from libs.helper import to_timestamp
from models.account import Account
from models.model import (
    AppModelConfigDict,
    MessageAgentThought,
    MessageAnnotation,
    MessageFeedback,
    MessageFileInfo,
)
from models.model import (
    Conversation as ConversationModel,
)
from models.model import (
    Message as MessageModel,
)

type AnyJSONValue = Any
UUIDString = Annotated[str, WithJsonSchema({"format": "uuid", "type": "string"})]
Int64 = Annotated[int, WithJsonSchema({"format": "int64", "type": "integer"})]
URLString = Annotated[str, WithJsonSchema({"format": "url", "type": "string"})]


class _FeedbackResponseSource(SessionResponseSource[MessageFeedback]):
    @property
    def from_account(self) -> Account | None:
        return self._source.from_account_with_session(session=self._session)


class _AnnotationResponseSource(SessionResponseSource[MessageAnnotation]):
    @property
    def account(self) -> Account | None:
        return self._source.account_with_session(session=self._session)

    @property
    def annotation_create_account(self) -> Account | None:
        return self._source.annotation_create_account_with_session(session=self._session)


class MessageResponseSource(SessionResponseSource[MessageModel]):
    @property
    def inputs(self) -> dict[str, Any]:
        return self._source.inputs_with_session(session=self._session)

    @property
    def feedbacks(self) -> list[_FeedbackResponseSource]:
        return [
            _FeedbackResponseSource(feedback, session=self._session)
            for feedback in self._source.feedbacks_with_session(session=self._session)
        ]

    @property
    def user_feedback(self) -> MessageFeedback | None:
        return self._source.user_feedback_with_session(session=self._session)

    @property
    def annotation(self) -> _AnnotationResponseSource | None:
        annotation = self._source.annotation_with_session(session=self._session)
        return _AnnotationResponseSource(annotation, session=self._session) if annotation else None

    @property
    def annotation_hit_history(self) -> _AnnotationResponseSource | None:
        annotation = self._source.annotation_hit_history_with_session(session=self._session)
        return _AnnotationResponseSource(annotation, session=self._session) if annotation else None

    @property
    def agent_thoughts(self) -> Sequence[MessageAgentThought]:
        return self._source.agent_thoughts_with_session(session=self._session)

    @property
    def message_files(self) -> list[MessageFileInfo]:
        return self._source.message_files_with_session(session=self._session)


class ConversationResponseSource(SessionResponseSource[ConversationModel]):
    @property
    def inputs(self) -> dict[str, Any]:
        return self._source.inputs_with_session(session=self._session)

    @property
    def model_config(self) -> AppModelConfigDict:
        return self._source.model_config_with_session(session=self._session)

    @property
    def summary_or_query(self) -> str:
        return self._source.summary_or_query_with_session(session=self._session)

    @property
    def annotated(self) -> bool:
        return self._source.annotated_with_session(session=self._session)

    @property
    def annotation(self) -> _AnnotationResponseSource | None:
        annotation = self._source.annotation_with_session(session=self._session)
        return _AnnotationResponseSource(annotation, session=self._session) if annotation else None

    @property
    def message_count(self) -> int:
        return self._source.message_count_with_session(session=self._session)

    @property
    def user_feedback_stats(self) -> dict[str, int]:
        return self._source.user_feedback_stats_with_session(session=self._session)

    @property
    def admin_feedback_stats(self) -> dict[str, int]:
        return self._source.admin_feedback_stats_with_session(session=self._session)

    @property
    def status_count(self) -> dict[str, int] | None:
        return self._source.status_count_with_session(session=self._session)

    @property
    def first_message(self) -> MessageResponseSource | None:
        message = self._source.first_message_with_session(session=self._session)
        return MessageResponseSource(message, session=self._session) if message else None

    @property
    def from_end_user_session_id(self) -> str | None:
        return self._source.from_end_user_session_id_with_session(session=self._session)

    @property
    def from_account_name(self) -> str | None:
        return self._source.from_account_name_with_session(session=self._session)


class MessageFile(ResponseModel):
    id: UUIDString
    filename: str
    type: str
    url: URLString | None = None
    mime_type: str | None = None
    size: int | None = None
    transfer_method: str
    belongs_to: str | None = None
    upload_file_id: UUIDString | None = None

    @field_validator("transfer_method", mode="before")
    @classmethod
    def _normalize_transfer_method(cls, value: object) -> str:
        if isinstance(value, str):
            return value
        return str(value)


class SimpleConversation(ResponseModel):
    id: UUIDString
    name: str
    inputs: dict[str, AnyJSONValue]
    status: str
    introduction: str | None = None
    created_at: Int64 | None = None
    updated_at: Int64 | None = None

    @field_validator("inputs", mode="before")
    @classmethod
    def _normalize_inputs(cls, value: AnyJSONValue) -> AnyJSONValue:
        return format_files_contained(value)

    @field_validator("created_at", "updated_at", mode="before")
    @classmethod
    def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class ConversationInfiniteScrollPagination(ResponseModel):
    limit: int
    has_more: bool
    data: list[SimpleConversation]


class ConversationDelete(ResponseModel):
    result: str


class ResultResponse(ResponseModel):
    result: str


class SimpleAccount(ResponseModel):
    id: str
    name: str
    email: str


class Feedback(ResponseModel):
    rating: str
    content: str | None = None
    from_source: str
    from_end_user_id: str | None = None
    from_account: SimpleAccount | None = None


class ConversationAnnotation(ResponseModel):
    id: str
    question: str | None = None
    content: str
    account: SimpleAccount | None = None
    created_at: int | None = None

    @field_validator("created_at", mode="before")
    @classmethod
    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class ConversationAnnotationHitHistory(ResponseModel):
    annotation_id: str = Field(validation_alias="id")
    annotation_create_account: SimpleAccount | None = None
    created_at: int | None = None

    @field_validator("created_at", mode="before")
    @classmethod
    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class AgentThought(ResponseModel):
    id: str
    chain_id: str | None = None
    message_chain_id: str | None = Field(default=None, exclude=True, validation_alias="message_chain_id")
    message_id: str
    position: int
    thought: str | None = None
    answer: str | None = None
    tool: str | None = None
    tool_labels: AnyJSONValue = Field(description="Labels for tools used.")
    tool_input: str | None = None
    created_at: int | None = None
    observation: str | None = None
    files: list[str]

    @field_validator("created_at", mode="before")
    @classmethod
    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)

    @model_validator(mode="after")
    def _fallback_chain_id(self):
        if self.chain_id is None and self.message_chain_id:
            self.chain_id = self.message_chain_id
        return self


class MessageDetail(ResponseModel):
    retriever_resources: list[RetrievalSourceMetadata] = Field(default_factory=list)
    id: str
    conversation_id: str
    inputs: dict[str, AnyJSONValue]
    query: str
    message: AnyJSONValue
    message_tokens: int
    answer: str = Field(validation_alias="re_sign_file_url_answer")
    answer_tokens: int
    provider_response_latency: float
    from_source: str
    from_end_user_id: str | None = None
    from_account_id: str | None = None
    feedbacks: list[Feedback]
    workflow_run_id: str | None = None
    annotation: ConversationAnnotation | None = None
    annotation_hit_history: ConversationAnnotationHitHistory | None = None
    created_at: int | None = None
    agent_thoughts: list[AgentThought]
    message_files: list[MessageFile]
    metadata: AnyJSONValue = Field(validation_alias="message_metadata_dict")
    status: str
    error: str | None = None
    parent_message_id: str | None = None

    @field_validator("inputs", mode="before")
    @classmethod
    def _normalize_inputs(cls, value: AnyJSONValue) -> AnyJSONValue:
        return format_files_contained(value)

    @field_validator("created_at", mode="before")
    @classmethod
    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class FeedbackStat(ResponseModel):
    like: int
    dislike: int


class StatusCount(ResponseModel):
    success: int
    failed: int
    partial_success: int
    paused: int


class ConversationModelConfig(ResponseModel):
    opening_statement: str | None = None
    suggested_questions: AnyJSONValue | None = Field(default=None)
    model: AnyJSONValue | None = Field(default=None)
    user_input_form: AnyJSONValue | None = Field(default=None)
    pre_prompt: str | None = None
    agent_mode: AnyJSONValue | None = Field(default=None)


class SimpleModelConfig(ResponseModel):
    model: AnyJSONValue | None = Field(
        default=None,
        validation_alias="model_dict",
    )
    pre_prompt: str | None = None


class SimpleMessageDetail(ResponseModel):
    inputs: dict[str, AnyJSONValue]
    query: str
    message: str
    answer: str

    @field_validator("inputs", mode="before")
    @classmethod
    def _normalize_inputs(cls, value: AnyJSONValue) -> AnyJSONValue:
        return format_files_contained(value)

    @field_validator("message", mode="before")
    @classmethod
    def _normalize_message(cls, value: AnyJSONValue) -> str:
        return message_text(value)


class Conversation(ResponseModel):
    id: str
    status: str
    from_source: str
    from_end_user_id: str | None = None
    from_end_user_session_id: str | None = None
    from_account_id: str | None = None
    from_account_name: str | None = None
    read_at: int | None = None
    created_at: int | None = None
    updated_at: int | None = None
    annotation: ConversationAnnotation | None = None
    model_config_: SimpleModelConfig | None = Field(default=None, alias="model_config")
    user_feedback_stats: FeedbackStat | None = None
    admin_feedback_stats: FeedbackStat | None = None
    message: SimpleMessageDetail | None = Field(default=None, validation_alias="first_message")

    @field_validator("read_at", "created_at", "updated_at", mode="before")
    @classmethod
    def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class ConversationPagination(ResponseModel):
    page: int
    limit: int = Field(validation_alias="per_page")
    total: int
    has_more: bool = Field(validation_alias="has_next")
    data: list[Conversation] = Field(validation_alias="items")


class ConversationMessageDetail(ResponseModel):
    id: str
    status: str
    from_source: str
    from_end_user_id: str | None = None
    from_account_id: str | None = None
    created_at: int | None = None
    model_config_: ConversationModelConfig | None = Field(default=None, alias="model_config")
    message: MessageDetail | None = Field(default=None, validation_alias="first_message")

    @field_validator("created_at", mode="before")
    @classmethod
    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class ConversationWithSummary(ResponseModel):
    id: str
    status: str
    from_source: str
    from_end_user_id: str | None = None
    from_end_user_session_id: str | None = None
    from_account_id: str | None = None
    from_account_name: str | None = None
    name: str
    summary: str = Field(validation_alias="summary_or_query")
    read_at: int | None = None
    created_at: int | None = None
    updated_at: int | None = None
    annotated: bool
    model_config_: SimpleModelConfig | None = Field(default=None, alias="model_config")
    message_count: int
    user_feedback_stats: FeedbackStat | None = None
    admin_feedback_stats: FeedbackStat | None = None
    status_count: StatusCount | None = None

    @field_validator("read_at", "created_at", "updated_at", mode="before")
    @classmethod
    def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


class ConversationWithSummaryPagination(ResponseModel):
    page: int
    limit: int = Field(validation_alias="per_page")
    total: int
    has_more: bool = Field(validation_alias="has_next")
    data: list[ConversationWithSummary] = Field(validation_alias="items")


class ConversationDetail(ResponseModel):
    id: str
    status: str
    from_source: str
    from_end_user_id: str | None = None
    from_account_id: str | None = None
    created_at: int | None = None
    updated_at: int | None = None
    annotated: bool
    introduction: str | None = None
    model_config_: ConversationModelConfig | None = Field(default=None, alias="model_config")
    message_count: int
    user_feedback_stats: FeedbackStat | None = None
    admin_feedback_stats: FeedbackStat | None = None

    @field_validator("created_at", "updated_at", mode="before")
    @classmethod
    def _normalize_timestamp(cls, value: datetime | int | None) -> int | None:
        return to_timestamp(value)


def format_files_contained(value: AnyJSONValue) -> AnyJSONValue:
    if isinstance(value, File):
        # Response payloads must preserve legacy file keys like `related_id`/`url`
        # while still exposing the new graph-layer `reference` field.
        return value.to_dict()
    if isinstance(value, dict):
        return {k: format_files_contained(v) for k, v in value.items()}
    if isinstance(value, list):
        return [format_files_contained(v) for v in value]
    return value


def message_text(value: AnyJSONValue) -> str:
    if isinstance(value, list) and value:
        first = value[0]
        if isinstance(first, dict):
            text = first.get("text")
            if isinstance(text, str):
                return text
    return ""


def extract_model_config(value: object | None) -> dict[str, AnyJSONValue]:
    if value is None:
        return {}
    if isinstance(value, dict):
        return value
    if hasattr(value, "to_dict"):
        return value.to_dict()
    return {}