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

2026-05-26 09:21:09 +00:00
parent 66ee72055f
commit 2387a64181
1 changed files with 158 additions and 0 deletions
@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+cli_client.py
+
+Интерактивный клиент для работы с агентом RAG‑памяти.
+Поддерживает команды:
+    /add     - добавить документ в базу знаний
+    /search  - выполнить семантический поиск по базе
+    /quit    - выйти из программы
+
+Команды реализованы через инструменты, объявленные в модуле tools.py:
+    add_to_knowledge_base(content: str, title: str) -> None
+    search_knowledge_base(query: str, max_results: int = 5) -> list[dict]
+
+Для красивого вывода используется библиотека rich.
+"""
+
+import sys
+from pathlib import Path
+
+try:
+    from rich.console import Console
+    from rich.table import Table
+except ImportError as exc:
+    print("Необходимо установить rich: pip install rich", file=sys.stderr)
+    raise exc
+
+# Импортируем инструменты из общего модуля. Предполагается, что они уже реализованы.
+try:
+    from tools import add_to_knowledge_base, search_knowledge_base
+except Exception as exc:
+    print("Не удалось импортировать инструменты из modules.tools", file=sys.stderr)
+    raise exc
+
+console = Console()
+
+
+def _print_help() -> None:
+    """Вывод справки по доступным командам."""
+    console.print(
+        """
+[bold cyan]/add[/]     – добавить документ в базу знаний
+[bold cyan]/search[/]  – выполнить поиск по базе
+[bold cyan]/quit[/]   – выйти из программы
+
+Для добавления можно указать путь к файлу или вводить текст вручную.
+"""
+    )
+
+
+def _handle_add() -> None:
+    """Обработчик команды /add."""
+    console.print("[green]Введите заголовок документа:[/]")
+    title = input("> ").strip()
+    if not title:
+        console.print("[red]Заголовок не может быть пустым.[/]")
+        return
+
+    console.print(
+        "[green]Выберите способ ввода содержимого:\n"
+        "1 – Ввести текст вручную\n"
+        "2 – Указать путь к файлу[/]"
+    )
+    choice = input("> ").strip()
+    if choice == "2":
+        path_str = input("[green]Введите путь к файлу:[/]").strip()
+        try:
+            content = Path(path_str).read_text(encoding="utf-8")
+        except Exception as exc:
+            console.print(f"[red]Не удалось прочитать файл: {exc}[/]")
+            return
+    else:
+        console.print("[green]Введите содержимое (конец ввода – пустая строка):[/]")
+        lines = []
+        while True:
+            line = input()
+            if line == "":
+                break
+            lines.append(line)
+        content = "\n".join(lines)
+
+    try:
+        add_to_knowledge_base(content=content, title=title)
+        console.print(f"[bold green]Документ '{title}' успешно добавлен.[/]")
+    except Exception as exc:
+        console.print(f"[red]Ошибка при добавлении документа: {exc}[/]")
+
+
+def _handle_search() -> None:
+    """Обработчик команды /search."""
+    query = input("[green]Введите запрос для поиска:[/]").strip()
+    if not query:
+        console.print("[red]Запрос не может быть пустым.[/]")
+        return
+
+    max_results_str = input(
+        "[green]Укажите количество результатов (по умолчанию 5):[/]"
+    ).strip()
+    try:
+        max_results = int(max_results_str) if max_results_str else 5
+    except ValueError:
+        console.print("[red]Неверное число. Будет использовано значение по умолчанию: 5[/]")
+        max_results = 5
+
+    try:
+        results = search_knowledge_base(query=query, max_results=max_results)
+    except Exception as exc:
+        console.print(f"[red]Ошибка при поиске: {exc}[/]")
+        return
+
+    if not results:
+        console.print("[yellow]Ничего не найдено.[/]")
+        return
+
+    table = Table(title="Результаты поиска", show_lines=True)
+    table.add_column("№", style="dim")
+    table.add_column("Заголовок", style="bold cyan")
+    table.add_column("Текст")
+
+    for idx, item in enumerate(results, start=1):
+        title = item.get("title", "Без заголовка")
+        snippet = item.get("content", "")[:200] + ("…" if len(item.get("content", "")) > 200 else "")
+        table.add_row(str(idx), title, snippet)
+
+    console.print(table)
+
+
+def main() -> None:
+    """Главный цикл программы."""
+    console.print("[bold magenta]=== RAG‑Память CLI ===[/]")
+    _print_help()
+
+    while True:
+        try:
+            command = input("\n[bold cyan]> [/]").strip()
+        except (KeyboardInterrupt, EOFError):
+            console.print("\n[bold red]Выход...[/]")
+            break
+
+        if not command:
+            continue
+
+        if command.lower() == "/quit":
+            console.print("[bold red]Завершение работы.[/]")
+            break
+        elif command.lower() == "/add":
+            _handle_add()
+        elif command.lower() == "/search":
+            _handle_search()
+        else:
+            console.print(f"[red]Неизвестная команда: {command}[/]")
+            _print_help()
+
+
+if __name__ == "__main__":
+    main()