| import os |
| import chromadb |
| from chromadb.utils import embedding_functions |
|
|
| |
| 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"): |
| |
| os.makedirs(PERSIST_DIR, exist_ok=True) |
| |
| self.client = chromadb.PersistentClient(path=PERSIST_DIR) |
| |
| |
| self.embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction( |
| model_name="paraphrase-multilingual-MiniLM-L12-v2" |
| ) |
| |
| 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) |
|
|