Update graph.py

graph.py

@@ -8,68 +8,154 @@ from agents.memory import Memory
 from agents.planner import PlannerAgent
 
 
-# Dummy LLM placeholder (replace with OpenAI/Groq later)
+# -------------------------------
+# Dummy LLM
+# -------------------------------
+
 class DummyLLM:
+
     def invoke(self, prompt):
-        class R:
+
+        class Response:
             content = "generate"
-        return R()
+
+        return Response()
 
 
 llm = DummyLLM()
 
 memory = Memory()
+
 router = RouterAgent(llm)
+
 planner = PlannerAgent()
+
 agent = ImageAgent(router, memory, planner)
 
 
-# ---------------- NODES ----------------
+# -------------------------------
+# ROUTE
+# -------------------------------
+
+def route_node(state):
 
-def route_node(state: AgentState):
     task = router.route(state["user_input"])
+
     state["task"] = task
+
     return state
 
 
-def plan_node(state: AgentState):
+# -------------------------------
+# PLAN
+# -------------------------------
+
+def plan_node(state):
+
     state["steps"] = planner.plan(state["task"])
+
     return state
 
 
-def execute_node(state: AgentState):
+# -------------------------------
+# LOAD UPLOADED IMAGE
+# -------------------------------
+
+def upload_node(state):
+
+    uploaded = state.get("uploaded_files", [])
+
+    if uploaded:
+
+        # Store uploaded image in memory
+        memory.uploaded_files = uploaded
+        memory.last_used_image = uploaded[0]
+
+    return state
+
+
+# -------------------------------
+# AUTO ANALYZE
+# -------------------------------
+
+def analyze_node(state):
+
+    image = memory.last_used_image
+
+    if image:
+
+        analysis = agent.auto_analyze(image)
+
+        state["analysis"] = analysis
+
+    return state
+
+
+# -------------------------------
+# EXECUTE
+# -------------------------------
+
+def execute_node(state):
+
     result = agent.run(state["user_input"])
+
     state["result"] = result
+
     return state
 
 
-# ---------------- GRAPH ----------------
+# -------------------------------
+# GRAPH
+# -------------------------------
 
 workflow = StateGraph(AgentState)
 
+workflow.add_node("upload", upload_node)
+
 workflow.add_node("route", route_node)
+
 workflow.add_node("plan", plan_node)
+
+workflow.add_node("analyze", analyze_node)
+
 workflow.add_node("execute", execute_node)
 
-workflow.set_entry_point("route")
+
+workflow.set_entry_point("upload")
+
+workflow.add_edge("upload", "route")
 
 workflow.add_edge("route", "plan")
-workflow.add_edge("plan", "execute")
+
+workflow.add_edge("plan", "analyze")
+
+workflow.add_edge("analyze", "execute")
+
 workflow.add_edge("execute", END)
 
 app = workflow.compile()
 
 
-# MAIN RUN FUNCTION
-def run_graph(user_input: str, uploaded_files=None):
+# -------------------------------
+# MAIN
+# -------------------------------
+
+def run_graph(user_input, uploaded_files=None):
 
     state = {
+
         "user_input": user_input,
+
         "task": "",
+
         "steps": [],
+
         "uploaded_files": uploaded_files or [],
-        "current_image": None,
+
+        "analysis": None,
+
         "result": None,
+
         "history": []
     }