24 lines
1.0 KiB
Python
24 lines
1.0 KiB
Python
import os
|
|
from rag_store import KnowledgeBase
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
from langchain_core.documents import Document
|
|
|
|
def load_documents_from_dir(directory: str, kb: KnowledgeBase):
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
for root, _, files in os.walk(directory):
|
|
for file in files:
|
|
if file.lower().endswith(('.txt', '.md')):
|
|
path = os.path.join(root, file)
|
|
try:
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
text = f.read()
|
|
docs = [Document(page_content=text, metadata={"title": file})]
|
|
chunks = splitter.split_documents(docs)
|
|
kb.add_documents(chunks)
|
|
print(f"Loaded {file} into KB.")
|
|
except Exception as e:
|
|
print(f"Failed to load {file}: {e}")
|
|
|
|
if __name__ == "__main__":
|
|
kb = KnowledgeBase()
|
|
load_documents_from_dir("./data", kb) |