62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
"""
|
|
init_knowledge_base.py
|
|
|
|
Script for loading documents from a directory into the vector store.
|
|
|
|
Usage:
|
|
python init_knowledge_base.py <documents_directory>
|
|
|
|
Supported file types: .txt, .md
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
from pathlib import Path
|
|
from vector_store import add_documents
|
|
|
|
SUPPORTED_EXTENSIONS = {".txt", ".md"}
|
|
|
|
|
|
def load_documents_from_directory(directory: str) -> None:
|
|
dir_path = Path(directory)
|
|
if not dir_path.exists() or not dir_path.is_dir():
|
|
print(f"Error: '{directory}' is not a valid directory.")
|
|
sys.exit(1)
|
|
|
|
files = [
|
|
f for f in dir_path.rglob("*")
|
|
if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS
|
|
]
|
|
|
|
if not files:
|
|
print(f"No supported files found in '{directory}'.")
|
|
return
|
|
|
|
print(f"Found {len(files)} file(s) to load.\n")
|
|
total_chunks = 0
|
|
|
|
for file_path in files:
|
|
try:
|
|
content = file_path.read_text(encoding="utf-8")
|
|
title = file_path.stem
|
|
num_chunks = add_documents(content=content, title=title)
|
|
total_chunks += num_chunks
|
|
print(f" [OK] '{file_path.name}' -> {num_chunks} chunk(s) added.")
|
|
except Exception as e:
|
|
print(f" [ERROR] '{file_path.name}': {e}")
|
|
|
|
print(f"\nDone. Total chunks added: {total_chunks}")
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python init_knowledge_base.py <documents_directory>")
|
|
sys.exit(1)
|
|
|
|
directory = sys.argv[1]
|
|
print(f"Loading documents from: {directory}\n")
|
|
load_documents_from_directory(directory)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |