107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
"""
|
||
Main agent implementation using LangChain.
|
||
"""
|
||
|
||
import os
|
||
from typing import Dict
|
||
|
||
from langchain import LLMChain, PromptTemplate
|
||
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 rag_tools import search_local_kb, web_search
|
||
from vectorstore import create_vectorstore, load_documents
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
MODEL_NAME = "llama3"
|
||
CHROMA_DIR = "./chroma_db"
|
||
DOCS_DIR = "documents"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Load or create vector store
|
||
# ---------------------------------------------------------------------------
|
||
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 _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.
|
||
"""
|
||
# 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}
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI loop
|
||
# ---------------------------------------------------------------------------
|
||
if __name__ == "__main__":
|
||
print("RAG Agent ready. Type 'exit' to quit.")
|
||
while True:
|
||
user_input = input("\nЗапрос: ")
|
||
if user_input.lower() in {"exit", "quit", "q"}:
|
||
break
|
||
output = decide_and_run(user_input)
|
||
print(f"\nОтвет:\n{output['answer']}")
|
||
print(f"Источник: {output['source']}")
|