diff --git a/main.py b/main.py index 831efee..615c756 100644 --- a/main.py +++ b/main.py @@ -1,34 +1,49 @@ """Интерактивный RAG-агент (Qdrant + Ollama + LangChain).""" from __future__ import annotations -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 +import os -__all__ = [ - "build_agent", - "build_llm", - "build_knowledge_base", - "get_knowledge_base", - "search_knowledge_base", - "add_to_knowledge_base", - "run_interactive_cli", -] +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-агент (Qdrant + Ollama). Команды: /add, /search, /quit") + print("RAG-агент. Команды: /add, /search, /quit") print(" /add <заголовок> | <текст>") print(" /search <запрос>") - print(" или задайте вопрос агенту\n") while True: - user_input = input("Вы: ").strip() + user_input = input("\n> ").strip() if not user_input: continue - if user_input.lower() in ("/quit", "exit", "quit", "q"): - print("До встречи!") + if user_input == "/quit": + print("Выход.") break if user_input.startswith("/search "): @@ -45,9 +60,11 @@ def run_interactive_cli() -> None: print(add_to_knowledge_base.invoke({"content": content, "title": title})) continue - result = agent.invoke({"messages": [{"role": "human", "content": user_input}]}) + result = agent.invoke( + {"messages": [{"role": "human", "content": user_input}]} + ) answer = result["messages"][-1] - print(f"\nАгент: {getattr(answer, 'content', answer)}\n") + print(getattr(answer, "content", answer)) if __name__ == "__main__":