Source code for superduperdb.vector_search.server.service

import math
import typing as t

from fastapi import Request

from superduperdb.base.datalayer import Datalayer
from superduperdb.ext.utils import superduperdecode
from superduperdb.vector_search.base import VectorItem

ListVectorType = t.Union[t.List[t.Union[float, int]], t.Dict]

VectorSearchResultType = t.Tuple[t.List[str], t.List[float]]


def _vector_search(
    x: t.Union[str, ListVectorType],
    n: int,
    vector_index: str,
    by_array: bool = True,
    db=None,
) -> VectorSearchResultType:
    vi = db.fast_vector_searchers[vector_index]
    if by_array:
        x = superduperdecode(x, db.datatypes)
        ids, scores = vi.searcher.find_nearest_from_array(x, n=n)
    else:
        ids, scores = vi.searcher.find_nearest_from_id(x, n=n)
    scores = [-1.0 if math.isnan(s) else s for s in scores]
    return ids, scores


[docs] def database(request: Request) -> Datalayer: return request.app.state.pool
[docs] def post_create(vector_index: str, db=None): ''' Performs post create method of vector searcher to incorporate remaining vectors to be added in cache. :param vector_index: Vector class to initiate :param db: Datalayer instance ''' vi = db.fast_vector_searchers[vector_index] vi.post_create()
[docs] def query_search_from_array( array: ListVectorType, vector_index: str, n: int = 100, db=None ) -> VectorSearchResultType: ''' Perform a vector search with an array :param array: Array to perform vector search on index. :param vector_index: Vector search index :param db: Datalayer instance ''' return _vector_search(array, n=n, vector_index=vector_index, db=db)
[docs] def query_search_from_id( id: str, vector_index: str, n: int = 100, db=None ) -> VectorSearchResultType: ''' Perform a vector search with an id :param id: Identifier for vector :param vector_index: Vector search index :param db: Datalayer instance ''' return _vector_search(id, n=n, vector_index=vector_index, db=db, by_array=False)