leocrm/app/plugins/builtins/unified_search/providers/ai_chat_provider.py

"""AI chat search provider — FTS search on ai_chat_messages table."""

from __future__ import annotations

import logging
import uuid
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.plugins.builtins.unified_search.base_provider import BaseSearchProvider

logger = logging.getLogger(__name__)


class AIChatSearchProvider(BaseSearchProvider):
    """Search provider for AI chat messages."""

    entity_type = "ai_chat"
    supports_fts = True
    supports_vector = False

    async def _search_fts_filtered(
        self,
        db: AsyncSession,
        tsquery: str,
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """Full-text search on ai_chat_messages.content, joined with sessions for title."""
        if visible_ids is not None:
            sql = text(
                """
                SELECT m.id, m.tenant_id, m.session_id, m.role, m.content,
                       m.model_used, m.tokens,
                       s.title AS session_title,
                       ts_rank(to_tsvector('pg_catalog.german', m.content),
                               to_tsquery('pg_catalog.german', :q)) AS rank
                FROM ai_chat_messages m
                JOIN ai_chat_sessions s ON s.id = m.session_id
                WHERE m.tenant_id = :tid
                  AND to_tsvector('pg_catalog.german', m.content) @@ to_tsquery('pg_catalog.german', :q)
                  AND m.id = ANY(:visible_ids)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {
                    "q": tsquery,
                    "tid": tenant_id,
                    "lim": limit,
                    "visible_ids": list(visible_ids),
                },
            )
        else:
            sql = text(
                """
                SELECT m.id, m.tenant_id, m.session_id, m.role, m.content,
                       m.model_used, m.tokens,
                       s.title AS session_title,
                       ts_rank(to_tsvector('pg_catalog.german', m.content),
                               to_tsquery('pg_catalog.german', :q)) AS rank
                FROM ai_chat_messages m
                JOIN ai_chat_sessions s ON s.id = m.session_id
                WHERE m.tenant_id = :tid
                  AND to_tsvector('pg_catalog.german', m.content) @@ to_tsquery('pg_catalog.german', :q)
                ORDER BY rank DESC
                LIMIT :lim
                """
            )
            result = await db.execute(
                sql,
                {"q": tsquery, "tid": tenant_id, "lim": limit},
            )
        rows = result.mappings().all()
        return [dict(r) for r in rows]

    async def _search_vector_filtered(
        self,
        db: AsyncSession,
        embedding: list[float],
        tenant_id: uuid.UUID,
        limit: int,
        visible_ids: set[uuid.UUID] | None,
    ) -> list[dict[str, Any]]:
        """No vector support for chat messages — return empty list."""
        return []

    async def get_embedding_text(
        self, db: AsyncSession, entity_id: uuid.UUID, tenant_id: uuid.UUID
    ) -> str:
        """Get text for embedding generation — returns message content."""
        sql = text(
            """
            SELECT content
            FROM ai_chat_messages
            WHERE id = :eid AND tenant_id = :tid
            """
        )
        result = await db.execute(sql, {"eid": entity_id, "tid": tenant_id})
        row = result.mappings().first()
        if not row:
            return ""
        return str(row.get("content", ""))

    def to_search_result(self, entity: object) -> dict[str, Any]:
        """Convert chat message to search result dict."""
        if isinstance(entity, dict):
            entity_id = str(entity.get("id", ""))
            content = entity.get("content", "")
            role = entity.get("role", "")
            session_title = entity.get("session_title", "")
            session_id = str(entity.get("session_id", ""))
            score = entity.get("rank", 0.0)
        else:
            entity_id = str(getattr(entity, "id", ""))
            content = getattr(entity, "content", "")
            role = getattr(entity, "role", "")
            session_title = getattr(entity, "session_title", "")
            session_id = str(getattr(entity, "session_id", ""))
            score = getattr(entity, "rank", 0.0)
        return {
            "entity_type": self.entity_type,
            "entity_id": entity_id,
            "title": session_title or content[:80] if content else "",
            "snippet": content,
            "score": float(score) if score else 0.0,
            "data": {
                "role": role,
                "session_id": session_id,
                "session_title": session_title,
            },
        }