52 lines
2.9 KiB
Python
52 lines
2.9 KiB
Python
"""
|
|
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 "Агент не вернул ответ." |