feat: ChatOllama(llama3) ReAct agent + CLI

This commit is contained in:
2026-06-05 14:55:43 +00:00
parent 96a8608aa9
commit 10994e48db
+64 -93
View File
@@ -1,105 +1,76 @@
import os
import asyncio
from pathlib import Path
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
"""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 environment variables
load_dotenv()
# ---------- LLM and Embeddings ----------
llm = ChatOpenAI(
model="openai/gpt-oss-20b:free",
base_url="https://openrouter.ai/api/v1",
api_key=os.getenv("OPENAI_API_KEY"),
temperature=0.0,
# ── 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,
)
embeddings = OpenAIEmbeddings(
model="text-embedding-3-small",
base_url="https://openrouter.ai/api/v1",
api_key=os.getenv("OPENAI_API_KEY"),
ROUTING_HINT = (
"Для вопросов по локальным документам/конспектам используй search_local_kb. "
"Для актуальных новостей и фактов из интернета — web_search. "
"В ответе всегда укажи источник: chromadb или tavily."
)
# ---------- Vector Store ----------
CHROMA_DIR = Path("./chroma_db")
vector_store = Chroma(
collection_name="knowledge",
embedding_function=embeddings,
persist_directory=str(CHROMA_DIR),
)
# ---------- Document Loader ----------
def ask(question: str) -> str:
"""Задать вопрос агенту."""
result = agent_executor.invoke({"input": f"{ROUTING_HINT}
def load_documents(directory: str, vectorstore: Chroma):
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
docs = []
for file_path in Path(directory).glob("**/*.*"):
if file_path.suffix.lower() not in {".txt", ".md"}:
continue
text = file_path.read_text(encoding="utf-8")
chunks = splitter.split_text(text)
docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks])
if docs:
vectorstore.add_documents(docs)
vectorstore.persist()
Вопрос: {question}"})
return result.get("output", str(result))
# Load initial documents if collection is empty
if not CHROMA_DIR.exists() or not list(CHROMA_DIR.iterdir()):
load_documents("documents", vector_store)
# ---------- Tools ----------
@tool
def search_local_kb(query: str, top_k: int = 3) -> str:
"""Semantic search in the local knowledge base."""
docs = vector_store.similarity_search(query, k=top_k)
if not docs:
return "No relevant information found in local knowledge base."
return "\n---\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs])
@tool
def web_search(query: str) -> str:
"""Web search using Tavily."""
from langchain_tavily import TavilySearchResults
tavily = TavilySearchResults()
results = tavily.run(query)
return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results])
# ---------- Backend ----------
backend = CompositeBackend([
LocalShellBackend(workspace_dir="./workspace"),
FilesystemBackend(),
])
# ---------- Agent ----------
agent = create_deep_agent(
model=llm,
tools=[search_local_kb, web_search],
backend=backend,
system_prompt="You are a helpful assistant that can search both a local knowledge base and the web.\nWhen answering, always indicate the source: either 'chromadb' or 'tavily'.\nUse the appropriate tool based on the query context.",
)
# ---------- CLI ----------
async def main():
print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.")
while True:
user_input = input("\nЗапрос: ")
if user_input.lower() in {"exit", "quit"}:
print("Goodbye!")
break
result = await agent.ainvoke(
{"messages": [{"role": "user", "content": user_input}]},
{"configurable": {"thread_id": "session-1"}},
)
# Extract last message content
content = result["messages"][-1]["content"]
print(content)
if __name__ == "__main__":
asyncio.run(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}")