File size: 3,703 Bytes
e454383 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 | 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 chunking if transcript_chunker is available, else treat as single doc
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:
# Fallback: treat as single document
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)
# Remove existing vector DB if present
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.")
|