diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..2c75eac --- /dev/null +++ b/agent.py @@ -0,0 +1,52 @@ +""" +agent.py — AI-агент с RAG-памятью на базе Ollama (llama3) и Qdrant. + +Использует create_react_agent из langgraph (рекомендуемый способ в LangChain v0.3+) +с системным промптом, инструктирующим агента работать с базой знаний. +""" + +from langchain_ollama import ChatOllama +from langgraph.prebuilt import create_react_agent +from langchain_core.messages import SystemMessage + +from tools import search_knowledge_base, add_to_knowledge_base + +# ── Модель ──────────────────────────────────────────────────────────────────── +LLM_MODEL = "llama3" + +llm = ChatOllama(model=LLM_MODEL, temperature=0.3) + +# ── Системный промпт ────────────────────────────────────────────────────────── +SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний. + +Правила работы: +1. ВСЕГДА начинай с поиска в базе знаний с помощью инструмента search_knowledge_base, + прежде чем отвечать на вопрос пользователя. +2. Если в базе найдена релевантная информация — используй её как основу ответа + и явно ссылайся на источник. +3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. +4. Если база знаний не содержит нужной информации — честно сообщи об этом + и постарайся помочь на основе своих знаний. +5. Отвечай на том же языке, на котором задан вопрос. +""" + +# ── Агент ───────────────────────────────────────────────────────────────────── +rag_tools = [search_knowledge_base, add_to_knowledge_base] + +agent = create_react_agent( + model=llm, + tools=rag_tools, + prompt=SYSTEM_PROMPT, +) + + +def run_agent(user_input: str) -> str: + """ + Запускает агента с переданным запросом и возвращает финальный ответ. + """ + result = agent.invoke({"messages": [("user", user_input)]}) + # Финальное сообщение — последнее от ассистента + for msg in reversed(result["messages"]): + if msg.type == "ai" and msg.content: + return msg.content + return "Агент не вернул ответ." \ No newline at end of file