# Versión simplificada (código real es 400+ líneas)
from typing import TypedDict, Literal
from langchain_aws import ChatBedrock
from langchain_pinecone import PineconeVectorStore
from langgraph.graph import StateGraph, END

# State
class SupportState(TypedDict):
    messages: list
    category: str
    confidence: float
    escalate_to_human: bool
    resolved: bool

# LLMs
llm_fast = ChatBedrock(model="anthropic.claude-3-haiku-20240307")
llm_smart = ChatBedrock(model="anthropic.claude-3-5-sonnet-20240620")

# Vector store
vectorstore = PineconeVectorStore.from_existing_index(
    index_name="mastersuiteai-kb",
    embedding=OpenAIEmbeddings()
)

# Agents
def triage_agent(state: SupportState) -> dict:
    """Clasifica ticket con Claude Haiku."""
    prompt = f"Clasifica este ticket en: billing, technical, product, general\\ \\ {state['messages'][-1]}"
    response = llm_fast.invoke(prompt)
    category = response.content.strip().lower()
    return {"category": category}

def knowledge_agent(state: SupportState) -> dict:
    """Retrieval + reranking en Pinecone."""
    query = state["messages"][-1]

    # Vector search (top 5 docs)
    docs = vectorstore.similarity_search(query, k=5)

    # Rerank con LLM
    context = "\\ \\ ".join([doc.page_content for doc in docs])
    prompt = f"""Responde basándote SOLO en este contexto:
{context}

Pregunta: {query}

Si el contexto no tiene info suficiente, di "No encontré información específica sobre esto"."""
    response = llm_smart.invoke(prompt)

    # Calculate confidence (simple heuristic)
    confidence = 0.9 if "No encontré" not in response.content else 0.3

    return {
        "messages": state["messages"] + [response],
        "confidence": confidence
    }

def action_agent(state: SupportState) -> dict:
    """Ejecuta acciones basadas en category."""
    category = state["category"]

    # Simular acciones (en producción: API calls reales)
    if category == "billing":
        action = "Created refund request in Stripe"
    elif category == "technical":
        action = "Escalated to engineering team in Zendesk"
    else:
        action = "Tagged user in Intercom for follow-up"

    return {
        "messages": state["messages"] + [{"role": "system", "content": action}],
        "resolved": True
    }

def supervisor(state: SupportState) -> dict:
    """Decide siguiente paso."""
    # Si no hay category, ir a triage
    if not state.get("category"):
        return {"next": "triage"}

    # Si no hay knowledge response, ir a knowledge
    if len(state["messages"]) < 3:
        return {"next": "knowledge"}

    # Si confidence baja O technical, escalar
    if state.get("confidence", 0) < 0.8 or state["category"] == "technical":
        return {"escalate_to_human": True, "next": "end"}

    # Sino, ejecutar action
    return {"next": "action"}

# Graph construction
workflow = StateGraph(SupportState)
workflow.add_node("supervisor", supervisor)
workflow.add_node("triage", triage_agent)
workflow.add_node("knowledge", knowledge_agent)
workflow.add_node("action", action_agent)
workflow.set_entry_point("supervisor")

def route(state: SupportState) -> Literal["triage", "knowledge", "action", "end"]:
    next_step = state.get("next", "triage")
    return END if next_step == "end" else next_step

workflow.add_conditional_edges("supervisor", route, {
    "triage": "triage",
    "knowledge": "knowledge",
    "action": "action",
    "end": END
})
workflow.add_edge("triage", "supervisor")
workflow.add_edge("knowledge", "supervisor")
workflow.add_edge("action", "supervisor")

# Compile con Postgres checkpointer
from langgraph.checkpoint.postgres import PostgresSaver
checkpointer = PostgresSaver.from_conn_string(os.environ["POSTGRES_URI"])
app = workflow.compile(checkpointer=checkpointer)

from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages

# Opción 1: TypedDict simple (más común)
class AgentState(TypedDict):
    """
    State compartido por todos los nodes del graph.
    'messages' usa el reducer add_messages para append automático.
    """
    messages: Annotated[Sequence[BaseMessage], add_messages]
    current_step: str
    user_input: str
    final_answer: str

# Opción 2: Pydantic model (validación strict)
from pydantic import BaseModel, Field

class AgentStatePydantic(BaseModel):
    """State con validación Pydantic para producción."""
    messages: list[BaseMessage] = Field(default_factory=list)
    current_step: str = "init"
    user_input: str = ""
    final_answer: str = ""
    metadata: dict = Field(default_factory=dict)

# Opción 3: State con reducers custom
def concat_strings(existing: str, new: str) -> str:
    """Custom reducer para concatenar strings."""
    return f"{existing}\\ {new}" if existing else new

class AdvancedState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    research_notes: Annotated[str, concat_strings]  # Acumula notas
    iterations: int  # Se sobreescribe (default behavior)

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage

# Node simple que llama al LLM
def call_llm_node(state: AgentState) -> dict:
    """
    Node que toma messages del state, llama a GPT-4,
    y retorna la respuesta actualizada.
    """
    llm = ChatOpenAI(model="gpt-4", temperature=0)

    # LangGraph hace merge automático con state existente
    response = llm.invoke(state["messages"])

    return {
        "messages": [response],  # add_messages reducer hace append
        "current_step": "llm_called"
    }

# Node con tool calling
from langchain_community.tools.tavily_search import TavilySearchResults

def search_node(state: AgentState) -> dict:
    """Node que ejecuta web search con Tavily."""
    search_tool = TavilySearchResults(max_results=3)

    # Extraer query del último mensaje
    last_message = state["messages"][-1]
    query = last_message.content

    # Ejecutar búsqueda
    results = search_tool.invoke({"query": query})

    # Agregar resultados como nuevo mensaje
    search_summary = "\\ ".join([r["content"] for r in results])

    return {
        "messages": [AIMessage(content=f"Encontré: {search_summary}")],
        "current_step": "search_completed"
    }

# Node con decisión condicional (usado en edge routing)
def should_continue(state: AgentState) -> str:
    """
    Decide si continuar o finalizar basándose en state.
    Retorna string que se usa en conditional_edges.
    """
    last_message = state["messages"][-1]

    # Si el mensaje tiene tool_calls, continuar con tools
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "continue"

    # Si tiene respuesta final, terminar
    if "RESPUESTA FINAL:" in last_message.content:
        return "end"

    # Default: continuar iterando
    return "continue"

from langgraph.graph import StateGraph, END

# Crear graph
graph = StateGraph(AgentState)

# Agregar nodes
graph.add_node("agent", call_llm_node)
graph.add_node("tools", search_node)

# Fixed edge: siempre search → agent
graph.add_edge("tools", "agent")

# Conditional edge: agent → decide dinámicamente
graph.add_conditional_edges(
    "agent",  # Source node
    should_continue,  # Función que retorna string
    {
        "continue": "tools",  # Si retorna "continue", ir a tools
        "end": END  # Si retorna "end", terminar workflow
    }
)

# Set entry point (primer node)
graph.set_entry_point("agent")

# Compilar graph
app = graph.compile()

from langgraph.checkpoint.memory import MemorySaver
from langchain_core.messages import HumanMessage, AIMessage

# 1. State con approval flag
class HITLState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    action_proposed: str
    approved: bool

# 2. Node que propone acción crítica
def propose_refund(state: HITLState) -> dict:
    """Propone refund, espera aprobación humana."""
    return {
        "messages": [AIMessage(content="Propongo refund de $150. ¿Aprobar?")],
        "action_proposed": "refund_150",
        "approved": False
    }

# 3. Node que ejecuta solo si aprobado
def execute_action(state: HITLState) -> dict:
    """Ejecuta acción solo si está aprobada."""
    if not state["approved"]:
        return {"messages": [AIMessage(content="❌ Acción cancelada por usuario.")]}

    action = state["action_proposed"]

    # Ejecutar acción real
    if action == "refund_150":
        result = "✅ Refund de $150 procesado. ID: REF-99999"
    else:
        result = f"✅ Acción '{action}' ejecutada."

    return {"messages": [AIMessage(content=result)]}

# 4. Construir graph con interrupt
workflow = StateGraph(HITLState)
workflow.add_node("propose", propose_refund)
workflow.add_node("execute", execute_action)
workflow.set_entry_point("propose")
workflow.add_edge("propose", "execute")

# CLAVE: Usar checkpointer + interrupt_before
checkpointer = MemorySaver()
app_hitl = workflow.compile(
    checkpointer=checkpointer,
    interrupt_before=["execute"]  # Pausa ANTES de ejecutar
)

# 5. Ejecutar workflow con thread_id (para reanudar después)
thread_config = {"configurable": {"thread_id": "1"}}

# Primera ejecución: pausa en execute
result = app_hitl.invoke(
    {"messages": [HumanMessage(content="Cobró el doble, necesito refund")]},
    config=thread_config
)

print("\\ === ESTADO DESPUÉS DE PROPUESTA ===")
print(result["messages"][-1].content)
print(f"Action propuesta: {result['action_proposed']}")
print(f"Aprobada: {result['approved']}")

# Simular aprobación humana (en producción: UI button click)
print("\\ [HUMANO APRUEBA LA ACCIÓN]")

# Actualizar state con aprobación
updated_state = app_hitl.get_state(thread_config)
updated_state.values["approved"] = True
app_hitl.update_state(thread_config, updated_state.values)

# Reanudar ejecución desde donde pausó
final_result = app_hitl.invoke(None, config=thread_config)

print("\\ === ESTADO FINAL DESPUÉS DE APROBACIÓN ===")
print(final_result["messages"][-1].content)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langgraph.checkpoint.postgres import PostgresSaver

app = FastAPI()

# PostgreSQL checkpointer para production
checkpointer = PostgresSaver.from_conn_string("postgresql://user:pass@localhost/langgraph")
app_hitl_prod = workflow.compile(
    checkpointer=checkpointer,
    interrupt_before=["execute"]
)

class StartWorkflowRequest(BaseModel):
    user_message: str
    thread_id: str

class ApproveActionRequest(BaseModel):
    thread_id: str
    approved: bool

@app.post("/workflow/start")
async def start_workflow(req: StartWorkflowRequest):
    """Inicia workflow, pausa en approval gate."""
    config = {"configurable": {"thread_id": req.thread_id}}
    result = app_hitl_prod.invoke(
        {"messages": [HumanMessage(content=req.user_message)]},
        config=config
    )
    return {
        "status": "waiting_approval",
        "message": result["messages"][-1].content,
        "action_proposed": result["action_proposed"],
        "thread_id": req.thread_id
    }

@app.post("/workflow/approve")
async def approve_action(req: ApproveActionRequest):
    """Aprueba/rechaza acción y reanuda workflow."""
    config = {"configurable": {"thread_id": req.thread_id}}

    # Obtener state actual
    try:
        state = app_hitl_prod.get_state(config)
    except Exception:
        raise HTTPException(status_code=404, detail="Thread not found")

    # Actualizar approval
    state.values["approved"] = req.approved
    app_hitl_prod.update_state(config, state.values)

    # Reanudar ejecución
    final_result = app_hitl_prod.invoke(None, config=config)

    return {
        "status": "completed",
        "message": final_result["messages"][-1].content
    }

# Frontend JavaScript (React hook example)
"""
const useHITLWorkflow = () => {
  const [workflowState, setWorkflowState] = useState(null);

  const startWorkflow = async (message) => {
    const threadId = uuidv4();
    const res = await fetch('/workflow/start', {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({user_message: message, thread_id: threadId})
    });
    const data = await res.json();
    setWorkflowState(data);
  };

  const approveAction = async (approved) => {
    const res = await fetch('/workflow/approve', {
      method: 'POST',
      headers: {'Content-Type': 'application/json'},
      body: JSON.stringify({thread_id: workflowState.thread_id, approved})
    });
    const data = await res.json();
    setWorkflowState(data);
  };

  return {workflowState, startWorkflow, approveAction};
};
"""

import os

# Enable tracing (environment variables)
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "ls__..."  # Tu API key de LangSmith
os.environ["LANGCHAIN_PROJECT"] = "langgraph-production"  # Proyecto en dashboard

# Ahora TODAS las ejecuciones se tracean automáticamente
result = app.invoke({"messages": [HumanMessage(content="Test")]})

# Ver en dashboard: https://smith.langchain.com

# ❌ ANTES (sequential): 10s total
workflow.add_edge("supervisor", "research_agent")  # 5s
workflow.add_edge("research_agent", "code_agent")  # 5s

# ✅ DESPUÉS (parallel): 5s total
from langgraph.graph import parallel

# Ejecutar research_agent y code_agent en paralelo
workflow.add_conditional_edges(
    "supervisor",
    lambda s: ["research_agent", "code_agent"],  # Lista = paralelo
    ["research_agent", "code_agent"]
)

# Ambos convergen en "merge_node"
workflow.add_edge("research_agent", "merge_node")
workflow.add_edge("code_agent", "merge_node")

from functools import lru_cache

# Opción 1: Python @lru_cache (simple)
@lru_cache(maxsize=1000)
def expensive_node(state_hash: str) -> dict:
    """Node cacheado en memoria Python."""
    # Computación costosa aquí
    return {"result": "..."}

# Opción 2: Redis caching (production)
import redis
import json
import hashlib

redis_client = redis.Redis.from_url(os.environ["REDIS_URL"])

def cached_search_node(state: AgentState) -> dict:
    """Node con Redis caching."""
    query = state["messages"][-1].content

    # Hash del input para cache key
    cache_key = f"search:{hashlib.md5(query.encode()).hexdigest()}"

    # Check cache
    cached = redis_client.get(cache_key)
    if cached:
        print("✅ Cache HIT")
        return json.loads(cached)

    # Cache MISS: ejecutar search
    print("❌ Cache MISS: ejecutando search")
    results = search_tool.invoke(query)

    # Guardar en cache (TTL 1 hora)
    output = {"messages": [AIMessage(content=results)]}
    redis_client.setex(cache_key, 3600, json.dumps(output))

    return output

from langchain_openai import ChatOpenAI

# Modelos con diferentes cost/performance
llm_fast = ChatOpenAI(model="gpt-4o-mini", temperature=0)  # $0.15/1M tokens
llm_smart = ChatOpenAI(model="gpt-4", temperature=0)  # $30/1M tokens

def router_node(state: AgentState) -> dict:
    """Decide qué modelo usar basándose en complejidad."""
    query = state["messages"][-1].content

    # Heurística simple: si query < 50 palabras y no menciona "complejo"
    # → usar modelo barato
    is_simple = (
        len(query.split()) < 50
        and "complejo" not in query.lower()
        and "análisis" not in query.lower()
    )

    llm = llm_fast if is_simple else llm_smart
    model_used = "gpt-4o-mini" if is_simple else "gpt-4"

    response = llm.invoke(state["messages"])

    return {
        "messages": [response],
        "model_used": model_used
    }

from langgraph.checkpoint.postgres import PostgresSaver
import psycopg

# 1. Crear connection string
DB_URI = "postgresql://langraph_user:secure_password@localhost:5432/langgraph_db"

# 2. Inicializar checkpointer
with psycopg.connect(DB_URI) as conn:
    # Setup automático de tablas (primera vez)
    checkpointer = PostgresSaver(conn)
    checkpointer.setup()

# 3. Compilar graph con Postgres checkpointer
app_prod = workflow.compile(checkpointer=checkpointer)

# 4. Uso idéntico a MemorySaver
config = {"configurable": {"thread_id": "user-123"}}
result = app_prod.invoke({"messages": [HumanMessage(content="Test")]}, config=config)

# El state se persiste en Postgres automáticamente
# Si la app crashea, puedes reanudar desde último checkpoint

-- Tablas creadas automáticamente por checkpointer.setup()
CREATE TABLE checkpoints (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    parent_checkpoint_id TEXT,
    type TEXT,
    checkpoint JSONB NOT NULL,
    metadata JSONB NOT NULL DEFAULT '{}',
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
);

CREATE TABLE checkpoint_writes (
    thread_id TEXT NOT NULL,
    checkpoint_ns TEXT NOT NULL DEFAULT '',
    checkpoint_id TEXT NOT NULL,
    task_id TEXT NOT NULL,
    idx INTEGER NOT NULL,
    channel TEXT NOT NULL,
    type TEXT,
    blob BYTEA,
    PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
);

-- Indexes para performance
CREATE INDEX idx_checkpoints_thread_id ON checkpoints(thread_id);
CREATE INDEX idx_checkpoint_writes_thread_id ON checkpoint_writes(thread_id);

FROM python:3.11-slim

WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \\
    postgresql-client \\
    && rm -rf /var/lib/apt/lists/*

# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY . .

# Environment variables (overridden by ECS task definition)
ENV OPENAI_API_KEY=""
ENV POSTGRES_URI=""
ENV LANGCHAIN_API_KEY=""

# Expose port
EXPOSE 8000

# Health check
HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \\
    CMD curl -f http://localhost:8000/health || exit 1

# Run FastAPI with uvicorn
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

version: '3.8'

services:
  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: langgraph_user
      POSTGRES_PASSWORD: secure_password
      POSTGRES_DB: langgraph_db
    ports:
      - "5432:5432"
    volumes:
      - postgres_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U langgraph_user"]
      interval: 10s
      timeout: 5s
      retries: 5

  redis:
    image: redis:7-alpine
    ports:
      - "6379:6379"
    healthcheck:
      test: ["CMD", "redis-cli", "ping"]
      interval: 10s
      timeout: 3s
      retries: 5

  app:
    build: .
    ports:
      - "8000:8000"
    environment:
      OPENAI_API_KEY: ${OPENAI_API_KEY}
      POSTGRES_URI: postgresql://langgraph_user:secure_password@postgres:5432/langgraph_db
      LANGCHAIN_API_KEY: ${LANGCHAIN_API_KEY}
      REDIS_URL: redis://redis:6379
    depends_on:
      postgres:
        condition: service_healthy
      redis:
        condition: service_healthy

volumes:
  postgres_data:

terraform {
  required_version = ">= 1.0"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.0"
    }
  }
}

provider "aws" {
  region = "eu-west-1"
}

# VPC y Networking
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  name    = "langgraph-vpc"
  cidr    = "10.0.0.0/16"

  azs             = ["eu-west-1a", "eu-west-1b"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]
  public_subnets  = ["10.0.101.0/24", "10.0.102.0/24"]

  enable_nat_gateway   = true
  enable_dns_hostnames = true
}

# RDS Postgres (Checkpointer DB)
resource "aws_db_instance" "langgraph" {
  identifier     = "langgraph-postgres"
  engine         = "postgres"
  engine_version = "16.3"
  instance_class = "db.t4g.micro"  # Free tier elegible

  allocated_storage = 20
  db_name          = "langgraph_db"
  username         = "langgraph_admin"
  password         = var.db_password  # Terraform variable

  vpc_security_group_ids = [aws_security_group.postgres.id]
  db_subnet_group_name   = aws_db_subnet_group.langgraph.name

  backup_retention_period  = 7
  skip_final_snapshot     = false
  final_snapshot_identifier = "langgraph-final-snapshot"

  tags = {
    Environment = "production"
    Service     = "langgraph"
  }
}

# ElastiCache Redis (Caching layer)
resource "aws_elasticache_cluster" "langgraph" {
  cluster_id           = "langgraph-cache"
  engine               = "redis"
  node_type            = "cache.t4g.micro"
  num_cache_nodes      = 1
  parameter_group_name = "default.redis7"
  port                 = 6379

  subnet_group_name  = aws_elasticache_subnet_group.langgraph.name
  security_group_ids = [aws_security_group.redis.id]
}

# ECS Cluster
resource "aws_ecs_cluster" "langgraph" {
  name = "langgraph-cluster"

  setting {
    name  = "containerInsights"
    value = "enabled"
  }
}

# ECR Repository (Docker images)
resource "aws_ecr_repository" "langgraph" {
  name                 = "langgraph-app"
  image_tag_mutability = "MUTABLE"

  image_scanning_configuration {
    scan_on_push = true
  }
}

# ECS Task Definition
resource "aws_ecs_task_definition" "langgraph" {
  family                   = "langgraph-app"
  network_mode             = "awsvpc"
  requires_compatibilities = ["FARGATE"]
  cpu                      = "512"   # 0.5 vCPU
  memory                   = "1024"  # 1GB

  execution_role_arn = aws_iam_role.ecs_execution.arn
  task_role_arn      = aws_iam_role.ecs_task.arn

  container_definitions = jsonencode([{
    name  = "langgraph-app"
    image = "${aws_ecr_repository.langgraph.repository_url}:latest"

    portMappings = [{
      containerPort = 8000
      protocol      = "tcp"
    }]

    environment = [
      {
        name  = "POSTGRES_URI"
        value = "postgresql://${aws_db_instance.langgraph.username}:${var.db_password}@${aws_db_instance.langgraph.endpoint}/${aws_db_instance.langgraph.db_name}"
      },
      {
        name  = "REDIS_URL"
        value = "redis://${aws_elasticache_cluster.langgraph.cache_nodes[0].address}:6379"
      }
    ]

    secrets = [
      {
        name      = "OPENAI_API_KEY"
        valueFrom = aws_secretsmanager_secret.openai_key.arn
      },
      {
        name      = "LANGCHAIN_API_KEY"
        valueFrom = aws_secretsmanager_secret.langchain_key.arn
      }
    ]

    logConfiguration = {
      logDriver = "awslogs"
      options = {
        "awslogs-group"         = aws_cloudwatch_log_group.langgraph.name
        "awslogs-region"        = "eu-west-1"
        "awslogs-stream-prefix" = "ecs"
      }
    }

    healthCheck = {
      command     = ["CMD-SHELL", "curl -f http://localhost:8000/health || exit 1"]
      interval    = 30
      timeout     = 5
      retries     = 3
      startPeriod = 60
    }
  }])
}

# ECS Service
resource "aws_ecs_service" "langgraph" {
  name            = "langgraph-service"
  cluster         = aws_ecs_cluster.langgraph.id
  task_definition = aws_ecs_task_definition.langgraph.arn
  desired_count   = 2  # 2 instancias para HA
  launch_type     = "FARGATE"

  network_configuration {
    subnets          = module.vpc.private_subnets
    security_groups  = [aws_security_group.ecs_tasks.id]
    assign_public_ip = false
  }

  load_balancer {
    target_group_arn = aws_lb_target_group.langgraph.arn
    container_name   = "langgraph-app"
    container_port   = 8000
  }

  depends_on = [aws_lb_listener.langgraph]
}

# Application Load Balancer
resource "aws_lb" "langgraph" {
  name               = "langgraph-alb"
  internal           = false
  load_balancer_type = "application"
  security_groups    = [aws_security_group.alb.id]
  subnets            = module.vpc.public_subnets
}

resource "aws_lb_target_group" "langgraph" {
  name        = "langgraph-tg"
  port        = 8000
  protocol    = "HTTP"
  vpc_id      = module.vpc.vpc_id
  target_type = "ip"

  health_check {
    path                = "/health"
    healthy_threshold   = 2
    unhealthy_threshold = 3
    timeout             = 5
    interval            = 30
  }
}

resource "aws_lb_listener" "langgraph" {
  load_balancer_arn = aws_lb.langgraph.arn
  port              = "443"
  protocol          = "HTTPS"
  ssl_policy        = "ELBSecurityPolicy-TLS13-1-2-2021-06"
  certificate_arn   = var.acm_certificate_arn

  default_action {
    type             = "forward"
    target_group_arn = aws_lb_target_group.langgraph.arn
  }
}

# Outputs
output "alb_dns_name" {
  description = "DNS name del ALB"
  value       = aws_lb.langgraph.dns_name
}

output "postgres_endpoint" {
  description = "Endpoint RDS Postgres"
  value       = aws_db_instance.langgraph.endpoint
  sensitive   = true
}

output "redis_endpoint" {
  description = "Endpoint ElastiCache Redis"
  value       = aws_elasticache_cluster.langgraph.cache_nodes[0].address
}

from typing import TypedDict, Annotated, Sequence, Literal
from langchain_core.messages import BaseMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages

# 1. State compartido por todos los agentes
class MultiAgentState(TypedDict):
    """State con metadata adicional para supervisor."""
    messages: Annotated[Sequence[BaseMessage], add_messages]
    next_agent: str  # Qué agente ejecutar siguiente
    ticket_category: str  # billing, technical, general
    escalate_to_human: bool  # Flag para human handoff

# 2. Definir workers especializados
llm = ChatOpenAI(model="gpt-4", temperature=0)

def triage_agent(state: MultiAgentState) -> dict:
    """
    Clasifica el ticket en: billing, technical, general.
    """
    system_prompt = """Eres un agente de triage para customer support.
Analiza el mensaje del usuario y clasifícalo en una de estas categorías:
- billing: Problemas de facturación, pagos, refunds
- technical: Bugs, errores, integraciones
- general: Preguntas generales, onboarding, features

Responde SOLO con la categoría, nada más."""

    messages = [SystemMessage(content=system_prompt)] + state["messages"]
    response = llm.invoke(messages)
    category = response.content.strip().lower()

    return {
        "messages": [response],
        "ticket_category": category
    }

def knowledge_agent(state: MultiAgentState) -> dict:
    """
    Busca en knowledge base y genera respuesta basada en docs.
    """
    # Simular búsqueda en knowledge base (en producción: vector search)
    knowledge_base = {
        "billing": "Para refunds, ve a Settings → Billing → Request Refund. Procesamos en 3-5 días.",
        "technical": "Para bugs, envía logs a support@ejemplo.com con descripción detallada.",
        "general": "Revisa nuestra documentación en docs.ejemplo.com para guías completas."
    }

    category = state["ticket_category"]
    knowledge = knowledge_base.get(category, "No encontré información relevante.")

    system_prompt = f"""Eres un agente de knowledge base.
Responde la pregunta del usuario usando esta información:
{knowledge}

Sé conciso pero amigable."""

    messages = [SystemMessage(content=system_prompt)] + state["messages"]
    response = llm.invoke(messages)

    return {"messages": [response]}

def action_agent(state: MultiAgentState) -> dict:
    """
    Ejecuta acciones: crear refund, escalar ticket, etc.
    """
    category = state["ticket_category"]

    # Lógica de acción basada en categoría
    if category == "billing":
        action_result = "✅ He creado un refund request (ID: REF-12345). Recibirás email de confirmación."
    elif category == "technical":
        action_result = "🎫 Ticket escalado a equipo técnico (ID: TECH-67890). Responderán en 24h."
    else:
        action_result = "📚 Te envié links relevantes por email."

    return {
        "messages": [HumanMessage(content=action_result)],
        "escalate_to_human": category == "technical"  # Escalar bugs a humano
    }

# 3. Supervisor: decide qué agente ejecutar siguiente
def supervisor_node(state: MultiAgentState) -> dict:
    """
    Supervisor analiza el state y decide el siguiente paso:
    - Si no hay categoría → triage
    - Si hay categoría pero no respuesta → knowledge
    - Si knowledge no resuelve → action
    - Si action completa O escalate → end
    """
    # Primera vez: ir a triage
    if not state.get("ticket_category"):
        return {"next_agent": "triage"}

    # Si triage completo pero no hay respuesta knowledge
    messages = state["messages"]
    has_knowledge_response = any(
        "knowledge base" in m.content.lower()
        for m in messages
        if hasattr(m, "content")
    )

    if not has_knowledge_response:
        return {"next_agent": "knowledge"}

    # Si knowledge no resuelve, ejecutar action
    has_action = any(
        "creado" in m.content.lower() or "escalado" in m.content.lower()
        for m in messages
        if hasattr(m, "content")
    )

    if not has_action:
        return {"next_agent": "action"}

    # Si action completa, terminar
    return {"next_agent": "end"}

# 4. Routing logic basada en supervisor
def route_supervisor(state: MultiAgentState) -> Literal["triage", "knowledge", "action", "end"]:
    """
    Lee next_agent del state y decide routing.
    """
    next_agent = state.get("next_agent", "triage")

    if next_agent == "end":
        return END

    return next_agent

# 5. Construir multi-agent graph
workflow = StateGraph(MultiAgentState)

# Agregar todos los agents como nodes
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("triage", triage_agent)
workflow.add_node("knowledge", knowledge_agent)
workflow.add_node("action", action_agent)

# Entry point siempre supervisor
workflow.set_entry_point("supervisor")

# Supervisor decide qué worker llamar
workflow.add_conditional_edges(
    "supervisor",
    route_supervisor,
    {
        "triage": "triage",
        "knowledge": "knowledge",
        "action": "action",
        "end": END
    }
)

# Después de cada worker, volver a supervisor para re-evaluar
workflow.add_edge("triage", "supervisor")
workflow.add_edge("knowledge", "supervisor")
workflow.add_edge("action", "supervisor")

# Compilar
multi_agent_app = workflow.compile()

# 6. Testear sistema completo
test_input = {
    "messages": [
        HumanMessage(content="Necesito un refund de mi último pago, cobró el doble")
    ],
    "ticket_category": "",
    "next_agent": "",
    "escalate_to_human": False
}

result = multi_agent_app.invoke(test_input)

print("\\ === FLUJO COMPLETO ===")
for i, msg in enumerate(result["messages"], 1):
    if hasattr(msg, "content"):
        print(f"{i}. {msg.content[:80]}...")

print(f"\\ Categoría: {result['ticket_category']}")
print(f"Escalado a humano: {result['escalate_to_human']}")

# Prevenir supervisor loops
class ImprovedState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    agents_called: list[str]  # Track qué agentes ya ejecutaron
    next_agent: str

def improved_supervisor(state: ImprovedState) -> dict:
    """Supervisor que previene loops."""
    agents_called = state.get("agents_called", [])

    # Si triage no ejecutado y no en lista
    if "triage" not in agents_called:
        return {
            "next_agent": "triage",
            "agents_called": agents_called + ["triage"]
        }

    # Si knowledge no ejecutado
    if "knowledge" not in agents_called:
        return {
            "next_agent": "knowledge",
            "agents_called": agents_called + ["knowledge"]
        }

    # Si action no ejecutado
    if "action" not in agents_called:
        return {
            "next_agent": "action",
            "agents_called": agents_called + ["action"]
        }

    # Todos ejecutados → end
    return {"next_agent": "end"}

import asyncio
from langchain_core.messages import HumanMessage

async def stream_agent_response():
    """Stream token-by-token del agente."""
    input_msg = {"messages": [HumanMessage(content="¿Qué es LangGraph?")]}

    # astream_events() es async, itera sobre cada evento
    async for event in app.astream_events(input_msg, version="v2"):
        kind = event["event"]

        # Filtrar eventos de streaming del LLM
        if kind == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                print(content, end="", flush=True)

# FastAPI endpoint con streaming
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

@app.post("/chat/stream")
async def chat_stream(message: str):
    """Endpoint que retorna SSE (Server-Sent Events)."""
    async def generate():
        input_msg = {"messages": [HumanMessage(content=message)]}

        async for event in app.astream_events(input_msg, version="v2"):
            if event["event"] == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                if content:
                    # Formato SSE
                    yield f"data: {content}\\ \\ "

        yield "data: [DONE]\\ \\ "

    return StreamingResponse(generate(), media_type="text/event-stream")

# Frontend JavaScript (fetch SSE)
"""
const streamChat = async (message) => {
  const response = await fetch('/chat/stream', {
    method: 'POST',
    headers: {'Content-Type': 'application/json'},
    body: JSON.stringify({message})
  });

  const reader = response.body.getReader();
  const decoder = new TextDecoder();

  while (true) {
    const {value, done} = await reader.read();
    if (done) break;

    const chunk = decoder.decode(value);
    const lines = chunk.split('\\ ');

    for (const line of lines) {
      if (line.startsWith('data: ')) {
        const data = line.slice(6);
        if (data === '[DONE]') return;

        // Append token al UI
        appendToChat(data);
      }
    }
  }
};
"""

# Stream solo updates (deltas)
for update in app.stream({"messages": [HumanMessage(content="Help")]}, stream_mode="updates"):
    for node_name, node_output in update.items():
        print(f"\\ [{node_name}] Update:")
        print(node_output)

# Output ejemplo:
# [agent] Update:
# {'messages': [AIMessage(content='Calling search tool...')]}
#
# [tools] Update:
# {'messages': [ToolMessage(content='Search results: ...')]}
#
# [agent] Update:
# {'messages': [AIMessage(content='Based on search: LangGraph is...')]}

from langsmith import Client
from langsmith.evaluation import evaluate

client = Client()

# 1. Crear dataset de test cases
dataset_name = "customer-support-eval"
examples = [
    {
        "inputs": {"messages": [{"role": "user", "content": "Necesito refund"}]},
        "outputs": {"category": "billing", "escalate": False}
    },
    {
        "inputs": {"messages": [{"role": "user", "content": "App crashea al login"}]},
        "outputs": {"category": "technical", "escalate": True}
    },
    # 50+ ejemplos representativos
]

dataset = client.create_dataset(dataset_name, examples=examples)

# 2. Definir evaluators (métricas)
def correctness_evaluator(run, example):
    """Verifica si categorización es correcta."""
    predicted_category = run.outputs.get("ticket_category")
    expected_category = example.outputs.get("category")

    return {
        "key": "correctness",
        "score": 1 if predicted_category == expected_category else 0
    }

def latency_evaluator(run, example):
    """Penaliza latency > 3s."""
    latency_seconds = (run.end_time - run.start_time).total_seconds()

    return {
        "key": "latency",
        "score": 1 if latency_seconds < 3 else 0.5 if latency_seconds < 5 else 0
    }

# 3. Ejecutar evaluation
results = evaluate(
    lambda inputs: app.invoke(inputs),
    data=dataset_name,
    evaluators=[correctness_evaluator, latency_evaluator],
    experiment_prefix="support-bot-v1"
)

# Ver resultados en LangSmith dashboard
print(f"Correctness: {results['correctness']:.2%}")
print(f"Latency: {results['latency']:.2%}")

from langsmith import Client

client = Client()

# Obtener traces de últimos 7 días con feedback positivo
runs = client.list_runs(
    project_name="langgraph-production",
    start_time="2025-11-05",
    filter='feedback_key="user_rating" AND feedback_score >= 4'
)

# Convertir traces a ejemplos
examples = []
for run in runs:
    examples.append({
        "inputs": run.inputs,
        "outputs": run.outputs,
        "metadata": {
            "run_id": str(run.id),
            "feedback_score": run.feedback_scores.get("user_rating")
        }
    })

# Crear dataset
client.create_dataset("production-golden-set", examples=examples[:100])

# ❌ INCORRECTO
def bad_node(state: AgentState) -> AgentState:
    state["messages"].append(new_message)  # Muta state directamente
    return state  # Retorna state completo

# ✅ CORRECTO
def good_node(state: AgentState) -> dict:
    return {
        "messages": [new_message]  # Solo updates, LangGraph hace merge
    }

# Solución 1: Aumentar limit
app = workflow.compile(
    checkpointer=checkpointer,
    recursion_limit=100  # Default es 25
)

# Solución 2: Agregar contador in state para prevenir loops
class StateWithCounter(TypedDict):
    messages: list
    iterations: int  # Counter

def node_with_counter(state: StateWithCounter) -> dict:
    return {
        "messages": [...],
        "iterations": state.get("iterations", 0) + 1
    }

def should_continue(state: StateWithCounter) -> str:
    # Terminar si > 10 iterations
    if state.get("iterations", 0) > 10:
        return "end"
    return "continue"

# ❌ INCORRECTO (no serializable)
class BadState(TypedDict):
    messages: list
    db_connection: psycopg.Connection  # ❌ No serializable!

# ✅ CORRECTO (solo JSON-serializable data)
class GoodState(TypedDict):
    messages: list
    db_connection_string: str  # String es serializable

# Reconstruir connection en cada node
def node_with_db(state: GoodState) -> dict:
    conn = psycopg.connect(state["db_connection_string"])
    # Usar connection aquí
    conn.close()
    return {"messages": [...]}

from langchain_core.messages import AIMessage

def robust_llm_node(state: AgentState) -> dict:
    """Node con error handling robusto."""
    try:
        response = llm.invoke(state["messages"], timeout=30)
        return {"messages": [response]}

    except TimeoutError:
        error_msg = AIMessage(content="Error: LLM timeout. Please retry.")
        return {"messages": [error_msg], "error": "timeout"}

    except Exception as e:
        error_msg = AIMessage(content=f"Error inesperado: {str(e)}")
        return {"messages": [error_msg], "error": str(e)}

# Crear virtual environment 
python3 -m venv venv source venv/bin/activate 
# Instalar dependencias core 
pip install langgraph langchain langchain-openai langchain-community 
# Tools (Tavily para web search) 
pip install tavily-python 
# Para debugging (opcional pero recomendado) 
pip install langsmith 

# OpenAI API key (requerido) 
OPENAI_API_KEY=sk-... 
# Tavily API key para web search (free tier: 1000 requests/mes) 
TAVILY_API_KEY=tvly-... 
# LangSmith para debugging (opcional) LANGCHAIN_TRACING_V2=true LANGCHAIN_API_KEY=ls__... 
LANGCHAIN_PROJECT=langgraph-tutorial 

import os 
from typing import TypedDict, Annotated, Sequence 
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage 
from langchain_openai import ChatOpenAI 
from langchain_community.tools.tavily_search import TavilySearchResults 
from langgraph.graph import StateGraph, END 
from langgraph.graph.message import add_messages 
from langgraph.prebuilt import ToolNode # 1. Definir State 

class AgentState(TypedDict): """State compartido del agente."""
     messages: Annotated[Sequence[BaseMessage], add_messages] 
# 2. Inicializar LLM con tools 
llm = ChatOpenAI(model="gpt-4", temperature=0) 
# 3. Definir tools disponibles 
tools = [ TavilySearchResults( max_results=3, description="Busca información actualizada en la web. Usa esto cuando necesites datos recientes o verificar hechos." ) ] 
# Bind tools al LLM (permite tool calling) 
llm_with_tools = llm.bind_tools(tools)
# 4. Crear ToolNode (ejecuta tools automáticamente)
tool_node = ToolNode(tool
s) 

# 5. Node que llama al LLM
def call_model(state: AgentState) -> dict:
    """
    Llama al LLM con el historial de messages.
    El LLM decide si necesita tools o puede responder directamente.
    """
    messages = state["messages"]
    response = llm_with_tools.invoke(messages)
    return {"messages": [response]}


# 6. Routing logic (decide siguiente node)
def should_continue(state: AgentState) -> str:
    """
    Decide si continuar con tools o terminar.

    Returns:
        "tools": Si el LLM hizo tool_calls
        "end": Si tiene respuesta final
    """
    messages = state["messages"]
    last_message = messages[-1]

    # Si el LLM pidió ejecutar tools, ir a tool_node
    if hasattr(last_message, "tool_calls") and last_message.tool_calls:
        return "tools"

    # Si no hay tool_calls, terminamos
    return "end"
  

# 7. Construir StateGraph 
workflow = StateGraph(AgentState) 
# Agregar nodes 
workflow.add_node("agent", call_model) workflow.add_node("tools", tool_node) 
# Set entry point 
workflow.set_entry_point("agent") 
# Agregar edges 
workflow.add_conditional_edges( "agent", should_continue, { "tools": "tools", "end": END } ) 
# Después de ejecutar tools, siempre volver a agent 
workflow.add_edge("tools", "agent")
# 8. Compilar el graph app = workflow.compile() 

from langchain_core.messages import HumanMessage

# Ejecutar agente con pregunta que requiere búsqueda web
input_messages = [
    HumanMessage(content="¿Cuáles son las últimas noticias sobre LangGraph en 2025?")
]

# Opción 1: Invoke (bloquea hasta terminar)
result = app.invoke({"messages": input_messages})
print("\n=== RESPUESTA FINAL ===")
print(result["messages"][-1].content)

# Opción 2: Stream (eventos en tiempo real)
print("\n=== STREAMING EVENTS ===")
for event in app.stream({"messages": input_messages}):
    for node_name, output in event.items():
        print(f"\n[{node_name}]")
        if "messages" in output:
            print(output["messages"][-1].content[:100] + "...")

from IPython.display import Image, display 
# Generar visualización del graph 
display(Image(app.get_graph().draw_mermaid_png()))