From 79f06a19eba97d924d6e1c980b945a0ec54a1e43 Mon Sep 17 00:00:00 2001 From: lonpatovaadelina Date: Thu, 4 Jun 2026 11:51:24 +0000 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D1=80=D0=B0=D0=B1=D0=BE=D1=82?= =?UTF-8?q?=D0=BA=D0=B0:=201.=20=D0=97=D0=B0=D0=B2=D0=B5=D1=80=D1=88=D0=B8?= =?UTF-8?q?=D1=82=D1=8C=20vectorstore.py:=20=D1=80=D0=B5=D0=B0=D0=BB=D0=B8?= =?UTF-8?q?=D0=B7=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20load=5Fdocuments=20(?= =?UTF-8?q?=D1=87=D0=B0=D0=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent.py | 69 ++++++++------------------------------------------------ 1 file changed, 9 insertions(+), 60 deletions(-) diff --git a/agent.py b/agent.py index aff3050..112085c 100644 --- a/agent.py +++ b/agent.py @@ -1,12 +1,10 @@ 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_ollama import OllamaLLM from langchain.agents import initialize_agent, AgentType +from langchain_community.tools.tavily_search import TavilySearchResults +from langchain.tools import tool +from vectorstore import get_vectorstore # Load environment variables load_dotenv() @@ -20,74 +18,25 @@ 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() +# Initialize or load vectorstore +vectorstore = get_vectorstore() # Define tools +@tool 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]) +@tool 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] +tools = [search_local_kb, web_search] # Initialize LLM llm = OllamaLLM(model=LLM_MODEL)