63 lines
2.0 KiB
Python
63 lines
2.0 KiB
Python
from abc import ABC, abstractmethod
|
|
from typing import List, TypedDict
|
|
|
|
|
|
class SearchResult(TypedDict):
|
|
id: str
|
|
distance: float
|
|
content: str
|
|
|
|
|
|
class BaseVectorSearch(ABC):
|
|
"""
|
|
Abstract base class for a vector search provider.
|
|
|
|
This class defines the standard interface for creating a vector search index
|
|
and running queries against it.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def create_index(self, name: str, content_path: str, **kwargs) -> None:
|
|
"""
|
|
Creates a new vector search index and populates it with the provided content.
|
|
|
|
Args:
|
|
name: The desired name for the new index.
|
|
content_path: The local file system path to the data that will be used to
|
|
populate the index. This is expected to be a JSON file
|
|
containing a list of objects, each with an 'id', 'name',
|
|
and 'embedding' key.
|
|
**kwargs: Additional provider-specific arguments for index creation.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def update_index(self, index_name: str, content_path: str, **kwargs) -> None:
|
|
"""
|
|
Updates an existing vector search index with new content.
|
|
|
|
Args:
|
|
index_name: The name of the index to update.
|
|
content_path: The local file system path to the data that will be used to
|
|
populate the index.
|
|
**kwargs: Additional provider-specific arguments for index update.
|
|
"""
|
|
...
|
|
|
|
@abstractmethod
|
|
def run_query(
|
|
self, index: str, query: List[float], limit: int
|
|
) -> List[SearchResult]:
|
|
"""
|
|
Runs a similarity search query against the index.
|
|
|
|
Args:
|
|
query: The embedding vector to use for the search query.
|
|
limit: The maximum number of nearest neighbors to return.
|
|
|
|
Returns:
|
|
A list of dictionaries, where each dictionary represents a matched item
|
|
and contains at least the item's 'id' and the search 'distance'.
|
|
"""
|
|
...
|