Доработка: Добавьте модуль с инструментом web_search, реализующим пои

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.
@@ -54,3 +61,65 @@ def load_documents(directory: str, vectorstore: Chroma) -> None:
     # Сохраняем изменения
     vectorstore.persist()
     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)