From a55a5e039e593bc5491cf4ac0596220937a65644 Mon Sep 17 00:00:00 2001 From: lonpatovaadelina Date: Thu, 4 Jun 2026 16:20:00 +0000 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D1=80=D0=B0=D0=B1=D0=BE=D1=82?= =?UTF-8?q?=D0=BA=D0=B0:=201)=20=D0=92=20agent.py=20=D1=80=D0=B5=D0=B0?= =?UTF-8?q?=D0=BB=D0=B8=D0=B7=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20=D1=84=D1=83?= =?UTF-8?q?=D0=BD=D0=BA=D1=86=D0=B8=D1=8E=20main()=20=D1=81=20=D1=86=D0=B8?= =?UTF-8?q?=D0=BA=D0=BB=D0=BE=D0=BC=20=D1=87=D1=82=D0=B5=D0=BD=D0=B8=D1=8F?= =?UTF-8?q?=20=D0=BF=D0=BE=D0=BB?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 162 +++++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 161 insertions(+), 1 deletion(-) diff --git a/main.py b/main.py index 8f8a025..53eca11 100644 --- a/main.py +++ b/main.py @@ -140,4 +140,164 @@ def main() -> None: print(f"\n{source_line}\n") if __name__ == "__main__": - main() \ No newline at end of file + 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" \ No newline at end of file