agentworkflow-viz 0.2.0__py3-none-any.whl

agentviz/__init__.py

@@ -0,0 +1,88 @@
+import os
+from typing import Optional
+
+from agentviz.config import (
+    VizConfig,
+    VizStyle,
+    NodeStyle,
+    SectionVisibility,
+    ThemeColors,
+    VizCustomization,
+    NodeCustomProps,
+    NodeRenderer,
+)
+from agentviz.client import VizClient
+from agentviz.server import _run_server, mount_app
+from agentviz.tracer import LangGraphTracer, VizTracer
+
+
+def _read_env_config():
+    enabled = os.environ.get("AGENT_VIZ_ENABLED", "true").lower() not in ("false", "0", "no")
+    port = int(os.environ.get("AGENT_VIZ_PORT", "8100"))
+    host = os.environ.get("AGENT_VIZ_HOST", "0.0.0.0")
+    mount = os.environ.get("AGENT_VIZ_MOUNT", "false").lower() in ("true", "1", "yes")
+    return VizConfig(enable=enabled, port=port, host=host, mount=mount)
+
+
+def _merge_config(user_cfg, env_cfg):
+    return VizConfig(
+        enable=user_cfg.enable if user_cfg.enable is not None else env_cfg.enable,
+        port=user_cfg.port if user_cfg.port is not None else env_cfg.port,
+        host=user_cfg.host if user_cfg.host is not None else env_cfg.host,
+        mount=user_cfg.mount if user_cfg.mount is not None else env_cfg.mount,
+        style=user_cfg.style if user_cfg.style is not None else env_cfg.style,
+        customization=user_cfg.customization if user_cfg.customization is not None else env_cfg.customization,
+    )
+
+
+class VizPlugin:
+    def __init__(self, config: Optional[VizConfig] = None):
+        env_cfg = _read_env_config()
+        if config is None:
+            self.config = env_cfg
+        else:
+            self.config = _merge_config(config, env_cfg)
+        self._ready = False
+        self._client = VizClient(
+            host=self.config.host,
+            port=self.config.port,
+            style=self.config.style,
+            customization=self.config.customization,
+        )
+        self._tracer = LangGraphTracer(client=self._client) if self.config.enable else None
+
+    def start(self):
+        if not self.config.enable:
+            self._ready = True
+            return
+        if self.config.mount:
+            self._ready = True
+            return
+        _run_server(self.config, client=self._client)
+        self._client.push_style()
+        self._client.push_customization()
+        self._ready = True
+
+    def stop(self):
+        self._client.flush()
+        self._client.close()
+
+    def wrap(self, graph):
+        if not self.config.enable:
+            return graph
+        return self._tracer.wrap(graph)
+
+
+__all__ = [
+    "VizPlugin",
+    "VizConfig",
+    "VizStyle",
+    "NodeStyle",
+    "SectionVisibility",
+    "ThemeColors",
+    "VizCustomization",
+    "NodeCustomProps",
+    "NodeRenderer",
+    "VizTracer",
+    "mount_app",
+]

agentviz/client.py

@@ -0,0 +1,151 @@
+import threading
+import time
+from collections import deque
+from dataclasses import asdict
+from typing import Optional
+
+import requests
+
+from agentviz.config import VizStyle, VizCustomization, NodeRenderer
+
+
+class VizClient:
+    MAX_BUFFER = 200
+    MAX_RETRIES = 3
+    BACKOFF_BASE = 0.1
+
+    def __init__(self, host: str = "0.0.0.0", port: int = 8100, style: Optional[VizStyle] = None, customization: Optional[VizCustomization] = None):
+        self._host = host
+        self._port = port
+        self._url = f"http://{host}:{port}"
+        self._lock = threading.Lock()
+        self._trace_log: deque[dict] = deque(maxlen=self.MAX_BUFFER)
+        self._graph_data: Optional[dict] = None
+        self._trace_id = 0
+        self._closed = False
+        self._style = style
+        self._customization = customization
+
+    @property
+    def trace_log(self):
+        return list(self._trace_log)
+
+    @property
+    def graph_data(self):
+        return self._graph_data
+
+    @graph_data.setter
+    def graph_data(self, value: dict):
+        self._graph_data = value
+
+    @property
+    def style_config(self) -> Optional[dict]:
+        if self._style is None:
+            return None
+        return self._style_to_dict(self._style)
+
+    @staticmethod
+    def _style_to_dict(style: VizStyle) -> dict:
+        out = {}
+        ns = {}
+        for node_id, nst in style.node_styles.items():
+            ns[node_id] = asdict(nst)
+        out["node_styles"] = ns
+        out["sections"] = asdict(style.sections)
+        tc = {}
+        for k, v in asdict(style.theme).items():
+            if v is not None:
+                tc[k] = v
+        out["theme"] = tc
+        out["title"] = style.title
+        out["show_groups"] = style.show_groups
+        return out
+
+    def _build_customization_payload(self) -> Optional[dict]:
+        if self._customization is None:
+            return None
+        cfg = self._customization
+        out = {}
+        out["title"] = cfg.title
+        out["show_groups"] = cfg.show_groups
+        out["custom_css"] = cfg.custom_css
+        out["custom_js"] = cfg.custom_js
+        out["custom_html_header"] = cfg.custom_html_header
+        out["custom_html_body"] = cfg.custom_html_body
+        if cfg.theme_colors is not None:
+            tc = {}
+            for k, v in asdict(cfg.theme_colors).items():
+                if v is not None:
+                    tc[k] = v
+            out["theme"] = tc
+        if cfg.section_visibility is not None:
+            out["sections"] = asdict(cfg.section_visibility)
+        node_props = {}
+        for node_id, ncp in cfg.node_props.items():
+            node_props[node_id] = asdict(ncp)
+        if cfg.node_renderer is not None:
+            node_props["_renderer"] = True
+        out["node_props"] = node_props
+        return out
+
+    def _apply_node_renderer(self, nodes: list) -> list:
+        if self._customization is None or self._customization.node_renderer is None:
+            return nodes
+        renderer: NodeRenderer = self._customization.node_renderer
+        result = []
+        for node in nodes:
+            n = dict(node)
+            nid = n.get("id", "")
+            rendered = renderer.render(nid, n)
+            if rendered:
+                existing = n.get("_custom_props") or {}
+                existing.update(rendered)
+                n["_custom_props"] = existing
+            result.append(n)
+        return result
+
+    def _next_id(self):
+        with self._lock:
+            self._trace_id += 1
+            return self._trace_id
+
+    def push_graph(self, nodes: list, edges: list):
+        enriched = self._apply_node_renderer(nodes)
+        self.graph_data = {"nodes": enriched, "edges": edges}
+        self._http_post("/api/graph/push", {"nodes": enriched, "edges": edges})
+
+    def push_style(self):
+        if self._style is not None:
+            payload = self.style_config
+            self._http_post("/api/style/push", payload)
+
+    def push_customization(self):
+        payload = self._build_customization_payload()
+        if payload is not None:
+            self._http_post("/api/customize/push", payload)
+
+    def push_trace(self, node: str, duration: Optional[float] = None):
+        tid = self._next_id()
+        entry = {"id": tid, "node": node, "duration": duration}
+        with self._lock:
+            self._trace_log.append(entry)
+        self._http_post("/api/trace/push", entry)
+
+    def _http_post(self, path: str, payload: dict):
+        for attempt in range(self.MAX_RETRIES):
+            try:
+                requests.post(
+                    f"{self._url}{path}",
+                    json=payload,
+                    timeout=2,
+                )
+                return
+            except Exception:
+                if attempt < self.MAX_RETRIES - 1:
+                    time.sleep(self.BACKOFF_BASE * (2 ** attempt))
+
+    def flush(self):
+        pass
+
+    def close(self):
+        self._closed = True

agentviz/config.py

@@ -0,0 +1,123 @@
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+
+@dataclass
+class NodeStyle:
+    fill: Optional[str] = None
+    stroke: Optional[str] = None
+    text_color: Optional[str] = None
+    stroke_width: Optional[float] = None
+    shape: Optional[str] = None
+
+
+@dataclass
+class SectionVisibility:
+    graph: bool = True
+    nodes_table: bool = True
+    trace_panel: bool = True
+    header: bool = True
+
+
+@dataclass
+class ThemeColors:
+    bg: Optional[str] = None
+    surface: Optional[str] = None
+    text: Optional[str] = None
+    accent: Optional[str] = None
+    border: Optional[str] = None
+
+
+@dataclass
+class VizStyle:
+    node_styles: dict[str, NodeStyle] = field(default_factory=dict)
+    sections: SectionVisibility = field(default_factory=SectionVisibility)
+    theme: ThemeColors = field(default_factory=ThemeColors)
+    title: str = "Agent Graph"
+    show_groups: bool = True
+
+
+@dataclass
+class NodeCustomProps:
+    """Rich per-node customization sent to the frontend renderer."""
+    width: Optional[int] = None
+    height: Optional[int] = None
+    icon: Optional[str] = None
+    css_class: Optional[str] = None
+    html_template: Optional[str] = None
+    metadata: dict = field(default_factory=dict)
+
+
+class NodeRenderer:
+    """Callable-based Python node renderer.
+
+    Users implement ``render(self, node_id, node_data)`` to return a dict
+    of custom properties that override the default rendering for that node.
+
+    Returns a dict with any subset of keys from :class:`NodeCustomProps`.
+    Example::
+
+        class MyRenderer(NodeRenderer):
+            def render(self, node_id, node_data):
+                if node_id == "agent":
+                    return {
+                        "icon": "🤖",
+                        "css_class": "custom-agent",
+                        "width": 160,
+                        "height": 60,
+                        "metadata": {"custom_key": "value"},
+                    }
+                return {}
+    """
+
+    def render(self, node_id: str, node_data: dict) -> dict:
+        return {}
+
+
+@dataclass
+class VizCustomization:
+    """Full customization config that bridges Python and the frontend.
+
+    Parameters
+    ----------
+    node_renderer : NodeRenderer | None
+        Python-side renderer called for each node to compute custom props.
+    node_props : dict[str, NodeCustomProps]
+        Static per-node overrides (merged with renderer output).
+    custom_css : str | None
+        Raw CSS injected into the page via a ``<style>`` tag.
+    custom_js : str | None
+        Raw JS injected into the page via a ``<script>`` tag.
+    custom_html_header : str | None
+        HTML injected into ``<head>`` (useful for external CSS/JS links).
+    custom_html_body : str | None
+        HTML injected at the end of ``<body>``.
+    theme_colors : ThemeColors | None
+        Theme overrides.
+    title : str
+        Page and header title.
+    show_groups : bool
+        Whether to draw group containers in the SVG graph.
+    section_visibility : SectionVisibility | None
+        Which UI sections to show/hide.
+    """
+    node_renderer: Optional[NodeRenderer] = None
+    node_props: dict[str, NodeCustomProps] = field(default_factory=dict)
+    custom_css: Optional[str] = None
+    custom_js: Optional[str] = None
+    custom_html_header: Optional[str] = None
+    custom_html_body: Optional[str] = None
+    theme_colors: Optional[ThemeColors] = None
+    title: str = "Agent Graph"
+    show_groups: bool = True
+    section_visibility: Optional[SectionVisibility] = None
+
+
+@dataclass
+class VizConfig:
+    enable: bool = True
+    port: int = 8100
+    host: str = "0.0.0.0"
+    mount: bool = False
+    style: Optional[VizStyle] = None
+    customization: Optional[VizCustomization] = None

agentviz/server.py

@@ -0,0 +1,104 @@
+import os
+import threading
+import time
+from typing import Optional
+
+from fastapi import FastAPI, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
+
+from agentviz.client import VizClient
+
+
+_static_dir = os.path.join(os.path.dirname(__file__), "static")
+
+
+def _create_app(client: Optional[VizClient] = None) -> FastAPI:
+    app = FastAPI(title="Agent Viz", version="0.1.0")
+    _style_cache: dict = {}
+    _customize_cache: dict = {}
+
+    @app.get("/health")
+    async def health():
+        return {"status": "ok"}
+
+    @app.get("/api/config")
+    async def api_config():
+        customize = _customize_cache if _customize_cache else None
+        return {
+            "style": _style_cache if _style_cache else None,
+            "customize": customize,
+        }
+
+    @app.post("/api/style/push")
+    async def style_push(data: dict):
+        nonlocal _style_cache
+        _style_cache = data
+        return {"status": "ok"}
+
+    @app.post("/api/customize/push")
+    async def customize_push(data: dict):
+        nonlocal _customize_cache
+        _customize_cache = data
+        return {"status": "ok"}
+
+    @app.get("/api/graph")
+    async def api_graph():
+        if client is None or client.graph_data is None:
+            raise HTTPException(status_code=503, detail="No graph data available")
+        return client.graph_data
+
+    @app.get("/api/trace")
+    async def api_trace(after: int = 0):
+        if client is None:
+            return []
+        return [t for t in client.trace_log if t["id"] > after]
+
+    @app.post("/api/trace/push")
+    async def trace_push(entry: dict):
+        if client is not None:
+            with client._lock:
+                client._trace_log.append(entry)
+        return {"status": "ok"}
+
+    @app.post("/api/graph/push")
+    async def graph_push(data: dict):
+        if client is not None:
+            client.graph_data = data
+        return {"status": "ok"}
+
+    app.mount("/static", StaticFiles(directory=_static_dir), name="static")
+
+    @app.get("/graph")
+    async def graph_page():
+        html_path = os.path.join(_static_dir, "graph.html")
+        return FileResponse(html_path)
+
+    return app
+
+
+def mount_app(client: Optional[VizClient] = None) -> FastAPI:
+    return _create_app(client)
+
+
+def _run_server(config, client: Optional[VizClient] = None):
+    import uvicorn
+
+    app = _create_app(client)
+
+    cfg = uvicorn.Config(
+        app,
+        host=config.host,
+        port=config.port,
+        log_level="warning",
+        loop="asyncio",
+    )
+    server = uvicorn.Server(cfg)
+
+    def _run():
+        server.run()
+
+    thread = threading.Thread(target=_run, daemon=True)
+    thread.start()
+    time.sleep(0.5)
+    return server