agentviz 0.3.0__py3-none-any.whl

agentviz/__init__.py

@@ -0,0 +1,53 @@
+"""
+AgentViz — Real-time 3D visualization for multi-agent AI systems.
+
+LangChain-style quick start
+----------------------------
+    import agentviz
+
+    # Option A: env vars (no code changes needed)
+    #   export AGENTVIZ_SERVER=https://agentviz.yourdomain.com
+    #   export AGENTVIZ_PROJECT=my-team
+
+    # Option B: one-line programmatic config
+    agentviz.init(server="https://agentviz.yourdomain.com", project="my-team")
+
+    # Decorate any function — async or sync
+    @agentviz.trace
+    async def fetch_data(query: str) -> str:
+        return await db.query(query)
+
+    @agentviz.trace(name="Planner", to="orchestrator", color="#9B59B6")
+    async def plan(goal: str) -> str:
+        return await llm.plan(goal)
+
+    # Auto-patch every installed framework at once
+    agentviz.instrument_all()
+
+Manual SDK (full control)
+--------------------------
+    from agentviz import AgentVizClient       # WebSocket — local / same-network
+    from agentviz import HttpAgentVizClient   # HTTP     — cloud / serverless
+"""
+
+# ── LangChain-style top-level API ─────────────────────────────────────────────
+from .config  import init, configure
+from .tracing import trace, Agent, instrument_all
+
+# ── Low-level SDK clients ─────────────────────────────────────────────────────
+from .sdk         import AgentVizClient
+from .http_client import HttpAgentVizClient
+
+__all__ = [
+    # LangChain-style
+    "init",
+    "configure",
+    "trace",
+    "Agent",
+    "instrument_all",
+    # SDK clients
+    "AgentVizClient",
+    "HttpAgentVizClient",
+]
+
+__version__ = "0.3.0"

agentviz/__main__.py

@@ -0,0 +1,4 @@
+# Allows: python -m agentviz serve
+from agentviz.cli import cli
+
+cli()

agentviz/_server_app.py

@@ -0,0 +1,322 @@
+"""
+Standalone AgentViz server.
+
+Used by `agentviz serve` (CLI) and can be imported directly:
+
+    from agentviz._server_app import create_app
+    app = create_app()
+
+Room isolation
+--------------
+Every WebSocket / HTTP endpoint accepts an optional ?room= query parameter.
+Agents and browser viewers in the same room share one 3D scene; different
+rooms are completely isolated from each other.
+
+This makes it safe to host one AgentViz instance on a public server (e.g.
+Linode) and have multiple teams / projects use it simultaneously.
+
+    # Team A
+    AgentVizClient(server="wss://agentviz.example.com/agent-ws?room=team-a")
+
+    # Team B — sees nothing from Team A
+    AgentVizClient(server="wss://agentviz.example.com/agent-ws?room=team-b")
+
+    # Browser
+    http://agentviz.example.com?room=team-a
+"""
+
+import asyncio
+import json
+import time
+from pathlib import Path
+from typing import Optional
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException, Query
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse, HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from agentviz.recorder import SessionRecorder
+
+STATIC_DIR = Path(__file__).parent / "static"
+
+_CORE_SRC   = Path(__file__).parent.parent / "packages" / "core" / "src"
+_CORE_BUILT = STATIC_DIR / "agentviz-core"
+CORE_DIR    = _CORE_SRC if _CORE_SRC.exists() else (_CORE_BUILT if _CORE_BUILT.exists() else None)
+
+DEFAULT_ROOM = "default"
+
+
+def create_app() -> FastAPI:
+    app = FastAPI(title="AgentViz", version="0.2.0")
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    # ── Per-room state ─────────────────────────────────────────────────────────
+
+    class Hub:
+        def __init__(self):
+            self._conns: list[WebSocket] = []
+
+        async def connect(self, ws: WebSocket):
+            await ws.accept()
+            self._conns.append(ws)
+
+        def disconnect(self, ws: WebSocket):
+            if ws in self._conns:
+                self._conns.remove(ws)
+
+        async def broadcast(self, data: dict):
+            dead = []
+            for ws in list(self._conns):
+                try:
+                    await ws.send_json(data)
+                except Exception:
+                    dead.append(ws)
+            for ws in dead:
+                self.disconnect(ws)
+
+        def __len__(self):
+            return len(self._conns)
+
+    class AgentRegistry:
+        def __init__(self):
+            self._agents: dict[str, dict] = {}
+
+        def register(self, info: dict):
+            aid = info["id"]
+            self._agents[aid] = {
+                "id":           aid,
+                "name":         info.get("name", aid),
+                "color":        info.get("color", "#4488ff"),
+                "role":         info.get("role", "worker"),
+                "capabilities": info.get("capabilities", []),
+                "connected_at": time.time(),
+            }
+
+        def deregister(self, aid: str):
+            self._agents.pop(aid, None)
+
+        def snapshot(self) -> list[dict]:
+            return list(self._agents.values())
+
+    class Room:
+        """All state for one isolated project / team room."""
+
+        def __init__(self, room_id: str):
+            self.id       = room_id
+            self.hub      = Hub()
+            self.registry = AgentRegistry()
+            self.recorder = SessionRecorder()
+            self._last_event_ts: float       = 0.0
+            self._end_task: Optional[asyncio.Task] = None
+
+        def touch_session(self) -> None:
+            self._last_event_ts = time.time()
+            if self.recorder._current_session_id is None:
+                self.recorder.start_session()
+            if self._end_task and not self._end_task.done():
+                self._end_task.cancel()
+            self._end_task = asyncio.ensure_future(self._idle_close())
+
+        async def _idle_close(self) -> None:
+            await asyncio.sleep(45)
+            if time.time() - self._last_event_ts >= 44:
+                self.recorder.end_session()
+
+    # ── Room manager ──────────────────────────────────────────────────────────
+
+    _rooms: dict[str, Room] = {}
+
+    def get_room(room_id: str) -> Room:
+        room_id = (room_id or DEFAULT_ROOM).strip() or DEFAULT_ROOM
+        if room_id not in _rooms:
+            _rooms[room_id] = Room(room_id)
+        return _rooms[room_id]
+
+    # ── WebSocket: browser viewer ─────────────────────────────────────────────
+
+    @app.websocket("/ws")
+    async def ws_frontend(ws: WebSocket, room: str = Query(default=DEFAULT_ROOM)):
+        r = get_room(room)
+        await r.hub.connect(ws)
+        if r.registry.snapshot():
+            await ws.send_json({"type": "registry_update", "agents": r.registry.snapshot()})
+        try:
+            while True:
+                await ws.receive_text()
+        except WebSocketDisconnect:
+            r.hub.disconnect(ws)
+
+    # ── WebSocket: SDK agents ─────────────────────────────────────────────────
+
+    @app.websocket("/agent-ws")
+    async def ws_agent(ws: WebSocket, room: str = Query(default=DEFAULT_ROOM)):
+        r = get_room(room)
+        await ws.accept()
+        agent_id: Optional[str] = None
+
+        try:
+            async for raw in ws.iter_text():
+                raw = raw.strip()
+                if not raw:
+                    continue
+                try:
+                    msg = json.loads(raw)
+                except json.JSONDecodeError:
+                    continue
+
+                etype = msg.get("type", "")
+
+                if etype == "register":
+                    agent_id = msg.get("id")
+                    if agent_id:
+                        r.registry.register(msg)
+                        await r.hub.broadcast({
+                            "type":   "registry_update",
+                            "agents": r.registry.snapshot(),
+                        })
+
+                elif etype in ("call", "response", "thinking", "error", "complete"):
+                    await r.hub.broadcast(msg)
+                    r.touch_session()
+                    r.recorder.record_event(msg)
+
+        except WebSocketDisconnect:
+            pass
+        finally:
+            if agent_id:
+                r.registry.deregister(agent_id)
+                await r.hub.broadcast({
+                    "type":   "registry_update",
+                    "agents": r.registry.snapshot(),
+                })
+
+    # ── HTTP ingest (serverless / Lambda / @trace decorator) ──────────────────
+
+    class IngestPayload(BaseModel):
+        events: list[dict]
+
+    @app.post("/ingest")
+    async def ingest(payload: IngestPayload, room: str = Query(default=DEFAULT_ROOM)):
+        """
+        Batch-ingest AgentViz events over plain HTTP.
+        Used by HttpAgentVizClient and the @trace decorator.
+
+        Pass ?room=<project> to route events to the right isolated scene.
+        """
+        r       = get_room(room)
+        ALLOWED = {"register", "call", "response", "thinking", "error", "complete", "token", "stream_end"}
+        accepted = 0
+
+        for msg in payload.events:
+            etype = msg.get("type", "")
+            if etype not in ALLOWED:
+                continue
+            if "timestamp" not in msg:
+                msg["timestamp"] = time.time()
+
+            if etype == "register":
+                agent_id = msg.get("id")
+                if agent_id:
+                    r.registry.register(msg)
+                    await r.hub.broadcast({
+                        "type":   "registry_update",
+                        "agents": r.registry.snapshot(),
+                    })
+
+            elif etype in ("call", "response", "thinking", "error", "complete", "token", "stream_end"):
+                await r.hub.broadcast(msg)
+                r.touch_session()
+                r.recorder.record_event(msg)
+
+            accepted += 1
+
+        return {"accepted": accepted, "total": len(payload.events), "room": room}
+
+    # ── Sessions ──────────────────────────────────────────────────────────────
+
+    @app.get("/sessions")
+    async def list_sessions(room: str = Query(default=DEFAULT_ROOM)):
+        return {"sessions": get_room(room).recorder.list_sessions(), "room": room}
+
+    @app.get("/sessions/{session_id}/events")
+    async def get_session_events(session_id: str, room: str = Query(default=DEFAULT_ROOM)):
+        return {"events": get_room(room).recorder.get_session_events(session_id)}
+
+    @app.post("/sessions/{session_id}/replay")
+    async def replay_session(session_id: str, room: str = Query(default=DEFAULT_ROOM)):
+        r      = get_room(room)
+        events = r.recorder.get_session_events(session_id)
+        if not events:
+            raise HTTPException(status_code=404, detail="Session not found or empty")
+        asyncio.create_task(_replay(events, r.hub))
+        return {"status": "replaying", "event_count": len(events)}
+
+    @app.delete("/sessions/{session_id}")
+    async def delete_session(session_id: str, room: str = Query(default=DEFAULT_ROOM)):
+        get_room(room).recorder.delete_session(session_id)
+        return {"status": "deleted"}
+
+    # ── Health ────────────────────────────────────────────────────────────────
+
+    @app.get("/health")
+    async def health():
+        return {
+            "status": "ok",
+            "rooms":  {
+                rid: {
+                    "agents":   len(r.registry.snapshot()),
+                    "browsers": len(r.hub),
+                }
+                for rid, r in _rooms.items()
+            },
+        }
+
+    @app.get("/rooms")
+    async def list_rooms():
+        return {
+            "rooms": [
+                {
+                    "id":      rid,
+                    "agents":  len(r.registry.snapshot()),
+                    "browsers": len(r.hub),
+                }
+                for rid, r in _rooms.items()
+            ]
+        }
+
+    # ── Static UI ─────────────────────────────────────────────────────────────
+
+    if CORE_DIR and CORE_DIR.exists():
+        app.mount("/static/agentviz-core", StaticFiles(directory=str(CORE_DIR)), name="core")
+
+    if STATIC_DIR.exists():
+        app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
+
+    @app.get("/")
+    async def index():
+        index_html = STATIC_DIR / "index.html"
+        if index_html.exists():
+            return FileResponse(str(index_html))
+        return HTMLResponse("<h1>AgentViz server running</h1><p>Static UI not found.</p>")
+
+    return app
+
+
+# ── Replay helper ──────────────────────────────────────────────────────────────
+
+async def _replay(events: list[dict], hub):
+    await hub.broadcast({"type": "replay_start", "event_count": len(events)})
+    for i, ev in enumerate(events):
+        if i > 0:
+            delay = min((ev["timestamp"] - events[i - 1]["timestamp"]), 3.0)
+            if delay > 0:
+                await asyncio.sleep(delay)
+        await hub.broadcast(ev["data"])
+    await hub.broadcast({"type": "replay_end"})

agentviz/cli.py

@@ -0,0 +1,116 @@
+"""
+agentviz CLI
+
+Usage:
+    agentviz serve [--host HOST] [--port PORT] [--no-browser]
+    agentviz --version
+    agentviz --help
+"""
+
+import argparse
+import sys
+import threading
+import time
+import webbrowser
+
+
+def cli():
+    parser = argparse.ArgumentParser(
+        prog="agentviz",
+        description="AgentViz — real-time 3D visualization for multi-agent AI systems",
+    )
+    parser.add_argument("--version", action="store_true", help="Print version and exit")
+
+    sub = parser.add_subparsers(dest="command")
+
+    serve_p = sub.add_parser("serve", help="Start the AgentViz visualization server")
+    serve_p.add_argument("--host", default="0.0.0.0", help="Bind host (default: 0.0.0.0)")
+    serve_p.add_argument("--port", type=int, default=8000, help="Bind port (default: 8000)")
+    serve_p.add_argument(
+        "--no-browser", action="store_true", help="Don't open browser automatically"
+    )
+    serve_p.add_argument(
+        "--demo", action="store_true",
+        help="Also start the demo orchestrator + agents on ports 8100-8103",
+    )
+
+    args = parser.parse_args()
+
+    if args.version:
+        from agentviz import __version__
+        print(f"agentviz {__version__}")
+        sys.exit(0)
+
+    if args.command == "serve":
+        _serve(args)
+    else:
+        parser.print_help()
+        sys.exit(1)
+
+
+def _serve(args):
+    import uvicorn
+    from agentviz._server_app import create_app
+
+    app = create_app()
+
+    url = f"http://localhost:{args.port}"
+
+    if not args.no_browser:
+        # Open browser after a short delay so the server has time to bind
+        def _open():
+            time.sleep(1.2)
+            webbrowser.open(url)
+
+        threading.Thread(target=_open, daemon=True).start()
+
+    print(f"\n  AgentViz  v{_version()}")
+    print(f"  Listening on  {url}")
+    print(f"\n  Connect your agents:")
+    print(f"    Python SDK  →  ws://localhost:{args.port}/agent-ws")
+    print(f"    JS SDK      →  ws://localhost:{args.port}/agent-ws")
+    print(f"    HTTP ingest →  POST http://localhost:{args.port}/ingest\n")
+
+    if args.demo:
+        _start_demo_processes()
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
+
+
+def _start_demo_processes():
+    """Launch the demo orchestrator + agents as subprocesses."""
+    import subprocess
+    from pathlib import Path
+
+    agents_dir = Path(__file__).parent.parent / "agents"
+    if not agents_dir.exists():
+        print("  [demo] agents/ directory not found — skipping demo agents")
+        return
+
+    scripts = [
+        ("orchestrator", agents_dir / "orchestrator.py"),
+        ("alpha",        agents_dir / "alpha_agent.py"),
+        ("beta",         agents_dir / "beta_agent.py"),
+        ("gamma",        agents_dir / "gamma_agent.py"),
+    ]
+
+    for name, script in scripts:
+        if script.exists():
+            subprocess.Popen(
+                [sys.executable, str(script)],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            print(f"  [demo] started {name}")
+
+
+def _version():
+    try:
+        from agentviz import __version__
+        return __version__
+    except Exception:
+        return "?"
+
+
+if __name__ == "__main__":
+    cli()

agentviz/config.py

@@ -0,0 +1,121 @@
+"""
+AgentViz global configuration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Reads from environment variables so no code changes are needed to enable
+visualization — just set the env vars before running your agents.
+
+Environment variables
+---------------------
+AGENTVIZ_SERVER   Base URL of the AgentViz server.
+                  Default: http://localhost:8000
+                  Example: https://agentviz.yourdomain.com
+
+AGENTVIZ_PROJECT  Project / room name.  Each project gets its own isolated
+                  3D scene.  Multiple agents in the same project share the
+                  same visualization.
+                  Default: "default"
+
+AGENTVIZ_ENABLED  Set to "false" to disable all telemetry without removing
+                  @trace decorators from your code.
+                  Default: "true"
+
+Usage
+-----
+    # Option A — env vars only (no code changes)
+    # export AGENTVIZ_SERVER=https://agentviz.yourdomain.com
+    # export AGENTVIZ_PROJECT=my-team
+
+    # Option B — programmatic (call before importing agents)
+    import agentviz
+    agentviz.init(server="https://agentviz.yourdomain.com", project="my-team")
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+
+# ── Defaults ─────────────────────────────────────────────────────────────────
+
+_DEFAULTS = {
+    "server":  "http://localhost:8000",
+    "project": "default",
+    "enabled": True,
+}
+
+# ── Live config dict (mutated by init / configure) ───────────────────────────
+
+_cfg: dict = {
+    "server":  os.getenv("AGENTVIZ_SERVER",  _DEFAULTS["server"]),
+    "project": os.getenv("AGENTVIZ_PROJECT", _DEFAULTS["project"]),
+    "enabled": os.getenv("AGENTVIZ_ENABLED", "true").strip().lower() != "false",
+}
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def init(
+    server:  Optional[str]  = None,
+    project: Optional[str]  = None,
+    enabled: Optional[bool] = None,
+) -> None:
+    """
+    Configure AgentViz once at the top of your program.
+
+    Parameters
+    ----------
+    server   Base URL of the AgentViz server.
+             e.g. "https://agentviz.yourdomain.com"
+    project  Project / room name.  Agents with the same project appear in
+             the same 3D scene.  Agents in different projects are isolated.
+    enabled  Set False to silence all telemetry without removing decorators.
+    """
+    if server  is not None: _cfg["server"]  = server.rstrip("/")
+    if project is not None: _cfg["project"] = project
+    if enabled is not None: _cfg["enabled"] = enabled
+
+
+# configure is an alias for init — same ergonomics as langsmith.configure()
+configure = init
+
+
+def get() -> dict:
+    """Return a snapshot of the current configuration."""
+    return dict(_cfg)
+
+
+def is_enabled() -> bool:
+    return bool(_cfg["enabled"])
+
+
+def server_url() -> str:
+    return _cfg["server"]
+
+
+def project() -> str:
+    return _cfg["project"]
+
+
+def ws_url() -> str:
+    """WebSocket URL derived from the base server URL."""
+    base = _cfg["server"]
+    # Normalise scheme
+    if base.startswith("https://"):
+        ws_base = "wss://" + base[len("https://"):]
+    elif base.startswith("http://"):
+        ws_base = "ws://" + base[len("http://"):]
+    else:
+        ws_base = base  # already ws:// or wss://
+    return f"{ws_base}/agent-ws?room={_cfg['project']}"
+
+
+def http_ingest_url() -> str:
+    """HTTP ingest URL derived from the base server URL."""
+    base = _cfg["server"]
+    # Normalise to http(s)
+    if base.startswith("wss://"):
+        base = "https://" + base[len("wss://"):]
+    elif base.startswith("ws://"):
+        base = "http://" + base[len("ws://"):]
+    return f"{base}/ingest?room={_cfg['project']}"