Агент с RAG‑памятью: cli_client.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/cli_client.py

@@ -4,7 +4,7 @@
 """
 cli_client.py
 
-Интерактивный клиент для работы с агентом RAG‑памяти.
+Интерактивный клиент для работы с агентом RAG.
 Поддерживает команды:
     /add <title> <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__":