fix: load_docs.py — Агент с RAG-памятью
This commit is contained in:
+14
-12
@@ -1,19 +1,21 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from rag_agent import add_to_knowledge_base
|
||||
|
||||
from utils import splitter, vector_store
|
||||
from langchain_core.documents import Document
|
||||
|
||||
def load_documents(dir_path: str):
|
||||
for file_path in Path(dir_path).rglob("*.txt"):
|
||||
content = file_path.read_text(encoding="utf-8")
|
||||
chunks = splitter.split_text(content)
|
||||
docs = [Document(page_content=chunk, metadata={"title": file_path.name}) for chunk in chunks]
|
||||
vector_store.add_documents(docs)
|
||||
print(f"Added {len(docs)} chunks from {file_path}")
|
||||
def load_docs_from_dir(directory: str):
|
||||
for file_path in Path(directory).rglob("*"):
|
||||
if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md", ".pdf"}:
|
||||
try:
|
||||
content = file_path.read_text(encoding="utf-8", errors="ignore")
|
||||
except Exception as e:
|
||||
print(f"Failed to read {file_path}: {e}")
|
||||
continue
|
||||
title = file_path.stem
|
||||
result = add_to_knowledge_base(content, title)
|
||||
print(f"Loaded {file_path}: {result}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) != 2:
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python load_docs.py <directory>")
|
||||
sys.exit(1)
|
||||
load_documents(sys.argv[1])
|
||||
load_docs_from_dir(sys.argv[1])
|
||||
Reference in New Issue
Block a user