Update agent.py
This commit is contained in:
@@ -1,15 +1,20 @@
|
||||
"""
|
||||
Main agent implementation using LangChain.
|
||||
"""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
|
||||
from typing import Dict, Any
|
||||
|
||||
from langchain import LLMChain, PromptTemplate
|
||||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain_core.messages import HumanMessage, SystemMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langchain_core.tools import BaseTool
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
|
||||
from rag_tools import search_local_kb, web_search
|
||||
from vectorstore import create_vectorstore, load_documents
|
||||
@@ -17,90 +22,76 @@ from vectorstore import create_vectorstore, load_documents
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# ---------------------------------------------------------------------------
|
||||
MODEL_NAME = "llama3"
|
||||
# Directory containing the documents to index
|
||||
DOCUMENTS_DIR = "./documents"
|
||||
# Persistence directory for Chroma
|
||||
CHROMA_DIR = "./chroma_db"
|
||||
DOCS_DIR = "documents"
|
||||
# The LLM used by the agent
|
||||
LLM = ChatOllama(model="llama3")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load or create vector store
|
||||
# Helper functions
|
||||
# ---------------------------------------------------------------------------
|
||||
vectorstore = create_vectorstore(CHROMA_DIR)
|
||||
# Load documents only if the store is empty
|
||||
if not vectorstore._collection.count():
|
||||
load_documents(DOCS_DIR, vectorstore)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Define tools
|
||||
# ---------------------------------------------------------------------------
|
||||
class LocalKBTool(BaseTool):
|
||||
name = "search_local_kb"
|
||||
description = "Perform a semantic search in the local knowledge base."
|
||||
def should_use_web(query: str) -> bool:
|
||||
"""Return True if the query looks like it needs up‑to‑date information.
|
||||
|
||||
def _run(self, query: str, top_k: int = 3) -> str: # pragma: no cover
|
||||
return search_local_kb(query, top_k, vectorstore)
|
||||
|
||||
class WebSearchTool(BaseTool):
|
||||
name = "web_search"
|
||||
description = "Search the web using Tavily."
|
||||
|
||||
def _run(self, query: str, top_k: int = 3) -> str: # pragma: no cover
|
||||
return web_search(query, top_k)
|
||||
|
||||
tools = [LocalKBTool(), WebSearchTool()]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt template
|
||||
# ---------------------------------------------------------------------------
|
||||
SYSTEM_PROMPT = """
|
||||
You are an AI assistant that answers user questions.
|
||||
- If the answer can be found in the local knowledge base, use the tool `search_local_kb`.
|
||||
- If the answer requires up‑to‑date information, use the tool `web_search`.
|
||||
After providing the answer, always state the source in the format:
|
||||
|
||||
Source: <chromadb|tavily>
|
||||
"""
|
||||
|
||||
PROMPT = PromptTemplate(
|
||||
input_variables=["input", "chat_history"],
|
||||
template="""
|
||||
{chat_history}
|
||||
User: {input}
|
||||
Assistant: """
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent chain
|
||||
# ---------------------------------------------------------------------------
|
||||
llm = ChatOllama(model=MODEL_NAME, temperature=0.2)
|
||||
chain = LLMChain(llm=llm, prompt=PROMPT)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper to decide which tool to use
|
||||
# ---------------------------------------------------------------------------
|
||||
def decide_and_run(query: str) -> Dict[str, str]:
|
||||
"""Use the LLM to decide whether to use local KB or web search.
|
||||
|
||||
Returns a dict with keys: answer, source.
|
||||
The heuristic looks for a few Russian keywords that usually indicate a
|
||||
request for recent news.
|
||||
"""
|
||||
# Simple heuristic: if the query contains words like "news", "latest", "today" use web
|
||||
web_keywords = {"news", "latest", "today", "current", "recent", "update"}
|
||||
if any(word in query.lower() for word in web_keywords):
|
||||
result = web_search(query)
|
||||
source = "tavily"
|
||||
else:
|
||||
result = search_local_kb(query, vectorstore=vectorstore)
|
||||
source = "chromadb"
|
||||
return {"answer": result, "source": source}
|
||||
keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
|
||||
lowered = query.lower()
|
||||
return any(k in lowered for k in keywords)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI loop
|
||||
# Agent setup
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
# 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_input = input("\nЗапрос: ")
|
||||
if user_input.lower() in {"exit", "quit", "q"}:
|
||||
user_query = input("\nЗапрос: ")
|
||||
if user_query.lower() in {"exit", "quit", "q"}:
|
||||
break
|
||||
output = decide_and_run(user_input)
|
||||
print(f"\nОтвет:\n{output['answer']}")
|
||||
print(f"Источник: {output['source']}")
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user