From cd030f4910cd8de08aeab3dbc0a662790f633c45 Mon Sep 17 00:00:00 2001 From: lonpatovaadelina Date: Thu, 28 May 2026 10:12:40 +0000 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=B8=D1=82?= =?UTF-8?q?=D1=8C=20agent.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- agent.py | 166 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 agent.py diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..9caba9a --- /dev/null +++ b/agent.py @@ -0,0 +1,166 @@ +# agent.py + +import os +from typing import Any, Dict, Iterable, Tuple + +from langchain_ollama import ChatOllama +from langchain_qdrant import QdrantStore +from langchain.embeddings.nomic import NomicEmbeddings +from langgraph.checkpoint.memory import MemorySaver +from langgraph.graph import StateGraph, END +from langgraph.prebuilt import create_chat_agent +from langchain.tools import tool +from rich.console import Console + +# ------------------------------------------------------------------ +# Конфигурация +# ------------------------------------------------------------------ + +CONSOLE = Console() + +LLM_MODEL = "llama3" +EMBEDDINGS_MODEL = "nomic-embed-text" + +QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") +QDRANT_COLLECTION = "knowledge_base" + +# ------------------------------------------------------------------ +# Векторный хранилище +# ------------------------------------------------------------------ + +embeddings = NomicEmbeddings(model=EMBEDDINGS_MODEL) +vector_store = QdrantStore( + url=QDRANT_URL, + collection_name=QDRANT_COLLECTION, + embeddings=embeddings, +) + +# ------------------------------------------------------------------ +# Инструменты RAG +# ------------------------------------------------------------------ + + +@tool("search_knowledge_base") +def search_knowledge_base(query: str) -> str: + """ + Поиск в базе знаний по запросу. + Возвращает найденный текст (или сообщение об отсутствии результата). + """ + results = vector_store.similarity_search_with_score(query, k=3) + if not results: + return "Ничего не найдено." + # Собираем топ‑k результатов + texts = "\n\n".join([f"Score: {score:.2f}\n{doc.page_content}" for doc, score in results]) + return f"Найденные документы:\n{texts}" + + +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str) -> str: + """ + Добавляет новый текст в базу знаний. + Возвращает подтверждение добавления. + """ + vector_store.add_texts([content]) + return "Текст успешно добавлен в базу знаний." + + +# ------------------------------------------------------------------ +# Создание агента +# ------------------------------------------------------------------ + +llm = ChatOllama(model=LLM_MODEL) + +memory = MemorySaver() + +def create_agent() -> StateGraph: + """ + Возвращает готовый граф агента с памятью и паузой перед вызовом инструментов. + """ + # Базовый агент от LangChain + base_agent = create_chat_agent( + llm=llm, + tools=[search_knowledge_base, add_to_knowledge_base], + system_prompt="Ты полезный ассистент. При необходимости используй инструменты.", + ) + + # Создаём граф с паузой перед инструментами + graph = StateGraph() + graph.add_node("agent", base_agent) + graph.set_entry_point("agent") + graph.add_edge("agent", END) + + # Включаем interrupt_before для паузы + graph.interrupt_before = ["tools"] + + return graph + +AGENT_GRAPH = create_agent() + + +# ------------------------------------------------------------------ +# Функция взаимодействия +# ------------------------------------------------------------------ + + +def ask_and_run(user_input: Dict[str, Any], config: Dict[str, Any]) -> None: + """ + Обрабатывает ввод пользователя и выводит потоковые ответы. + При паузе спрашивает подтверждение перед вызовом инструмента. + """ + for chunk in AGENT_GRAPH.stream( + user_input, + config=config, + stream_mode=["messages", "updates"], + ): + state = AGENT_GRAPH.get_state(config) + chunk_type, chunk_data = chunk + + # Печать токенов + if chunk_type == "messages": + CONSOLE.print(chunk_data["content"], end="") + + # Печать вызова инструмента + if chunk_type == "updates" and "tool_calls" in chunk_data: + for call in chunk_data["tool_calls"]: + name = call["name"] + args = call.get("arguments", {}) + CONSOLE.print(f"\n[bold cyan]Инструмент:{name}({args})[/bold cyan]") + + # Обработка паузы + if "__interrupt__" in chunk_data and state.next == ("tools",): + # Получаем последний вызов инструмента + last_msg = state.values["messages"][-1] + tool_call = last_msg.tool_calls[0] + name, args = tool_call["name"], tool_call.get("arguments", {}) + CONSOLE.print(f"\n[bold magenta]Агент хочет вызвать инструмент {name}({args})[/bold magenta]") + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() in ("y", ""): + # Возобновляем без новых сообщений + ask_and_run(None, config) + else: + CONSOLE.print("[red]Отменено[/red]") + break + + +# ------------------------------------------------------------------ +# Основной цикл чата +# ------------------------------------------------------------------ + + +def main() -> None: + thread_id = "thread-1" + config = {"configurable": {"thread_id": thread_id}} + + while True: + user_input = input("\nВы: ") + if user_input.lower().strip() == "exit": + break + + ask_and_run( + {"messages": [{"role": "human", "content": user_input}]}, + config, + ) + + +if __name__ == "__main__": + main() \ No newline at end of file