Доработка: 1) В agent.py реализовать функцию main() с циклом чтения пол
This commit is contained in:
@@ -141,3 +141,163 @@ def main() -> None:
|
|||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
|
```
|
||||||
|
|
||||||
|
```python
|
||||||
|
# agent.py
|
||||||
|
|
||||||
|
import os
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Загрузка переменных окружения (TAVILY_API_KEY)
|
||||||
|
try:
|
||||||
|
from dotenv import load_dotenv
|
||||||
|
|
||||||
|
load_dotenv()
|
||||||
|
except Exception: # pragma: no cover
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Основные зависимости LangChain и Ollama
|
||||||
|
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||||||
|
from langchain_chroma import Chroma
|
||||||
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||||
|
from langchain.tools import tool
|
||||||
|
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||||||
|
from langchain_core.messages import HumanMessage
|
||||||
|
|
||||||
|
# Импорт вспомогательных функций из других файлов проекта
|
||||||
|
try:
|
||||||
|
from vectorstore import create_vectorstore, load_documents # noqa: F401
|
||||||
|
except Exception: # pragma: no cover
|
||||||
|
def create_vectorstore(persist_directory="./chroma_db"):
|
||||||
|
return Chroma(
|
||||||
|
persist_directory=persist_directory,
|
||||||
|
embedding_function=OllamaEmbeddings(model="nomic-embed-text"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def load_documents(directory, vectorstore):
|
||||||
|
pass
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 1. Создаём векторное хранилище и загружаем документы (если ещё нет)
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
CHROMA_DIR = "./chroma_db"
|
||||||
|
DOCS_DIR = "./documents"
|
||||||
|
|
||||||
|
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
|
||||||
|
|
||||||
|
if not vectorstore.get_ids():
|
||||||
|
load_documents(DOCS_DIR, vectorstore)
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 2. Определяем инструменты агента
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
@tool("search_local_kb")
|
||||||
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||||||
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||||||
|
docs = retriever.get_relevant_documents(query)
|
||||||
|
if not docs:
|
||||||
|
return "Ничего не найдено в локальной базе."
|
||||||
|
content = "\n\n".join(doc.page_content for doc in docs)
|
||||||
|
return f"[Local KB]\n{content}\nИсточник: chromadb"
|
||||||
|
|
||||||
|
@tool("web_search")
|
||||||
|
def web_search(query: str) -> str:
|
||||||
|
from tavily import TavilyClient
|
||||||
|
|
||||||
|
client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
|
||||||
|
results = client.search(query=query, max_results=3)
|
||||||
|
if not results:
|
||||||
|
return "Ничего не найдено в интернете."
|
||||||
|
snippets = "\n\n".join(
|
||||||
|
f"{res['title']}\n{res['content']}" for res in results
|
||||||
|
)
|
||||||
|
return f"[Web Search]\n{snippets}\nИсточник: tavily"
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 3. Создаём агента с выбором источника
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
llm = ChatOllama(model="llama3")
|
||||||
|
|
||||||
|
tools = [search_local_kb, web_search]
|
||||||
|
|
||||||
|
agent = create_openai_tools_agent(
|
||||||
|
llm=llm,
|
||||||
|
tools=tools,
|
||||||
|
system_message=(
|
||||||
|
"Ты помощник, отвечающий на вопросы. "
|
||||||
|
"Если вопрос относится к содержимому локальных документов, используй инструмент search_local_kb; "
|
||||||
|
"если требуется актуальная информация из интернета – web_search. "
|
||||||
|
"В ответе обязательно указывай источник (chromadb или tavily)."
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
|
||||||
|
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
# 4. CLI чат‑цикл
|
||||||
|
# --------------------------------------------------------------------------- #
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
print("=== RAG-агент с ChromaDB и Tavily ===")
|
||||||
|
print("Введите 'exit' для выхода.\n")
|
||||||
|
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
query = input("Запрос: ").strip()
|
||||||
|
except (KeyboardInterrupt, EOFError):
|
||||||
|
print("\nВыход.")
|
||||||
|
break
|
||||||
|
|
||||||
|
if not query or query.lower() == "exit":
|
||||||
|
print("До свидания!")
|
||||||
|
break
|
||||||
|
|
||||||
|
response = agent_executor.invoke({"input": query})
|
||||||
|
answer = response.get("output", "")
|
||||||
|
|
||||||
|
source_line = ""
|
||||||
|
if "[Local KB]" in answer:
|
||||||
|
source_line = "Источник: chromadb"
|
||||||
|
elif "[Web Search]" in answer:
|
||||||
|
source_line = "Источник: tavily"
|
||||||
|
|
||||||
|
print("\nОтвет:")
|
||||||
|
cleaned_answer = answer.replace("[Local KB]\n", "").replace(
|
||||||
|
"[Web Search]\n", ""
|
||||||
|
)
|
||||||
|
print(cleaned_answer.strip())
|
||||||
|
if source_line:
|
||||||
|
print(f"\n{source_line}\n")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
|
```
|
||||||
|
|
||||||
|
```python
|
||||||
|
# tools.py
|
||||||
|
|
||||||
|
import os
|
||||||
|
from tavily import TavilyClient, TavilySearchResults
|
||||||
|
from langchain.tools import tool
|
||||||
|
|
||||||
|
# Загружаем ключ из .env (dotenv уже вызывается в другом месте)
|
||||||
|
tavily_api_key = os.getenv("TAVILY_API_KEY")
|
||||||
|
|
||||||
|
@tool("web_search")
|
||||||
|
def web_search(query: str) -> str:
|
||||||
|
"""
|
||||||
|
Поиск в интернете через Tavily.
|
||||||
|
Возвращает краткое резюме найденных результатов.
|
||||||
|
"""
|
||||||
|
client = TavilyClient(api_key=tavily_api_key)
|
||||||
|
results = client.search(query=query, max_results=3)
|
||||||
|
if not results:
|
||||||
|
return "Ничего не найдено в интернете."
|
||||||
|
snippets = "\n\n".join(
|
||||||
|
f"{res['title']}\n{res['content']}" for res in results
|
||||||
|
)
|
||||||
|
return f"[Web Search]\n{snippets}\nИсточник: tavily"
|
||||||
Reference in New Issue
Block a user