98 lines
3.8 KiB
Python
98 lines
3.8 KiB
Python
"""Core logic for the RAG agent.
|
||
|
||
The agent decides whether to use the local knowledge base or Tavily based on simple heuristics:
|
||
|
||
- If the query contains words that usually refer to recent events (e.g. "новости", "актуальные", "сегодня", "сейчас"), we route to Tavily.
|
||
- Otherwise we assume the answer can be found in the local KB.
|
||
|
||
The agent uses LangChain's :class:`langchain.agents.AgentExecutor` with a custom prompt that
|
||
instructs the LLM to specify the source in the response.
|
||
"""
|
||
|
||
import os
|
||
from typing import Dict, Any
|
||
|
||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||
from langchain_ollama import ChatOllama
|
||
from langchain_core.prompts import ChatPromptTemplate
|
||
|
||
from rag_tools import search_local_kb, web_search
|
||
from vectorstore import create_vectorstore, load_documents
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
# Directory containing the documents to index
|
||
DOCUMENTS_DIR = "./documents"
|
||
# Persistence directory for Chroma
|
||
CHROMA_DIR = "./chroma_db"
|
||
# The LLM used by the agent
|
||
LLM = ChatOllama(model="llama3")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper functions
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def should_use_web(query: str) -> bool:
|
||
"""Return True if the query looks like it needs up‑to‑date information.
|
||
|
||
The heuristic looks for a few Russian keywords that usually indicate a
|
||
request for recent news.
|
||
"""
|
||
keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
|
||
lowered = query.lower()
|
||
return any(k in lowered for k in keywords)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent setup
|
||
# ---------------------------------------------------------------------------
|
||
# Load or create the vector store
|
||
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
|
||
# If the store is empty, load documents from the documents directory
|
||
if not vectorstore._collection.count(): # type: ignore[attr-defined]
|
||
load_documents(DOCUMENTS_DIR, vectorstore)
|
||
|
||
# Tools – we pass the vectorstore instance to the local search tool via a closure
|
||
local_kb_tool = search_local_kb
|
||
local_kb_tool.__globals__["vectorstore"] = vectorstore # inject vectorstore
|
||
|
||
TOOLS = [local_kb_tool, web_search]
|
||
|
||
# Prompt template – the LLM is instructed to specify the source.
|
||
prompt = ChatPromptTemplate.from_messages([
|
||
("system", "You are an AI assistant that answers user questions.") ,
|
||
("user", "{input}"),
|
||
])
|
||
|
||
# Agent – we use the simple tool‑calling agent
|
||
agent = create_openai_tools_agent(llm=LLM, tools=TOOLS, prompt=prompt)
|
||
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def main() -> None:
|
||
print("RAG Agent ready. Type 'exit' to quit.")
|
||
while True:
|
||
user_query = input("\nЗапрос: ")
|
||
if user_query.lower() in {"exit", "quit", "q"}:
|
||
break
|
||
# Decide which tool to use
|
||
if should_use_web(user_query):
|
||
# Explicitly call the web search tool
|
||
result = web_search(user_query)
|
||
source = "tavily"
|
||
else:
|
||
result = search_local_kb(user_query, vectorstore=vectorstore)
|
||
source = "chromadb"
|
||
# Ask the LLM to format the answer
|
||
formatted = LLM.invoke({"input": user_query + "\n\nAnswer: " + result})
|
||
print("\nОтвет:", formatted.content)
|
||
print("Источник:", source)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
# End of agent.py
|