70 lines
1.9 KiB
Python
70 lines
1.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Initialization client - loads documents from a directory into the vector store."""
|
|
|
|
import os
|
|
from typing import List
|
|
|
|
from langchain_core.documents import Document
|
|
from langchain_community.documenloaders import TextLoader
|
|
|
|
from vector_store import get_vector_store, add_documents_to_store
|
|
|
|
|
|
def load_documents_from_dir(dir_path: str) -> List[Document]:
|
|
"""Load all .txt files from a directory as LangChain Documents.
|
|
|
|
Args:
|
|
dir_path: Path to directory containing .txt files.
|
|
|
|
Returns:
|
|
List of Document objects.
|
|
"""
|
|
documents = []
|
|
for filename in os.listdir(dir_path):
|
|
if filename.endswith(".txt"):
|
|
file_path = os.path.join(dir_path, filename)
|
|
loader = TextLoader(file_path, encoding="utf-8")
|
|
docs = loader.load()
|
|
# Add title metadata from filename
|
|
title = os.path.splitext(filename)[0]
|
|
for doc in docs:
|
|
doc.metadata["title"] = title
|
|
documents.extend(docs)
|
|
return documents
|
|
|
|
|
|
def init_knowledge_base(dir_path: str) -> int:
|
|
"""Initialize the knowledge base by loading documents from a directory.
|
|
|
|
Args:
|
|
dir_path: Path to directory containing .txt files.
|
|
|
|
Returns:
|
|
Number of chunks added to the store.
|
|
"""
|
|
store = get_vector_store()
|
|
documents = load_documents_from_dir(dir_path)
|
|
|
|
if not documents:
|
|
print(f"No .txt files found in {dir_path}")
|
|
return 0
|
|
|
|
ids = add_documents_to_store(store, documents)
|
|
print(f"Loaded {len(documents)} document(s) into knowledge base ({len(ids)} chunks).")
|
|
return len(ids)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python init_client.py <directory_path>")
|
|
sys.exit(1)
|
|
|
|
dir_path = sys.argv[1]
|
|
if not os.path.isdir(dir_path):
|
|
print(f"Error: {dir_path} is not a directory")
|
|
sys.exit(1)
|
|
|
|
init_knowledge_base(dir_path)
|