diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/cli_client.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/cli_client.py index 069eeb2..bc5e72f 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/cli_client.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/cli_client.py @@ -4,7 +4,7 @@ """ cli_client.py -Интерактивный клиент для работы с агентом RAG‑памяти. +Интерактивный клиент для работы с агентом RAG. Поддерживает команды: /add <content> – добавить документ в базу знаний /search <query> – выполнить семантический поиск @@ -14,110 +14,122 @@ cli_client.py import sys from typing import List -# Предполагается, что модули с инструментами уже реализованы и находятся в пакете `rag_agent`. -# Импортируем необходимые функции. +# Rich для красивого вывода try: - from rag_agent.tools import add_to_knowledge_base, search_knowledge_base # noqa: F401 -except Exception as exc: - print(f"Ошибка при импорте инструментов: {exc}", file=sys.stderr) + 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) -def parse_command(line: str) -> (str, List[str]): +console = Console() + + +def _parse_command(line: str) -> List[str]: """ - Разбирает строку команды и возвращает имя команды и список аргументов. - Команда начинается с символа '/'. + Разбирает строку ввода пользователя. + Возвращает список токенов, разделённых пробелами, + но сохраняет аргументы в кавычках как один элемент. """ - if not line.startswith("/"): - raise ValueError("Команды должны начинаться со знака '/'.") - parts = line.strip().split(maxsplit=1) - cmd = parts[0] - args_str = parts[1] if len(parts) > 1 else "" - # Для /add аргументы разделены пробелом: title content - # Чтобы сохранить пробелы в содержимом, делаем split только один раз. - return cmd, [args_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(args: List[str]) -> None: +def _handle_add(tokens: List[str]) -> None: """ Обработчик команды /add. - Ожидает строку вида "<title> <content>". + Ожидает два аргумента: title и content (в кавычках). """ - if not args or not args[0]: - print("Ошибка: нужно указать заголовок и содержимое.") + if len(tokens) < 3: + console.print("[red]Неверное количество аргументов. Используйте:/add \"title\" \"content\"[/red]") return - # Разделяем на заголовок и контент. Заголовок – первое слово, остальное – контент. - parts = args[0].split(maxsplit=1) - if len(parts) < 2: - print("Ошибка: нужно указать как минимум заголовок и содержимое.") - return - title, content = parts[0], parts[1] + title = tokens[1] + content = " ".join(tokens[2:]) try: add_to_knowledge_base(content=content, title=title) - print(f"Документ '{title}' успешно добавлен в базу знаний.") - except Exception as exc: - print(f"Ошибка при добавлении документа: {exc}") + console.print(f"[green]Документ '{title}' успешно добавлен.[/green]") + except Exception as exc: # pragma: no cover + console.print(f"[red]Ошибка при добавлении:[/red] {exc}") -def handle_search(args: List[str]) -> None: +def _handle_search(tokens: List[str]) -> None: """ Обработчик команды /search. - Ожидает строку запроса. + Ожидает один аргумент – запрос (в кавычках). """ - if not args or not args[0]: - print("Ошибка: нужно указать поисковый запрос.") + if len(tokens) < 2: + console.print("[red]Неверное количество аргументов. Используйте:/search \"query\"[/red]") return - query = args[0] + query = " ".join(tokens[1:]) try: results = search_knowledge_base(query=query, max_results=5) if not results: - print("Ничего не найдено.") + console.print("[yellow]Результаты не найдены.[/yellow]") return - print(f"Найдено {len(results)} результатов:") - for i, res in enumerate(results, 1): - # Предполагаем, что результат – словарь с полями 'title' и 'content'. - title = res.get("metadata", {}).get("title", "Без названия") - snippet = res.get("content")[:200] + ("..." if len(res.get("content", "")) > 200 else "") - print(f"\n{i}. {title}\n {snippet}") - except Exception as exc: - print(f"Ошибка при поиске: {exc}") + + 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. """ - print("=== RAG‑память CLI ===") - print("Доступные команды:") - print("/add <title> <content> – добавить документ в базу знаний") - print("/search <query> – выполнить семантический поиск") - print("/quit – выйти из программы\n") + 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 = input("> ").strip() - except (EOFError, KeyboardInterrupt): - print("\nВыход.") + 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: + if not line.strip(): continue - try: - cmd, args = parse_command(line) - except ValueError as ve: - print(ve) + tokens = _parse_command(line) + if not tokens: continue - if cmd == "/add": - handle_add(args) - elif cmd == "/search": - handle_search(args) - elif cmd == "/quit": - print("До свидания!") + 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: - print(f"Неизвестная команда: {cmd}") + console.print(f"[red]Неизвестная команда: {command}[/red]") if __name__ == "__main__":