Доработка: 1. Завершить vectorstore.py: реализовать load_documents (чан
This commit is contained in:
+4
-72
@@ -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}")
|
||||
Reference in New Issue
Block a user