77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
"""RAG-агент с ChromaDB (Ollama) и веб-поиском (Tavily). CLI."""
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
from langchain_ollama import ChatOllama
|
||
from langchain.agents import AgentExecutor, create_react_agent
|
||
from langchain import hub
|
||
|
||
from vectorstore import create_vectorstore, load_documents
|
||
from tools import init_tools, search_local_kb, web_search
|
||
|
||
load_dotenv()
|
||
|
||
# ── LLM: Ollama llama3 ────────────────────────────────────────────────────
|
||
llm = ChatOllama(model="llama3", temperature=0)
|
||
|
||
# ── Векторное хранилище ───────────────────────────────────────────────────
|
||
CHROMA_DIR = "./chroma_db"
|
||
vectorstore = create_vectorstore(CHROMA_DIR)
|
||
init_tools(vectorstore)
|
||
|
||
# Автозагрузка документов из documents/ (если есть и БД пустая)
|
||
docs_dir = Path("documents")
|
||
if docs_dir.exists() and not any(Path(CHROMA_DIR).rglob("*.sqlite3")):
|
||
count = load_documents(str(docs_dir), vectorstore)
|
||
print(f"[init] Загружено {count} чанков из {docs_dir}/")
|
||
|
||
# ── ReAct-агент ───────────────────────────────────────────────────────────
|
||
tools = [search_local_kb, web_search]
|
||
prompt = hub.pull("hwchase17/react")
|
||
agent = create_react_agent(llm, tools, prompt)
|
||
agent_executor = AgentExecutor(
|
||
agent=agent,
|
||
tools=tools,
|
||
verbose=True,
|
||
max_iterations=10,
|
||
handle_parsing_errors=True,
|
||
)
|
||
|
||
ROUTING_HINT = (
|
||
"Для вопросов по локальным документам/конспектам используй search_local_kb. "
|
||
"Для актуальных новостей и фактов из интернета — web_search. "
|
||
"В ответе всегда укажи источник: chromadb или tavily."
|
||
)
|
||
|
||
|
||
def ask(question: str) -> str:
|
||
"""Задать вопрос агенту."""
|
||
result = agent_executor.invoke({"input": f"{ROUTING_HINT}
|
||
|
||
Вопрос: {question}"})
|
||
return result.get("output", str(result))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
print("RAG-агент: Ollama llama3 + nomic-embed-text + ChromaDB + Tavily")
|
||
print("Введите 'exit' для выхода.")
|
||
print("-" * 60)
|
||
while True:
|
||
try:
|
||
user_input = input("
|
||
Запрос: ").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("
|
||
Завершение.")
|
||
break
|
||
if user_input.lower() in ("exit", "quit", "выход"):
|
||
print("До свидания!")
|
||
break
|
||
if not user_input:
|
||
continue
|
||
answer = ask(user_input)
|
||
print(f"
|
||
{answer}")
|