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