Files
task-6a02e23da6fe2e4ac16acf65/cli.py
T
2026-05-28 09:15:22 +00:00

44 lines
1.4 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
Commandline interface for the RAG agent.
Commands:
add <file> add document from file to knowledge base
search <q> search query in knowledge base
quit exit
"""
import argparse
import asyncio
from pathlib import Path
from langchain_core.messages import HumanMessage
from agent import agent
async def run_cli():
parser = argparse.ArgumentParser(description="RAG Agent CLI")
subparsers = parser.add_subparsers(dest="cmd", required=True)
add_parser = subparsers.add_parser("add", help="Add document from file")
add_parser.add_argument("file", type=Path, help="Path to text file")
search_parser = subparsers.add_parser("search", help="Search query in knowledge base")
search_parser.add_argument("query", type=str, help="Search string")
args = parser.parse_args()
if args.cmd == "add":
content = args.file.read_text(encoding="utf-8")
title = args.file.stem
await agent.ainvoke(
{"messages": [HumanMessage(content=f"Add document: {title}")]},
{"configurable": {"thread_id": "cli-add"}},
)
elif args.cmd == "search":
result = await agent.ainvoke(
{"messages": [HumanMessage(content=args.query)]},
{"configurable": {"thread_id": "cli-search"}},
)
print(result["messages"][-1].content)
if __name__ == "__main__": # pragma: no cover
asyncio.run(run_cli())