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

This commit is contained in:
2026-06-04 11:49:40 +00:00
parent d0cb373baf
commit 2aba338633
+4 -72
View File
@@ -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)
print(f"Загружено {len(chunks)} чанков из {len(raw_docs)} документов в {vectorstore.persist_directory}")