Update src/cli.py

This commit is contained in:
2026-06-05 11:43:52 +00:00
parent f6cf962208
commit ae2e0a49c6
+36 -3
View File
@@ -1,7 +1,18 @@
"""CLI entry point for the RAG agent.
This module provides a :func:`run_cli` function that loads documents from a
directory into the knowledge base and then starts an interactive chat loop.
folder into the knowledge base and then starts an interactive chat loop.
The loop now supports the following commands:
* ``/add`` add a new document to the knowledge base. The user will be
prompted for a title and content.
* ``/search`` perform a semantic search in the knowledge base and display
the results.
* ``/quit`` exit the program.
Any other input is treated as a normal user query and is forwarded to the
agent.
"""
from __future__ import annotations
@@ -11,6 +22,7 @@ from typing import Iterable
from src.vector_store import kb
from src.agent import run_query
from src.tools import search_knowledge_base, add_to_knowledge_base
def load_documents_from_dir(directory: str | Path) -> None:
"""Load all ``.txt`` files from *directory* into the knowledge base.
@@ -40,11 +52,32 @@ def run_cli(docs_dir: str | Path) -> None:
print("\n--- RAG Agent ready. Type your question (or /quit to exit). ---\n")
while True:
user_input = input("You: ")
if user_input.strip().lower() == "/quit":
if not user_input:
continue
cmd = user_input.strip().split(" ", 1)
if cmd[0].lower() == "/quit":
print("Bye!")
break
if user_input.strip() == "":
if cmd[0].lower() == "/add":
# Prompt for title and content
title = input("Enter document title: ")
print("Enter document content. Finish with an empty line.")
lines = []
while True:
line = input()
if line == "":
break
lines.append(line)
content = "\n".join(lines)
response = add_to_knowledge_base(content=content, title=title)
print(f"Agent: {response}\n")
continue
if cmd[0].lower() == "/search":
query = cmd[1] if len(cmd) > 1 else input("Enter search query: ")
response = search_knowledge_base(query=query, max_results=5)
print(f"Agent: {response}\n")
continue
# Default: forward to agent
response = run_query(user_input)
print(f"Agent: {response}\n")