Экзамен: RAG-агент с ChromaDB и веб-поиском: main.py

main.py

@@ -0,0 +1,139 @@
+#!/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()