restore own solution: main.py

This commit is contained in:
2026-05-27 13:27:58 +00:00
parent 4415a06167
commit 548e41fa4a
+36 -19
View File
@@ -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__":