85 lines
3.1 KiB
Python
85 lines
3.1 KiB
Python
"""Main agent logic.
|
|
|
|
Creates a Chroma vector store, loads documents from the ``documents`` directory,
|
|
and runs a simple chat loop. The agent decides whether to use the local KB
|
|
or perform a web search based on the presence of the word "news" or
|
|
"latest" in the query.
|
|
"""
|
|
|
|
import os
|
|
from typing import List
|
|
|
|
from langchain_ollama import ChatOllama
|
|
from langchain.agents import tool, AgentExecutor, ZeroShotAgent
|
|
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
|
|
|
|
from vectorstore import create_vectorstore, load_documents
|
|
from rag_tools import web_search
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 1. Setup vector store
|
|
# ---------------------------------------------------------------------------
|
|
VECTORSTORE_DIR = "./chroma_db"
|
|
DOCS_DIR = "./documents"
|
|
|
|
vectorstore = create_vectorstore(VECTORSTORE_DIR)
|
|
load_documents(DOCS_DIR, vectorstore)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 2. Define local search tool (needs the vectorstore)
|
|
# ---------------------------------------------------------------------------
|
|
@tool
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""Semantic search over the local ChromaDB collection.
|
|
|
|
Returns a formatted string containing the top_k snippets and a source tag.
|
|
"""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.invoke(query)
|
|
snippets = "\n".join([f"{idx+1}. {doc.page_content[:200]}" for idx, doc in enumerate(docs)])
|
|
return f"[Local KB]\n{snippets}\nSource: chromadb"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 3. Agent prompt and execution
|
|
# ---------------------------------------------------------------------------
|
|
# The agent will be given two tools: search_local_kb and web_search.
|
|
# We provide a simple instruction to choose the appropriate tool.
|
|
|
|
agent_prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
HumanMessagePromptTemplate.from_template(
|
|
"You are an assistant that can search a local knowledge base or the web. "
|
|
"If the question is about recent events or news, use web_search. "
|
|
"Otherwise, use search_local_kb. "
|
|
"Respond with the answer and the source (chromadb or tavily)."
|
|
),
|
|
]
|
|
)
|
|
|
|
# Create the agent with the two tools
|
|
tools = [search_local_kb, web_search]
|
|
|
|
agent = ZeroShotAgent.from_llm_and_tools(
|
|
llm=ChatOllama(model="llama3", temperature=0),
|
|
tools=tools,
|
|
prompt=agent_prompt,
|
|
)
|
|
|
|
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 4. Chat loop
|
|
# ---------------------------------------------------------------------------
|
|
if __name__ == "__main__":
|
|
print("RAG Agent ready. Type 'exit' to quit.")
|
|
while True:
|
|
user_input = input("Запрос: ")
|
|
if user_input.lower() in {"exit", "quit", "q"}:
|
|
print("Bye!")
|
|
break
|
|
try:
|
|
result = agent_executor.invoke({"input": user_input})
|
|
print(result["output"])
|
|
except Exception as e:
|
|
print(f"Error: {e}")
|