Update graph.py

graph.py

@@ -1,35 +1,101 @@
-from typing import Annotated, Sequence, TypedDict
+import os
+from typing import Annotated, Literal, TypedDict
+from pydantic import BaseModel
+
+from langchain_core.messages import BaseMessage, SystemMessage
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, START, END
-from langchain_core.messages import BaseMessage
-import os
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import create_react_agent
 
-# Состояние нашего чата
+# Импорты из соседних файлов нашего проекта
+from tools import ads_tools
+from prompts import SUPERVISOR_PROMPT, ADS_ANALYST_PROMPT
+from database import checkpointer
+
+# ==========================================
+# 1. СТРУКТУРА СОСТОЯНИЯ (STATE)
+# ==========================================
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[BaseMessage], "add_messages"]
-    next_node: str # Кто ходит следующим
+    # Указываем, что новые сообщения будут добавляться к старым (склеиваться)
+    messages: Annotated[list[BaseMessage], add_messages]
+    # Переменная для хранения решения Супервайзера (кому передать ход)
+    next_node: str
 
-# 1. Настройка модели через OpenRouter
+# ==========================================
+# 2. ИНИЦИАЛИЗАЦИЯ LLM (OpenRouter)
+# ==========================================
 llm = ChatOpenAI(
-    model="google/gemini-2.0-flash-001", # Быстрая и дешевая для логики
+    model="google/gemini-2.0-flash-001", # Можно поменять на claude-3-5-sonnet
     openai_api_key=os.getenv("OPENROUTER_API_KEY"),
-    openai_api_base="https://openrouter.ai/api/v1"
+    openai_api_base="https://openrouter.ai/api/v1",
+    temperature=0.2 # Низкая температура для аналитики, чтобы модель меньше фантазировала
+)
+
+# ==========================================
+# 3. УЗЕЛ: АНАЛИТИК РЕКЛАМЫ (Суб-агент)
+# ==========================================
+# Используем готовый шаблон ReAct-агента, который умеет вызывать Tools
+ads_agent_runnable = create_react_agent(
+    model=llm,
+    tools=ads_tools,
+    state_modifier=ADS_ANALYST_PROMPT # Передаем огромный промпт из файла
 )
 
-# 2. Узел Супервайзера (Директор)
-def supervisor_node(state: AgentState):
-    # Промпт: реши, кому передать задачу: SALES, ADV или закончить (FINISH)
-    # Здесь логика выбора на основе текста пользователя
-    return {"next_node": "SALES_AGENT"} 
+async def ads_node(state: AgentState):
+    """
+    Агент по рекламе берет историю диалога, вызывает нужное API (tools)
+    и формирует итоговый аналитический ответ.
+    """
+    # Вызываем агента и передаем ему текущую историю сообщений
+    result = await ads_agent_runnable.ainvoke({"messages": state["messages"]})
+    
+    # Возвращаем только последнее сгенерированное сообщение, 
+    # чтобы функция add_messages корректно добавила его в State
+    return {"messages": [result["messages"][-1]]}
+
+# ==========================================
+# 4. УЗЕЛ: СУПЕРВАЙЗЕР (Оркестратор)
+# ==========================================
+class RouterOutput(BaseModel):
+    # Строго ограничиваем варианты ответа Супервайзера
+    next_node: Literal["Ads_Analyst", "FINISH"]
+
+async def supervisor_node(state: AgentState):
+    """
+    Супервайзер читает весь диалог и решает, кому передать задачу.
+    Он возвращает только JSON с названием следующего узла.
+    """
+    # Собираем промпт + историю
+    messages = [SystemMessage(content=SUPERVISOR_PROMPT)] + state["messages"]
+    
+    # Заставляем LLM ответить строго по структуре RouterOutput
+    response = await llm.with_structured_output(RouterOutput).ainvoke(messages)
+    
+    return {"next_node": response.next_node}
 
-# Собираем граф
+# ==========================================
+# 5. СБОРКА ГРАФА СОСТОЯНИЙ (LangGraph)
+# ==========================================
 workflow = StateGraph(AgentState)
-# Добавляем узлы (Nodes)
-workflow.add_node("SUPERVISOR", supervisor_node)
-# Добавим сюда узлы для SALES и ADV...
 
-workflow.add_edge(START, "SUPERVISOR")
-# ... логика переходов
+# Добавляем узлы в граф
+workflow.add_node("Supervisor", supervisor_node)
+workflow.add_node("Ads_Analyst", ads_node)
+
+# Логика условного перехода: куда идти после Супервайзера?
+def router(state: AgentState) -> str:
+    if state["next_node"] == "FINISH":
+        return END
+    return state["next_node"]
+
+# Строим связи (ребра графа)
+workflow.add_edge(START, "Supervisor") # Старт всегда направляется к Супервайзеру
+workflow.add_conditional_edges("Supervisor", router) # От Супервайзера - по условию
+workflow.add_edge("Ads_Analyst", "Supervisor") # После того как Аналитик ответил, возвращаем ход Супервайзеру
 
-# Компилируем с нашей базой данных!
+# ==========================================
+# 6. КОМПИЛЯЦИЯ И ПОДКЛЮЧЕНИЕ ПАМЯТИ
+# ==========================================
+# Компилируем граф, передавая checkpointer (Supabase PostgresSaver)
 graph = workflow.compile(checkpointer=checkpointer)