Добавить client.py
This commit is contained in:
@@ -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()
|
||||||
Reference in New Issue
Block a user