Удалить main.py

This commit is contained in:
2026-05-28 13:32:46 +00:00
parent 6c1ddf8bdd
commit 6d81aced8c
-82
View File
@@ -1,82 +0,0 @@
"""
CLI for interacting with the RAG agent.
Commands:
/add <file_path> add document to knowledge base
/search <query> search knowledge base
/quit exit
Optional argument: --load-dir DIR load all .txt files from directory at start.
"""
import argparse
import sys
from pathlib import Path
from rag_agent import agent, add_to_knowledge_base, search_knowledge_base
from langchain_core.messages import HumanMessage
def load_directory(dir_path: str):
"""Load all .txt files from a directory into the knowledge base."""
p = Path(dir_path)
if not p.is_dir():
print(f"{dir_path} is not a directory")
return
for txt_file in p.rglob("*.txt"):
try:
content = txt_file.read_text(encoding="utf-8")
title = txt_file.stem
add_to_knowledge_base(content=content, title=title)
print(f"Loaded {txt_file}")
except Exception as e:
print(f"Failed to load {txt_file}: {e}")
def main():
parser = argparse.ArgumentParser(description="RAG agent CLI")
parser.add_argument("--load-dir", type=str, help="Directory with .txt files to preload")
args = parser.parse_args()
if args.load_dir:
load_directory(args.load_dir)
print("Enter commands (/add <file>, /search <query>, /quit). Press Ctrl+C to exit.")
while True:
try:
line = input(">>> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nExiting.")
break
if not line:
continue
if line.lower() == "/quit":
print("Bye!")
break
if line.startswith("/add "):
path = line[5:].strip()
try:
content = Path(path).read_text(encoding="utf-8")
title = Path(path).stem
res = add_to_knowledge_base(content=content, title=title)
print(res)
except Exception as e:
print(f"Error adding file: {e}")
elif line.startswith("/search "):
query = line[8:].strip()
try:
res = search_knowledge_base(query=query, max_results=5)
print(res)
except Exception as e:
print(f"Error searching: {e}")
else:
# Treat any other input as a message to the agent
try:
result = agent.ainvoke({"messages": [HumanMessage(content=line)]})
print(result["messages"][-1].content)
except Exception as e:
print(f"Agent error: {e}")
if __name__ == "__main__":
import asyncio
asyncio.run(main())