Доработка: 1. Завершить vectorstore.py: реализовать load_documents (чан

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)