diff --git a/main.py b/main.py new file mode 100644 index 0000000..615c756 --- /dev/null +++ b/main.py @@ -0,0 +1,71 @@ +"""Интерактивный RAG-агент (Qdrant + Ollama + LangChain).""" +from __future__ import annotations + +import os + +from langchain.agents import create_agent +from langchain_ollama import ChatOllama + +from rag_tools import add_to_knowledge_base, search_knowledge_base + +LLM_MODEL = os.getenv("OLLAMA_LLM_MODEL", "llama3") +OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") + +SYSTEM_PROMPT = """\ +Ты ассистент с доступом к базе знаний. +Перед ответом на фактические вопросы используй search_knowledge_base. +Если пользователь просит сохранить информацию — вызывай add_to_knowledge_base. +Отвечай на русском, кратко и по существу. +""" + + +def build_agent(): + llm = ChatOllama( + model=LLM_MODEL, + base_url=OLLAMA_BASE_URL, + temperature=0.2, + ) + return create_agent( + model=llm, + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt=SYSTEM_PROMPT, + ) + + +def run_interactive_cli() -> None: + agent = build_agent() + print("RAG-агент. Команды: /add, /search, /quit") + print(" /add <заголовок> | <текст>") + print(" /search <запрос>") + + while True: + user_input = input("\n> ").strip() + if not user_input: + continue + if user_input == "/quit": + print("Выход.") + break + + if user_input.startswith("/search "): + 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] + print(getattr(answer, "content", answer)) + + +if __name__ == "__main__": + run_interactive_cli()