initial release: deploy code + split scripts

0e41b61 verified 7 months ago

26.1 kB

#!/usr/bin/env python3
"""
mac-tensor ui — Web chat UI for the distributed agent.
Serves a single-page HTML chat interface and a Server-Sent Events
endpoint that streams agent events (steps, tool calls, results, final answer).
Usage:
    mac-tensor ui --model gemma4 --nodes http://mac2:8401,http://mac3:8401
    # Then open http://localhost:8500 in your browser
"""

import json
import os
import sys
import time
import threading
from queue import Queue, Empty


def run_server(model_key, node_urls=None, host="0.0.0.0", port=8500, allow_write=False,
               vision=False, stream_dir=None, source_dir=None,
               falcon=False, falcon_model=None,
               swarm_leader=False,
               turbo_url=None):
    """Start the FastAPI server with the agent backend pre-loaded.
    Modes:
      - Distributed text-only: pass node_urls
      - Single-machine vision (Gemma 4 only): vision=True
      - Vision + Falcon Perception (segmentation): vision=True, falcon=True
      - Swarm leader: swarm_leader=True (peer registry + dynamic coordinator)
    """
    from fastapi import FastAPI, Request, UploadFile, File, Form
    from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
    from .agent import AgentBackend, run_agent_turn_stream

    vision_engine = None
    falcon_tools = None
    swarm_registry = None

    if swarm_leader:
        from .swarm import SwarmRegistry, reaper_loop
        swarm_registry = SwarmRegistry(model_key=model_key)
        threading.Thread(target=reaper_loop, args=(swarm_registry,), daemon=True).start()
        print(f"Swarm registry started (model={model_key})")

    if swarm_leader:
        # Leader mode: no LLM backend loaded yet — peers register dynamically.
        # The leader can ALSO load a backend on demand (later) when peers exist.
        backend = None
        print("Running as swarm leader. Workers join via mac-tensor join.")
    elif vision:
        print(f"Loading vision Gemma 4 sniper (single-machine)...")
        from .vision_engine import VisionGemma4Sniper
        vision_engine = VisionGemma4Sniper(
            stream_dir=stream_dir or "~/models/gemma4-stream",
            source_dir=source_dir or "~/models/gemma4-26b-4bit",
        )
        vision_engine.load()
        print("Vision engine ready.")

        if falcon:
            print(f"Loading Falcon Perception...")
            from .falcon_perception import FalconPerceptionTools
            falcon_tools = FalconPerceptionTools.load(
                model_path=falcon_model or "/Users/bigneek/models/falcon-perception"
            )
            print("Falcon Perception ready.")

        backend = None  # Not used in vision mode
    elif falcon:
        # Falcon-only mode: no Gemma, no distributed nodes. Used by the
        # data labeling factory where we only need /api/falcon.
        print("Loading Falcon Perception (standalone, no Gemma)...")
        from .falcon_perception import FalconPerceptionTools
        falcon_tools = FalconPerceptionTools.load(
            model_path=falcon_model or "/Users/bigneek/models/falcon-perception"
        )
        print("Falcon Perception ready. (~1.5 GB resident, no Gemma loaded)")
        backend = None
    else:
        print(f"Loading {model_key} distributed engine...")
        backend = AgentBackend(model_key=model_key, node_urls=node_urls)
        backend.load()
        print(f"Backend ready. Connected to {len(node_urls)} expert nodes.")

    app = FastAPI(title="mac-tensor agent UI")

    # Read the static HTML file shipped alongside this server
    static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
    html_path = os.path.join(static_dir, "chat.html")
    with open(html_path) as f:
        chat_html = f.read()

    # Inject backend info into the HTML so the UI can show it
    if vision:
        model_label = "Gemma 4-26B-A4B (Vision)"
        node_count_label = "single Mac · vision enabled"
    elif falcon_tools is not None and vision_engine is None:
        model_label = "Falcon Perception (labeling factory)"
        node_count_label = "single Mac · Falcon-only"
    elif swarm_leader:
        model_label = {"gemma4": "Gemma 4-26B-A4B",
                       "qwen35": "Qwen 3.5-35B-A3B"}.get(model_key, model_key)
        node_count_label = "swarm leader · waiting for peers"
    else:
        model_label = {"gemma4": "Gemma 4-26B-A4B",
                       "qwen35": "Qwen 3.5-35B-A3B"}.get(model_key, model_key)
        node_count_label = f"{len(node_urls)} expert nodes"

    chat_html = chat_html.replace("{{MODEL_NAME}}", model_label) \
                          .replace("{{NODE_COUNT}}", node_count_label) \
                          .replace("{{VISION_ENABLED}}", "true" if vision else "false") \
                          .replace("{{FALCON_ENABLED}}", "true" if falcon_tools is not None else "false")

    # Lock so only one chat request runs at a time (single MoE engine)
    lock = threading.Lock()

    @app.get("/")
    async def index():
        return HTMLResponse(chat_html)

    @app.get("/api/info")
    async def info():
        return {
            "model": model_key,
            "nodes": node_urls,
            "allow_write": allow_write,
            "vision": vision,
            "falcon": falcon_tools is not None,
            "swarm_leader": swarm_registry is not None,
        }

    # ============================================================
    # Swarm endpoints (only when running as leader)
    # ============================================================

    if swarm_registry is not None:

        @app.post("/swarm/register")
        async def swarm_register(request: Request):
            body = await request.json()
            url = body.get("url")
            mem_gb = body.get("mem_gb", 0)
            meta = body.get("meta", {})
            if not url:
                return JSONResponse({"error": "url required"}, status_code=400)
            peer_id, partition = swarm_registry.register(url, mem_gb, meta)
            print(f"[swarm] +peer {peer_id} at {url} → partition {partition}")
            return {
                "peer_id": peer_id,
                "partition": partition,
                "model": model_key,
                "partition_version": swarm_registry.partition_version,
            }

        @app.post("/swarm/heartbeat")
        async def swarm_heartbeat(request: Request):
            body = await request.json()
            peer_id = body.get("peer_id")
            if not peer_id:
                return JSONResponse({"error": "peer_id required"}, status_code=400)
            ok, version = swarm_registry.heartbeat(peer_id)
            if not ok:
                return JSONResponse({"error": "unknown peer"}, status_code=404)

            # Tell the peer if its partition has been reassigned
            current_partition = None
            with swarm_registry.lock:
                if peer_id in swarm_registry.peers:
                    current_partition = swarm_registry.peers[peer_id]["partition"]

            return {
                "ok": True,
                "partition_version": version,
                "partition": current_partition,
            }

        @app.post("/swarm/leave")
        async def swarm_leave(request: Request):
            body = await request.json()
            peer_id = body.get("peer_id")
            if not peer_id:
                return JSONResponse({"error": "peer_id required"}, status_code=400)
            swarm_registry.leave(peer_id)
            print(f"[swarm] -peer {peer_id} (graceful leave)")
            return {"ok": True}

        @app.get("/swarm/peers")
        async def swarm_peers():
            return swarm_registry.status()

    @app.post("/api/reset")
    async def reset():
        with lock:
            if vision_engine:
                vision_engine.sniper.reset_cache()
            elif backend:
                backend.reset()
        return {"ok": True}

    # When running as leader, lazily build a backend from the swarm registry.
    # We cache it and rebuild when the partition_version changes.
    leader_backend_state = {"backend": None, "version": -1}

    def get_swarm_backend():
        """Lazy backend that uses live peers from the swarm registry.
        Rebuilds when partition_version changes. Returns None if no live peers.
        """
        if swarm_registry is None:
            return backend  # static mode

        live_peers = [p for p in swarm_registry.get_live_peers() if p.get("alive")]
        if not live_peers:
            return None

        # Check if registry version changed → rebuild backend
        current_version = swarm_registry.partition_version
        if leader_backend_state["version"] != current_version or leader_backend_state["backend"] is None:
            print(f"[swarm] building backend for {len(live_peers)} peers (v{current_version})")
            from .agent import AgentBackend
            node_urls = [p["url"] for p in live_peers]
            try:
                bk = AgentBackend(model_key=model_key, node_urls=node_urls)
                bk.load()
                leader_backend_state["backend"] = bk
                leader_backend_state["version"] = current_version
            except Exception as e:
                print(f"[swarm] backend load failed: {e}")
                return None

        return leader_backend_state["backend"]

    @app.post("/api/chat")
    async def chat(request: Request):
        # Use the static backend if not in swarm mode, else build from registry
        active = get_swarm_backend() if swarm_registry is not None else backend

        if active is None:
            return JSONResponse({
                "error": "no live peers in swarm yet — run `mac-tensor join "
                         f"http://{_local_ip()}:{port}` on a worker Mac first"
            }, status_code=503)

        body = await request.json()
        message = body.get("message", "").strip()
        if not message:
            return JSONResponse({"error": "empty message"}, status_code=400)

        max_iterations = int(body.get("max_iterations", 5))
        max_tokens = int(body.get("max_tokens", 300))

        def event_stream():
            with lock:
                try:
                    for event in run_agent_turn_stream(
                        active, message,
                        max_iterations=max_iterations,
                        max_tokens=max_tokens,
                        allow_write=allow_write,
                    ):
                        yield f"data: {json.dumps(event)}\n\n"
                except Exception as e:
                    yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
                     "Connection": "keep-alive"},
        )

    # Build a vision agent backend if Falcon is loaded
    vision_agent = None
    if vision_engine is not None and falcon_tools is not None:
        from .agent import VisionAgentBackend
        vision_agent = VisionAgentBackend(
            vision_engine=vision_engine, falcon_tools=falcon_tools
        )
        print("Vision agent ready (Gemma 4 + Falcon Perception chained).")

    @app.post("/api/chat_vision")
    async def chat_vision(
        message: str = Form(...),
        max_tokens: int = Form(300),
        image: UploadFile = File(None),
    ):
        """Vision chat endpoint — accepts an optional image upload.
        If Falcon Perception is loaded, uses the vision agent loop with
        tool calling. Otherwise falls back to plain Gemma 4 vision.
        """
        if vision_engine is None:
            return JSONResponse({"error": "vision mode not enabled"}, status_code=400)

        image_path = None
        if image is not None and image.filename:
            import tempfile
            tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
            tmp.write(await image.read())
            tmp.close()
            image_path = tmp.name

        def event_stream():
            with lock:
                try:
                    if vision_agent is not None and image_path:
                        # Chained mode: Gemma 4 + Falcon tool calls
                        from .agent import run_vision_agent_turn_stream
                        for event in run_vision_agent_turn_stream(
                            vision_agent, message, image_path,
                            max_iterations=4,
                            max_tokens=max_tokens,
                        ):
                            yield f"data: {json.dumps(event)}\n\n"
                    else:
                        # Simple mode: just Gemma 4 vision (no tools)
                        yield f"data: {json.dumps({'type': 'step_start', 'step': 1, 'max': 1})}\n\n"
                        chunks = []
                        def on_chunk(text):
                            chunks.append(text)
                        output = vision_engine.generate(
                            message, image_path=image_path,
                            max_tokens=max_tokens, temperature=0.6,
                            on_chunk=on_chunk,
                        )
                        for chunk in chunks:
                            yield f"data: {json.dumps({'type': 'token', 'text': chunk})}\n\n"
                        yield f"data: {json.dumps({'type': 'final', 'text': output.strip()})}\n\n"
                        yield f"data: {json.dumps({'type': 'done'})}\n\n"
                except Exception as e:
                    import traceback
                    traceback.print_exc()
                    yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
                finally:
                    if image_path and os.path.exists(image_path):
                        try:
                            os.unlink(image_path)
                        except Exception:
                            pass

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
                     "Connection": "keep-alive"},
        )

    @app.post("/api/turbo_chat")
    async def turbo_chat(
        message: str = Form(...),
        max_tokens: int = Form(300),
        image: UploadFile = File(None),
    ):
        """Turbo mode: Gemma 4 encodes image once, then a fast small LLM
        (Qwen3-1.7B 4-bit via mlx-lm) handles the reasoning loop.
        Streams SSE events compatible with the existing chat UI.
        """
        if vision_engine is None:
            return JSONResponse({"error": "vision mode not enabled"}, status_code=400)
        if not turbo_url:
            return JSONResponse(
                {"error": "turbo brain not configured (start server with --turbo-url)"},
                status_code=400,
            )

        image_path = None
        if image is not None and image.filename:
            import tempfile
            tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
            tmp.write(await image.read())
            tmp.close()
            image_path = tmp.name

        def event_stream():
            with lock:
                try:
                    # Step 1: Gemma 4 vision encodes the image ONCE
                    yield f"data: {json.dumps({'type': 'step_start', 'step': 1, 'max': 2})}\n\n"
                    yield f"data: {json.dumps({'type': 'token', 'text': '🔍 Encoding image with Gemma 4 vision...\\n'})}\n\n"

                    description = ""
                    if image_path:
                        try:
                            description = vision_engine.generate(
                                "Describe this image briefly: what's in it, where is it, what colors and notable details. Be concise (3-4 sentences max).",
                                image_path=image_path,
                                max_tokens=150,
                                temperature=0.5,
                            ).strip()
                            yield f"data: {json.dumps({'type': 'tool_call', 'tool': 'vision_describe', 'args': 'image'})}\n\n"
                            yield f"data: {json.dumps({'type': 'tool_result', 'result': description})}\n\n"
                        except Exception as e:
                            yield f"data: {json.dumps({'type': 'error', 'message': f'vision encode failed: {e}'})}\n\n"
                            return

                    # Step 2: fast turbo brain reasons over the description
                    yield f"data: {json.dumps({'type': 'step_start', 'step': 2, 'max': 2})}\n\n"

                    # Build the prompt for the turbo brain
                    if description:
                        system_msg = (
                            "You are a fast assistant that answers questions about images. "
                            "A vision model has already described the image for you. "
                            "Give a SHORT, direct answer (1-3 sentences). Do not show your reasoning."
                        )
                        user_msg = (
                            f"Image description: {description}\n\n"
                            f"User's question: {message}"
                        )
                    else:
                        system_msg = "You are a fast helpful assistant. Give SHORT, direct answers."
                        user_msg = message

                    # Stream from the turbo brain (mlx-lm OpenAI-compatible endpoint)
                    import urllib.request
                    payload = {
                        "model": "qwen3-1.7b",
                        "messages": [
                            {"role": "system", "content": system_msg},
                            {"role": "user", "content": user_msg},
                        ],
                        "max_tokens": max_tokens,
                        "temperature": 0.5,
                        "stream": True,
                    }
                    req = urllib.request.Request(
                        f"{turbo_url}/v1/chat/completions",
                        data=json.dumps(payload).encode(),
                        headers={"Content-Type": "application/json"},
                    )
                    final_text = ""
                    in_think = False
                    try:
                        with urllib.request.urlopen(req, timeout=120) as resp:
                            buffer = ""
                            for chunk in iter(lambda: resp.read(512), b""):
                                buffer += chunk.decode("utf-8", errors="ignore")
                                while "\n" in buffer:
                                    line, buffer = buffer.split("\n", 1)
                                    line = line.strip()
                                    if not line.startswith("data: "):
                                        continue
                                    data = line[6:]
                                    if data == "[DONE]":
                                        break
                                    try:
                                        ev = json.loads(data)
                                        delta = ev.get("choices", [{}])[0].get("delta", {})
                                        text = delta.get("content", "")
                                        if not text:
                                            continue
                                        # Strip out <think>...</think> reasoning blocks
                                        for ch in text:
                                            if "<think>" in (final_text + ch)[-7:]:
                                                in_think = True
                                                continue
                                            if "</think>" in (final_text + ch)[-8:]:
                                                in_think = False
                                                final_text = ""
                                                continue
                                            if not in_think:
                                                final_text += ch
                                                yield f"data: {json.dumps({'type': 'token', 'text': ch})}\n\n"
                                    except Exception:
                                        continue
                    except Exception as e:
                        yield f"data: {json.dumps({'type': 'error', 'message': f'turbo brain failed: {e}'})}\n\n"
                        return

                    yield f"data: {json.dumps({'type': 'final', 'text': final_text.strip()})}\n\n"
                    yield f"data: {json.dumps({'type': 'done'})}\n\n"
                except Exception as e:
                    import traceback
                    traceback.print_exc()
                    yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
                finally:
                    if image_path and os.path.exists(image_path):
                        try:
                            os.unlink(image_path)
                        except Exception:
                            pass

        return StreamingResponse(
            event_stream(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
                     "Connection": "keep-alive"},
        )

    @app.post("/api/falcon")
    async def falcon_ground(
        query: str = Form(...),
        image: UploadFile = File(...),
    ):
        """Run Falcon Perception on uploaded image with text query.
        Returns JSON with detected masks (count, IDs, metadata) plus a
        base64-encoded annotated image showing bounding boxes + labels.
        """
        if falcon_tools is None:
            return JSONResponse({"error": "Falcon Perception not loaded"}, status_code=400)

        # Save uploaded image
        import tempfile, base64, io
        tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
        tmp.write(await image.read())
        tmp.close()

        try:
            with lock:
                # Set image in Falcon session
                falcon_tools.set_image(tmp.name)
                # Run grounding
                t0 = time.time()
                result = falcon_tools.ground(query, slot=query.replace(" ", "_")[:32])
                elapsed = time.time() - t0

                if "error" in result:
                    return JSONResponse(result, status_code=500)

                # Annotate the image with bounding boxes
                annotated = falcon_tools.annotate_image(mask_ids=result["mask_ids"])

                # Encode annotated image as base64 PNG
                buf = io.BytesIO()
                annotated.save(buf, format="PNG")
                annotated_b64 = base64.b64encode(buf.getvalue()).decode()

                return JSONResponse({
                    "query": query,
                    "count": result["count"],
                    "mask_ids": result["mask_ids"],
                    "masks": result["masks"],
                    "annotated_image": f"data:image/png;base64,{annotated_b64}",
                    "elapsed_seconds": round(elapsed, 2),
                })
        except Exception as e:
            import traceback
            traceback.print_exc()
            return JSONResponse({"error": str(e)}, status_code=500)
        finally:
            try:
                os.unlink(tmp.name)
            except Exception:
                pass

    print()
    print("=" * 60)
    print(f"  mac-tensor UI ready")
    print(f"  Open: http://localhost:{port}")
    print(f"        http://{_local_ip()}:{port}  (LAN access)")
    print("=" * 60)
    print()

    import uvicorn
    uvicorn.run(app, host=host, port=port, log_level="warning")


def _local_ip():
    """Best-effort detection of the LAN IP."""
    import socket
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except Exception:
        return "localhost"


def main(args):
    vision = getattr(args, "vision", False)
    falcon_only = getattr(args, "falcon_only", False)

    if falcon_only:
        # Falcon-only mode: load Falcon Perception, skip Gemma entirely.
        # ~1.5 GB resident. Used by the data labeling factory.
        run_server(
            model_key="falcon",
            node_urls=None,
            host=args.host or "0.0.0.0",
            port=args.port or 8500,
            allow_write=False,
            vision=False,
            falcon=True,
            falcon_model=getattr(args, "falcon_model", None),
        )
        return

    if vision:
        # Vision mode: single-machine, no distributed nodes needed
        run_server(
            model_key="gemma4",
            node_urls=None,
            host=args.host or "0.0.0.0",
            port=args.port or 8500,
            allow_write=getattr(args, "write", False),
            vision=True,
            stream_dir=getattr(args, "stream_dir", None),
            source_dir=getattr(args, "source_dir", None),
            falcon=getattr(args, "falcon", False),
            falcon_model=getattr(args, "falcon_model", None),
        )
    else:
        if not args.nodes:
            print("Error: --nodes is required (or pass --vision for single-machine mode)")
            sys.exit(1)
        node_urls = [u.strip() for u in args.nodes.split(",")]
        run_server(
            model_key=args.model or "gemma4",
            node_urls=node_urls,
            host=args.host or "0.0.0.0",
            port=args.port or 8500,
            allow_write=args.write,
        )

	#!/usr/bin/env python3
	"""
	mac-tensor ui — Web chat UI for the distributed agent.

	Serves a single-page HTML chat interface and a Server-Sent Events
	endpoint that streams agent events (steps, tool calls, results, final answer).

	Usage:
	mac-tensor ui --model gemma4 --nodes http://mac2:8401,http://mac3:8401
	# Then open http://localhost:8500 in your browser
	"""

	import json
	import os
	import sys
	import time
	import threading
	from queue import Queue, Empty


	def run_server(model_key, node_urls=None, host="0.0.0.0", port=8500, allow_write=False,
	vision=False, stream_dir=None, source_dir=None,
	falcon=False, falcon_model=None,
	swarm_leader=False,
	turbo_url=None):
	"""Start the FastAPI server with the agent backend pre-loaded.

	Modes:
	- Distributed text-only: pass node_urls
	- Single-machine vision (Gemma 4 only): vision=True
	- Vision + Falcon Perception (segmentation): vision=True, falcon=True
	- Swarm leader: swarm_leader=True (peer registry + dynamic coordinator)
	"""
	from fastapi import FastAPI, Request, UploadFile, File, Form
	from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
	from .agent import AgentBackend, run_agent_turn_stream

	vision_engine = None
	falcon_tools = None
	swarm_registry = None

	if swarm_leader:
	from .swarm import SwarmRegistry, reaper_loop
	swarm_registry = SwarmRegistry(model_key=model_key)
	threading.Thread(target=reaper_loop, args=(swarm_registry,), daemon=True).start()
	print(f"Swarm registry started (model={model_key})")

	if swarm_leader:
	# Leader mode: no LLM backend loaded yet — peers register dynamically.
	# The leader can ALSO load a backend on demand (later) when peers exist.
	backend = None
	print("Running as swarm leader. Workers join via mac-tensor join.")
	elif vision:
	print(f"Loading vision Gemma 4 sniper (single-machine)...")
	from .vision_engine import VisionGemma4Sniper
	vision_engine = VisionGemma4Sniper(
	stream_dir=stream_dir or "~/models/gemma4-stream",
	source_dir=source_dir or "~/models/gemma4-26b-4bit",
	)
	vision_engine.load()
	print("Vision engine ready.")

	if falcon:
	print(f"Loading Falcon Perception...")
	from .falcon_perception import FalconPerceptionTools
	falcon_tools = FalconPerceptionTools.load(
	model_path=falcon_model or "/Users/bigneek/models/falcon-perception"
	)
	print("Falcon Perception ready.")

	backend = None # Not used in vision mode
	elif falcon:
	# Falcon-only mode: no Gemma, no distributed nodes. Used by the
	# data labeling factory where we only need /api/falcon.
	print("Loading Falcon Perception (standalone, no Gemma)...")
	from .falcon_perception import FalconPerceptionTools
	falcon_tools = FalconPerceptionTools.load(
	model_path=falcon_model or "/Users/bigneek/models/falcon-perception"
	)
	print("Falcon Perception ready. (~1.5 GB resident, no Gemma loaded)")
	backend = None
	else:
	print(f"Loading {model_key} distributed engine...")
	backend = AgentBackend(model_key=model_key, node_urls=node_urls)
	backend.load()
	print(f"Backend ready. Connected to {len(node_urls)} expert nodes.")

	app = FastAPI(title="mac-tensor agent UI")

	# Read the static HTML file shipped alongside this server
	static_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")
	html_path = os.path.join(static_dir, "chat.html")
	with open(html_path) as f:
	chat_html = f.read()

	# Inject backend info into the HTML so the UI can show it
	if vision:
	model_label = "Gemma 4-26B-A4B (Vision)"
	node_count_label = "single Mac · vision enabled"
	elif falcon_tools is not None and vision_engine is None:
	model_label = "Falcon Perception (labeling factory)"
	node_count_label = "single Mac · Falcon-only"
	elif swarm_leader:
	model_label = {"gemma4": "Gemma 4-26B-A4B",
	"qwen35": "Qwen 3.5-35B-A3B"}.get(model_key, model_key)
	node_count_label = "swarm leader · waiting for peers"
	else:
	model_label = {"gemma4": "Gemma 4-26B-A4B",
	"qwen35": "Qwen 3.5-35B-A3B"}.get(model_key, model_key)
	node_count_label = f"{len(node_urls)} expert nodes"

	chat_html = chat_html.replace("{{MODEL_NAME}}", model_label) \
	.replace("{{NODE_COUNT}}", node_count_label) \
	.replace("{{VISION_ENABLED}}", "true" if vision else "false") \
	.replace("{{FALCON_ENABLED}}", "true" if falcon_tools is not None else "false")

	# Lock so only one chat request runs at a time (single MoE engine)
	lock = threading.Lock()

	@app.get("/")
	async def index():
	return HTMLResponse(chat_html)

	@app.get("/api/info")
	async def info():
	return {
	"model": model_key,
	"nodes": node_urls,
	"allow_write": allow_write,
	"vision": vision,
	"falcon": falcon_tools is not None,
	"swarm_leader": swarm_registry is not None,
	}

	# ============================================================
	# Swarm endpoints (only when running as leader)
	# ============================================================

	if swarm_registry is not None:

	@app.post("/swarm/register")
	async def swarm_register(request: Request):
	body = await request.json()
	url = body.get("url")
	mem_gb = body.get("mem_gb", 0)
	meta = body.get("meta", {})
	if not url:
	return JSONResponse({"error": "url required"}, status_code=400)
	peer_id, partition = swarm_registry.register(url, mem_gb, meta)
	print(f"[swarm] +peer {peer_id} at {url} → partition {partition}")
	return {
	"peer_id": peer_id,
	"partition": partition,
	"model": model_key,
	"partition_version": swarm_registry.partition_version,
	}

	@app.post("/swarm/heartbeat")
	async def swarm_heartbeat(request: Request):
	body = await request.json()
	peer_id = body.get("peer_id")
	if not peer_id:
	return JSONResponse({"error": "peer_id required"}, status_code=400)
	ok, version = swarm_registry.heartbeat(peer_id)
	if not ok:
	return JSONResponse({"error": "unknown peer"}, status_code=404)

	# Tell the peer if its partition has been reassigned
	current_partition = None
	with swarm_registry.lock:
	if peer_id in swarm_registry.peers:
	current_partition = swarm_registry.peers[peer_id]["partition"]

	return {
	"ok": True,
	"partition_version": version,
	"partition": current_partition,
	}

	@app.post("/swarm/leave")
	async def swarm_leave(request: Request):
	body = await request.json()
	peer_id = body.get("peer_id")
	if not peer_id:
	return JSONResponse({"error": "peer_id required"}, status_code=400)
	swarm_registry.leave(peer_id)
	print(f"[swarm] -peer {peer_id} (graceful leave)")
	return {"ok": True}

	@app.get("/swarm/peers")
	async def swarm_peers():
	return swarm_registry.status()

	@app.post("/api/reset")
	async def reset():
	with lock:
	if vision_engine:
	vision_engine.sniper.reset_cache()
	elif backend:
	backend.reset()
	return {"ok": True}

	# When running as leader, lazily build a backend from the swarm registry.
	# We cache it and rebuild when the partition_version changes.
	leader_backend_state = {"backend": None, "version": -1}

	def get_swarm_backend():
	"""Lazy backend that uses live peers from the swarm registry.

	Rebuilds when partition_version changes. Returns None if no live peers.
	"""
	if swarm_registry is None:
	return backend # static mode

	live_peers = [p for p in swarm_registry.get_live_peers() if p.get("alive")]
	if not live_peers:
	return None

	# Check if registry version changed → rebuild backend
	current_version = swarm_registry.partition_version
	if leader_backend_state["version"] != current_version or leader_backend_state["backend"] is None:
	print(f"[swarm] building backend for {len(live_peers)} peers (v{current_version})")
	from .agent import AgentBackend
	node_urls = [p["url"] for p in live_peers]
	try:
	bk = AgentBackend(model_key=model_key, node_urls=node_urls)
	bk.load()
	leader_backend_state["backend"] = bk
	leader_backend_state["version"] = current_version
	except Exception as e:
	print(f"[swarm] backend load failed: {e}")
	return None

	return leader_backend_state["backend"]

	@app.post("/api/chat")
	async def chat(request: Request):
	# Use the static backend if not in swarm mode, else build from registry
	active = get_swarm_backend() if swarm_registry is not None else backend

	if active is None:
	return JSONResponse({
	"error": "no live peers in swarm yet — run `mac-tensor join "
	f"http://{_local_ip()}:{port}` on a worker Mac first"
	}, status_code=503)

	body = await request.json()
	message = body.get("message", "").strip()
	if not message:
	return JSONResponse({"error": "empty message"}, status_code=400)

	max_iterations = int(body.get("max_iterations", 5))
	max_tokens = int(body.get("max_tokens", 300))

	def event_stream():
	with lock:
	try:
	for event in run_agent_turn_stream(
	active, message,
	max_iterations=max_iterations,
	max_tokens=max_tokens,
	allow_write=allow_write,
	):
	yield f"data: {json.dumps(event)}\n\n"
	except Exception as e:
	yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"

	return StreamingResponse(
	event_stream(),
	media_type="text/event-stream",
	headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
	"Connection": "keep-alive"},
	)

	# Build a vision agent backend if Falcon is loaded
	vision_agent = None
	if vision_engine is not None and falcon_tools is not None:
	from .agent import VisionAgentBackend
	vision_agent = VisionAgentBackend(
	vision_engine=vision_engine, falcon_tools=falcon_tools
	)
	print("Vision agent ready (Gemma 4 + Falcon Perception chained).")

	@app.post("/api/chat_vision")
	async def chat_vision(
	message: str = Form(...),
	max_tokens: int = Form(300),
	image: UploadFile = File(None),
	):
	"""Vision chat endpoint — accepts an optional image upload.

	If Falcon Perception is loaded, uses the vision agent loop with
	tool calling. Otherwise falls back to plain Gemma 4 vision.
	"""
	if vision_engine is None:
	return JSONResponse({"error": "vision mode not enabled"}, status_code=400)

	image_path = None
	if image is not None and image.filename:
	import tempfile
	tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
	tmp.write(await image.read())
	tmp.close()
	image_path = tmp.name

	def event_stream():
	with lock:
	try:
	if vision_agent is not None and image_path:
	# Chained mode: Gemma 4 + Falcon tool calls
	from .agent import run_vision_agent_turn_stream
	for event in run_vision_agent_turn_stream(
	vision_agent, message, image_path,
	max_iterations=4,
	max_tokens=max_tokens,
	):
	yield f"data: {json.dumps(event)}\n\n"
	else:
	# Simple mode: just Gemma 4 vision (no tools)
	yield f"data: {json.dumps({'type': 'step_start', 'step': 1, 'max': 1})}\n\n"
	chunks = []
	def on_chunk(text):
	chunks.append(text)
	output = vision_engine.generate(
	message, image_path=image_path,
	max_tokens=max_tokens, temperature=0.6,
	on_chunk=on_chunk,
	)
	for chunk in chunks:
	yield f"data: {json.dumps({'type': 'token', 'text': chunk})}\n\n"
	yield f"data: {json.dumps({'type': 'final', 'text': output.strip()})}\n\n"
	yield f"data: {json.dumps({'type': 'done'})}\n\n"
	except Exception as e:
	import traceback
	traceback.print_exc()
	yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
	finally:
	if image_path and os.path.exists(image_path):
	try:
	os.unlink(image_path)
	except Exception:
	pass

	return StreamingResponse(
	event_stream(),
	media_type="text/event-stream",
	headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
	"Connection": "keep-alive"},
	)

	@app.post("/api/turbo_chat")
	async def turbo_chat(
	message: str = Form(...),
	max_tokens: int = Form(300),
	image: UploadFile = File(None),
	):
	"""Turbo mode: Gemma 4 encodes image once, then a fast small LLM
	(Qwen3-1.7B 4-bit via mlx-lm) handles the reasoning loop.

	Streams SSE events compatible with the existing chat UI.
	"""
	if vision_engine is None:
	return JSONResponse({"error": "vision mode not enabled"}, status_code=400)
	if not turbo_url:
	return JSONResponse(
	{"error": "turbo brain not configured (start server with --turbo-url)"},
	status_code=400,
	)

	image_path = None
	if image is not None and image.filename:
	import tempfile
	tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
	tmp.write(await image.read())
	tmp.close()
	image_path = tmp.name

	def event_stream():
	with lock:
	try:
	# Step 1: Gemma 4 vision encodes the image ONCE
	yield f"data: {json.dumps({'type': 'step_start', 'step': 1, 'max': 2})}\n\n"
	yield f"data: {json.dumps({'type': 'token', 'text': '🔍 Encoding image with Gemma 4 vision...\\n'})}\n\n"

	description = ""
	if image_path:
	try:
	description = vision_engine.generate(
	"Describe this image briefly: what's in it, where is it, what colors and notable details. Be concise (3-4 sentences max).",
	image_path=image_path,
	max_tokens=150,
	temperature=0.5,
	).strip()
	yield f"data: {json.dumps({'type': 'tool_call', 'tool': 'vision_describe', 'args': 'image'})}\n\n"
	yield f"data: {json.dumps({'type': 'tool_result', 'result': description})}\n\n"
	except Exception as e:
	yield f"data: {json.dumps({'type': 'error', 'message': f'vision encode failed: {e}'})}\n\n"
	return

	# Step 2: fast turbo brain reasons over the description
	yield f"data: {json.dumps({'type': 'step_start', 'step': 2, 'max': 2})}\n\n"

	# Build the prompt for the turbo brain
	if description:
	system_msg = (
	"You are a fast assistant that answers questions about images. "
	"A vision model has already described the image for you. "
	"Give a SHORT, direct answer (1-3 sentences). Do not show your reasoning."
	)
	user_msg = (
	f"Image description: {description}\n\n"
	f"User's question: {message}"
	)
	else:
	system_msg = "You are a fast helpful assistant. Give SHORT, direct answers."
	user_msg = message

	# Stream from the turbo brain (mlx-lm OpenAI-compatible endpoint)
	import urllib.request
	payload = {
	"model": "qwen3-1.7b",
	"messages": [
	{"role": "system", "content": system_msg},
	{"role": "user", "content": user_msg},
	],
	"max_tokens": max_tokens,
	"temperature": 0.5,
	"stream": True,
	}
	req = urllib.request.Request(
	f"{turbo_url}/v1/chat/completions",
	data=json.dumps(payload).encode(),
	headers={"Content-Type": "application/json"},
	)
	final_text = ""
	in_think = False
	try:
	with urllib.request.urlopen(req, timeout=120) as resp:
	buffer = ""
	for chunk in iter(lambda: resp.read(512), b""):
	buffer += chunk.decode("utf-8", errors="ignore")
	while "\n" in buffer:
	line, buffer = buffer.split("\n", 1)
	line = line.strip()
	if not line.startswith("data: "):
	continue
	data = line[6:]
	if data == "[DONE]":
	break
	try:
	ev = json.loads(data)
	delta = ev.get("choices", [{}])[0].get("delta", {})
	text = delta.get("content", "")
	if not text:
	continue
	# Strip out <think>...</think> reasoning blocks
	for ch in text:
	if "<think>" in (final_text + ch)[-7:]:
	in_think = True
	continue
	if "</think>" in (final_text + ch)[-8:]:
	in_think = False
	final_text = ""
	continue
	if not in_think:
	final_text += ch
	yield f"data: {json.dumps({'type': 'token', 'text': ch})}\n\n"
	except Exception:
	continue
	except Exception as e:
	yield f"data: {json.dumps({'type': 'error', 'message': f'turbo brain failed: {e}'})}\n\n"
	return

	yield f"data: {json.dumps({'type': 'final', 'text': final_text.strip()})}\n\n"
	yield f"data: {json.dumps({'type': 'done'})}\n\n"
	except Exception as e:
	import traceback
	traceback.print_exc()
	yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
	finally:
	if image_path and os.path.exists(image_path):
	try:
	os.unlink(image_path)
	except Exception:
	pass

	return StreamingResponse(
	event_stream(),
	media_type="text/event-stream",
	headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no",
	"Connection": "keep-alive"},
	)

	@app.post("/api/falcon")
	async def falcon_ground(
	query: str = Form(...),
	image: UploadFile = File(...),
	):
	"""Run Falcon Perception on uploaded image with text query.

	Returns JSON with detected masks (count, IDs, metadata) plus a
	base64-encoded annotated image showing bounding boxes + labels.
	"""
	if falcon_tools is None:
	return JSONResponse({"error": "Falcon Perception not loaded"}, status_code=400)

	# Save uploaded image
	import tempfile, base64, io
	tmp = tempfile.NamedTemporaryFile(suffix="_" + image.filename, delete=False)
	tmp.write(await image.read())
	tmp.close()

	try:
	with lock:
	# Set image in Falcon session
	falcon_tools.set_image(tmp.name)
	# Run grounding
	t0 = time.time()
	result = falcon_tools.ground(query, slot=query.replace(" ", "_")[:32])
	elapsed = time.time() - t0

	if "error" in result:
	return JSONResponse(result, status_code=500)

	# Annotate the image with bounding boxes
	annotated = falcon_tools.annotate_image(mask_ids=result["mask_ids"])

	# Encode annotated image as base64 PNG
	buf = io.BytesIO()
	annotated.save(buf, format="PNG")
	annotated_b64 = base64.b64encode(buf.getvalue()).decode()

	return JSONResponse({
	"query": query,
	"count": result["count"],
	"mask_ids": result["mask_ids"],
	"masks": result["masks"],
	"annotated_image": f"data:image/png;base64,{annotated_b64}",
	"elapsed_seconds": round(elapsed, 2),
	})
	except Exception as e:
	import traceback
	traceback.print_exc()
	return JSONResponse({"error": str(e)}, status_code=500)
	finally:
	try:
	os.unlink(tmp.name)
	except Exception:
	pass

	print()
	print("=" * 60)
	print(f" mac-tensor UI ready")
	print(f" Open: http://localhost:{port}")
	print(f" http://{_local_ip()}:{port} (LAN access)")
	print("=" * 60)
	print()

	import uvicorn
	uvicorn.run(app, host=host, port=port, log_level="warning")


	def _local_ip():
	"""Best-effort detection of the LAN IP."""
	import socket
	try:
	s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	s.connect(("8.8.8.8", 80))
	ip = s.getsockname()[0]
	s.close()
	return ip
	except Exception:
	return "localhost"


	def main(args):
	vision = getattr(args, "vision", False)
	falcon_only = getattr(args, "falcon_only", False)

	if falcon_only:
	# Falcon-only mode: load Falcon Perception, skip Gemma entirely.
	# ~1.5 GB resident. Used by the data labeling factory.
	run_server(
	model_key="falcon",
	node_urls=None,
	host=args.host or "0.0.0.0",
	port=args.port or 8500,
	allow_write=False,
	vision=False,
	falcon=True,
	falcon_model=getattr(args, "falcon_model", None),
	)
	return

	if vision:
	# Vision mode: single-machine, no distributed nodes needed
	run_server(
	model_key="gemma4",
	node_urls=None,
	host=args.host or "0.0.0.0",
	port=args.port or 8500,
	allow_write=getattr(args, "write", False),
	vision=True,
	stream_dir=getattr(args, "stream_dir", None),
	source_dir=getattr(args, "source_dir", None),
	falcon=getattr(args, "falcon", False),
	falcon_model=getattr(args, "falcon_model", None),
	)
	else:
	if not args.nodes:
	print("Error: --nodes is required (or pass --vision for single-machine mode)")
	sys.exit(1)
	node_urls = [u.strip() for u in args.nodes.split(",")]
	run_server(
	model_key=args.model or "gemma4",
	node_urls=node_urls,
	host=args.host or "0.0.0.0",
	port=args.port or 8500,
	allow_write=args.write,
	)