diff --git a/client.py b/client.py index 53ee41c..51c3b97 100644 --- a/client.py +++ b/client.py @@ -1,153 +1,79 @@ """ -client.py — интерактивный CLI-клиент для демонстрации работы RAG-агента. +client.py -Команды: - /add <текст> — добавить текст прямо в базу знаний - /add-file <путь> — загрузить файл в базу знаний - /search <запрос> — прямой семантический поиск (без агента) - /quit или /exit — выйти из программы - <любой другой текст> — отправить запрос агенту +Interactive CLI client for the RAG agent. + +Commands: + /add | <content> — Add a document to the knowledge base + /search <query> — Semantic search in the knowledge base + /quit — Exit the client + <any other input> — Send query to the agent """ -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 tools import search_knowledge_base, add_to_knowledge_base from agent import run_agent -console = Console() - -BANNER = """ -╔══════════════════════════════════════════╗ -║ 🤖 RAG-Agent • Qdrant+Ollama ║ -║ /add <текст> — добавить в базу ║ -║ /add-file <путь> — загрузить файл ║ -║ /search <запрос> — поиск в базе ║ -║ /quit — выйти ║ -╚══════════════════════════════════════════╝ +HELP_TEXT = """ +Commands: + /add <title> | <content> Add a document to the knowledge base + /search <query> Search the knowledge base directly + /quit Exit + <any text> Ask the agent a question """ -def handle_add(text: str) -> None: - """Добавляет текст в базу знаний через инструмент.""" - if not text.strip(): - console.print("[yellow]⚠ Укажите текст после /add[/yellow]") +def handle_add(args: str) -> None: + if "|" not in args: + print("[Error] Usage: /add <title> | <content>") 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]") + title, _, content = args.partition("|") + title = title.strip() + content = content.strip() + if not title or not content: + print("[Error] Both title and content are required.") 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]") + result = add_to_knowledge_base.invoke({"content": content, "title": title}) + print(result) def handle_search(query: str) -> None: - """Прямой семантический поиск — без агента.""" if not query.strip(): - console.print("[yellow]⚠ Укажите запрос после /search[/yellow]") + print("[Error] Please provide a search query.") 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", - ) - ) + result = search_knowledge_base.invoke({"query": query.strip(), "max_results": 5}) + print(result) -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) +def main(): + print("=== RAG Agent Interactive Client ===") + print(HELP_TEXT) while True: try: - user_input = Prompt.ask("\n[bold cyan]Вы[/bold cyan]").strip() - except (KeyboardInterrupt, EOFError): - console.print("\n[yellow]До свидания![/yellow]") + user_input = input("You: ").strip() + except (EOFError, KeyboardInterrupt): + print("\nGoodbye!") break if not user_input: continue - # ── Команды ─────────────────────────────────────────────────────────── - if user_input.lower() in ("/quit", "/exit", "/выход"): - console.print("[yellow]До свидания![/yellow]") + if user_input.lower() == "/quit": + print("Goodbye!") 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)) - + handle_add(user_input[5:]) 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]") - - # ── Запрос к агенту ─────────────────────────────────────────────────── + handle_search(user_input[8:]) + elif user_input.lower() == "/help": + print(HELP_TEXT) else: - handle_agent(user_input) + print("Agent: thinking...\n") + try: + response = run_agent(user_input) + print(f"Agent: {response}\n") + except Exception as e: + print(f"[Error] Agent failed: {e}\n") if __name__ == "__main__":