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

2026-05-27 06:34:20 +00:00
parent 3d3cb640a9
commit fb9bc02371
1 changed files with 124 additions and 0 deletions
@@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+cli_client.py
+
+Интерактивный клиент для работы с агентом RAG‑памяти.
+Поддерживает команды:
+    /add <title> <content>   – добавить документ в базу знаний
+    /search <query>          – выполнить семантический поиск
+    /quit                    – выйти из программы
+"""
+
+import sys
+from typing import List
+
+# Предполагается, что модули с инструментами уже реализованы и находятся в пакете `rag_agent`.
+# Импортируем необходимые функции.
+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)
+    sys.exit(1)
+
+
+def parse_command(line: str) -> (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]
+
+
+def handle_add(args: List[str]) -> None:
+    """
+    Обработчик команды /add.
+    Ожидает строку вида "<title> <content>".
+    """
+    if not args or not args[0]:
+        print("Ошибка: нужно указать заголовок и содержимое.")
+        return
+    # Разделяем на заголовок и контент. Заголовок – первое слово, остальное – контент.
+    parts = args[0].split(maxsplit=1)
+    if len(parts) < 2:
+        print("Ошибка: нужно указать как минимум заголовок и содержимое.")
+        return
+    title, content = parts[0], parts[1]
+    try:
+        add_to_knowledge_base(content=content, title=title)
+        print(f"Документ '{title}' успешно добавлен в базу знаний.")
+    except Exception as exc:
+        print(f"Ошибка при добавлении документа: {exc}")
+
+
+def handle_search(args: List[str]) -> None:
+    """
+    Обработчик команды /search.
+    Ожидает строку запроса.
+    """
+    if not args or not args[0]:
+        print("Ошибка: нужно указать поисковый запрос.")
+        return
+    query = args[0]
+    try:
+        results = search_knowledge_base(query=query, max_results=5)
+        if not results:
+            print("Ничего не найдено.")
+            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}")
+
+
+def main() -> None:
+    """
+    Основной цикл интерактивного клиента.
+    """
+    print("=== RAG‑память CLI ===")
+    print("Доступные команды:")
+    print("/add <title> <content>   – добавить документ в базу знаний")
+    print("/search <query>          – выполнить семантический поиск")
+    print("/quit                    – выйти из программы\n")
+
+    while True:
+        try:
+            line = input("> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nВыход.")
+            break
+
+        if not line:
+            continue
+
+        try:
+            cmd, args = parse_command(line)
+        except ValueError as ve:
+            print(ve)
+            continue
+
+        if cmd == "/add":
+            handle_add(args)
+        elif cmd == "/search":
+            handle_search(args)
+        elif cmd == "/quit":
+            print("До свидания!")
+            break
+        else:
+            print(f"Неизвестная команда: {cmd}")
+
+
+if __name__ == "__main__":
+    main()