53 lines
1.7 KiB
Python
53 lines
1.7 KiB
Python
"""Интерактивный клиент: агент с RAG-памятью."""
|
|
from __future__ import annotations
|
|
|
|
from agent_app import build_agent
|
|
from rag_tools import get_knowledge_base, search_knowledge_base
|
|
|
|
|
|
def _cmd_add(raw: str) -> None:
|
|
if "|" not in raw:
|
|
print("Формат: /add заголовок | текст документа")
|
|
return
|
|
title, content = raw.split("|", maxsplit=1)
|
|
result = get_knowledge_base().add_document(content.strip(), title.strip())
|
|
print(f"Добавлено: «{title.strip()}» ({result} чанк(ов))")
|
|
|
|
|
|
def _cmd_search(raw: str) -> None:
|
|
query = raw.strip()
|
|
if not query:
|
|
print("Формат: /search ваш запрос")
|
|
return
|
|
print(search_knowledge_base.invoke({"query": query, "max_results": 5}))
|
|
|
|
|
|
def main() -> None:
|
|
agent = build_agent()
|
|
print("RAG-агент (Qdrant + Ollama). Команды: /add, /search, /quit")
|
|
print(" /add заголовок | текст")
|
|
print(" /search запрос")
|
|
print(" или просто задайте вопрос агенту\n")
|
|
|
|
while True:
|
|
user_input = input("Вы: ").strip()
|
|
if not user_input:
|
|
continue
|
|
if user_input.lower() in ("/quit", "exit", "quit", "q"):
|
|
print("До встречи!")
|
|
break
|
|
if user_input.startswith("/add "):
|
|
_cmd_add(user_input[5:])
|
|
continue
|
|
if user_input.startswith("/search "):
|
|
_cmd_search(user_input[8:])
|
|
continue
|
|
|
|
result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
|
|
answer = result["messages"][-1].content
|
|
print(f"\nАгент: {answer}\n")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|