File size: 983 Bytes
e68d535
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
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