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]