87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
# -------------------- vectorstore.py --------------------
|
|
from pathlib import Path
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
|
|
def create_vectorstore(persist_directory: str = "./chroma_db"):
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
vector_store = Chroma(
|
|
collection_name="rag_collection",
|
|
embedding_function=embeddings,
|
|
persist_directory=persist_directory,
|
|
)
|
|
return vector_store
|
|
|
|
def load_documents(directory: str, vectorstore):
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
docs = []
|
|
for file_path in Path(directory).glob("*.txt"):
|
|
text = file_path.read_text(encoding="utf-8")
|
|
docs.extend(splitter.create_documents([text]))
|
|
for file_path in Path(directory).glob("*.md"):
|
|
text = file_path.read_text(encoding="utf-8")
|
|
docs.extend(splitter.create_documents([text]))
|
|
vectorstore.add_documents(docs)
|
|
|
|
# -------------------- tools.py --------------------
|
|
from langchain.tools import tool
|
|
from langchain_ollama import ChatOllama
|
|
|
|
@tool
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""Semantic search in the local ChromaDB knowledge base."""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.invoke({"query": query})["documents"]
|
|
return "\n".join(doc.page_content for doc in docs)
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Web search using Tavily."""
|
|
from langchain_tavily import TavilySearchResults
|
|
tavily = TavilySearchResults(api_key=__import__("os").environ["TAVILY_API_KEY"])
|
|
results = tavily.invoke({"query": query})
|
|
return "\n".join(f"{r['title']}: {r['url']}" for r in results)
|
|
|
|
# -------------------- agent.py --------------------
|
|
from langchain.agents import create_agent
|
|
from langchain_ollama import ChatOllama
|
|
|
|
llm = ChatOllama(model="llama3", temperature=0.2)
|
|
|
|
system_prompt = """
|
|
You are an assistant that answers user questions.
|
|
If the answer can be found in the local knowledge base, use `search_local_kb`.
|
|
Otherwise, use `web_search`.
|
|
Always indicate the source of your answer: either "chromadb" or "tavily".
|
|
"""
|
|
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[search_local_kb, web_search],
|
|
system_prompt=system_prompt,
|
|
)
|
|
|
|
# -------------------- main.py --------------------
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
if __name__ == "__main__":
|
|
# Initialize vectorstore and load documents if not already loaded
|
|
vectorstore = create_vectorstore()
|
|
if not vectorstore.get_collection().count():
|
|
load_documents("documents", vectorstore)
|
|
vectorstore.persist()
|
|
|
|
print("Chat started. Type 'exit' to quit.")
|
|
while True:
|
|
user_input = input("\nЗапрос: ").strip()
|
|
if user_input.lower() in ("exit", "quit", "выход"):
|
|
break
|
|
result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
|
|
ai_msg = result["messages"][-1]
|
|
print(f"[{ai_msg.tool_calls[0]['name'].capitalize()}] {ai_msg.content}")
|
|
source = "chromadb" if ai_msg.tool_calls[0]["name"] == "search_local_kb" else "tavily"
|
|
print(f"Источник: {source}") |