31 lines
1.1 KiB
Python
31 lines
1.1 KiB
Python
import argparse
|
|
from tools import add_to_knowledge_base, search_knowledge_base
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="RAG Agent CLI")
|
|
subparsers = parser.add_subparsers(dest='command')
|
|
|
|
add_parser = subparsers.add_parser('add', help='Add document to knowledge base')
|
|
add_parser.add_argument('--title', required=True, help='Title of the document')
|
|
add_parser.add_argument('--content', required=True, help='Content of the document')
|
|
|
|
search_parser = subparsers.add_parser('search', help='Search knowledge base')
|
|
search_parser.add_argument('--query', required=True, help='Query string')
|
|
search_parser.add_argument('--max', type=int, default=5, help='Max results')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.command == 'add':
|
|
result = add_to_knowledge_base(args.content, args.title)
|
|
print(result)
|
|
elif args.command == 'search':
|
|
results = search_knowledge_base(args.query, args.max)
|
|
for i, r in enumerate(results, 1):
|
|
print(f"{i}. {r}")
|
|
else:
|
|
parser.print_help()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|