133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_ollama import OllamaLLM, OllamaEmbeddings
|
|
from langchain_chroma import Chroma
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
from langchain_community.document_loaders import DirectoryLoader, TextLoader, UnstructuredMarkdownLoader
|
|
from langchain.tools import Tool
|
|
from langchain_community.tools.tavily_search import TavilySearchResults
|
|
from langchain.agents import initialize_agent, AgentType
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
|
if not TAVILY_API_KEY:
|
|
raise ValueError("TAVILY_API_KEY not found in .env file")
|
|
|
|
# Configuration
|
|
PERSIST_DIRECTORY = "./chroma_db"
|
|
DOCUMENTS_DIR = "./documents"
|
|
EMBEDDING_MODEL = "nomic-embed-text"
|
|
LLM_MODEL = "llama3"
|
|
|
|
def setup_vectorstore():
|
|
"""Initialize or load ChromaDB vectorstore with Ollama embeddings."""
|
|
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
|
|
|
if os.path.exists(PERSIST_DIRECTORY) and os.listdir(PERSIST_DIRECTORY):
|
|
vectorstore = Chroma(
|
|
persist_directory=PERSIST_DIRECTORY,
|
|
embedding_function=embeddings
|
|
)
|
|
else:
|
|
# Load and process documents
|
|
loader = DirectoryLoader(
|
|
DOCUMENTS_DIR,
|
|
glob="**/*",
|
|
loader_cls=lambda path: TextLoader(path, encoding="utf-8") if path.endswith(".txt")
|
|
else UnstructuredMarkdownLoader(path) if path.endswith(".md")
|
|
else None,
|
|
show_progress=True,
|
|
use_multithreading=True
|
|
)
|
|
documents = loader.load()
|
|
|
|
text_splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
length_function=len
|
|
)
|
|
texts = text_splitter.split_documents(documents)
|
|
|
|
vectorstore = Chroma.from_documents(
|
|
documents=texts,
|
|
embedding=embeddings,
|
|
persist_directory=PERSIST_DIRECTORY
|
|
)
|
|
vectorstore.persist()
|
|
|
|
return vectorstore
|
|
|
|
# Initialize vectorstore
|
|
vectorstore = setup_vectorstore()
|
|
|
|
# Define tools
|
|
def search_local_kb(query: str) -> str:
|
|
"""Search local knowledge base using ChromaDB."""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return "\n\n".join([doc.page_content for doc in docs])
|
|
|
|
def web_search(query: str) -> str:
|
|
"""Search the web using Tavily."""
|
|
search = TavilySearchResults(tavily_api_key=TAVILY_API_KEY, max_results=3)
|
|
results = search.run(query)
|
|
return "\n\n".join([result["content"] for result in results])
|
|
|
|
# Create LangChain tools
|
|
local_tool = Tool(
|
|
name="search_local_kb",
|
|
func=search_local_kb,
|
|
description="Useful for answering questions about local documents stored in the knowledge base."
|
|
)
|
|
|
|
web_tool = Tool(
|
|
name="web_search",
|
|
func=web_search,
|
|
description="Useful for answering questions about current events, news, or general knowledge from the internet."
|
|
)
|
|
|
|
tools = [local_tool, web_tool]
|
|
|
|
# Initialize LLM
|
|
llm = OllamaLLM(model=LLM_MODEL)
|
|
|
|
# Create agent with routing logic
|
|
agent = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
system_message="""You are an AI agent that must choose between two tools:
|
|
1. search_local_kb: For questions about local documents (concepts, notes, stored information)
|
|
2. web_search: For questions requiring up-to-date information from the internet
|
|
|
|
When answering, ALWAYS specify your information source:
|
|
- If using local knowledge base: [Source: chromadb]
|
|
- If using web search: [Source: tavily]
|
|
|
|
Be concise and accurate in your responses."""
|
|
)
|
|
|
|
# Chat interface
|
|
def main():
|
|
print("RAG Agent with ChromaDB and Tavily Web Search")
|
|
print("Type 'exit' to quit\n")
|
|
|
|
while True:
|
|
query = input("Запрос: ").strip()
|
|
if query.lower() == "exit":
|
|
break
|
|
|
|
if not query:
|
|
continue
|
|
|
|
try:
|
|
response = agent.run(query)
|
|
print(f"\n{response}\n")
|
|
except Exception as e:
|
|
print(f"Error: {str(e)}\n")
|
|
|
|
if __name__ == "__main__":
|
|
main() |