95 lines
3.0 KiB
Python
95 lines
3.0 KiB
Python
"""CLI entry point for the RAG agent.
|
||
|
||
This module provides a :func:`run_cli` function that loads documents from a
|
||
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
|
||
|
||
from pathlib import Path
|
||
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.
|
||
|
||
Parameters
|
||
----------
|
||
directory:
|
||
Path to the folder containing documents.
|
||
"""
|
||
directory = Path(directory)
|
||
for file_path in directory.rglob("*.txt"):
|
||
title = file_path.stem
|
||
content = file_path.read_text(encoding="utf-8")
|
||
kb.add_document(title=title, content=content)
|
||
print(f"Loaded {title}")
|
||
|
||
def run_cli(docs_dir: str | Path) -> None:
|
||
"""Run an interactive CLI.
|
||
|
||
Parameters
|
||
----------
|
||
docs_dir:
|
||
Directory with documents to load into the knowledge base.
|
||
"""
|
||
load_documents_from_dir(docs_dir)
|
||
|
||
print("\n--- RAG Agent ready. Type your question (or /quit to exit). ---\n")
|
||
while True:
|
||
user_input = input("You: ")
|
||
if not user_input:
|
||
continue
|
||
cmd = user_input.strip().split(" ", 1)
|
||
if cmd[0].lower() == "/quit":
|
||
print("Bye!")
|
||
break
|
||
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")
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="RAG agent CLI")
|
||
parser.add_argument(
|
||
"--docs",
|
||
type=str,
|
||
default="docs",
|
||
help="Directory with documents to load into the knowledge base.",
|
||
)
|
||
args = parser.parse_args()
|
||
run_cli(args.docs) |