Add initialization client to load documents from directory
This commit is contained in:
@@ -0,0 +1,69 @@
|
||||
#!/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)
|
||||
Reference in New Issue
Block a user