Агент с 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()