arc53--docsgpt/application/services/search_service.py

"""Shared retrieval service used by the HTTP search route and the MCP tool.

Flask-free. Raises domain exceptions (``InvalidAPIKey``, ``SearchFailed``)
that callers translate into their own wire protocol (HTTP status codes,
MCP error responses, etc.).
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List

from application.core.settings import settings
from application.storage.db.repositories.agents import AgentsRepository
from application.storage.db.session import db_readonly
from application.vectorstore.vector_creator import VectorCreator

logger = logging.getLogger(__name__)


class InvalidAPIKey(Exception):
    """The supplied ``api_key`` does not resolve to an agent."""


class SearchFailed(Exception):
    """Unexpected error during retrieval (e.g. DB outage). Caller maps to 5xx."""


def _collect_source_ids(agent: Dict[str, Any]) -> List[str]:
    """Extract the ordered list of source UUIDs to search.

    Prefers ``extra_source_ids`` (PG ARRAY(UUID) of multi-source agents);
    falls back to the legacy single ``source_id`` field.
    """
    source_ids: List[str] = []
    extra = agent.get("extra_source_ids") or []
    for src in extra:
        if src:
            source_ids.append(str(src))
    if not source_ids:
        single = agent.get("source_id")
        if single:
            source_ids.append(str(single))
    return source_ids


def _search_sources(
    query: str, source_ids: List[str], chunks: int
) -> List[Dict[str, Any]]:
    """Search across each source's vectorstore and return up to ``chunks`` hits.

    Per-source errors are logged and skipped so one broken index doesn't
    take down the whole search. Results are de-duplicated by content hash.
    """
    if chunks <= 0 or not source_ids:
        return []

    results: List[Dict[str, Any]] = []
    chunks_per_source = max(1, chunks // len(source_ids))
    seen_texts: set[int] = set()

    for source_id in source_ids:
        if not source_id or not source_id.strip():
            continue

        try:
            docsearch = VectorCreator.create_vectorstore(
                settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
            )
            docs = docsearch.search(query, k=chunks_per_source * 2)

            for doc in docs:
                if len(results) >= chunks:
                    break

                if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
                    page_content = doc.page_content
                    metadata = doc.metadata
                else:
                    page_content = doc.get("text", doc.get("page_content", ""))
                    metadata = doc.get("metadata", {})

                text_hash = hash(page_content[:200])
                if text_hash in seen_texts:
                    continue
                seen_texts.add(text_hash)

                title = metadata.get("title", metadata.get("post_title", ""))
                if not isinstance(title, str):
                    title = str(title) if title else ""

                if title:
                    title = title.split("/")[-1]
                else:
                    title = metadata.get("filename", page_content[:50] + "...")

                source = metadata.get("source", source_id)

                results.append(
                    {
                        "text": page_content,
                        "title": title,
                        "source": source,
                    }
                )

            if len(results) >= chunks:
                break

        except Exception as e:
            logger.error(
                f"Error searching vectorstore {source_id}: {e}",
                exc_info=True,
            )
            continue

    return results[:chunks]


def search(api_key: str, query: str, chunks: int = 5) -> List[Dict[str, Any]]:
    """Resolve an agent by API key and search its sources.

    Args:
        api_key: Agent API key (the opaque string stored on
            ``agents.key`` in Postgres).
        query: Free-text search query.
        chunks: Max number of hits to return.

    Returns:
        List of hit dicts with ``text``, ``title``, ``source`` keys.
        Empty list if the agent has no sources configured.

    Raises:
        InvalidAPIKey: if ``api_key`` does not resolve to an agent.
        SearchFailed: on unexpected DB / infrastructure errors.
    """
    if chunks <= 0:
        return []

    try:
        with db_readonly() as conn:
            agent = AgentsRepository(conn).find_by_key(api_key)
    except Exception as e:
        raise SearchFailed("agent lookup failed") from e

    if not agent:
        raise InvalidAPIKey()

    source_ids = _collect_source_ids(agent)
    if not source_ids:
        return []

    return _search_sources(query, source_ids, chunks)