124 lines
5.2 KiB
Python
124 lines
5.2 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from langchain_ollama import OllamaEmbeddings, OllamaLLM
|
|
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:
|
|
"""
|
|
Создаёт и возвращает объект Chroma, подключённый к указанной директории.
|
|
Если директория не существует, она будет создана.
|
|
"""
|
|
Path(persist_directory).mkdir(parents=True, exist_ok=True)
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
return Chroma(
|
|
persist_directory=persist_directory,
|
|
embedding_function=embeddings,
|
|
)
|
|
|
|
def _load_text_files(directory: str) -> List[Document]:
|
|
"""
|
|
Читает все .txt и .md файлы из указанной директории и возвращает список Document.
|
|
"""
|
|
docs: List[Document] = []
|
|
for file_path in Path(directory).rglob("*"):
|
|
if file_path.suffix.lower() in {".txt", ".md"}:
|
|
text = file_path.read_text(encoding="utf-8")
|
|
docs.append(Document(page_content=text, metadata={"source": str(file_path)}))
|
|
return docs
|
|
|
|
def load_documents(directory: str, vectorstore: Chroma) -> None:
|
|
"""
|
|
Загружает документы из указанной директории в ChromaDB.
|
|
Делает чанкинг с помощью RecursiveCharacterTextSplitter и добавляет в коллекцию.
|
|
"""
|
|
# Читаем файлы
|
|
raw_docs = _load_text_files(directory)
|
|
|
|
# Разбиваем на чанки
|
|
splitter = RecursiveCharacterTextSplitter(
|
|
chunk_size=1000,
|
|
chunk_overlap=200,
|
|
separators=["\n\n", "\n", " ", ""],
|
|
)
|
|
chunks = splitter.split_documents(raw_docs)
|
|
|
|
# Добавляем в коллекцию
|
|
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) |