letxinet / backend /vector_store.py
C2MV's picture
Initial upload for Build Small Hackathon
68fb5e2 verified
Raw
History Blame Contribute Delete
1.94 kB
import os
import chromadb
from chromadb.utils import embedding_functions
# Configurar persistencia en la carpeta raíz del proyecto
PERSIST_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".chroma_db")
class VectorStore:
"""Implementación nativa de base de datos vectorial usando ChromaDB para búsquedas locales."""
def __init__(self, collection_name: str = "letxipu_docs"):
# Asegurar que el directorio de persistencia exista
os.makedirs(PERSIST_DIR, exist_ok=True)
self.client = chromadb.PersistentClient(path=PERSIST_DIR)
# Usar un modelo ligero local multilingüe o estándar
self.embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
model_name="paraphrase-multilingual-MiniLM-L12-v2" # Soporta español e inglés muy bien
)
self.collection = self.client.get_or_create_collection(
name=collection_name,
embedding_function=self.embedding_fn
)
def add_documents(self, documents: list[str], metadatas: list[dict], ids: list[str]):
"""Añade documentos (chunks) a la base vectorial."""
if not documents:
return
self.collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
def search(self, query: str, n_results: int = 5, filter_dict: dict = None) -> dict:
"""Busca los fragmentos semánticamente más similares a la consulta."""
results = self.collection.query(
query_texts=[query],
n_results=n_results,
where=filter_dict
)
return results
def clear(self):
"""Elimina todos los documentos de la colección actual."""
all_ids = self.collection.get().get("ids", [])
if all_ids:
self.collection.delete(ids=all_ids)