107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""Core logic for the RAG agent.
|
||
|
||
The agent decides whether to use the local knowledge base or Tavily based on a
|
||
very simple heuristic: if the query contains words like ``news``, ``latest``
|
||
or ``today`` it is routed to the web search; otherwise the local KB is used.
|
||
|
||
The decision logic can be replaced with a more sophisticated router if
|
||
desired.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Tuple
|
||
|
||
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
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Prompt template
|
||
# ---------------------------------------------------------------------------
|
||
SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search.
|
||
|
||
When answering, always include the source of the information:
|
||
- "chromadb" for local knowledge base results.
|
||
- "tavily" for web search results.
|
||
|
||
If you are uncertain, say "I don't know" but still mention the source you used.
|
||
"""
|
||
|
||
USER_PROMPT = """Question: {question}\n
|
||
When you respond, first state the source (chromadb or tavily) and then provide the answer.
|
||
"""
|
||
|
||
prompt = ChatPromptTemplate.from_messages([
|
||
("system", SYSTEM_PROMPT),
|
||
("user", USER_PROMPT),
|
||
])
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Decision logic
|
||
# ---------------------------------------------------------------------------
|
||
WEB_KEYWORDS = {"news", "latest", "today", "current", "recent"}
|
||
|
||
|
||
def choose_tool(question: str) -> Tuple[str, callable]:
|
||
"""Return the name of the tool and the function to call.
|
||
|
||
Parameters
|
||
----------
|
||
question: str
|
||
The user query.
|
||
|
||
Returns
|
||
-------
|
||
Tuple[str, callable]
|
||
The tool name and the corresponding function.
|
||
"""
|
||
lowered = question.lower()
|
||
if any(word in lowered for word in WEB_KEYWORDS):
|
||
return "web_search", web_search
|
||
return "search_local_kb", search_local_kb
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent creation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def create_agent() -> AgentExecutor:
|
||
"""Instantiate the agent with the two tools.
|
||
|
||
The LLM used is Ollama's ``llama3``.
|
||
"""
|
||
tools = [search_local_kb, web_search]
|
||
llm = ChatOllama(model="llama3", temperature=0)
|
||
# Build an agent that knows about the tools and uses the custom prompt
|
||
agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt)
|
||
return AgentExecutor(agent=agent, tools=tools, verbose=True)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI loop
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
# Ensure the vector store is loaded once
|
||
store = create_vectorstore()
|
||
# Load documents if the store is empty
|
||
if not store.get_index_info():
|
||
from vectorstore import load_documents
|
||
load_documents("documents", store)
|
||
|
||
agent = create_agent()
|
||
print("RAG Agent ready. Type 'exit' to quit.")
|
||
while True:
|
||
try:
|
||
question = input("\nЗапрос: ")
|
||
except EOFError:
|
||
break
|
||
if question.strip().lower() in {"exit", "quit"}:
|
||
break
|
||
# The agent will automatically call the chosen tool via the prompt.
|
||
# We simply pass the question to the agent.
|
||
result = agent.invoke({"input": question})
|
||
# The agent's output already contains the source.
|
||
print("Ответ:", result["output"]) # noqa: T201
|