File size: 974 Bytes
732b14f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
"""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,
        ),
    )