N8N-Inspired AI Automation Workflow Builder using Streamlit and LangGraph

N8N-Streamlit-Langgraph

# -*- coding: utf-8 -*-
"""
N8N-Inspired AI Automation Workflow Builder using Streamlit and LangGraph
"""
# Subscribe to the Deep Charts YouTube Channel (https://www.youtube.com/@DeepCharts)

import streamlit as st
import uuid
import os
import re
from typing import List, Dict, Any, Optional, Sequence, TypedDict as TypingTypedDict # Use typing's TypedDict
from datetime import datetime
import json

# --- LangGraph & Tool Imports ---
try:
    from langgraph.graph import StateGraph, START, END
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
    from langchain_core.tools import BaseTool

    LANGGRAPH_AVAILABLE = True
except ImportError as e:
    st.error(
        f"Required libraries not found: {e}. Please install them: \n"
        "`pip install streamlit streamlit-agraph langchain-openai langgraph typing-extensions regex langchain-core`"
    )
    LANGGRAPH_AVAILABLE = False
    # --- DUMMY DEFINITIONS ---
    class BaseMessage: pass
    class HumanMessage(BaseMessage): pass
    class AIMessage(BaseMessage): pass
    class SystemMessage(BaseMessage): pass
    class TypedDict(TypingTypedDict): pass
    def StateGraph(state): return None
    class BaseTool:
        name: str = "dummy_tool"; description: str = "A dummy tool"
        def _run(self, *args: Any, **kwargs: Any) -> Any: pass
        async def _arun(self, *args: Any, **kwargs: Any) -> Any: pass
    START, END = "START", "END"
    # --- END DUMMY DEFINITIONS ---
    st.stop()

# --- Streamlit Agraph Imports ---
try:
    from streamlit_agraph import agraph, Node, Edge, Config
    AGRAPH_AVAILABLE = True
except ImportError:
    st.warning("`streamlit-agraph` not found. Install: `pip install streamlit-agraph`")
    AGRAPH_AVAILABLE = False
    def agraph(*args, **kwargs): pass
    class Node:
        def __init__(self, id, label, **kwargs): self.id = id; self.label = label
    class Edge:
        def __init__(self, source, target, **kwargs): self.source = source; self.target = target
    class Config:
        def __init__(self, **kwargs): pass

# --- Constants ---
ROUTING_KEY_MARKER = "ROUTING_KEY:"
DEFAULT_ROUTING_KEY = "__DEFAULT__"
START_NODE_ID = "__START__"
END_NODE_ID = END

# --- Tool Definition (Using web_search_preview dictionary) ---
web_search_tool_dict = {"type": "web_search_preview"}
tools_list_for_binding = [web_search_tool_dict]

# --- LangGraph State Definition (Dictionary Based) ---
class WorkflowState(TypingTypedDict):
    input: str
    node_outputs: Dict[str, str]
    last_response_content: str
    current_node_id: str

# --- LLM Initialization ---
llm: Optional[ChatOpenAI] = None
llm_with_search: Optional[Any] = None

def initialize_llm() -> bool:
    global llm, llm_with_search
    openai_api_key = os.environ.get("OPENAI_API_KEY")
    openai_ready = bool(openai_api_key)
    if openai_ready and (llm is None or getattr(llm, 'openai_api_key', None) != openai_api_key):
        try:
            base_llm = ChatOpenAI(model_name="gpt-4o", openai_api_key=openai_api_key, temperature=0.2)
            llm = base_llm; llm_with_search = base_llm.bind_tools(tools_list_for_binding)
            print(f"LLM initialized (gpt-4o). Bound: {web_search_tool_dict}")
            return True
        except Exception as e: st.error(f"LLM Init Error: {e}", icon="🔥"); llm = None; llm_with_search = None; return False
    elif not openai_ready:
        if llm is not None: print("Clearing LLM."); llm = None; llm_with_search = None
        return False
    return llm is not None

# --- Helper Function ---
def get_node_display_name(node_id: str) -> str:
    if node_id == START_NODE_ID: return "⏹️ START"
    if node_id == END_NODE_ID: return "🏁 END"
    if "nodes" in st.session_state and isinstance(st.session_state.nodes, list):
        for i, node in enumerate(st.session_state.nodes):
            if isinstance(node, dict) and node.get("id") == node_id: return f"{i+1}. {node.get('name', f'Unk ({node_id})')}"
    return f"Unknown ({node_id})"


# --- Node Execution Function Factory (FIXED CONTENT EXTRACTION) ---
def create_agent_node_function(node_id: str, node_name: str, node_prompt: str, possible_keys: List[str]):
    """ Node function (dict state) with corrected content extraction. """
    def agent_node_function(state: WorkflowState) -> WorkflowState:
        print(f"\n--- Executing Node: {node_name} ({node_id}) ---")
        if "execution_log" not in st.session_state: st.session_state.execution_log = []
        st.session_state.execution_log.append(f"⚙️ Executing Node: **{node_name}**")

        if not llm_with_search: # Error handling for LLM init
            error_msg = f"ERROR: LLM not initialized."; st.session_state.execution_log.append(f"  -> ❌ Error: {error_msg}")
            updated_state = state.copy(); updated_state["last_response_content"] = f"{error_msg} {ROUTING_KEY_MARKER} error"; updated_state["current_node_id"] = node_id; return updated_state

        # --- Prepare Context ---
        context_input = ""; is_first_node = not state.get("last_response_content")
        if is_first_node:
            context_input = state.get("input", ""); print(f"DEBUG: First node, input: '{context_input[:100]}...'")
        else:
            prev_content = state.get("last_response_content", ""); print(f"DEBUG: Prev content type: {type(prev_content)}")
            if isinstance(prev_content, str): context_input = re.sub(rf"\s*{ROUTING_KEY_MARKER}\s*\w+\s*$", "", prev_content).strip()
            else: print("ERROR: Prev content not str"); context_input = str(prev_content) # Fallback
            print(f"DEBUG: Context input: '{context_input[:100]}...'")

        # --- Prepare Prompt ---
        prompt_with_context = node_prompt
        if '{input_text}' in node_prompt: prompt_with_context = node_prompt.replace('{input_text}', context_input)
        elif context_input: prompt_with_context += f"\n\nInput Context:\n{context_input}"
        current_task_prompt = f"Current Task ({node_name}):\n{prompt_with_context}\n(Search web if needed)."
        key_options_text = ", ".join(f"'{k}'" for k in possible_keys if k and k != DEFAULT_ROUTING_KEY)
        routing_instruction = f"\n\n--- ROUTING ---\nAfter response, MUST end with '{ROUTING_KEY_MARKER} <key>' (e.g., from [{key_options_text}]).\n--- END ROUTING ---"
        full_prompt = current_task_prompt + routing_instruction
        print(f"DEBUG: Full prompt type: {type(full_prompt)}"); print(f"Node '{node_name}' sending prompt: {full_prompt[:300]}..."); st.session_state.execution_log.append(f"  📝 Prompt Snippet: {prompt_with_context[:100]}...")

        try:
            # --- Invoke LLM ---
            result = llm_with_search.invoke(full_prompt)
            response_content = "" # Default empty

            # --- *** CORRECTED CONTENT EXTRACTION *** ---
            if hasattr(result, 'content'):
                raw_content = result.content
                if isinstance(raw_content, str):
                    response_content = raw_content # It was already a string
                elif isinstance(raw_content, list) and len(raw_content) > 0 and isinstance(raw_content[0], dict) and 'text' in raw_content[0]:
                    # Standard case for gpt-4o list output: extract text from first block
                    response_content = raw_content[0].get('text', '') # Use .get for safety
                    print("DEBUG: Extracted text from result.content[0]['text']")
                else: # Handle unexpected formats
                    print(f"WARNING: Unexpected format for result.content: {type(raw_content)}. Trying str().")
                    response_content = str(raw_content) # Fallback
            else: print("WARNING: result has no 'content' attribute.")
            # --- *** END CORRECTION *** ---

            # --- Debugging & Logging ---
            print(f"\nDEBUG: Raw result type: {type(result)}")
            # print(f"DEBUG: Raw result value: {result}") # Can be very verbose
            print(f"DEBUG: response_content type: {type(response_content)}") # Should be str now
            print(f"DEBUG: response_content value: '{str(response_content)[:500]}...'")
            log_snippet = str(response_content)[:100] if response_content is not None else "[No text content]"
            st.session_state.execution_log.append(f"  🤖 LLM Response Snippet: {log_snippet}...")

            # --- Tool Call Warning (if they still appear) ---
            if hasattr(result, 'tool_calls') and result.tool_calls:
                 calls = result.tool_calls; call_details = [f"{call.get('name', '?')}({call.get('args', {})})" for call in calls]
                 warning_msg = f"  ⚠️ WARNING: LLM response included 'tool_calls' ({call_details}). Graph not handling them."; print(warning_msg); st.session_state.execution_log.append(warning_msg); st.warning(warning_msg)

            # --- Routing Key Check (Should work on extracted string now) ---
            if isinstance(response_content, str):
                match = re.search(rf"{ROUTING_KEY_MARKER}\s*(\w+)\s*$", response_content)
                if match: st.session_state.execution_log.append(f"  🔑 Detected key: '{match.group(1)}'.")
                else: st.session_state.execution_log.append(f"  ⚠️ WARNING: No routing key found."); response_content += f" {ROUTING_KEY_MARKER} {DEFAULT_ROUTING_KEY}"; st.session_state.execution_log.append(f"  🔧 Appended default key.")
            else: # Should not happen now, but keep as safeguard
                 st.session_state.execution_log.append(f"  ⚠️ ERROR: Extracted content not str ({type(response_content)}).")
                 response_content = f"Error: Invalid type {ROUTING_KEY_MARKER} {DEFAULT_ROUTING_KEY}"

            # --- Update State ---
            updated_state = state.copy()
            if "node_outputs" not in updated_state or not isinstance(updated_state["node_outputs"], dict): updated_state["node_outputs"] = {}
            updated_state["node_outputs"][node_id] = str(response_content) # Store string
            updated_state["last_response_content"] = str(response_content) # Store string
            updated_state["current_node_id"] = node_id
            return updated_state

        # --- Exception Handling ---
        except Exception as e: # Catch any other errors during invoke/processing
            error_msg = f"Error in node {node_name} ({node_id}): {e}"
            print(error_msg); st.session_state.execution_log.append(f"  -> ❌ Error: {e}")
            updated_state = state.copy(); updated_state["last_response_content"] = f"ERROR: {error_msg} {ROUTING_KEY_MARKER} error"; updated_state["current_node_id"] = node_id
            import traceback; traceback.print_exc(); return updated_state

    return agent_node_function


# --- Generic Router Function (Dictionary State Version - CORRECTED FORMATTING) ---
def generic_router(state: WorkflowState) -> str:
    """ Determines route based on key in state['last_response_content']. """
    print("\n--- Routing Check ---")
    routing_key = DEFAULT_ROUTING_KEY # Default if no key found
    last_content = state.get("last_response_content", "")

    if last_content:
        # Ensure last_content is treated as a string
        if isinstance(last_content, str):
             match = re.search(rf"{ROUTING_KEY_MARKER}\s*(\w+)\s*$", last_content)
             if match:
                 # Key found, use it
                 routing_key = match.group(1).strip()
                 print(f"  Extracted key: '{routing_key}'")
             else:
                 # String content, but no key found at the end
                 print(f"  No routing key found in last response: '...{last_content[-50:]}'")
                 print(f"  -> Using default routing ('{DEFAULT_ROUTING_KEY}').")
        else:
             # Content is not a string (shouldn't happen often with current node logic, but good practice)
             print(f"  Last response content type is {type(last_content)}, not string.")
             print(f"  -> Using default routing ('{DEFAULT_ROUTING_KEY}').")
    else:
        # No previous response content found in state
        print(f"  No previous response content found.")
        print(f"  -> Using default routing ('{DEFAULT_ROUTING_KEY}').")

    # Log and return the determined routing key
    log_decision_msg = f"🚦 Routing decision: '{routing_key}'"
    print(f"  -> {log_decision_msg}")
    if "execution_log" not in st.session_state: st.session_state.execution_log = []
    st.session_state.execution_log.append(log_decision_msg)
    return routing_key

# --- Graph Compilation Function (Dictionary State Version - No Change) ---
def compile_graph() -> bool:
    if not llm: st.error("LLM not initialized.", icon="🔥"); return False
    if not st.session_state.nodes: st.warning("No nodes defined.", icon="⚠️"); return False
    print("\n--- Compiling Graph (Dictionary State) ---")
    try:
        graph_builder = StateGraph(WorkflowState); valid_nodes = [n for n in st.session_state.nodes if isinstance(n, dict) and all(k in n for k in ["id", "name", "prompt"])];
        if not valid_nodes: st.warning("No valid nodes.", icon="⚠️"); return False
        node_ids = {node['id'] for node in valid_nodes}; start_node_id_actual = valid_nodes[0]["id"]; print("  Adding Nodes:")
        possible_keys_per_node = {}
        for node_data in valid_nodes:
            node_id, node_name, node_prompt = node_data["id"], node_data["name"], node_data["prompt"]; print(f"    - ID: {node_id}, Name: '{node_name}'")
            routing_rules = node_data.get("routing_rules", {}); cond_keys = {rule.get("output_key", "").strip() for rule in routing_rules.get("conditional_targets", []) if rule.get("output_key")}
            all_keys = cond_keys.union({DEFAULT_ROUTING_KEY, "error"}); possible_keys_per_node[node_id] = list(all_keys); agent_func = create_agent_node_function(node_id, node_name, node_prompt, possible_keys_per_node[node_id])
            graph_builder.add_node(node_id, agent_func)
        print("  Adding Edges:"); graph_builder.add_edge(START, start_node_id_actual); print(f"    - START -> {get_node_display_name(start_node_id_actual)}")
        all_targets_valid = True
        for node_data in valid_nodes:
            node_id, node_name = node_data["id"], node_data["name"]; routing_rules = node_data.get("routing_rules", {}); default_target = routing_rules.get("default_target", END_NODE_ID); conditional_targets = routing_rules.get("conditional_targets", [])
            path_map = {}; print(f"    - Edges from '{node_name}' ({node_id}):"); seen_keys_for_node = set(); node_targets_valid = True
            for rule_idx, rule in enumerate(conditional_targets):
                key, target_id = rule.get("output_key", "").strip(), rule.get("target_node_id")
                if key and target_id:
                    if target_id != END_NODE_ID and target_id not in node_ids: st.error(f"❌ Invalid Target: {node_name} '{key}'->'{get_node_display_name(target_id)}'", icon="🔥"); all_targets_valid=False; node_targets_valid=False; continue
                    if key in seen_keys_for_node: st.warning(f"⚠️ Duplicate key '{key}' in '{node_name}'.", icon="⚠️")
                    path_map[key] = target_id; seen_keys_for_node.add(key); print(f"      - If key '{key}' -> {get_node_display_name(target_id)}")
                elif key or target_id: st.warning(f"Node '{node_name}' incomplete rule #{rule_idx+1}. Ignored.", icon="⚠️")
            if DEFAULT_ROUTING_KEY not in path_map:
                 if default_target != END_NODE_ID and default_target not in node_ids: st.error(f"❌ Invalid Default Target: {node_name}->'{get_node_display_name(default_target)}'", icon="🔥"); all_targets_valid=False; node_targets_valid=False
                 else: path_map[DEFAULT_ROUTING_KEY] = default_target; print(f"      - If key '{DEFAULT_ROUTING_KEY}' -> {get_node_display_name(default_target)}")
            if "error" not in path_map: path_map["error"] = END_NODE_ID; print(f"      - If key 'error' -> {get_node_display_name(END_NODE_ID)} (Implicit)")
            if node_targets_valid: graph_builder.add_conditional_edges(node_id, generic_router, path_map)
            else: print(f"      -> Skipping edges for '{node_name}'.")
        if not all_targets_valid: st.error("Compilation failed.", icon="🔥"); return False
        recursion_limit = len(valid_nodes) * 3 + 10; print(f"  Setting recursion limit to: {recursion_limit}")
        st.session_state.compiled_graph = graph_builder.compile(checkpointer=None); st.session_state.recursion_limit = recursion_limit
        print("✅ Graph compiled successfully!"); st.toast("Workflow compiled!", icon="✅"); return True
    except Exception as e: st.error(f"Compile error: {e}", icon="🔥"); print(f"❌ Compile Error: {e}"); import traceback; traceback.print_exc(); st.session_state.compiled_graph = None; return False

# --- Graph Visualization Data Generation (No Change) ---
def generate_agraph_data(nodes_data: List[Dict[str, Any]]) -> tuple[List[Node], List[Edge]]:
    agraph_nodes: List[Node] = []; agraph_edges: List[Edge] = []
    agraph_nodes.append(Node(id=START_NODE_ID, label="START", shape="ellipse", color="#4CAF50", title="Workflow Entry Point"))
    valid_nodes_vis = [node for node in nodes_data if isinstance(node, dict) and 'id' in node]
    node_ids_vis = {node['id'] for node in valid_nodes_vis}; node_indices = {node['id']: i for i, node in enumerate(valid_nodes_vis)}
    for i, node in enumerate(valid_nodes_vis):
        node_id = node['id']; node_name = node.get('name', 'Unnamed'); node_prompt_snippet = node.get('prompt', '')[:100] + "..."
        is_selected = st.session_state.get('selected_node_id') == node_id; border_width = 3 if is_selected else 1; node_color = "#FFC107" if is_selected else "#90CAF9"
        agraph_nodes.append(Node(id=node_id, label=f"{i+1}. {node_name}", shape="box", color=node_color, borderWidth=border_width, title=f"ID: {node_id}\nPrompt: {node_prompt_snippet}"))
        if i == 0: agraph_edges.append(Edge(source=START_NODE_ID, target=node_id, label="Start Flow", color="#4CAF50", width=2))
        routing_rules = node.get("routing_rules", {}); default_target = routing_rules.get("default_target", END_NODE_ID); conditional_targets = routing_rules.get("conditional_targets", [])
        added_vis_edges = set()
        if default_target == END_NODE_ID or default_target in node_ids_vis:
            is_overridden = any(r.get("output_key") == DEFAULT_ROUTING_KEY for r in conditional_targets); edge_id = (node_id, default_target, DEFAULT_ROUTING_KEY)
            if not is_overridden and edge_id not in added_vis_edges: agraph_edges.append(Edge(source=node_id, target=default_target, label=DEFAULT_ROUTING_KEY, color="#9E9E9E", dashes=True, arrows="to", font={'align': 'middle'})); added_vis_edges.add(edge_id)
        for rule in conditional_targets:
            key, target_id = rule.get("output_key", "").strip(), rule.get("target_node_id")
            if key and target_id and (target_id == END_NODE_ID or target_id in node_ids_vis):
                edge_id = (node_id, target_id, key)
                if edge_id not in added_vis_edges:
                    is_loopback = False
                    if target_id in node_indices and node_id in node_indices:
                         if node_indices[target_id] < node_indices[node_id]: is_loopback = True
                    edge_color = "#FF5722" if is_loopback else "#2196F3"; edge_label = f"LOOP: {key}" if is_loopback else key
                    agraph_edges.append(Edge(source=node_id, target=target_id, label=edge_label, color=edge_color, arrows="to", font={'align': 'middle'})); added_vis_edges.add(edge_id)
    if any(hasattr(edge, 'target') and edge.target == END_NODE_ID for edge in agraph_edges):
        agraph_nodes.append(Node(id=END_NODE_ID, label="END", shape="ellipse", color="#F44336", title="Workflow End"))
    return agraph_nodes, agraph_edges

# --- Example Workflow Definitions (No Change) ---
def get_simple_summarizer_workflow(): node_id = f"summarize_node_{uuid.uuid4().hex[:4]}"; return [ { "id": node_id, "name": "Summarize Input", "prompt": f"Summarize the input text concisely (1-2 sentences).", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]} } ]
def get_sentiment_workflow(): ids = {name: f"sentiment_node_{uuid.uuid4().hex[:4]}" for name in ["analyzer", "positive", "negative", "neutral"]}; nodes = [ { "id": ids["analyzer"], "name": "Analyze Sentiment", "prompt": f"Analyze sentiment (Positive, Negative, Neutral) of the input. Respond ONLY with the result.", "type": "llm_call", "routing_rules": { "default_target": ids["neutral"], "conditional_targets": [ {"output_key": "Positive", "target_node_id": ids["positive"]}, {"output_key": "Negative", "target_node_id": ids["negative"]}, {"output_key": "Neutral", "target_node_id": ids["neutral"]} ]}}, { "id": ids["positive"], "name": "Handle Positive", "prompt": f"Sentiment was positive. Respond cheerfully.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}}, { "id": ids["negative"], "name": "Handle Negative", "prompt": f"Sentiment was negative. Respond empathetically.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}}, { "id": ids["neutral"], "name": "Handle Neutral", "prompt": f"Sentiment was neutral. Acknowledge receipt.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}} ]; return nodes
def get_classification_workflow(): ids = {name: f"classify_node_{uuid.uuid4().hex[:4]}" for name in ["classify", "complaint", "query", "compliment", "general"]}; nodes = [ { "id": ids["classify"], "name": "Extract & Classify Intent", "prompt": f"From input, extract product/person (or 'None'). Classify intent: Complaint, Query, Compliment. Respond STRICTLY:\nExtracted Info: [Info]\nIntent: [IntentWord]", "type": "llm_call", "routing_rules": { "default_target": ids["general"], "conditional_targets": [ {"output_key": "Complaint", "target_node_id": ids["complaint"]}, {"output_key": "Query", "target_node_id": ids["query"]}, {"output_key": "Compliment", "target_node_id": ids["compliment"]} ]}}, { "id": ids["complaint"], "name": "Handle Complaint", "prompt": f"Complaint received. Respond empathetically, use extracted info.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}}, { "id": ids["query"], "name": "Answer Query", "prompt": f"Query received. Answer based on context/info. Use web search if needed.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}}, { "id": ids["compliment"], "name": "Handle Compliment", "prompt": f"Compliment received. Respond thankfully.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}}, { "id": ids["general"], "name": "General Response", "prompt": f"Intent unclear. Provide generic response.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"done", "target_node_id": END_NODE_ID}]}} ]; return nodes
def get_deep_research_workflow(): ids = {name: f"dr_node_{uuid.uuid4().hex[:4]}" for name in [ "planner", "search_A", "search_B", "cross_reference", "synthesize", "final_report" ]}; nodes = [ { "id": ids["planner"], "name": "📝 Research Planner", "prompt": f"Analyze the research goal. Break into 1-3 angles (A, B). State plan. Decide if single/multi angle needed.", "type": "llm_call", "routing_rules": { "default_target": ids["search_A"], "conditional_targets": [ {"output_key": "single_angle", "target_node_id": ids["search_A"]}, {"output_key": "multi_angle", "target_node_id": ids["search_A"]} ] } }, { "id": ids["search_A"], "name": "🔍 Research Angle A", "prompt": f"Research Angle A based on plan. Use web search if necessary. Summarize findings for Angle A.", "type": "llm_call", "routing_rules": { "default_target": ids["search_B"], "conditional_targets": [ {"output_key": "synthesize_direct", "target_node_id": ids["synthesize"]}, {"output_key": "next_angle", "target_node_id": ids["search_B"]} ] } }, { "id": ids["search_B"], "name": "🔎 Research Angle B", "prompt": f"Research Angle B based on plan. Use web search if necessary. Summarize findings for Angle B.", "type": "llm_call", "routing_rules": { "default_target": ids["cross_reference"], "conditional_targets": [ {"output_key": "cross_reference", "target_node_id": ids["cross_reference"]} ] } }, { "id": ids["cross_reference"], "name": "🔄 Cross-Reference & Validate", "prompt": f"Review findings from Angle A & B. Identify agreements/contradictions/gaps. Use web search to verify if needed.", "type": "llm_call", "routing_rules": { "default_target": ids["synthesize"], "conditional_targets": [ {"output_key": "revisit_A", "target_node_id": ids["search_A"]}, {"output_key": "revisit_B", "target_node_id": ids["search_B"]}, {"output_key": "synthesize", "target_node_id": ids["synthesize"]} ] } }, { "id": ids["synthesize"], "name": "🧩 Synthesize Findings", "prompt": f"Combine validated findings from research angles. Create concise summary for original goal.", "type": "llm_call", "routing_rules": { "default_target": ids["final_report"], "conditional_targets": [ {"output_key": "refine_synthesis", "target_node_id": ids["cross_reference"]}, {"output_key": "final_report", "target_node_id": ids["final_report"]} ] } }, { "id": ids["final_report"], "name": "📄 Generate Final Report", "prompt": f"Format synthesized findings into clear final report.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key": "done", "target_node_id": END_NODE_ID}]} } ]; return nodes
def get_enhanced_hedge_fund_workflow(): ids = {name: f"ehf_node_{uuid.uuid4().hex[:4]}" for name in [ "goal_risk", "planner", "macro", "sector", "company", "risk_assess", "strategist" ]}; nodes = [ { "id": ids["goal_risk"], "name": "🎯 Goal & Risk Profiler", "prompt": f"Analyze investment goal from input. Clarify timeframe, expectations, infer risk tolerance (Conservative, Balanced, Aggressive). State profile.", "type": "llm_call", "routing_rules": { "default_target": ids["planner"], "conditional_targets": [ {"output_key": "plan_research", "target_node_id": ids["planner"]} ] } }, { "id": ids["planner"], "name": "🗺️ Advanced Research Planner", "prompt": f"Based on Goal/Risk, determine necessary research steps & sequence (Macro, Sector, Company). State plan. Decide FIRST step.", "type": "llm_call", "routing_rules": { "default_target": ids["risk_assess"], "conditional_targets": [ {"output_key": "research_macro", "target_node_id": ids["macro"]}, {"output_key": "research_sector", "target_node_id": ids["sector"]}, {"output_key": "research_company", "target_node_id": ids["company"]}, {"output_key": "go_strategy", "target_node_id": ids["strategist"]} ] } }, { "id": ids["macro"], "name": "📈 Macro Researcher", "prompt": f"Perform macro analysis relevant to Goal/Risk. Use web search if needed (GDP, rates, inflation). Summarize findings & impacts. Determine NEXT planned step.", "type": "llm_call", "routing_rules": { "default_target": ids["risk_assess"], "conditional_targets": [ {"output_key": "research_sector", "target_node_id": ids["sector"]}, {"output_key": "research_company", "target_node_id": ids["company"]}, {"output_key": "assess_risk", "target_node_id": ids["risk_assess"]} ] } }, { "id": ids["sector"], "name": "🏭 Sector Researcher", "prompt": f"Perform sector analysis relevant to Goal/Risk (guided by macro). Use web search if needed (trends/competitors). Summarize findings & impacts. Determine NEXT planned step.", "type": "llm_call", "routing_rules": { "default_target": ids["risk_assess"], "conditional_targets": [ {"output_key": "research_macro", "target_node_id": ids["macro"]}, {"output_key": "research_company", "target_node_id": ids["company"]}, {"output_key": "assess_risk", "target_node_id": ids["risk_assess"]} ] } }, { "id": ids["company"], "name": "🏢 Company Researcher", "prompt": f"Perform company analysis relevant to Goal/Risk (guided by macro/sector). Use web search if needed (news, financials). Summarize findings (valuation, risks) & impacts. Determine NEXT planned step.", "type": "llm_call", "routing_rules": { "default_target": ids["risk_assess"], "conditional_targets": [ {"output_key": "research_macro", "target_node_id": ids["macro"]}, {"output_key": "research_sector", "target_node_id": ids["sector"]}, {"output_key": "assess_risk", "target_node_id": ids["risk_assess"]} ] } }, { "id": ids["risk_assess"], "name": "⚠️ Risk Assessor & Validator", "prompt": f"Review ALL research findings against Goal/Risk. Identify key risks, inconsistencies, gaps. Use web search to verify if needed.", "type": "llm_call", "routing_rules": { "default_target": ids["strategist"], "conditional_targets": [ {"output_key": "revisit_macro", "target_node_id": ids["macro"]}, {"output_key": "revisit_sector", "target_node_id": ids["sector"]}, {"output_key": "revisit_company", "target_node_id": ids["company"]}, {"output_key": "create_strategy", "target_node_id": ids["strategist"]} ] } }, { "id": ids["strategist"], "name": "💰 Portfolio Strategist", "prompt": f"Synthesize validated research. Develop specific, diversified portfolio allocation strategy tailored to Goal/Risk. Justify strategy.", "type": "llm_call", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": [{"output_key":"portfolio_ready", "target_node_id": END_NODE_ID}]} } ]; return nodes

# --- Helper Function to Load Workflow (No Change) ---
def load_workflow(workflow_func):
     try: st.session_state.nodes = workflow_func(); st.session_state.compiled_graph = None; st.session_state.selected_node_id = st.session_state.nodes[0]['id'] if st.session_state.nodes else None; st.session_state.execution_log = []; st.session_state.final_state = None; st.session_state.recursion_limit = None; wf_name = workflow_func.__name__.replace('get_','').replace('_workflow','').replace('_', ' ').title(); st.toast(f"{wf_name} loaded!", icon="📄"); st.rerun()
     except Exception as e: st.error(f"Load Error: {e}"); print(f"Load Error: {e}"); import traceback; traceback.print_exc()

# --- Session State Initialization (WITH SYNTAX FIX) ---
default_values = { "compiled_graph": None, "execution_log": [], "final_state": None, "selected_node_id": None, "recursion_limit": None, "openai_api_key_provided": False }
for key, value in default_values.items(): st.session_state.setdefault(key, value)
if "nodes" not in st.session_state or not st.session_state.nodes:
    try:
        print("Loading default workflow"); st.session_state.nodes = get_enhanced_hedge_fund_workflow()
        if st.session_state.nodes and isinstance(st.session_state.nodes, list): st.session_state.selected_node_id = st.session_state.nodes[0].get('id')
    except Exception as e: # Corrected Block
        print(f"Default load error: {e}"); st.session_state.nodes = []
    # --- END CORRECTION ---

# --- Helper UI Functions (Readable Versions with Syntax Fix) ---
def get_node_options_for_select(include_end=True, exclude_node_id: Optional[str] = None) -> List[tuple[str, str]]:
    options = []; node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
        options.extend([(node.get("id"), f"{i+1}. {node.get('name', 'Unnamed')}") for i, node in enumerate(node_list) if isinstance(node, dict) and node.get("id") != exclude_node_id])
    if include_end and END_NODE_ID != exclude_node_id: options.append((END_NODE_ID, get_node_display_name(END_NODE_ID)))
    return options
def select_node(node_id: Optional[str]): st.session_state.selected_node_id = node_id
def update_node_data_ui(node_id: str, data: Dict[str, Any]):
    updated = False; node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
        for i, node in enumerate(node_list):
            if isinstance(node, dict) and node.get("id") == node_id:
                try: new_rules = json.loads(json.dumps(data.get("routing_rules", {})))
                except Exception: st.error("Serialize Error."); return
                node["name"] = data.get("name", node.get("name")); node["prompt"] = data.get("prompt", node.get("prompt")); node["routing_rules"] = new_rules; updated = True; break
    if updated: st.session_state.compiled_graph = None; st.toast(f"Node '{data.get('name', node_id)}' updated.", icon="💾"); st.session_state.selected_node_id = node_id
    else: st.error(f"Node {node_id} not found.")
def delete_node_ui(node_id_to_delete: str):
    node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
        original_len = len(node_list); node_name_deleted = get_node_display_name(node_id_to_delete)
        st.session_state.nodes = [n for n in node_list if not (isinstance(n, dict) and n.get("id") == node_id_to_delete)]
        if len(st.session_state.nodes) < original_len:
            st.session_state.compiled_graph = None
            if st.session_state.selected_node_id == node_id_to_delete: st.session_state.selected_node_id = st.session_state.nodes[0]['id'] if st.session_state.nodes else None
            st.toast(f"Deleted {node_name_deleted}", icon="🗑️"); st.rerun()
        else: st.warning(f"Node {node_id_to_delete} not found.")
def move_node_ui(node_id_to_move: str, direction: int): # Readable Version with Syntax Fix
    node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
        try:
            index = next(i for i, n in enumerate(node_list) if isinstance(n, dict) and n.get("id") == node_id_to_move)
            new_index = index + direction
            if 0 <= new_index < len(node_list):
                node = node_list.pop(index)
                node_list.insert(new_index, node)
                st.session_state.compiled_graph = None
                st.session_state.selected_node_id = node_id_to_move
                st.rerun()
        except StopIteration: # Correctly placed except block
            st.warning(f"Node {node_id_to_move} not found for moving.")
def add_conditional_rule_ui(node_id: str):
    node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
        for i, node in enumerate(node_list):
            if isinstance(node, dict) and node.get("id") == node_id:
                if not isinstance(node.get("routing_rules"), dict): node["routing_rules"] = {}
                if not isinstance(node["routing_rules"].get("conditional_targets"), list): node["routing_rules"]["conditional_targets"] = []
                node["routing_rules"]["conditional_targets"].append({"output_key": "", "target_node_id": END_NODE_ID}); st.session_state.compiled_graph = None; st.session_state.selected_node_id = node_id; st.rerun(); return
    st.error(f"Node {node_id} not found.")
def delete_conditional_rule_ui(node_id: str, rule_index: int):
    node_list = st.session_state.get("nodes", [])
    if isinstance(node_list, list):
         for i, node in enumerate(node_list):
             if isinstance(node, dict) and node.get("id") == node_id:
                 rules_list = node.get("routing_rules", {}).get("conditional_targets")
                 if isinstance(rules_list, list) and 0 <= rule_index < len(rules_list):
                     del node["routing_rules"]["conditional_targets"][rule_index]; st.session_state.compiled_graph = None; st.session_state.selected_node_id = node_id; st.rerun(); return
                 else: st.warning(f"Invalid rule index {rule_index}."); return
         st.error(f"Node {node_id} not found.")

# --- Streamlit App UI (No Change) ---
st.set_page_config(layout="wide", page_title="Visual AI Automation Builder ", page_icon="🤖")
st.title("🤖🧠 Visual AI Automation Builder")
now = datetime.now(); current_timestamp = now.strftime("%Y-%m-%d %H:%M:%S")
st.caption(f"Uses Dictionary State. LLM: `gpt-4o` (bound: `{web_search_tool_dict}`). Refresh: {current_timestamp}")
with st.sidebar: # Sidebar UI (No Change)
    st.header("🔑 OpenAI API Key"); openai_api_key_input = st.text_input("API Key", type="password", value=os.environ.get("OPENAI_API_KEY", ""), help="Needed for LLM nodes.", key="api_key_input_sidebar"); api_key_in_env = os.environ.get("OPENAI_API_KEY"); api_key_changed = (openai_api_key_input != api_key_in_env if api_key_in_env else bool(openai_api_key_input))
    if api_key_changed:
        if openai_api_key_input: os.environ["OPENAI_API_KEY"] = openai_api_key_input
        elif "OPENAI_API_KEY" in os.environ: del os.environ["OPENAI_API_KEY"]
        llm = None; llm_with_search = None; llm_ready = initialize_llm(); st.session_state.openai_api_key_provided = bool(openai_api_key_input)
        if api_key_changed: st.session_state.compiled_graph = None; st.rerun()
    else: llm_ready = initialize_llm(); st.session_state.openai_api_key_provided = bool(os.environ.get("OPENAI_API_KEY"))
    if llm_ready: st.success("LLM Initialized.", icon="✅")
    elif st.session_state.get("openai_api_key_provided"): st.error("LLM Init Failed.", icon="🔥")
    else: st.warning("LLM needs API Key.", icon="⚠️")
    st.caption(f"Tool binding: `{web_search_tool_dict}`"); st.divider()
    st.header("🧩 Node Palette");
    with st.form("add_node_form"):
        new_node_name_input = st.text_input("New Node Name", placeholder="e.g., 'Summarize Input'", key="new_node_name_palette_main"); add_node_submitted = st.form_submit_button("➕ Add LLM Node", use_container_width=True)
        if add_node_submitted:
            if new_node_name_input:
                if any(isinstance(n, dict) and n.get('name') == new_node_name_input for n in st.session_state.nodes): st.warning(f"Name exists.", icon="⚠️")
                else: node_id = f"node_{uuid.uuid4().hex[:6]}"; st.session_state.nodes.append({"id": node_id, "name": new_node_name_input, "type": "llm_call", "prompt": f"Task: {new_node_name_input}\nInput: {{input_text}}\n(Use web search if needed.)", "routing_rules": {"default_target": END_NODE_ID, "conditional_targets": []}}); st.session_state.compiled_graph = None; st.session_state.selected_node_id = node_id; st.toast(f"Added: {new_node_name_input}", icon="➕"); st.rerun()
            else: st.error("Node Name required.", icon="❗")
    st.divider(); st.header("📜 Example Workflows"); ex_cols = st.columns(2)
    with ex_cols[0]: st.button("📄 Summarizer", use_container_width=True, key="load_summarizer_btn", on_click=load_workflow, args=(get_simple_summarizer_workflow,)); st.button("🎭 Sentiment", use_container_width=True, key="load_sentiment_btn", on_click=load_workflow, args=(get_sentiment_workflow,)); st.button("🔬 Deep Research", use_container_width=True, key="load_deep_research_btn", on_click=load_workflow, args=(get_deep_research_workflow,))
    with ex_cols[1]: st.button("🏷️ Classify", use_container_width=True, key="load_classify_btn", on_click=load_workflow, args=(get_classification_workflow,)); st.button("📈 Adv. Hedge Fund", use_container_width=True, key="load_adv_hedge_btn", on_click=load_workflow, args=(get_enhanced_hedge_fund_workflow,))
    st.divider(); st.header("⚙️ Workflow Control"); compile_disabled = not llm_ready or not st.session_state.nodes; tooltip_compile = "Requires API Key & nodes." if compile_disabled else "Compile workflow.";
    if st.button("🔄 Compile Workflow", type="primary", use_container_width=True, disabled=compile_disabled, help=tooltip_compile, key="compile_workflow_btn"):
        if compile_graph(): st.rerun()
    tooltip_reset = "Clear all nodes and reset."; st.button("🗑️ Reset Workflow", use_container_width=True, help=tooltip_reset, key="reset_workflow_btn", on_click=lambda: setattr(st.session_state, 'nodes', []) or st.rerun())
top_cols = st.columns([0.6, 0.4])
with top_cols[0]: # Graph Vis (No Change)
    st.subheader("📊 Workflow Graph")
    if not st.session_state.nodes: st.info("Add nodes or load an example.")
    elif not AGRAPH_AVAILABLE: st.warning("Install `streamlit-agraph` for visualization.")
    else:
        try:
            agraph_nodes, agraph_edges = generate_agraph_data(st.session_state.nodes); agraph_config = Config(width='100%', height=500, directed=True, physics={'enabled': True, 'solver': 'forceAtlas2Based', 'forceAtlas2Based': {'gravitationalConstant': -60, 'centralGravity': 0.01, 'springLength': 120, 'springConstant': 0.1, 'damping': 0.3}}, interaction={'navigationButtons': True, 'tooltipDelay': 300, 'hover': True}, nodes={'font': {'size': 14}}, edges={'font': {'size': 12, 'align': 'middle'}}, layout={'hierarchical': False}, manipulation=False )
            clicked_node_id = agraph(nodes=agraph_nodes, edges=agraph_edges, config=agraph_config); valid_node_ids = {n['id'] for n in st.session_state.nodes if isinstance(n, dict)}
            if clicked_node_id and clicked_node_id in valid_node_ids and clicked_node_id != st.session_state.selected_node_id: select_node(clicked_node_id); st.rerun()
        except Exception as e: st.error(f"Graph Error: {e}", icon="🔥"); print(f"Graph Error: {e}"); import traceback; traceback.print_exc()
with top_cols[1]: # Node Config (No Change)
    st.subheader("⚙️ Node Configuration"); valid_node_ids = [n.get('id') for n in st.session_state.nodes if isinstance(n, dict)]
    if not st.session_state.selected_node_id or st.session_state.selected_node_id not in valid_node_ids: st.session_state.selected_node_id = valid_node_ids[0] if valid_node_ids else None
    if not st.session_state.nodes: st.info("Add nodes or load example.")
    else:
        node_options_display = get_node_options_for_select(include_end=False); node_ids_only = [opt[0] for opt in node_options_display]; current_selection_index=0
        try: current_selection_index = node_ids_only.index(st.session_state.selected_node_id) if st.session_state.selected_node_id in node_ids_only else 0
        except: pass
        if not node_ids_only: current_selection_index = 0
        selected_id_from_dropdown = st.selectbox( "Select Node:", options=node_options_display, index=current_selection_index, format_func=lambda x: x[1], key="node_selector_config", label_visibility="collapsed" )
        if selected_id_from_dropdown and selected_id_from_dropdown[0] != st.session_state.selected_node_id: select_node(selected_id_from_dropdown[0]); st.rerun()
        st.divider()
    selected_node_data = next((n for n in st.session_state.nodes if isinstance(n, dict) and n.get("id") == st.session_state.selected_node_id), None)
    if selected_node_data:
         node_id, node_name, node_prompt = selected_node_data["id"], selected_node_data.get("name", ""), selected_node_data.get("prompt", ""); routing_rules = selected_node_data.get("routing_rules", {}); default_target = routing_rules.get("default_target", END_NODE_ID); conditional_targets = routing_rules.get("conditional_targets", [])
         if not isinstance(routing_rules, dict): routing_rules = {"default_target": END_NODE_ID, "conditional_targets": []}
         if not isinstance(conditional_targets, list): conditional_targets = []
         with st.container(border=True):
             form_key = f"config_form_{node_id}"
             with st.form(key=form_key):
                 st.markdown(f"**Editing: {node_name}** (`{node_id}`)"); edited_name = st.text_input("Node Name", value=node_name, key=f"cfg_name_{node_id}"); edited_prompt = st.text_area("LLM Prompt", value=node_prompt, height=150, key=f"cfg_prompt_{node_id}", help=f"Define task. Use '{{input_text}}' if needed. End response with '{ROUTING_KEY_MARKER} <key>'.")
                 st.markdown("**🚦 Routing Rules**"); node_options = get_node_options_for_select(include_end=True, exclude_node_id=node_id); current_default_idx=0;
                 try: current_default_idx = [i for i, (opt_id, _) in enumerate(node_options) if opt_id == default_target][0]
                 except: current_default_idx = next((i for i, (opt_id, _) in enumerate(node_options) if opt_id == END_NODE_ID), 0)
                 selected_default_option = st.selectbox("Default Target", options=node_options, index=current_default_idx, format_func=lambda x: x[1], key=f"cfg_default_{node_id}"); edited_default_target_id = selected_default_option[0] if selected_default_option else END_NODE_ID
                 st.markdown(f"**Conditional Targets (based on {ROUTING_KEY_MARKER} output):**"); edited_conditional_targets = []; current_conditional_targets = list(conditional_targets)
                 for rule_idx, rule in enumerate(current_conditional_targets):
                     st.caption(f"Rule {rule_idx+1}"); rule_cols = st.columns([0.5, 0.5]);
                     with rule_cols[0]: output_key = st.text_input(f"If Key Is", value=rule.get("output_key", ""), placeholder="e.g., success", key=f"cfg_key_{node_id}_{rule_idx}", label_visibility="collapsed")
                     with rule_cols[1]:
                          current_target_idx=0;
                          try: current_target_idx = [i for i, (opt_id, _) in enumerate(node_options) if opt_id == rule.get("target_node_id")][0]
                          except: current_target_idx = next((i for i, (opt_id, _) in enumerate(node_options) if opt_id == END_NODE_ID), 0)
                          selected_target_option = st.selectbox(f"Then Go To", options=node_options, index=current_target_idx, format_func=lambda x: x[1], key=f"cfg_target_{node_id}_{rule_idx}", label_visibility="collapsed"); target_node_id = selected_target_option[0] if selected_target_option else END_NODE_ID
                     edited_conditional_targets.append({"output_key": output_key.strip(), "target_node_id": target_node_id})
                 st.divider(); submitted = st.form_submit_button("💾 Save Changes", type="primary", use_container_width=True)
                 if submitted:
                     final_conditional_targets = [r for r in edited_conditional_targets if r.get("output_key")]; new_data = {"name": edited_name.strip(), "prompt": edited_prompt, "routing_rules": {"default_target": edited_default_target_id, "conditional_targets": final_conditional_targets}}
                     if not edited_name.strip(): st.warning("Node Name cannot be empty.")
                     elif any(isinstance(n, dict) and n.get('id') != node_id and n.get('name','').strip() == edited_name.strip() for n in st.session_state.nodes): st.warning(f"Name exists.")
                     else: update_node_data_ui(node_id, new_data); st.rerun()
             st.markdown("**Manage Rules & Node:**"); action_cols = st.columns([0.5, 0.5])
             with action_cols[0]: st.button("➕ Add Rule", key=f"add_rule_btn_{node_id}", on_click=add_conditional_rule_ui, args=(node_id,), use_container_width=True)
             if conditional_targets: rule_opts = [(idx, f"Rule {idx+1}: '{rule.get('output_key', '')[:10]}...'") for idx, rule in enumerate(conditional_targets)]; rule_opts.insert(0, (-1, "Delete Rule...")); selected_rule_idx_to_del = st.selectbox("Delete Rule", options=rule_opts, format_func=lambda x: x[1], label_visibility="collapsed", index=0, key=f"del_rule_sel_{node_id}")
             if 'selected_rule_idx_to_del' in locals() and selected_rule_idx_to_del and selected_rule_idx_to_del[0] != -1: rule_idx_del = selected_rule_idx_to_del[0]; st.button(f"🗑️ Confirm Del Rule {rule_idx_del+1}", key=f"del_rule_btn_{node_id}_{rule_idx_del}", on_click=delete_conditional_rule_ui, args=(node_id, rule_idx_del), use_container_width=True, type="secondary")
             with action_cols[1]: current_node_index = next((i for i, n in enumerate(st.session_state.nodes) if isinstance(n, dict) and n.get("id") == node_id), -1); st.button("⬆️ Up", key=f"mv_up_{node_id}", on_click=move_node_ui, args=(node_id, -1), disabled=(current_node_index <= 0), use_container_width=True); st.button("⬇️ Down", key=f"mv_dn_{node_id}", on_click=move_node_ui, args=(node_id, 1), disabled=(current_node_index < 0 or current_node_index >= len(st.session_state.nodes) - 1), use_container_width=True); st.button("❌ Delete Node", key=f"del_nd_{node_id}", on_click=delete_node_ui, args=(node_id,), use_container_width=True, type="secondary", help="Delete this node.")
    elif st.session_state.nodes: st.info("Select node.");
    if st.session_state.selected_node_id and st.session_state.nodes and st.session_state.selected_node_id not in [n.get('id') for n in st.session_state.nodes if isinstance(n, dict)]: st.session_state.selected_node_id = st.session_state.nodes[0].get('id'); st.rerun()
    elif not st.session_state.nodes: st.session_state.selected_node_id = None

# --- Execution Section (Adapted for Dict State - No Change) ---
st.divider(); st.header("🚀 Execute Workflow")
run_tooltip = ""; run_disabled = True
if st.session_state.compiled_graph and llm_ready: st.success("Workflow compiled.", icon="✅"); run_tooltip = "Run workflow."; run_disabled = False
elif not llm_ready: st.warning("LLM not ready.", icon="⚠️"); run_tooltip = "LLM needs API Key."
elif st.session_state.nodes: st.warning("Compile Workflow first.", icon="⚠️"); run_tooltip = "Compile first."
else: st.info("Workflow empty."); run_tooltip = "Workflow empty/not compiled."
initial_message = st.text_area( "Enter initial message:", height=80, key="initial_input_exec", value="I want moderate capital appreciation over 5-7 years, willing to accept some market volatility but avoid highly speculative assets. Focus on tech and renewable energy sectors.", help="Input for the first node." )
if not initial_message.strip(): run_disabled = True; run_tooltip += " Initial message required."
if st.button("▶️ Run Workflow", disabled=run_disabled, type="primary", help=run_tooltip):
    st.session_state.execution_log = ["**🚀 Starting Workflow...**"]; st.session_state.final_state = None
    initial_state = WorkflowState(input=initial_message, node_outputs={}, last_response_content="", current_node_id="")
    st.session_state.execution_log.append(f"📥 Input: {initial_message[:150]}{'...' if len(initial_message)>150 else ''}")
    log_placeholder = st.empty(); log_placeholder.info("⏳ Running workflow...")
    with st.spinner("Executing workflow..."):
        try:
            rec_limit = st.session_state.get('recursion_limit', 25); print(f"Invoking graph with limit: {rec_limit}")
            final_state_result: WorkflowState = st.session_state.compiled_graph.invoke(initial_state, config={"recursion_limit": rec_limit})
            st.session_state.final_state = final_state_result; st.session_state.execution_log.append("**🏁 Workflow Finished**"); st.toast("Finished!", icon="🏁")
        except Exception as e:
            err_msg = f"{e}"
            if "must be followed by tool messages" in str(e): err_msg = f"LLM generated 'tool_calls' which graph cannot handle. Error: {e}."; st.error(f"{err_msg}", icon="🔥")
            elif isinstance(e, RecursionError) or "recursion limit" in str(e).lower(): err_msg = f"Recursion limit ({rec_limit}) reached. Error: {e}"; st.error(f"{err_msg}", icon="🔥")
            else: st.error(f"Exec failed: {err_msg}", icon="🔥");
            st.session_state.execution_log.append(f"**💥 WORKFLOW ERROR:** {err_msg}"); print(f"Exec Error: {e}"); import traceback; traceback.print_exc(); st.toast("Failed!", icon="❌")
    log_placeholder.empty(); st.rerun()

# --- Results Display Section (Adapted for Dict State - No Change) ---
st.subheader("📊 Execution Results"); results_cols = st.columns(2)
with results_cols[0]:
    st.markdown("**Execution Log**");
    if st.session_state.execution_log: log_text = "\n".join(st.session_state.execution_log); st.text_area("Log Details:", value=log_text, height=300, disabled=True, key="log_display_final_main")
    else: st.caption("Run workflow to see log.")
with results_cols[1]:
    st.markdown("**Final Output Message**"); final_message_content = None
    if st.session_state.final_state and isinstance(st.session_state.final_state, dict):
        final_content_raw = st.session_state.final_state.get('last_response_content', '')
        if final_content_raw and isinstance(final_content_raw, str): final_message_content = re.sub(rf"\s*{ROUTING_KEY_MARKER}\s*\w+\s*$", "", final_content_raw).strip()
    if final_message_content:
        message_container = st.container(height=300, border=False); avatar = "🏁";
        if any("WORKFLOW ERROR" in log for log in st.session_state.execution_log): avatar = "💥"
        elif "ERROR:" in final_message_content: avatar = "⚠️"
        message_container.chat_message("assistant", avatar=avatar).write(final_message_content)
    elif any("WORKFLOW ERROR" in log for log in st.session_state.execution_log): st.caption("Ended with error.")
    else: st.caption("Run workflow.")

# --- End of App ---

Source : https://deepcharts.substack.com/p/the-end-of-n8n-vibe-coding-an-ai?subscribe_prompt=free