Обновить client.py
This commit is contained in:
@@ -1,153 +1,79 @@
|
|||||||
"""
|
"""
|
||||||
client.py — интерактивный CLI-клиент для демонстрации работы RAG-агента.
|
client.py
|
||||||
|
|
||||||
Команды:
|
Interactive CLI client for the RAG agent.
|
||||||
/add <текст> — добавить текст прямо в базу знаний
|
|
||||||
/add-file <путь> — загрузить файл в базу знаний
|
Commands:
|
||||||
/search <запрос> — прямой семантический поиск (без агента)
|
/add <title> | <content> — Add a document to the knowledge base
|
||||||
/quit или /exit — выйти из программы
|
/search <query> — Semantic search in the knowledge base
|
||||||
<любой другой текст> — отправить запрос агенту
|
/quit — Exit the client
|
||||||
|
<any other input> — Send query to the agent
|
||||||
"""
|
"""
|
||||||
|
|
||||||
import sys
|
from tools import search_knowledge_base, add_to_knowledge_base
|
||||||
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
|
from agent import run_agent
|
||||||
|
|
||||||
console = Console()
|
HELP_TEXT = """
|
||||||
|
Commands:
|
||||||
BANNER = """
|
/add <title> | <content> Add a document to the knowledge base
|
||||||
╔══════════════════════════════════════════╗
|
/search <query> Search the knowledge base directly
|
||||||
║ 🤖 RAG-Agent • Qdrant+Ollama ║
|
/quit Exit
|
||||||
║ /add <текст> — добавить в базу ║
|
<any text> Ask the agent a question
|
||||||
║ /add-file <путь> — загрузить файл ║
|
|
||||||
║ /search <запрос> — поиск в базе ║
|
|
||||||
║ /quit — выйти ║
|
|
||||||
╚══════════════════════════════════════════╝
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
def handle_add(text: str) -> None:
|
def handle_add(args: str) -> None:
|
||||||
"""Добавляет текст в базу знаний через инструмент."""
|
if "|" not in args:
|
||||||
if not text.strip():
|
print("[Error] Usage: /add <title> | <content>")
|
||||||
console.print("[yellow]⚠ Укажите текст после /add[/yellow]")
|
|
||||||
return
|
return
|
||||||
title = Prompt.ask(" Название документа (Enter — пропустить)", default="")
|
title, _, content = args.partition("|")
|
||||||
n = add_documents(text.strip(), title=title)
|
title = title.strip()
|
||||||
console.print(
|
content = content.strip()
|
||||||
f"[green]✓ Добавлено {n} чанков в базу знаний"
|
if not title or not content:
|
||||||
+ (f" (источник: «{title}»)" if title else "")
|
print("[Error] Both title and content are required.")
|
||||||
+ "[/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
|
return
|
||||||
content = p.read_text(encoding="utf-8", errors="ignore")
|
result = add_to_knowledge_base.invoke({"content": content, "title": title})
|
||||||
n = add_documents(content, title=p.name)
|
print(result)
|
||||||
console.print(f"[green]✓ Файл «{p.name}» загружен, создано {n} чанков[/green]")
|
|
||||||
|
|
||||||
|
|
||||||
def handle_search(query: str) -> None:
|
def handle_search(query: str) -> None:
|
||||||
"""Прямой семантический поиск — без агента."""
|
|
||||||
if not query.strip():
|
if not query.strip():
|
||||||
console.print("[yellow]⚠ Укажите запрос после /search[/yellow]")
|
print("[Error] Please provide a search query.")
|
||||||
return
|
return
|
||||||
console.print(f"[dim]Поиск: «{query}»...[/dim]")
|
result = search_knowledge_base.invoke({"query": query.strip(), "max_results": 5})
|
||||||
results = search(query, max_results=5)
|
print(result)
|
||||||
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:
|
def main():
|
||||||
"""Отправляет запрос агенту и выводит ответ."""
|
print("=== RAG Agent Interactive Client ===")
|
||||||
console.print("[dim]Агент думает...[/dim]")
|
print(HELP_TEXT)
|
||||||
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:
|
while True:
|
||||||
try:
|
try:
|
||||||
user_input = Prompt.ask("\n[bold cyan]Вы[/bold cyan]").strip()
|
user_input = input("You: ").strip()
|
||||||
except (KeyboardInterrupt, EOFError):
|
except (EOFError, KeyboardInterrupt):
|
||||||
console.print("\n[yellow]До свидания![/yellow]")
|
print("\nGoodbye!")
|
||||||
break
|
break
|
||||||
|
|
||||||
if not user_input:
|
if not user_input:
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# ── Команды ───────────────────────────────────────────────────────────
|
if user_input.lower() == "/quit":
|
||||||
if user_input.lower() in ("/quit", "/exit", "/выход"):
|
print("Goodbye!")
|
||||||
console.print("[yellow]До свидания![/yellow]")
|
|
||||||
break
|
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 "):
|
elif user_input.lower().startswith("/add "):
|
||||||
text_part = user_input[len("/add "):]
|
handle_add(user_input[5:])
|
||||||
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 "):
|
elif user_input.lower().startswith("/search "):
|
||||||
query_part = user_input[len("/search "):]
|
handle_search(user_input[8:])
|
||||||
handle_search(query_part)
|
elif user_input.lower() == "/help":
|
||||||
|
print(HELP_TEXT)
|
||||||
elif user_input.lower() == "/search":
|
|
||||||
console.print("[yellow]⚠ Укажите запрос: /search <текст>[/yellow]")
|
|
||||||
|
|
||||||
# ── Запрос к агенту ───────────────────────────────────────────────────
|
|
||||||
else:
|
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__":
|
if __name__ == "__main__":
|
||||||
|
|||||||
Reference in New Issue
Block a user