Demos / backend /classes /vector_database /base_vector_database.py
nikhile-galileo's picture
Adding finance protect demo
e68d535
raw
history blame
983 Bytes
from pydantic import BaseModel
import pandas as pd
class VectorDatabaseConfig(BaseModel):
"""Base configuration for vector databases."""
collection_name: str
class Config:
arbitrary_types_allowed = True
class VectorDatabase:
"""Abstract base class for vector databases."""
def __init__(self, config: VectorDatabaseConfig):
self.config = config
def add_texts(self, df: pd.DataFrame, embeddings: list):
"""Add texts and their embeddings to the collection."""
raise NotImplementedError
def search_similar_texts(self, query_embedding: list, limit: int = 5):
"""Search for similar texts based on embeddings."""
raise NotImplementedError
def drop_collection(self):
"""Drop the collection."""
raise NotImplementedError
def count_entities(self) -> int:
"""Get the number of entities in the collection."""
raise NotImplementedError