Blame for app.py

from __future__ import annotations

import json
import os
import re
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import streamlit as st

ROOT = Path(__file__).resolve().parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from qa_langgraph.export import build_run_export
from qa_langgraph.graph import build_graph
from qa_langgraph.registry import AGENTS, default_agent_configs
from qa_langgraph.run_logger import read_stage_traces_jsonl, write_run_export
from qa_langgraph.sample_scenarios import CUSTOM_SCENARIO, load_sample_scenario, scenario_options
from qa_langgraph.service_client import DEFAULT_SERVICE_URL, normalize_service_url, validate_agent_service


def init_session() -> None:
    default_service_url = normalize_service_url(os.getenv("QA_AGENT_SERVICE_URL", DEFAULT_SERVICE_URL))
    st.session_state.setdefault("agent_configs", default_agent_configs())
    st.session_state.setdefault("run_state", {})
    st.session_state.setdefault("run_export", None)
    st.session_state.setdefault("run_status", "not run")
    st.session_state.setdefault("current_run_id", "")
    st.session_state.setdefault("stage_traces_jsonl", "")
    st.session_state.setdefault("selected_scenario", CUSTOM_SCENARIO)
    st.session_state.setdefault("title", "")
    st.session_state.setdefault("requirements_text", "")
    st.session_state.setdefault("qa_agent_service_url", default_service_url)


def agent_status_rows(stage_traces: list[dict]) -> list[dict]:
    latest: dict[str, dict] = {}
    for trace in stage_traces:
        latest[trace.get("agent_key", "")] = trace
    rows = []
    for key, agent in AGENTS.items():
        trace = latest.get(key, {})
        rows.append(
            {
                "agent": agent.name,
                "status": trace.get("status", "not started"),
                "iteration": trace.get("iteration", ""),
                "duration_ms": trace.get("duration_ms", ""),
            }
        )
    return rows


def safe_filename_part(value: str, fallback: str = "scenario") -> str:
    cleaned = re.sub(r"[^A-Za-z0-9]+", "_", value.strip()).strip("_")
    return cleaned or fallback


def download_basename(export: dict) -> str:
    title = safe_filename_part(export.get("title", "scenario"))
    created_at = export.get("run_metadata", {}).get("created_at", "")
    try:
        timestamp = datetime.fromisoformat(created_at.replace("Z", "+00:00")).strftime("%Y%m%d-%H%M%S")
    except ValueError:
        timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    return f"LangGraph_{title}_{timestamp}"


def _dot_label(label: str) -> str:
    return label.replace('"', '\\"')


def runtime_path_edges(stage_traces: list[dict], orchestration_decision: dict) -> set[tuple[str, str]]:
    edges: set[tuple[str, str]] = set()
    ordered_nodes = [trace.get("agent_key") for trace in stage_traces if trace.get("agent_key")]
    if ordered_nodes:
        edges.add(("START", ordered_nodes[0]))
    for source, target in zip(ordered_nodes, ordered_nodes[1:]):
        edges.add((source, target))
    decision = orchestration_decision.get("decision")
    target_agent = orchestration_decision.get("target_agent")
    if ordered_nodes and ordered_nodes[-1] == "orchestrator":
        if decision == "finish":
            edges.add(("orchestrator", "END"))
        elif decision == "fail":
            edges.add(("orchestrator", "END_FAIL"))
        elif decision == "backtrack" and target_agent in {"requirements_analyst", "test_designer", "reviewer"}:
            edges.add(("orchestrator", target_agent))
    return edges


def build_workflow_dot(stage_traces: list[dict] | None = None, orchestration_decision: dict | None = None) -> str:
    active_edges = runtime_path_edges(stage_traces or [], orchestration_decision or {})
    completed_nodes = {trace.get("agent_key") for trace in (stage_traces or []) if trace.get("status") == "completed"}
    failed_nodes = {trace.get("agent_key") for trace in (stage_traces or []) if trace.get("status") == "failed"}

    nodes = {
        "START": {"label": "START", "shape": "circle", "fillcolor": "#1e293b", "fontcolor": "#f8fafc"},
        "requirements_analyst": {
            "label": "Requirements Analyst\\n(service)",
            "shape": "box",
            "fillcolor": "#0b5fff" if "requirements_analyst" in completed_nodes else "#f8fafc",
            "fontcolor": "#ffffff" if "requirements_analyst" in completed_nodes else "#0f172a",
        },
        "test_designer": {
            "label": "Test Designer\\n(service)",
            "shape": "box",
            "fillcolor": "#0b5fff" if "test_designer" in completed_nodes else "#f8fafc",
            "fontcolor": "#ffffff" if "test_designer" in completed_nodes else "#0f172a",
        },
        "reviewer": {
            "label": "Reviewer\\n(service)",
            "shape": "box",
            "fillcolor": "#0b5fff" if "reviewer" in completed_nodes else "#f8fafc",
            "fontcolor": "#ffffff" if "reviewer" in completed_nodes else "#0f172a",
        },
        "orchestrator": {
            "label": "Orchestrator\\n(local deterministic)",
            "shape": "diamond",
            "fillcolor": "#f59e0b" if "orchestrator" in completed_nodes else "#fef3c7",
            "fontcolor": "#0f172a",
        },
        "END": {"label": "END\\nfinish", "shape": "doublecircle", "fillcolor": "#16a34a", "fontcolor": "#ffffff"},
        "END_FAIL": {"label": "END\\nfail", "shape": "doublecircle", "fillcolor": "#dc2626", "fontcolor": "#ffffff"},
    }

    for node_key in failed_nodes:
        if node_key in nodes:
            nodes[node_key]["fillcolor"] = "#dc2626"
            nodes[node_key]["fontcolor"] = "#ffffff"

    edges = [
        ("START", "requirements_analyst", ""),
        ("requirements_analyst", "test_designer", ""),
        ("test_designer", "reviewer", ""),
        ("reviewer", "orchestrator", ""),
        ("orchestrator", "END", "finish"),
        ("orchestrator", "requirements_analyst", "backtrack"),
        ("orchestrator", "test_designer", "backtrack"),
        ("orchestrator", "reviewer", "backtrack"),
        ("orchestrator", "END_FAIL", "fail"),
    ]

    lines = [
        "digraph LangGraphWorkflow {",
        '  graph [rankdir=LR, bgcolor="#f8fafc", fontname="Helvetica", nodesep="0.55", ranksep="0.75", pad="0.25"];',
        '  node [style="rounded,filled", color="#0f172a", penwidth=1.5, fontname="Helvetica"];',
        '  edge [color="#334155", penwidth=1.9, fontname="Helvetica", fontsize=11, fontcolor="#0f172a"];',
    ]

    for node_key, attrs in nodes.items():
        lines.append(
            f'  {node_key} [label="{_dot_label(attrs["label"])}", shape="{attrs["shape"]}", fillcolor="{attrs["fillcolor"]}", fontcolor="{attrs["fontcolor"]}"];'
        )

    for source, target, label in edges:
        is_active = (source, target) in active_edges
        color = "#0b5fff" if is_active else "#334155"
        penwidth = "3.4" if is_active else "1.9"
        fontcolor = "#0b5fff" if is_active else "#0f172a"
        label_attr = f', label="{label}", fontcolor="{fontcolor}"' if label else ""
        lines.append(f'  {source} -> {target} [color="{color}", penwidth={penwidth}{label_attr}];')

    lines.append("}")
    return "\n".join(lines)


def runtime_path_summary(stage_traces: list[dict], orchestration_decision: dict) -> str:
    ordered_names = [AGENTS[trace["agent_key"]].name for trace in stage_traces if trace.get("agent_key") in AGENTS]
    if not ordered_names:
        return "No run has been executed yet."
    summary = " -> ".join(ordered_names)
    decision = orchestration_decision.get("decision")
    if decision == "finish":
        return f"{summary} -> END (finish)"
    if decision == "fail":
        return f"{summary} -> END (fail)"
    if decision == "backtrack" and orchestration_decision.get("target_agent") in AGENTS:
        target_name = AGENTS[orchestration_decision["target_agent"]].name
        return f"{summary} -> backtrack to {target_name}"
    return summary


def render_service_help(service_url: str) -> None:
    rest_url = normalize_service_url(service_url)
    st.markdown("**Use REST here, not MCP.**")
    st.caption("LangGraph calls FastAPI REST endpoints. MCP on port 8008 is for MCP clients and does not expose Swagger/OpenAPI.")
    st.table(
        [
            {
                "server": "Hosted REST/FastAPI",
                "default": "https://helmfridsson-qa-agent-service.hf.space/api/v1",
                "paths": "/health, /docs, /openapi.json, /agents/{agent_key}/invoke",
            },
            {
                "server": "Local REST/FastAPI",
                "default": "http://127.0.0.1:8009/api/v1",
                "paths": "/health, /docs, /openapi.json, /agents/{agent_key}/invoke",
            },
            {"server": "Local MCP", "default": "http://127.0.0.1:8008/mcp", "paths": "/mcp only, no Swagger/OpenAPI"},
        ]
    )
    st.caption("You can paste the full REST base URL, the service root, `/docs`, `/health`, or the Hugging Face Space page URL. The app normalizes it to `/api/v1`.")
    st.markdown(
        f"""
REST links for the configured URL:

- Health: [{rest_url}/health]({rest_url}/health)
- Swagger: [{rest_url}/docs]({rest_url}/docs)
- OpenAPI JSON: [{rest_url}/openapi.json]({rest_url}/openapi.json)
"""
    )
    st.code(
        """cd /Users/tomashelmfridsson/workspace/qa-agent-service
PYTHONPATH=src uvicorn api:app --reload --port 8009""",
        language="bash",
    )
    st.caption("The PYTHONPATH part fixes ModuleNotFoundError when qa-agent-service is not installed in the active Python environment.")
    st.caption("If the hosted /api/v1 health check returns Streamlit HTML, the Space is running the monitor UI but not exposing the REST API surface yet.")


def run_workflow(title: str, requirements_text: str, global_config: dict, agent_configs: list[dict], max_rounds: int) -> None:
    if not title.strip():
        st.error("Scenario title is required.")
        return
    if not requirements_text.strip():
        st.error("Requirement text is required.")
        return
    if any(config.get("execution_mode") not in {"Deterministic", "Disabled"} for config in agent_configs):
        validate_agent_service(global_config)

    run_id = f"run_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
    state = {
        "run_metadata": {
            "run_id": run_id,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "started_monotonic_ms": int(time.time() * 1000),
            "app_version": "langgraph-qa-compare-0.1.0",
        },
        "title": title.strip(),
        "source_requirements": requirements_text.strip(),
        "requirements": [],
        "test_designs": [],
        "generated_artifacts": [],
        "iterations": 1,
        "run_controls": {
            "max_rounds": max_rounds,
            "max_feedback_messages": 0,
            "max_feedback_per_agent_pair": 0,
        },
        "global_llm_config": global_config,
        "agent_configs": agent_configs,
        "stage_traces": [],
        "working_memory": {"shared": {}, "agent_private": {}, "timeline": []},
        "errors": [],
    }
    st.session_state.current_run_id = run_id
    st.session_state.run_status = "running"
    st.session_state.run_state = state
    st.session_state.run_export = None
    st.session_state.stage_traces_jsonl = ""

    graph = build_graph()
    final_state = state
    status_slot = st.empty()
    traces_slot = st.empty()
    memory_slot = st.empty()
    for snapshot in graph.stream(state, {"recursion_limit": max(8, max_rounds * 6)}, stream_mode="values"):
        final_state = snapshot
        st.session_state.run_state = final_state
        st.session_state.run_status = "running"
        traces = final_state.get("stage_traces", [])
        memory = final_state.get("working_memory", {"shared": {}, "agent_private": {}, "timeline": []})
        status_slot.info(f"Running... {len(traces)} stage trace(s), iteration {final_state.get('iterations', 1)}")
        traces_slot.dataframe(agent_status_rows(traces), use_container_width=True)
        memory_slot.json(
            {
                "shared": memory.get("shared", {}),
                "timeline": memory.get("timeline", []),
            }
        )
    export = build_run_export(final_state)
    export_dict = export.model_dump(mode="json")
    write_run_export(run_id, export_dict)
    st.session_state.run_state = final_state
    st.session_state.run_export = export_dict
    st.session_state.run_status = export_dict.get("run_metadata", {}).get("status", "partial")
    st.session_state.stage_traces_jsonl = read_stage_traces_jsonl(run_id)


init_session()
st.set_page_config(page_title="LangGraph QA Compare", layout="wide")
st.markdown(
    """
    <style>
    .stApp {
        background-color: #0b1220;
        color: #e5e7eb;
    }

    .stApp h1,
    .stApp h2,
    .stApp h3,
    .stApp h4,
    .stApp h5,
    .stApp h6 {
        color: #f8fafc !important;
    }

    .stApp [data-testid="stHeadingWithActionElements"],
    .stApp [data-testid="stHeadingWithActionElements"] *,
    .stApp [data-testid="stCaptionContainer"],
    .stApp [data-testid="stAlertContent"],
    .stApp [data-testid="stJson"] {
        color: #dbe4f0 !important;
        opacity: 1 !important;
    }

    .stApp .block-container [data-testid="stMarkdownContainer"] p,
    .stApp .block-container [data-testid="stMarkdownContainer"] li,
    .stApp .block-container [data-testid="stMarkdownContainer"] span,
    .stApp .block-container [data-testid="stCaptionContainer"] p {
        color: #dbe4f0 !important;
        opacity: 1 !important;
    }

    .stApp [data-testid="stCodeBlock"],
    .stApp pre,
    .stApp code {
        color: #f8fafc !important;
    }

    .stApp [data-testid="stCodeBlock"] pre,
    .stApp pre {
        background-color: #111827 !important;
    }

    .stApp [data-testid="stMetricValue"],
    .stApp [data-testid="stMetricLabel"] {
        color: #f8fafc !important;
    }

    .stApp [data-testid="stWidgetLabel"],
    .stApp [data-testid="stExpander"] summary p,
    .stApp [data-testid="stTextInput"] label,
    .stApp [data-testid="stTextArea"] label,
    .stApp [data-testid="stSelectbox"] label,
    .stApp [data-testid="stNumberInput"] label {
        color: #e2e8f0 !important;
        opacity: 1 !important;
    }

    .stApp input,
    .stApp textarea,
    .stApp [data-baseweb="select"] > div {
        background-color: #111827 !important;
        border-color: #374151 !important;
        color: #f8fafc !important;
    }

    .stApp [data-baseweb="select"] input,
    .stApp input,
    .stApp textarea {
        color: #f8fafc !important;
    }

    .stApp input::placeholder,
    .stApp textarea::placeholder {
        color: #cbd5e1 !important;
        opacity: 1 !important;
    }

    .agent-directives {
        white-space: pre-wrap;
        color: #f8fafc;
        background: #111827;
        border: 1px solid #374151;
        border-radius: 0.5rem;
        padding: 0.75rem;
        line-height: 1.45;
    }
    </style>
    """,
    unsafe_allow_html=True,
)
st.title("LangGraph QA Compare")

with st.sidebar:
    st.header("QA Agent Service")
    service_url = st.text_input("REST API URL", value=st.session_state.get("qa_agent_service_url", DEFAULT_SERVICE_URL))
    st.session_state.qa_agent_service_url = normalize_service_url(service_url)
    service_timeout_seconds = st.number_input("Service timeout seconds", min_value=10, max_value=600, value=150, step=10)
    max_rounds = st.number_input("Max rounds", min_value=1, max_value=8, value=2, step=1)
    global_config = {
        "qa_agent_service_url": st.session_state.qa_agent_service_url,
        "service_timeout_seconds": int(service_timeout_seconds),
    }
    if st.button("Validate QA Agent Service"):
        try:
            validate_agent_service(global_config)
            st.success("QA Agent Service is reachable.")
        except Exception as exc:
            st.error(str(exc))
    with st.expander("REST vs MCP endpoints", expanded=False):
        render_service_help(service_url)
    if st.button("Clear session"):
        st.session_state.clear()
        st.rerun()

tab_input, tab_agents, tab_flow, tab_runtime, tab_memory, tab_results, tab_downloads = st.tabs(
    ["Scenario Input", "Agent Configuration", "Flow Map", "Runtime Activity", "Working Memory", "Results", "Downloads"]
)

with tab_input:
    selected_scenario = st.selectbox(
        "Predefined scenario",
        scenario_options(),
        index=scenario_options().index(st.session_state.get("selected_scenario", CUSTOM_SCENARIO)),
    )
    if st.button("Load selected scenario"):
        title_value, requirements_value = load_sample_scenario(selected_scenario)
        st.session_state.selected_scenario = selected_scenario
        st.session_state.title = title_value
        st.session_state.requirements_text = requirements_value
        st.rerun()

    title = st.text_input("Scenario title", value=st.session_state.get("title", ""))
    requirements_text = st.text_area(
        "Raw requirements",
        value=st.session_state.get("requirements_text", ""),
        height=300,
    )
    if st.button("Run workflow", type="primary"):
        st.session_state.title = title
        st.session_state.requirements_text = requirements_text
        try:
            run_workflow(title, requirements_text, global_config, st.session_state.agent_configs, int(max_rounds))
            st.success("Workflow finished. Check Runtime Activity and Results.")
        except Exception as exc:
            st.session_state.run_status = "failed"
            st.error(str(exc))

with tab_agents:
    updated_configs = []
    for config in st.session_state.agent_configs:
        key = config["agent_key"]
        agent_definition = AGENTS[key]
        with st.expander(config["agent_name"], expanded=False):
            mode_options = ["Deterministic", "Disabled"] if key == "orchestrator" else ["Service-backed", "Deterministic", "Disabled"]
            current_mode = config.get("execution_mode", "Deterministic" if key == "orchestrator" else "Service-backed")
            if current_mode == "LLM-backed":
                current_mode = "Service-backed"
            if current_mode not in mode_options:
                current_mode = mode_options[0]
            execution_mode = st.selectbox(
                "Execution mode",
                mode_options,
                index=mode_options.index(current_mode),
                key=f"{key}_mode",
            )
            provider_strategy = st.selectbox(
                "Caller provider override",
                ["global", "ollama", "openai"],
                index=["global", "ollama", "openai"].index(config.get("provider_strategy", "global")),
                key=f"{key}_provider",
            )
            model_override = st.text_input("Caller model override", value=config.get("model_override", ""), key=f"{key}_model")
            agent_timeout = st.number_input(
                "Service call timeout seconds",
                min_value=10,
                max_value=600,
                value=int(config.get("timeout_seconds", 150)),
                step=10,
                key=f"{key}_timeout",
            )
            st.markdown("Default directives")
            st.markdown(
                f"<div class='agent-directives'>{agent_definition.default_directives}</div>",
                unsafe_allow_html=True,
            )
            directives = st.text_area("Additional directives", value=config.get("directives", ""), key=f"{key}_directives")
            updated = dict(config)
            updated.update(
                {
                    "execution_mode": execution_mode,
                    "provider_strategy": provider_strategy,
                    "model_override": model_override,
                    "timeout_seconds": int(agent_timeout),
                    "directives": directives,
                }
            )
            updated_configs.append(updated)
    st.session_state.agent_configs = updated_configs

run_state = st.session_state.get("run_state", {})
stage_traces = run_state.get("stage_traces", [])
orchestration_decision = run_state.get("orchestration_decision", {})

with tab_flow:
    col_static, col_runtime = st.columns(2)
    with col_static:
        st.subheader("Workflow structure")
        st.caption("Static LangGraph routing, including the orchestrator's conditional backtracking paths.")
        st.graphviz_chart(build_workflow_dot(), use_container_width=True)
    with col_runtime:
        st.subheader("Latest run path")
        st.caption("Highlighted edges show the path taken by the current run state.")
        st.graphviz_chart(build_workflow_dot(stage_traces, orchestration_decision), use_container_width=True)
    st.code(runtime_path_summary(stage_traces, orchestration_decision), language="text")
    if orchestration_decision:
        st.json(
            {
                "decision": orchestration_decision.get("decision"),
                "target_agent": orchestration_decision.get("target_agent"),
                "reason": orchestration_decision.get("reason"),
                "feedback": orchestration_decision.get("feedback"),
            }
        )

with tab_runtime:
    export = st.session_state.get("run_export")
    status = export.get("run_metadata", {}).get("status") if export else st.session_state.get("run_status", "not run")
    st.metric("Current run status", status)
    if st.session_state.get("current_run_id"):
        st.caption(f"Run ID: {st.session_state.current_run_id}")
    st.metric("Iteration count", run_state.get("iterations", 0))
    st.dataframe(agent_status_rows(stage_traces), use_container_width=True)
    st.caption(runtime_path_summary(stage_traces, orchestration_decision))
    st.json(stage_traces)

with tab_memory:
    memory = run_state.get("working_memory", {"shared": {}, "timeline": []})
    st.metric("Current run status", st.session_state.get("run_status", "not run"))
    if st.session_state.get("current_run_id"):
        st.caption(f"Run ID: {st.session_state.current_run_id}")
    st.subheader("Shared")
    st.json(memory.get("shared", {}))
    st.subheader("Timeline")
    st.json(memory.get("timeline", []))

with tab_results:
    st.subheader("Requirements")
    st.dataframe(run_state.get("requirements", []), use_container_width=True)
    st.subheader("Test cases")
    st.dataframe(run_state.get("test_designs", []), use_container_width=True)
    st.subheader("Review")
    st.json(run_state.get("review", {}))
    st.subheader("Orchestration decision")
    st.json(orchestration_decision)
    st.subheader("Final JSON preview")
    st.json(st.session_state.get("run_export") or {})

with tab_downloads:
    export = st.session_state.get("run_export")
    if export:
        basename = download_basename(export)
        st.download_button(
            "Download run_export.json",
            data=json.dumps(export, indent=2, ensure_ascii=False),
            file_name=f"{basename}_run_export.json",
            mime="application/json",
        )
        st.download_button(
            "Download stage_traces.jsonl",
            data=st.session_state.get("stage_traces_jsonl", ""),
            file_name=f"{basename}_stage_traces.jsonl",
            mime="application/x-ndjson",
        )
    else:
        st.info("Run the workflow to generate downloads.")