112 lines
3.5 KiB
Python
112 lines
3.5 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") -> Chroma:
|
|
"""Create a persistent Chroma vector store with Ollama embeddings."""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(
|
|
collection_name="rag_collection",
|
|
embedding_function=embeddings,
|
|
persist_directory=persist_directory,
|
|
)
|
|
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""Load .txt and .md files from `directory`, chunk them, and add to the store."""
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
docs = []
|
|
for file in Path(directory).glob("*.txt") | Path(directory).glob("*.md"):
|
|
text = file.read_text(encoding="utf-8")
|
|
docs.extend(splitter.create_documents([text]))
|
|
vectorstore.add_documents(docs)
|
|
|
|
|
|
# tools.py
|
|
from langchain.tools import tool
|
|
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
|
|
|
|
@tool("search_local_kb")
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""Semantic search in local knowledge base."""
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
vectorstore = Chroma(
|
|
collection_name="rag_collection",
|
|
embedding_function=embeddings,
|
|
persist_directory="./chroma_db",
|
|
)
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.invoke(query)
|
|
return "\n".join(doc.page_content for doc in docs)
|
|
|
|
|
|
@tool("web_search")
|
|
def web_search(query: str) -> str:
|
|
"""Search the web using Tavily."""
|
|
from langchain_tavily import TavilyAPIWrapper
|
|
|
|
tavily = TavilyAPIWrapper()
|
|
results = tavily.run(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
|
|
|
|
from tools import search_local_kb, web_search
|
|
|
|
|
|
def create_rag_agent():
|
|
llm = ChatOllama(model="llama3", temperature=0.2)
|
|
tools = [search_local_kb, web_search]
|
|
system_prompt = (
|
|
"You are an assistant that answers user questions.\n"
|
|
"If the answer can be found in local documents, use search_local_kb.\n"
|
|
"Otherwise, use web_search. Return the answer followed by a line\n"
|
|
"\"Source: chromadb\" or \"Source: tavily\"."
|
|
)
|
|
return create_agent(model=llm, tools=tools, system_prompt=system_prompt)
|
|
|
|
|
|
# main.py
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
from vectorstore import create_vectorstore, load_documents
|
|
from agent import create_rag_agent
|
|
|
|
|
|
def init_db():
|
|
"""Create and populate the Chroma DB if it does not exist."""
|
|
db_path = Path("./chroma_db")
|
|
if not db_path.exists():
|
|
store = create_vectorstore()
|
|
load_documents("documents", store)
|
|
store.persist()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
load_dotenv() # Loads TAVILY_API_KEY and any other env vars
|
|
init_db()
|
|
agent = create_rag_agent()
|
|
|
|
while True:
|
|
user_input = input("Запрос: ").strip()
|
|
if not user_input or user_input.lower() in ("exit", "quit", "выход"):
|
|
break
|
|
|
|
result = agent.invoke({"messages": [{"role": "human", "content": user_input}]})
|
|
for msg in result["messages"]:
|
|
# `msg` is a Pydantic model; use the `.content` attribute
|
|
if hasattr(msg, "content"):
|
|
print(msg.content) |