Добавить cli_client.py

2026-05-28 09:26:47 +00:00
parent ae4622db14
commit 8c08255a98
1 changed files with 136 additions and 0 deletions
@@ -0,0 +1,136 @@
+#!/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()