Добавить agent.py

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()