app.py

import asyncio
import json
import os
import time
import threading
import contextvars
from pathlib import Path

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles

app = FastAPI()

SYSTEM_PROMPT = Path("system.md").read_text()
MODEL = os.environ.get("SYNC_PLANNER_MODEL", "zai-org/GLM-5.1")
BASE_URL = os.environ.get("SYNC_BASE_URL", "https://router.huggingface.co/v1")
API_KEY = os.environ.get("SYNC_API_KEY", "")

if not API_KEY:
    try:
        from huggingface_hub import get_token
        API_KEY = get_token() or ""
    except ImportError:
        pass

from openai import AsyncOpenAI
client = AsyncOpenAI(base_url=BASE_URL, api_key=API_KEY)


# --- Minimal sync runtime (inlined) ---

_current_ctx = contextvars.ContextVar("ctx")


class Result:
    def __init__(self, messages):
        self.messages = messages

    @property
    def last(self):
        return self.messages[-1]["content"]


class _Context:
    def __init__(self):
        self._out = threading.Event()
        self._in = threading.Event()
        self._batch = None
        self._results = None
        self._done = False
        self._done_value = None

    def _call_many(self, calls):
        self._batch = ("calls", calls)
        self._out.set()
        self._in.wait()
        self._in.clear()
        return self._results

    def _update(self, message):
        self._batch = ("update", message)
        self._out.set()
        self._in.wait()
        self._in.clear()

    def _finish(self, value):
        self._done = True
        self._done_value = value
        self._out.set()


def agent(name, message, **_kw):
    ctx = _current_ctx.get()
    if isinstance(message, list):
        return ctx._call_many([(name, m) for m in message])
    return ctx._call_many([(name, message)])[0]


def update(message):
    ctx = _current_ctx.get()
    ctx._update(message)


def load_plan(source, task=""):
    ns = {"agent": agent, "update": update, "TASK": task}
    exec(source, ns)
    return ns.get("AGENTS", []), ns.get("plan")


# --- SSE endpoint ---

@app.get("/", response_class=HTMLResponse)
async def index():
    return Path("static/index.html").read_text()


@app.get("/generate")
async def generate(task: str):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": task},
    ]
    resp = await client.chat.completions.create(
        model=MODEL, messages=messages)
    code = resp.choices[0].message.content
    code = (code.strip()
            .removeprefix("```python")
            .removeprefix("```")
            .removesuffix("```")
            .strip())
    return {"code": code}


@app.get("/execute")
async def execute(task: str, code: str):
    async def stream():
        def send(event, data):
            return f"event: {event}\ndata: {json.dumps(data)}\n\n"

        try:
            agent_defs, plan_fn = load_plan(code, task=task)
        except Exception as e:
            yield send("error", {"msg": f"Plan load error: {e}"})
            return

        if not plan_fn:
            yield send("error", {"msg": "No plan() found"})
            return

        ctx = _Context()

        def run_plan():
            _current_ctx.set(ctx)
            try:
                result = plan_fn()
                ctx._finish(result)
            except Exception as e:
                ctx._finish(f"Error: {e}")

        thread = threading.Thread(target=run_plan, daemon=True)
        thread.start()
        t0 = time.time()

        while True:
            await asyncio.get_event_loop().run_in_executor(None, ctx._out.wait)
            ctx._out.clear()

            if ctx._done:
                elapsed = time.time() - t0
                yield send("log", {"msg": f"done — {elapsed:.1f}s total"})
                yield send("result", {"text": str(ctx._done_value or "")})
                return

            kind, payload = ctx._batch

            if kind == "update":
                elapsed = time.time() - t0
                yield send("log", {"msg": f"[{elapsed:.1f}s] {payload}"})
                ctx._in.set()
                continue

            if kind == "calls":
                names = {}
                for name, msg in payload:
                    names[name] = names.get(name, 0) + 1
                parts = []
                for n, c in names.items():
                    if c > 1:
                        parts.append(f"[{c}x] {n}")
                    else:
                        parts.append(n)
                desc = ", ".join(parts)
                elapsed = time.time() - t0
                yield send("log", {"msg": f"[{elapsed:.1f}s] ▶ {desc}"})

                results = []
                tasks = []
                for name, msg in payload:
                    ad = next((a for a in agent_defs if a["name"] == name), {})
                    tasks.append(_call_agent(ad, msg))

                agent_results = await asyncio.gather(*tasks)
                for (name, msg), messages in zip(payload, agent_results):
                    if isinstance(messages, list):
                        results.append(Result(messages))
                    else:
                        results.append(Result([
                            {"role": "user", "content": msg},
                            {"role": "assistant", "content": str(messages)},
                        ]))

                elapsed2 = time.time() - t0
                yield send("log", {"msg": f"[{elapsed2:.1f}s] ✓ {desc}"})
                ctx._results = results
                ctx._in.set()

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


async def _call_agent(agent_def, message):
    model = agent_def.get("model", MODEL)
    system = agent_def.get("system", "")
    tools = {fn.__name__: fn for fn in agent_def.get("tools", []) if callable(fn)}

    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": message})

    kwargs = {"model": model, "messages": messages}
    if tools:
        kwargs["tools"] = list(tools.values())

    try:
        while True:
            resp = await client.chat.completions.create(**kwargs)
            choice = resp.choices[0]

            if choice.finish_reason == "tool_calls":
                messages.append(choice.message.model_dump())
                for tc in choice.message.tool_calls:
                    fn = tools.get(tc.function.name)
                    if fn:
                        import json as j
                        output = fn(**j.loads(tc.function.arguments))
                    else:
                        output = f"Unknown tool: {tc.function.name}"
                    messages.append({
                        "role": "tool",
                        "tool_call_id": tc.id,
                        "content": str(output),
                    })
                continue

            messages.append({
                "role": "assistant",
                "content": choice.message.content,
            })
            return messages
    except Exception as e:
        return [
            {"role": "user", "content": message},
            {"role": "assistant", "content": f"Error: {e}"},
        ]