14 lines
444 B
Python
14 lines
444 B
Python
|
import numpy as np
|
||
|
|
import faiss
|
||
|
|
from backend.config import Config
|
||
|
|
|
||
|
|
class VectorDB:
|
||
|
|
def __init__(self):
|
||
|
|
self.index = faiss.IndexFlatL2(768) # Cohere embedding dim
|
||
|
|
|
||
|
|
def add_vectors(self, ids, embeddings):
|
||
|
|
self.index.add(np.array(embeddings).astype('float32'))
|
||
|
|
|
||
|
|
def search(self, query_embedding, k=5):
|
||
|
|
distances, indices = self.index.search(np.array([query_embedding]), k)
|
||
|
|
return indices[0]
|