from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from pydantic import BaseModel, Field
from typing import List, Optional
from google.colab import userdata

# Улучшенная модель состояния
class DebateState(BaseModel):
    topic: str
    round: int
    optimist_args: List[str] = Field(default_factory=list)
    pessimist_args: List[str] = Field(default_factory=list)
    verdict: Optional[str] = Field(default=None, description="Решение судьи")

    def add_optimist_arg(self, arg: str):
        self.optimist_args.append(arg)
        
    def add_pessimist_arg(self, arg: str):
        self.pessimist_args.append(arg)
        
    def increment_round(self):
        self.round += 1

# Инициализация LLM
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0.0,
    openai_api_base="https://openrouter.ai/api/v1",
    openai_api_key=userdata.get('API_KEY')
)

# Узлы графа
def optimist(state: DebateState):
    # Получаем последний аргумент оппонента (если есть)
    last_pessimist_arg = state.pessimist_args[-1] if state.pessimist_args else "Пока нет аргументов"
    
    response = llm.invoke(
        f"Ты оптимист. Раунд: {state.round}\n"
        f"Тема: {state.topic}\n"
        f"Последний аргумент оппонента: {last_pessimist_arg}\n"
        "Приведи один новый аргумент ЗА. Ответь только аргументом."
    )
    
    state.add_optimist_arg(response.content)
    print(f"\n[Раунд {state.round}] Оптимист: {response.content}")
    return state

def pessimist(state: DebateState):
    # Получаем последний аргумент оппонента
    last_optimist_arg = state.optimist_args[-1] if state.optimist_args else "Пока нет аргументов"
    
    response = llm.invoke(
        f"Ты пессимист. Раунд: {state.round}\n"
        f"Тема: {state.topic}\n"
        f"Последний аргумент оппонента: {last_optimist_arg}\n"
        "Приведи один новый аргумент ПРОТИВ. Ответь только аргументом."
    )
    
    state.add_pessimist_arg(response.content)
    print(f"[Раунд {state.round}] Пессимист: {response.content}")
    return state

def judge(state: DebateState):
    response = llm.invoke(
        f"Ты судья. Определи победителя дебатов на тему: '{state.topic}'\n\n"
        f"Все аргументы ЗА ({len(state.optimist_args)}):\n- " + "\n- ".join(state.optimist_args) + "\n\n"
        f"Все аргументы ПРОТИВ ({len(state.pessimist_args)}):\n- " + "\n- ".join(state.pessimist_args) + "\n\n"
        "Ответь строго в формате: 'Победил оптимист' или 'Победил пессимист'"
    )
    state.verdict = response.content
    print(f"\n[Судья] {state.verdict}")
    return state

# Логика переходов
def check_round(state: DebateState):
    state.increment_round()
    print(f"\n--- Начат раунд {state.round} ---")
    return "optimist" if state.round <= 3 else judge

# Строим граф
graph_builder = StateGraph(DebateState)
graph_builder.add_node("optimist", optimist)
graph_builder.add_node("pessimist", pessimist)
graph_builder.add_node("judge", judge)

graph_builder.set_entry_point("optimist")
graph_builder.add_edge("optimist", "pessimist")
graph_builder.add_conditional_edges("pessimist", check_round)
graph_builder.add_edge("judge", END)

graph = graph_builder.compile()

# Запускаем дебаты с правильной инициализацией
topic = "роботы заменят людей на большинстве работ"
initial_state = DebateState(topic=topic, round=0)  # Правильная инициализация
result = graph.invoke(initial_state)

print(f"\nФинальное решение: {result.verdict}")

Почему не меняется state.round?