diff --git a/agent.py b/agent.py index 0e08605..c92d20e 100644 --- a/agent.py +++ b/agent.py @@ -1,97 +1,106 @@ """Core logic for the RAG agent. -The agent decides whether to use the local knowledge base or Tavily based on simple heuristics: +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. -- 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. +The decision logic can be replaced with a more sophisticated router if +desired. """ -import os -from typing import Dict, Any +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, load_documents +from vectorstore import create_vectorstore # --------------------------------------------------------------------------- -# Configuration +# Prompt template # --------------------------------------------------------------------------- -# 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") +SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search. -# --------------------------------------------------------------------------- -# Helper functions -# --------------------------------------------------------------------------- +When answering, always include the source of the information: +- "chromadb" for local knowledge base results. +- "tavily" for web search results. -def should_use_web(query: str) -> bool: - """Return True if the query looks like it needs up‑to‑date information. +If you are uncertain, say "I don't know" but still mention the source you used. +""" - 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) +USER_PROMPT = """Question: {question}\n +When you respond, first state the source (chromadb or tavily) and then provide the answer. +""" -# --------------------------------------------------------------------------- -# 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}"), + ("system", SYSTEM_PROMPT), + ("user", USER_PROMPT), ]) -# 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) +# --------------------------------------------------------------------------- +# 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 # --------------------------------------------------------------------------- -# CLI +# Agent creation # --------------------------------------------------------------------------- -def main() -> None: +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: - user_query = input("\nЗапрос: ") - if user_query.lower() in {"exit", "quit", "q"}: + try: + question = input("\nЗапрос: ") + except EOFError: 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 + 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