diff --git a/main.py b/main.py index d4b859e..831efee 100644 --- a/main.py +++ b/main.py @@ -1,33 +1,27 @@ -"""Интерактивный клиент: агент с RAG-памятью.""" +"""Интерактивный RAG-агент (Qdrant + Ollama + LangChain).""" from __future__ import annotations -from agent_app import build_agent -from rag_tools import get_knowledge_base, search_knowledge_base +from agent_app import build_agent, build_llm +from knowledge_base import build_knowledge_base +from rag_tools import add_to_knowledge_base, get_knowledge_base, search_knowledge_base + +__all__ = [ + "build_agent", + "build_llm", + "build_knowledge_base", + "get_knowledge_base", + "search_knowledge_base", + "add_to_knowledge_base", + "run_interactive_cli", +] -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: +def run_interactive_cli() -> None: agent = build_agent() print("RAG-агент (Qdrant + Ollama). Команды: /add, /search, /quit") - print(" /add заголовок | текст") - print(" /search запрос") - print(" или просто задайте вопрос агенту\n") + print(" /add <заголовок> | <текст>") + print(" /search <запрос>") + print(" или задайте вопрос агенту\n") while True: user_input = input("Вы: ").strip() @@ -36,17 +30,25 @@ def main() -> None: 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:]) + query = user_input[len("/search ") :].strip() + print(search_knowledge_base.invoke({"query": query, "max_results": 5})) + continue + + if user_input.startswith("/add "): + payload = user_input[len("/add ") :].strip() + if "|" not in payload: + print("Формат: /add <заголовок> | <текст>") + continue + title, content = (part.strip() for part in payload.split("|", 1)) + print(add_to_knowledge_base.invoke({"content": content, "title": title})) continue result = agent.invoke({"messages": [{"role": "human", "content": user_input}]}) - answer = result["messages"][-1].content - print(f"\nАгент: {answer}\n") + answer = result["messages"][-1] + print(f"\nАгент: {getattr(answer, 'content', answer)}\n") if __name__ == "__main__": - main() + run_interactive_cli()