diff --git a/vectorstore.py b/vectorstore.py index 572bb76..4827b73 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -2,19 +2,11 @@ import os from pathlib import Path from typing import List -from langchain_ollama import OllamaEmbeddings, OllamaLLM +from langchain_ollama import OllamaEmbeddings from langchain_chroma import Chroma from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain.docstore.document import Document -from langchain_tavily import TavilySearchResults -from langchain.agents import Tool, AgentExecutor, initialize_agent -from dotenv import load_dotenv -# Загрузка переменных окружения из .env -load_dotenv() -TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") -if not TAVILY_API_KEY: - raise ValueError("Требуется переменная окружения TAVILY_API_KEY") def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: """ @@ -28,6 +20,7 @@ def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: embedding_function=embeddings, ) + def _load_text_files(directory: str) -> List[Document]: """ Читает все .txt и .md файлы из указанной директории и возвращает список Document. @@ -39,6 +32,7 @@ def _load_text_files(directory: str) -> List[Document]: docs.append(Document(page_content=text, metadata={"source": str(file_path)})) return docs + def load_documents(directory: str, vectorstore: Chroma) -> None: """ Загружает документы из указанной директории в ChromaDB. @@ -59,66 +53,4 @@ def load_documents(directory: str, vectorstore: Chroma) -> None: vectorstore.add_documents(chunks) # Сохраняем изменения vectorstore.persist() - print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") - -# ------------------------------------------------------------------ -# Инструмент web_search, использующий TavilySearchResults -def web_search(query: str) -> str: - """ - Выполняет поиск по интернету через Tavily и возвращает строку с результатами. - """ - tavily = TavilySearchResults(api_key=TAVILY_API_KEY) - results = tavily.run(query) - return "\n".join([f"{r['title']}: {r['url']}" for r in results]) - -# ------------------------------------------------------------------ -# Инструмент search_local_kb, использующий локальный Chroma -def search_local_kb(query: str, vectorstore: Chroma) -> List[dict]: - """ - Выполняет поиск по локальной базе знаний (Chroma) и возвращает список документов. - """ - docs = vectorstore.similarity_search_with_score(query, k=5) - return [{"content": doc.page_content, "score": score} for doc, score in docs] - -# ------------------------------------------------------------------ -# Создание агента, выбирающего между локальным поиском и web search -def create_agent(vectorstore: Chroma) -> AgentExecutor: - """ - Инициализирует агент с двумя инструментами: - - search_local_kb: поиск в локальной базе знаний - - web_search: поиск через Tavily - Агент выводит источник ответа. - """ - tools = [ - Tool( - name="search_local_kb", - func=lambda q: search_local_kb(q, vectorstore), - description="Search the local knowledge base stored in Chroma." - ), - Tool( - name="web_search", - func=web_search, - description="Search the web using Tavily. Useful for up-to-date information." - ) - ] - - llm = OllamaLLM(model="llama3") - agent_executor = initialize_agent( - tools, - llm, - agent_type="zero-shot-react-description", - verbose=True, - return_intermediate_steps=False - ) - return agent_executor - -# ------------------------------------------------------------------ -# Пример использования агента (необязательно, можно удалить) -if __name__ == "__main__": - vs = create_vectorstore() - # Предположим, что документы уже загружены ранее - agent = create_agent(vs) - query = "What is the capital of France?" - response = agent.run(query) - print("\nОтвет агента:") - print(response) \ No newline at end of file + print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") \ No newline at end of file