| import os
|
| from dotenv import load_dotenv
|
| from utils.document_management import load_document, transcript_chunker, merge_chunks
|
| from utils.vector_db import VectorDB
|
| from langchain.schema import Document
|
| from uuid import uuid4
|
| from pathlib import Path
|
| from typing import Optional
|
|
|
| def build_vector_db_for_dir(
|
| data_dir: str,
|
| collection_name: str = "custom_collection",
|
| embedding_length: int = 1536,
|
| file_suffix: Optional[str] = None,
|
| openai_api_key: Optional[str] = None,
|
| verbose: bool = True,
|
| ):
|
| """
|
| Build a new vector database for all files in a directory.
|
| Args:
|
| data_dir: Directory containing files to process.
|
| collection_name: Name for the vector DB collection.
|
| embedding_length: Embedding vector length.
|
| file_suffix: Only process files with this suffix (e.g. '.srt', '.md'). If None, process all files.
|
| openai_api_key: API key for embedding (if required by backend).
|
| verbose: Print progress.
|
| """
|
| load_dotenv()
|
| working_dir = os.path.abspath(data_dir)
|
| if not os.path.isdir(working_dir):
|
| raise ValueError(f"Directory not found: {working_dir}")
|
|
|
| all_documents = []
|
| all_files = []
|
| for root, dirs, files in os.walk(working_dir):
|
| for file in files:
|
| if file_suffix is None or file.endswith(file_suffix):
|
| file_path = os.path.join(root, file)
|
| relative_path = os.path.relpath(file_path, working_dir)
|
| all_files.append(relative_path)
|
|
|
| for file_path in all_files:
|
| path = os.path.join(working_dir, file_path)
|
| doc = load_document(path)
|
| if doc is None:
|
| if verbose:
|
| print(f"Skipping file {file_path} as it could not be loaded.")
|
| continue
|
|
|
| try:
|
| chunks = transcript_chunker(doc, filename=file_path)
|
| merged_chunks = merge_chunks(chunks)
|
| documents = [
|
| Document(
|
| page_content=chunk["text"],
|
| metadata={**chunk, "chunk_id": i, "id": str(uuid4()), "filename": file_path}
|
| )
|
| for i, chunk in enumerate(merged_chunks)
|
| ]
|
| except Exception:
|
|
|
| documents = [
|
| Document(
|
| page_content=doc if isinstance(doc, str) else str(doc),
|
| metadata={"id": str(uuid4()), "filename": file_path}
|
| )
|
| ]
|
| if verbose:
|
| print(f"Processing {file_path} with {len(documents)} chunks")
|
| all_documents.extend(documents)
|
|
|
|
|
| chroma_db_path = os.path.join(working_dir, "chroma_db")
|
| if all_documents:
|
| if os.path.exists(chroma_db_path):
|
| import shutil
|
| shutil.rmtree(chroma_db_path)
|
| if verbose:
|
| print("Removed existing vector database directory.")
|
| vector_db = VectorDB(
|
| collection_name=collection_name,
|
| embedding_length=embedding_length,
|
| open_ai_api_key=openai_api_key or os.getenv("OPENAI_API_KEY"),
|
| working_dir=working_dir,
|
| documents=all_documents
|
| )
|
| if verbose:
|
| print(f"Created new vector database with {len(all_documents)} documents")
|
| else:
|
| if verbose:
|
| print("No documents found to process")
|
| if verbose:
|
| print("Vector database reset and data upload complete with ChromaDB.")
|
|
|
