136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
cli_client.py
|
|
|
|
Интерактивный клиент для работы с агентом RAG.
|
|
Поддерживает команды:
|
|
/add <title> <content> – добавить документ в базу знаний
|
|
/search <query> – выполнить семантический поиск
|
|
/quit – выйти из программы
|
|
"""
|
|
|
|
import sys
|
|
from typing import List
|
|
|
|
# Rich для красивого вывода
|
|
try:
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
except ImportError: # pragma: no cover
|
|
print("Необходимо установить пакет 'rich'")
|
|
sys.exit(1)
|
|
|
|
# Инструменты агента (предполагается, что они реализованы в другом модуле)
|
|
try:
|
|
from rag_tools import search_knowledge_base, add_to_knowledge_base
|
|
except Exception as exc: # pragma: no cover
|
|
print(f"Ошибка при импорте инструментов: {exc}")
|
|
sys.exit(1)
|
|
|
|
console = Console()
|
|
|
|
|
|
def _parse_command(line: str) -> List[str]:
|
|
"""
|
|
Разбирает строку ввода пользователя.
|
|
Возвращает список токенов, разделённых пробелами,
|
|
но сохраняет аргументы в кавычках как один элемент.
|
|
"""
|
|
import shlex
|
|
|
|
try:
|
|
return shlex.split(line)
|
|
except ValueError as exc: # pragma: no cover
|
|
console.print(f"[red]Ошибка разбора команды:[/red] {exc}")
|
|
return []
|
|
|
|
|
|
def _handle_add(tokens: List[str]) -> None:
|
|
"""
|
|
Обработчик команды /add.
|
|
Ожидает два аргумента: title и content (в кавычках).
|
|
"""
|
|
if len(tokens) < 3:
|
|
console.print("[red]Неверное количество аргументов. Используйте:/add \"title\" \"content\"[/red]")
|
|
return
|
|
title = tokens[1]
|
|
content = " ".join(tokens[2:])
|
|
try:
|
|
add_to_knowledge_base(content=content, title=title)
|
|
console.print(f"[green]Документ '{title}' успешно добавлен.[/green]")
|
|
except Exception as exc: # pragma: no cover
|
|
console.print(f"[red]Ошибка при добавлении:[/red] {exc}")
|
|
|
|
|
|
def _handle_search(tokens: List[str]) -> None:
|
|
"""
|
|
Обработчик команды /search.
|
|
Ожидает один аргумент – запрос (в кавычках).
|
|
"""
|
|
if len(tokens) < 2:
|
|
console.print("[red]Неверное количество аргументов. Используйте:/search \"query\"[/red]")
|
|
return
|
|
query = " ".join(tokens[1:])
|
|
try:
|
|
results = search_knowledge_base(query=query, max_results=5)
|
|
if not results:
|
|
console.print("[yellow]Результаты не найдены.[/yellow]")
|
|
return
|
|
|
|
table = Table(title="Поиск", show_lines=True)
|
|
table.add_column("Номер", style="cyan")
|
|
table.add_column("Заголовок", style="magenta")
|
|
table.add_column("Содержание")
|
|
|
|
for idx, res in enumerate(results, start=1):
|
|
# Предполагается, что результат – словарь с ключами title и content
|
|
title = res.get("title", "Без заголовка")
|
|
content = res.get("content", "")
|
|
snippet = (content[:200] + "...") if len(content) > 200 else content
|
|
table.add_row(str(idx), title, snippet)
|
|
|
|
console.print(table)
|
|
except Exception as exc: # pragma: no cover
|
|
console.print(f"[red]Ошибка при поиске:[/red] {exc}")
|
|
|
|
|
|
def main() -> None:
|
|
"""
|
|
Основной цикл CLI.
|
|
"""
|
|
console.print("[bold green]Добро пожаловать в RAG-агент![/bold green]")
|
|
console.print("Доступные команды:")
|
|
console.print("/add \"title\" \"content\" – добавить документ")
|
|
console.print("/search \"query\" – выполнить поиск")
|
|
console.print("/quit – выйти из программы")
|
|
|
|
while True:
|
|
try:
|
|
line = console.input("\n[bold blue]Введите команду[/bold blue]> ")
|
|
except (KeyboardInterrupt, EOFError): # pragma: no cover
|
|
console.print("\n[bold red]Выход...[/bold red]")
|
|
break
|
|
|
|
if not line.strip():
|
|
continue
|
|
|
|
tokens = _parse_command(line)
|
|
if not tokens:
|
|
continue
|
|
|
|
command = tokens[0].lower()
|
|
if command == "/add":
|
|
_handle_add(tokens)
|
|
elif command == "/search":
|
|
_handle_search(tokens)
|
|
elif command in {"/quit", "exit"}:
|
|
console.print("[bold red]Выход...[/bold red]")
|
|
break
|
|
else:
|
|
console.print(f"[red]Неизвестная команда: {command}[/red]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |