Добавить client.py

2026-05-28 10:09:52 +00:00
parent 3109ea9a43
commit 43424eedb4
1 changed files with 154 additions and 0 deletions
@@ -0,0 +1,154 @@
+"""
+client.py — интерактивный CLI-клиент для демонстрации работы RAG-агента.
+
+Команды:
+  /add <текст>          — добавить текст прямо в базу знаний
+  /add-file <путь>      — загрузить файл в базу знаний
+  /search <запрос>      — прямой семантический поиск (без агента)
+  /quit или /exit       — выйти из программы
+  <любой другой текст>  — отправить запрос агенту
+"""
+
+import sys
+from pathlib import Path
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.markdown import Markdown
+from rich.prompt import Prompt
+from rich.rule import Rule
+
+from vector_store import init_collection, add_documents, search
+from agent import run_agent
+
+console = Console()
+
+BANNER = """
+╔══════════════════════════════════════════╗
+║       🤖  RAG-Agent  •  Qdrant+Ollama    ║
+║  /add <текст>   — добавить в базу        ║
+║  /add-file <путь> — загрузить файл       ║
+║  /search <запрос> — поиск в базе         ║
+║  /quit          — выйти                  ║
+╚══════════════════════════════════════════╝
+"""
+
+
+def handle_add(text: str) -> None:
+    """Добавляет текст в базу знаний через инструмент."""
+    if not text.strip():
+        console.print("[yellow]⚠  Укажите текст после /add[/yellow]")
+        return
+    title = Prompt.ask("  Название документа (Enter — пропустить)", default="")
+    n = add_documents(text.strip(), title=title)
+    console.print(
+        f"[green]✓ Добавлено {n} чанков в базу знаний"
+        + (f" (источник: «{title}»)" if title else "")
+        + "[/green]"
+    )
+
+
+def handle_add_file(path_str: str) -> None:
+    """Загружает файл в базу знаний."""
+    p = Path(path_str.strip())
+    if not p.exists():
+        console.print(f"[red]✗ Файл не найден: {p}[/red]")
+        return
+    content = p.read_text(encoding="utf-8", errors="ignore")
+    n = add_documents(content, title=p.name)
+    console.print(f"[green]✓ Файл «{p.name}» загружен, создано {n} чанков[/green]")
+
+
+def handle_search(query: str) -> None:
+    """Прямой семантический поиск — без агента."""
+    if not query.strip():
+        console.print("[yellow]⚠  Укажите запрос после /search[/yellow]")
+        return
+    console.print(f"[dim]Поиск: «{query}»...[/dim]")
+    results = search(query, max_results=5)
+    if not results:
+        console.print("[yellow]Ничего не найдено.[/yellow]")
+        return
+    console.print(Rule("Результаты поиска"))
+    for i, (doc, score) in enumerate(results, 1):
+        title = doc.metadata.get("title", "—")
+        chunk = doc.metadata.get("chunk_index", 0)
+        console.print(
+            Panel(
+                doc.page_content,
+                title=f"[cyan]#{i}  «{title}»  chunk={chunk}  score={score:.3f}[/cyan]",
+                border_style="dim",
+            )
+        )
+
+
+def handle_agent(user_input: str) -> None:
+    """Отправляет запрос агенту и выводит ответ."""
+    console.print("[dim]Агент думает...[/dim]")
+    try:
+        answer = run_agent(user_input)
+        console.print(Panel(Markdown(answer), title="[green]Агент[/green]", border_style="green"))
+    except Exception as e:
+        console.print(f"[red]✗ Ошибка агента: {e}[/red]")
+
+
+def main() -> None:
+    console.print(BANNER)
+
+    # Инициализация коллекции при старте
+    try:
+        init_collection()
+        console.print("[green]✓ Подключение к Qdrant установлено[/green]\n")
+    except Exception as e:
+        console.print(f"[red]✗ Не удалось подключиться к Qdrant: {e}[/red]")
+        console.print("[yellow]Убедитесь, что Qdrant запущен: docker run -p 6333:6333 qdrant/qdrant[/yellow]")
+        sys.exit(1)
+
+    while True:
+        try:
+            user_input = Prompt.ask("\n[bold cyan]Вы[/bold cyan]").strip()
+        except (KeyboardInterrupt, EOFError):
+            console.print("\n[yellow]До свидания![/yellow]")
+            break
+
+        if not user_input:
+            continue
+
+        # ── Команды ───────────────────────────────────────────────────────────
+        if user_input.lower() in ("/quit", "/exit", "/выход"):
+            console.print("[yellow]До свидания![/yellow]")
+            break
+
+        elif user_input.lower().startswith("/add-file "):
+            path_part = user_input[len("/add-file "):]
+            handle_add_file(path_part)
+
+        elif user_input.lower().startswith("/add "):
+            text_part = user_input[len("/add "):]
+            handle_add(text_part)
+
+        elif user_input.lower() == "/add":
+            # /add без аргументов — многострочный ввод
+            console.print("[dim]Введите текст (пустая строка для завершения):[/dim]")
+            lines = []
+            while True:
+                line = input()
+                if line == "":
+                    break
+                lines.append(line)
+            handle_add("\n".join(lines))
+
+        elif user_input.lower().startswith("/search "):
+            query_part = user_input[len("/search "):]
+            handle_search(query_part)
+
+        elif user_input.lower() == "/search":
+            console.print("[yellow]⚠  Укажите запрос: /search <текст>[/yellow]")
+
+        # ── Запрос к агенту ───────────────────────────────────────────────────
+        else:
+            handle_agent(user_input)
+
+
+if __name__ == "__main__":
+    main()