diff --git a/vectorstore.py b/vectorstore.py index 4827b73..bcc74c4 100644 --- a/vectorstore.py +++ b/vectorstore.py @@ -6,7 +6,16 @@ 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 langchain.llms import Ollama +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: """ @@ -20,7 +29,6 @@ def create_vectorstore(persist_directory: str = "./chroma_db") -> Chroma: embedding_function=embeddings, ) - def _load_text_files(directory: str) -> List[Document]: """ Читает все .txt и .md файлы из указанной директории и возвращает список Document. @@ -32,7 +40,6 @@ 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. @@ -53,4 +60,66 @@ def load_documents(directory: str, vectorstore: Chroma) -> None: vectorstore.add_documents(chunks) # Сохраняем изменения vectorstore.persist() - print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") \ No newline at end of file + print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}") + +# ------------------------------------------------------------------ +# Инструмент web_search, использующий TavilySearchResults +def web_search(query: str) -> List[dict]: + """ + Выполняет поиск по интернету через Tavily и возвращает список результатов. + """ + tavily = TavilySearchResults(api_key=TAVILY_API_KEY) + results = tavily.run(query) + return 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 = Ollama(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