139 lines
5.5 KiB
Python
139 lines
5.5 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
main.py – CLI чат‑цикл для RAG‑агента с ChromaDB и веб‑поиском через Tavily.
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
|
||
from dotenv import load_dotenv
|
||
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
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 1. Загрузка переменных окружения (TAVILY_API_KEY)
|
||
# --------------------------------------------------------------------------- #
|
||
load_dotenv()
|
||
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
|
||
if not TAVILY_API_KEY:
|
||
raise RuntimeError("Не найдено переменной окружения TAVILY_API_KEY")
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 2. Векторное хранилище ChromaDB
|
||
# --------------------------------------------------------------------------- #
|
||
CHROMA_DIR = Path("./chroma_db")
|
||
CHROMA_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||
vectorstore = Chroma(
|
||
persist_directory=str(CHROMA_DIR),
|
||
embedding_function=embeddings,
|
||
)
|
||
# --------------------------------------------------------------------------- #
|
||
# 3. Инструменты агента
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
@tool("search_local_kb")
|
||
def search_local_kb(query: str, top_k: int = 5) -> str:
|
||
"""
|
||
Семантический поиск в локальной базе знаний (ChromaDB).
|
||
Возвращает объединённый текст найденных документов.
|
||
"""
|
||
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
||
docs = retriever.invoke(query)
|
||
if not docs:
|
||
return "Ничего не найдено."
|
||
# Собираем содержимое всех документов
|
||
return "\n\n".join(doc.page_content for doc in docs)
|
||
|
||
|
||
@tool("web_search")
|
||
def web_search(query: str) -> str:
|
||
"""
|
||
Поиск в интернете через Tavily.
|
||
Возвращает краткое резюме результатов поиска.
|
||
"""
|
||
from tavily import TavilyClient
|
||
|
||
client = TavilyClient(api_key=TAVILY_API_KEY)
|
||
results = client.search(query=query, max_results=3)
|
||
if not results:
|
||
return "Ничего не найдено в интернете."
|
||
# Формируем строку с заголовками и ссылками
|
||
snippets = []
|
||
for r in results:
|
||
snippets.append(f"{r['title']}\n{r['url']}")
|
||
return "\n\n".join(snippets)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 4. Создание агента
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
llm = ChatOllama(model="llama3", temperature=0)
|
||
|
||
tools = [search_local_kb, web_search]
|
||
|
||
agent = create_openai_tools_agent(
|
||
llm,
|
||
tools,
|
||
system_message="""
|
||
Вы – интеллектуальный агент, который отвечает на вопросы пользователя.
|
||
Если вопрос относится к локальным документам (например, конспектам, заметкам),
|
||
используйте инструмент «search_local_kb».
|
||
Если требуется актуальная информация из интернета (новости, факты), используйте
|
||
инструмент «web_search».
|
||
|
||
После выполнения запроса возвращайте ответ в формате:
|
||
[Источник] <текст ответа>
|
||
|
||
Где источник – либо «chromadb», либо «tavily».
|
||
""",
|
||
)
|
||
|
||
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False)
|
||
|
||
|
||
# --------------------------------------------------------------------------- #
|
||
# 5. Чат‑цикл CLI
|
||
# --------------------------------------------------------------------------- #
|
||
|
||
def main() -> None:
|
||
print("=== RAG‑агент с ChromaDB и Tavily ===")
|
||
print("Введите запрос (или 'exit' для выхода):")
|
||
|
||
while True:
|
||
user_input = input("\n> ").strip()
|
||
if not user_input or user_input.lower() == "exit":
|
||
print("До свидания!")
|
||
break
|
||
|
||
# Запускаем агента
|
||
try:
|
||
result = agent_executor.invoke({"input": user_input})
|
||
except Exception as e:
|
||
print(f"Ошибка при обработке запроса: {e}")
|
||
continue
|
||
|
||
# Результат – строка вида "[Источник] ...". Если формат не соответствует,
|
||
# просто выводим как есть.
|
||
if isinstance(result, str) and result.startswith("["):
|
||
try:
|
||
source_part, answer = result.split("]", 1)
|
||
source = source_part.strip("[")
|
||
print(f"\nОтвет: {answer.strip()}")
|
||
print(f"Источник: {source.lower()}")
|
||
except ValueError:
|
||
# Не удалось распарсить – выводим полностью
|
||
print("\nОтвет:", result)
|
||
else:
|
||
print("\nОтвет:", result)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |