diff --git a/client.py b/client.py new file mode 100644 index 0000000..53ee41c --- /dev/null +++ b/client.py @@ -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() \ No newline at end of file