Update src/cli.py
This commit is contained in:
+36
-3
@@ -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")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user