company-ai-platform/app/modules/personalization/services/conversations.py

from datetime import timedelta
from hashlib import sha256
from typing import Any

from fastapi import HTTPException, status
from sqlalchemy import delete, exists, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm import Session

from app.core.utils.time import utc_now
from app.modules.personalization.constants import (
    CONVERSATION_MAX_CONTENT_LENGTH,
    CONVERSATION_MAX_MESSAGES,
    CONVERSATION_RETENTION_DAYS,
    UNAVAILABLE_AI_PROVIDERS,
    ConversationChatType,
    ConversationRole,
)
from app.modules.personalization.models import AIConversation, AIConversationMessage


class ConversationService:
    """Persist isolated Feishu conversations with bounded history."""

    def __init__(self, db: Session):
        self.db = db

    def history(
        self,
        owner_id: int,
        chat_type: str | ConversationChatType,
        chat_key: str,
    ) -> list[dict[str, Any]]:
        owner_id, chat_type_value, chat_key_value = _conversation_identity(
            owner_id,
            chat_type,
            chat_key,
        )
        self.cleanup_expired(owner_id=owner_id)
        conversation = self._find(owner_id, chat_type_value, chat_key_value)
        if conversation is None:
            self.db.commit()
            return []
        messages = list(
            self.db.execute(
                select(AIConversationMessage)
                .where(AIConversationMessage.conversation_id == conversation.id)
                .order_by(
                    AIConversationMessage.created_at.desc(),
                    AIConversationMessage.id.desc(),
                )
                .limit(CONVERSATION_MAX_MESSAGES)
            ).scalars()
        )
        self.db.commit()
        messages.reverse()
        return [_serialize_message(message) for message in messages]

    def record_turn(
        self,
        owner_id: int,
        chat_type: str | ConversationChatType,
        chat_key: str,
        *,
        user_content: str,
        assistant_content: str,
        provider_name: str,
        ai_available: bool = True,
    ) -> bool:
        """Record one complete turn only after a real AI answer succeeds."""

        if not ai_available or provider_name.strip().lower() in UNAVAILABLE_AI_PROVIDERS:
            return False
        owner_id, chat_type_value, chat_key_value = _conversation_identity(
            owner_id,
            chat_type,
            chat_key,
        )
        user_text = _message_content(user_content)
        assistant_text = _message_content(assistant_content)
        self.cleanup_expired(owner_id=owner_id)
        conversation = self._get_or_create(
            owner_id,
            chat_type_value,
            chat_key_value,
        )
        now = utc_now()
        self.db.add_all(
            [
                AIConversationMessage(
                    conversation_id=conversation.id,
                    role=ConversationRole.USER,
                    content=user_text,
                    created_at=now,
                ),
                AIConversationMessage(
                    conversation_id=conversation.id,
                    role=ConversationRole.ASSISTANT,
                    content=assistant_text,
                    created_at=now,
                ),
            ]
        )
        conversation.updated_at = now
        self.db.flush()
        self._trim(conversation.id)
        self.db.commit()
        return True

    def reset(
        self,
        owner_id: int,
        chat_type: str | ConversationChatType,
        chat_key: str,
    ) -> bool:
        owner_id, chat_type_value, chat_key_value = _conversation_identity(
            owner_id,
            chat_type,
            chat_key,
        )
        conversation = self._find(owner_id, chat_type_value, chat_key_value)
        if conversation is None:
            return False
        self.db.execute(
            delete(AIConversationMessage).where(
                AIConversationMessage.conversation_id == conversation.id
            )
        )
        self.db.delete(conversation)
        self.db.commit()
        return True

    def cleanup_expired(self, owner_id: int | None = None) -> int:
        """Stage retention cleanup for one owner or all owners.

        The surrounding operation owns the transaction. Interactive history and
        write paths use this method before completing their own commit.
        """

        return self._cleanup_expired(owner_id)["conversation_messages"]

    def cleanup_expired_globally(self) -> dict[str, int]:
        """Delete expired messages for every owner and commit the maintenance run."""

        deleted = self._cleanup_expired(owner_id=None)
        self.db.commit()
        return deleted

    def _cleanup_expired(self, owner_id: int | None) -> dict[str, int]:
        cutoff = utc_now() - timedelta(days=CONVERSATION_RETENTION_DAYS)
        conversation_ids = select(AIConversation.id)
        if owner_id is not None:
            conversation_ids = conversation_ids.where(AIConversation.owner_id == owner_id)
        message_result = self.db.execute(
            delete(AIConversationMessage).where(
                AIConversationMessage.conversation_id.in_(conversation_ids),
                AIConversationMessage.created_at < cutoff,
            )
        )
        empty_conversations = delete(AIConversation).where(
            ~exists(
                select(AIConversationMessage.id).where(
                    AIConversationMessage.conversation_id == AIConversation.id
                )
            )
        )
        if owner_id is not None:
            empty_conversations = empty_conversations.where(
                AIConversation.owner_id == owner_id
            )
        conversation_result = self.db.execute(empty_conversations)
        return {
            "conversation_messages": max(0, int(message_result.rowcount or 0)),
            "conversations": max(0, int(conversation_result.rowcount or 0)),
        }

    def delete_owner_conversations(self, owner_id: int) -> dict[str, int]:
        conversation_ids = select(AIConversation.id).where(
            AIConversation.owner_id == owner_id
        )
        message_result = self.db.execute(
            delete(AIConversationMessage).where(
                AIConversationMessage.conversation_id.in_(conversation_ids)
            )
        )
        conversation_result = self.db.execute(
            delete(AIConversation).where(AIConversation.owner_id == owner_id)
        )
        return {
            "conversation_messages": max(0, int(message_result.rowcount or 0)),
            "conversations": max(0, int(conversation_result.rowcount or 0)),
        }

    @staticmethod
    def provider_session_id(
        owner_id: int,
        chat_type: str | ConversationChatType,
        chat_key: str,
    ) -> str:
        owner_id, chat_type_value, chat_key_value = _conversation_identity(
            owner_id,
            chat_type,
            chat_key,
        )
        digest = sha256(
            f"{owner_id}\0{chat_type_value}\0{chat_key_value}".encode("utf-8")
        ).hexdigest()
        return f"feishu-{digest}"

    def _find(
        self,
        owner_id: int,
        chat_type: str,
        chat_key: str,
    ) -> AIConversation | None:
        return self.db.execute(
            select(AIConversation).where(
                AIConversation.owner_id == owner_id,
                AIConversation.chat_type == chat_type,
                AIConversation.chat_key == chat_key,
            )
        ).scalar_one_or_none()

    def _get_or_create(
        self,
        owner_id: int,
        chat_type: str,
        chat_key: str,
    ) -> AIConversation:
        existing = self._find(owner_id, chat_type, chat_key)
        if existing is not None:
            return existing
        conversation = AIConversation(
            owner_id=owner_id,
            chat_type=chat_type,
            chat_key=chat_key,
        )
        try:
            with self.db.begin_nested():
                self.db.add(conversation)
                self.db.flush()
        except IntegrityError:
            conversation = self._find(owner_id, chat_type, chat_key)
            if conversation is None:
                raise
        return conversation

    def _trim(self, conversation_id: int) -> int:
        stale_ids = list(
            self.db.execute(
                select(AIConversationMessage.id)
                .where(AIConversationMessage.conversation_id == conversation_id)
                .order_by(
                    AIConversationMessage.created_at.desc(),
                    AIConversationMessage.id.desc(),
                )
                .offset(CONVERSATION_MAX_MESSAGES)
            ).scalars()
        )
        if not stale_ids:
            return 0
        result = self.db.execute(
            delete(AIConversationMessage).where(
                AIConversationMessage.id.in_(stale_ids)
            )
        )
        return max(0, int(result.rowcount or 0))


def _conversation_identity(
    owner_id: int,
    chat_type: str | ConversationChatType,
    chat_key: str,
) -> tuple[int, str, str]:
    if owner_id <= 0:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Valid conversation owner is required",
        )
    raw_chat_type = str(chat_type).strip().lower()
    aliases = {
        "p2p": ConversationChatType.PRIVATE,
        "private": ConversationChatType.PRIVATE,
        "group": ConversationChatType.GROUP,
        "group_chat": ConversationChatType.GROUP,
    }
    try:
        chat_type_value = str(aliases[raw_chat_type])
    except KeyError as exc:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Unsupported conversation chat type",
        ) from exc
    chat_key_value = str(chat_key).strip()
    if not chat_key_value or len(chat_key_value) > 256:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Valid conversation chat key is required",
        )
    return owner_id, chat_type_value, chat_key_value


def _message_content(value: str) -> str:
    content = str(value).strip()
    if not content:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Conversation message is required",
        )
    if len(content) > CONVERSATION_MAX_CONTENT_LENGTH:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Conversation message is too long",
        )
    return content


def _serialize_message(message: AIConversationMessage) -> dict[str, Any]:
    return {
        "role": message.role,
        "content": message.content,
        "created_at": message.created_at.isoformat(),
    }