app.py

import os
from dotenv import load_dotenv
from typing import TypedDict, List, Literal

from langchain_groq import ChatGroq
from langgraph.graph import StateGraph, END

# Load Environment
load_dotenv()

# LLM Configuration
llm = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.3,
    max_tokens=800,
    timeout=60,
)

# State Definition-
class SolverState(TypedDict):
    problem: str
    history: List[str]
    round: int
    max_rounds: int
    last_speaker: Literal["creative", "critic", "none"]

# System Prompts
CREATIVE_SYS = """You are the Creative Brainstormer.
Generate structured, concrete ideas.
Be concise (<= 120 words).
End with:
QUESTION_FOR_CRITIC: <one sharp question>
"""

CRITIC_SYS = """You are the Harsh Critic.
Identify flaws, risks, assumptions, and suggest improvements.
Be concise (<= 120 words).
End with:
QUESTION_FOR_CREATIVE: <one sharp question>
"""

EDITOR_SYS = """You are the Editor.
You will receive a problem and a short debate log.

Produce the final result in this structure:

1) Final Answer (concise)
2) Step-by-step execution plan
3) Key risks + mitigations
4) Open questions (max 5 bullets)

Be specific and actionable.
"""

# Helper: Safe LLM Call
def safe_llm_call(prompt: str) -> str:
    try:
        response = llm.invoke(prompt)
        return response.content.strip()
    except Exception as e:
        return f"[LLM ERROR]: {str(e)}"

# Agent Nodes
def creative_node(state: SolverState) -> SolverState:
    prompt = (
        f"{CREATIVE_SYS}\n\n"
        f"PROBLEM:\n{state['problem']}\n\n"
        f"HISTORY:\n" + "\n".join(state["history"][-6:])
    )

    resp = safe_llm_call(prompt)

    return {
        **state,
        "history": state["history"] + [f"CREATIVE:\n{resp}"],
        "round": state["round"] + 1,
        "last_speaker": "creative",
    }


def critic_node(state: SolverState) -> SolverState:
    prompt = (
        f"{CRITIC_SYS}\n\n"
        f"PROBLEM:\n{state['problem']}\n\n"
        f"HISTORY:\n" + "\n".join(state["history"][-6:])
    )

    resp = safe_llm_call(prompt)

    return {
        **state,
        "history": state["history"] + [f"CRITIC:\n{resp}"],
        "round": state["round"] + 1,
        "last_speaker": "critic",
    }


def editor_node(state: SolverState) -> SolverState:
    prompt = (
        f"{EDITOR_SYS}\n\n"
        f"PROBLEM:\n{state['problem']}\n\n"
        f"DEBATE LOG:\n" + "\n".join(state["history"])
    )

    resp = safe_llm_call(prompt)

    return {
        **state,
        "history": state["history"] + [f"EDITOR_FINAL:\n{resp}"],
    }

# Routing Logic
def route_next(state: SolverState):
    if state["round"] >= state["max_rounds"]:
        return "editor"

    if state["last_speaker"] in ("none", "critic"):
        return "creative"

    return "critic"

# Build Graph
graph = StateGraph(SolverState)

graph.add_node("creative", creative_node)
graph.add_node("critic", critic_node)
graph.add_node("editor", editor_node)

graph.set_entry_point("creative")

graph.add_conditional_edges(
    "creative",
    route_next,
    {
        "creative": "creative",
        "critic": "critic",
        "editor": "editor",
    },
)

graph.add_conditional_edges(
    "critic",
    route_next,
    {
        "creative": "creative",
        "critic": "critic",
        "editor": "editor",
    },
)

graph.add_edge("editor", END)

app = graph.compile()

# Main Runner
if __name__ == "__main__":
    problem = input("\nPaste your complex problem:\n> ").strip()

    initial_state: SolverState = {
        "problem": problem,
        "history": [],
        "round": 0,
        "max_rounds": 6,  # 3 full rounds (Creative + Critic)
        "last_speaker": "none",
    }

    result = app.invoke(initial_state)

    print("\n\n================= FULL RUN LOG =================\n")

    for msg in result["history"]:
        print(msg)
        print("\n" + "-" * 70 + "\n")

    print("\n================= FINAL ANSWER =================\n")

    final = next(
        (x for x in result["history"] if x.startswith("EDITOR_FINAL")),
        "No final output generated."
    )

    print(final)
    print("\n✅ Done.\n")