Spaces:
Runtime error
Runtime error
File size: 1,642 Bytes
5a67683 |
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 41 42 43 44 45 46 47 48 49 50 51 52 |
import logging
import typing
from llama_index import VectorStoreIndex
from llama_index.indices.vector_store import VectorIndexRetriever
from llama_index.vector_stores.types import VectorStore
from app._config import settings
from app.enums import WEAVIATE_INDEX_NAME, VectorDatabase
logger = logging.getLogger(__name__)
class VectorStoreComponent:
vector_store: VectorStore
def __init__(self) -> None:
match settings.VECTOR_DATABASE:
case VectorDatabase.WEAVIATE:
import weaviate
from llama_index.vector_stores import WeaviateVectorStore
client = weaviate.Client(settings.WEAVIATE_CLIENT_URL)
self.vector_store = typing.cast(
VectorStore,
WeaviateVectorStore(
weaviate_client=client, index_name=WEAVIATE_INDEX_NAME
),
)
case _:
# Should be unreachable
# The settings validator should have caught this
raise ValueError(
f"Vectorstore database {settings.VECTOR_DATABASE} not supported"
)
@staticmethod
def get_retriever(
index: VectorStoreIndex,
doc_ids: list[str] | None = None,
similarity_top_k: int = 2,
) -> VectorIndexRetriever:
return VectorIndexRetriever(
index=index,
similarity_top_k=similarity_top_k,
doc_ids=doc_ids,
)
def close(self) -> None:
if hasattr(self.vector_store.client, "close"):
self.vector_store.client.close()
|