Удалить main.py
This commit is contained in:
@@ -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())
|
|
||||||
Reference in New Issue
Block a user