app/retrieval/vector_search.py

"""Async wrappers for vector store search."""

from __future__ import annotations

from functools import partial

from app.async_executor import run_sync_in_executor
from app.models.schemas import SearchResult
from app.vectorstore.base import VectorStore


async def async_vs_search(
    vs: VectorStore,
    query: str,
    tenant_id: str,
    k: int = 10,
    *,
    hierarchy_level: str | None = None,
    doc_id_in: frozenset[str] | None = None,
) -> list[SearchResult]:
    search_async = getattr(vs, "search_async", None)
    if callable(search_async):
        return await search_async(
            query,
            tenant_id,
            k=k,
            hierarchy_level=hierarchy_level,
            doc_id_in=doc_id_in,
        )
    return await run_sync_in_executor(
        partial(
            vs.search,
            query,
            tenant_id,
            k=k,
            hierarchy_level=hierarchy_level,
            doc_id_in=doc_id_in,
        ),
    )