example.py

"""
Exemplo simples de uso do LangGraph
"""
from pydantic import BaseModel
from typing import List, Optional
from langgraph.graph import START, END, StateGraph

# Definimos o modelo de estado que será usado no grafo
class EstadoCalculadora(BaseModel):
    numeros: List[int] = []
    operacao: Optional[str] = None
    resultado: Optional[float] = None
    historico: List[str] = []

# Função para adicionar números ao estado
def adicionar_numeros(state: EstadoCalculadora, numeros: List[int]):
    """Adiciona números ao estado atual."""
    return {"numeros": state.numeros + numeros}

# Função para definir a operação
def definir_operacao(state: EstadoCalculadora):
    """Define qual operação realizar com base nos números disponíveis."""
    # Simples exemplo: se temos mais de 2 números, multiplicamos, senão somamos
    if len(state.numeros) > 2:
        operacao = "multiplicar"
    else:
        operacao = "somar"
    
    return {"operacao": operacao}

# Função para calcular o resultado
def calcular(state: EstadoCalculadora):
    """Realiza o cálculo conforme a operação definida."""
    if state.operacao == "somar":
        resultado = sum(state.numeros)
        descricao = f"Soma de {state.numeros} = {resultado}"
    elif state.operacao == "multiplicar":
        resultado = 1
        for num in state.numeros:
            resultado *= num
        descricao = f"Multiplicação de {state.numeros} = {resultado}"
    else:
        resultado = 0
        descricao = "Operação não definida"
    
    # Adiciona o resultado e atualiza o histórico
    return {
        "resultado": resultado,
        "historico": state.historico + [descricao]
    }

# Função que determina se precisamos fazer mais cálculos
def verificar_fluxo(state: EstadoCalculadora):
    """Decide qual deve ser o próximo passo no fluxo."""
    if state.resultado is None:
        # Se não temos resultado, precisamos calcular
        return "calcular"
    else:
        # Se já temos um resultado, terminamos
        return END

# Construção do grafo
def construir_grafo():
    """Constrói e retorna o grafo de estado da calculadora."""
    builder = StateGraph(EstadoCalculadora)
    
    # Adicionamos os nós ao grafo
    builder.add_node("definir_operacao", definir_operacao)
    builder.add_node("calcular", calcular)
    
    # Configuramos as arestas no grafo
    builder.add_edge(START, "definir_operacao")
    builder.add_edge("definir_operacao", "calcular")
    builder.add_conditional_edges("calcular", verificar_fluxo, {
        "calcular": "calcular",  # Loop para mais cálculos se necessário
        END: END                 # Finalizar o fluxo
    })
    
    return builder.compile()

# Exemplo de uso
if __name__ == "__main__":
    # Construímos o grafo
    grafo = construir_grafo()
    
    # Definimos o estado inicial
    estado_inicial = {"numeros": [5, 10, 15]}
    
    # Executamos o grafo com o estado inicial
    resultado = grafo.invoke(estado_inicial)
    
    print(f"Estado final: {resultado}")
    print(f"Resultado do cálculo: {resultado.resultado}")
    print(f"Histórico de operações:")
    for entrada in resultado.historico:
        print(f"  - {entrada}")
    
    # Exemplo com execução passo a passo (streaming)
    print("\nExecutando novamente com streaming:")
    for passo in grafo.stream({"numeros": [2, 3]}):
        print(f"Passo: {passo}")