Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| from pathlib import Path | |
| from typing import List, Optional | |
| from langchain_core.documents import Document | |
| from src.config import settings | |
| from src.logger import logger | |
| def get_vectorstore( | |
| documents: List[Document], | |
| embedding_model, | |
| collection_name: str, | |
| persist_directory: str, | |
| store_type: Optional[str] = None, | |
| ): | |
| """ | |
| Create and populate a vector store. | |
| store_type: "faiss" | "chroma" — if omitted, falls back to settings.vector_db_type. | |
| """ | |
| store_type = (store_type or settings.vector_db_type or "faiss").lower().strip() | |
| logger.info( | |
| "Building vector store", | |
| context={ | |
| "type": store_type, | |
| "docs": len(documents), | |
| "collection": collection_name, | |
| "persist_dir": persist_directory, | |
| }, | |
| ) | |
| Path(persist_directory).mkdir(parents=True, exist_ok=True) | |
| if store_type == "chroma": | |
| from langchain_chroma import Chroma | |
| store = Chroma( | |
| collection_name=collection_name, | |
| persist_directory=persist_directory, | |
| embedding_function=embedding_model, | |
| ) | |
| batch_size = int(os.getenv("VECTORSTORE_BATCH_SIZE", "50")) | |
| total = len(documents) | |
| for i in range(0, total, batch_size): | |
| end = min(i + batch_size, total) | |
| logger.info( | |
| f"Ingesting batch {i // batch_size + 1}/{(total - 1) // batch_size + 1}" | |
| ) | |
| store.add_documents(documents[i:end]) | |
| logger.info("Chroma ingestion complete", context={"path": persist_directory}) | |
| return store | |
| if store_type == "faiss": | |
| from langchain_community.vectorstores import FAISS | |
| # Prefer Embeddings object; ST wrappers should implement embed_documents/embed_query | |
| store = FAISS.from_documents(documents, embedding_model) | |
| store.save_local(persist_directory) | |
| logger.info("FAISS store saved", context={"path": persist_directory}) | |
| return store | |
| raise ValueError(f"Unsupported vector_db_type: {store_type}") |