graph.py

from langgraph.graph import StateGraph, END

from state import AgentState

from agents.routing_agent import RouterAgent
from agents.image_agent import ImageAgent
from agents.memory import Memory
from agents.planner import PlannerAgent


# Dummy LLM placeholder (replace with OpenAI/Groq later)
class DummyLLM:
    def invoke(self, prompt):
        class R:
            content = "generate"
        return R()


llm = DummyLLM()

memory = Memory()
router = RouterAgent(llm)
planner = PlannerAgent()
agent = ImageAgent(router, memory, planner)


# ---------------- NODES ----------------

def route_node(state: AgentState):
    task = router.route(state["user_input"])
    state["task"] = task
    return state


def plan_node(state: AgentState):
    state["steps"] = planner.plan(state["task"])
    return state


def execute_node(state: AgentState):
    result = agent.run(state["user_input"])
    state["result"] = result
    return state


# ---------------- GRAPH ----------------

workflow = StateGraph(AgentState)

workflow.add_node("route", route_node)
workflow.add_node("plan", plan_node)
workflow.add_node("execute", execute_node)

workflow.set_entry_point("route")

workflow.add_edge("route", "plan")
workflow.add_edge("plan", "execute")
workflow.add_edge("execute", END)

app = workflow.compile()


# MAIN RUN FUNCTION
def run_graph(user_input: str, uploaded_files=None):

    state = {
        "user_input": user_input,
        "task": "",
        "steps": [],
        "uploaded_files": uploaded_files or [],
        "current_image": None,
        "result": None,
        "history": []
    }

    return app.invoke(state)