vrex-flow-engine 0.2.0__py3-none-any.whl

flow_engine/openai/videos.py

@@ -0,0 +1,289 @@
+"""Video generation — OpenAI-style job + poll, dynamic model catalog.
+
+  POST /v1/videos/generations  → create job, return {id, status:"in_progress"}
+  GET  /v1/videos/{id}         → {id, status, data:[{url}], ...}
+
+Flow video is asynchronous with no webhook, so flow_engine dispatches + polls in
+a background task and the client polls us.
+
+Only ``veo_3_1_lite_low_priority`` is accepted; any other ``model`` value
+returns HTTP 400.
+
+The generation MODE is inferred from the inputs you provide:
+  - ``ref_media_ids`` / ``ref_images``     → reference-to-video (r2v, "Ingredients")
+  - ``start_media_ids`` / ``start_images`` → image-to-video    (i2v, "Frames")
+  - neither                                → text-to-video     (t2v, prompt only)
+Optional ``reference_audio`` (voice ids from GET /v1/audio/voices) and
+``seconds`` (4/6/8 where the family supports it) refine the request.
+"""
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from pydantic import BaseModel
+
+from flow_engine import ingest, jobs
+from flow_engine.catalog import catalog
+from flow_engine.openai._util import media_url, require_api_key, video_aspect_from_size
+from flow_engine.posthog_client import capture_engine, capture_exception, elapsed_ms, start_time
+from flow_engine.session import ensure_catalog, get_default_project, get_sdk, require_ready
+
+router = APIRouter(prefix="/v1/videos", tags=["videos"])
+
+# The only accepted video model family. Any other value → HTTP 400.
+ALLOWED_VIDEO_MODEL = "veo_3_1_lite_low_priority"
+
+
+class VideoRequest(BaseModel):
+    prompt: str
+    model: str = ALLOWED_VIDEO_MODEL
+    size: Optional[str] = None
+    seconds: Optional[int] = 8
+    # inputs — presence selects the generation mode
+    ref_media_ids: Optional[list[str]] = None       # → r2v (Ingredients), pre-uploaded
+    start_media_ids: Optional[list[str]] = None     # → i2v (Frames), pre-uploaded
+    ref_images: Optional[list[str]] = None          # → r2v, raw (uploaded for you)
+    start_images: Optional[list[str]] = None        # → i2v, raw (uploaded for you)
+    reference_audio: Optional[list[str]] = None     # voice ids (r2v)
+    project: Optional[str] = None
+    # extend / edit — operate on a previously generated Veo video
+    mode: Optional[str] = None                      # "extend" | "edit" (else inferred)
+    source_media_id: Optional[str] = None           # the clip to extend/edit
+    # explicit context — only needed when the source isn't a flow_engine-generated
+    # video tracked this run (the registry fills these in automatically).
+    workflow_id: Optional[str] = None               # extend+edit: source's workflow
+    scene_id: Optional[str] = None                  # extend: override resolved scene
+    position: Optional[int] = None                  # extend: override timeline slot
+    source_seconds: Optional[int] = None            # edit: source clip length (→ frames)
+
+
+# Flow video runs at 30fps; edit selects a frame slice of the source clip.
+_VIDEO_FPS = 30
+
+
+async def _extend_kwargs(
+    sdk: Any,
+    req: "VideoRequest",
+    project_id: str,
+    tier: str,
+    aspect: str,
+    model_key: str,
+    vc_store: Any,
+) -> dict[str, Any]:
+    """Build gen_video_extend kwargs."""
+    if not req.source_media_id:
+        raise HTTPException(status_code=400, detail="source_media_id is required for extend")
+    ctx = vc_store.get(req.source_media_id)
+    workflow_id = req.workflow_id or (ctx.workflow_id if ctx else None)
+    if not workflow_id:
+        raise HTTPException(
+            status_code=400,
+            detail="workflow_id unknown for this source — pass workflow_id explicitly "
+            "(only flow_engine-generated videos from this run are auto-resolved)",
+        )
+    scene_id = req.scene_id
+    position = req.position
+    if not scene_id:
+        scene = await sdk.resolve_scene(project_id, [workflow_id])
+        if scene.get("error"):
+            raise HTTPException(
+                status_code=502,
+                detail=f"could not resolve scene for extend: {scene['error']}",
+            )
+        scene_id = scene["scene_id"]
+        if position is None:
+            position = scene["position"]
+    source_ref = ctx.operation_id if (ctx and ctx.operation_id) else req.source_media_id
+    return dict(
+        prompt=req.prompt.strip(),
+        project_id=project_id,
+        source_media_id=source_ref,
+        scene_id=scene_id,
+        position=position if position is not None else 1,
+        aspect_ratio=aspect,
+        paygate_tier=tier,
+        model_key=model_key,
+    )
+
+
+def _edit_kwargs(
+    req: "VideoRequest",
+    project_id: str,
+    tier: str,
+    aspect: str,
+    model_key: str,
+    vc_store: Any,
+) -> dict[str, Any]:
+    """Build gen_video_edit kwargs."""
+    if not req.source_media_id:
+        raise HTTPException(status_code=400, detail="source_media_id is required for edit")
+    ctx = vc_store.get(req.source_media_id)
+    workflow_id = req.workflow_id or (ctx.workflow_id if ctx else None)
+    if not workflow_id:
+        raise HTTPException(
+            status_code=400,
+            detail="workflow_id unknown for this source — pass workflow_id explicitly "
+            "(only flow_engine-generated videos from this run are auto-resolved)",
+        )
+    duration = req.source_seconds or (ctx.duration_sec if ctx else None) or 8
+    end_frame = int(round(float(duration) * _VIDEO_FPS))
+    return dict(
+        prompt=req.prompt.strip(),
+        project_id=project_id,
+        source_media_id=req.source_media_id,
+        workflow_id=workflow_id,
+        end_frame_index=end_frame,
+        aspect_ratio=aspect,
+        paygate_tier=tier,
+        model_key=model_key,
+    )
+
+
+def _job_view(job: jobs.Job) -> dict[str, Any]:
+    out: dict[str, Any] = {
+        "id": job.id,
+        "object": "video.generation",
+        "status": job.status,
+        "created_at": job.created_at,
+    }
+    if job.status == "completed":
+        out["data"] = [
+            {"url": media_url(mid)} for mid in job.media_ids if isinstance(mid, str)
+        ]
+        if job.partial_error:
+            out["partial_error"] = job.partial_error
+            out["slot_errors"] = job.slot_errors
+    elif job.status == "failed":
+        out["error"] = {"message": job.error, "slot_errors": job.slot_errors}
+    return out
+
+
+@router.post("/generations")
+async def create_video(
+    req: VideoRequest,
+    request: Request,
+    api_key: Optional[str] = Depends(require_api_key),
+) -> dict[str, Any]:
+    if not req.prompt or not req.prompt.strip():
+        raise HTTPException(status_code=400, detail="prompt is required")
+
+    # Enforce pinned model — reject any other family id.
+    if req.model != ALLOWED_VIDEO_MODEL:
+        raise HTTPException(
+            status_code=400,
+            detail=f"only model={ALLOWED_VIDEO_MODEL!r} is supported",
+        )
+
+    t0 = start_time()
+    inst, tier = require_ready(api_key)
+    await ensure_catalog(api_key)
+    project_id = req.project or await get_default_project(api_key)
+    sdk = get_sdk(api_key)
+    aspect = video_aspect_from_size(req.size)
+
+    ref_ids = list(req.ref_media_ids or [])
+    ref_ids += await ingest.resolve_image_inputs(sdk, project_id, req.ref_images)
+    start_ids = list(req.start_media_ids or [])
+    start_ids += await ingest.resolve_image_inputs(sdk, project_id, req.start_images)
+
+    explicit = (req.mode or "").lower().strip()
+    if explicit in ("extend", "extension"):
+        mode = "extension"
+    elif explicit == "edit":
+        mode = "edit"
+    elif ref_ids:
+        mode = "r2v"
+    elif start_ids:
+        mode = "i2v"
+    else:
+        mode = "t2v"
+
+    resolved = catalog.resolve_video(req.model, mode, aspect, req.seconds)
+    if resolved.get("error"):
+        raise HTTPException(status_code=400, detail=resolved["error"])
+    model_key = resolved["key"]
+
+    if mode == "r2v":
+        kwargs: dict[str, Any] = dict(
+            prompt=req.prompt.strip(),
+            project_id=project_id,
+            ref_media_ids=ref_ids,
+            aspect_ratio=aspect,
+            paygate_tier=tier,
+            reference_audio_ids=req.reference_audio,
+            model_key=model_key,
+        )
+    elif mode == "i2v":
+        kwargs = dict(
+            prompt=req.prompt.strip(),
+            project_id=project_id,
+            start_media_ids=start_ids,
+            aspect_ratio=aspect,
+            paygate_tier=tier,
+            model_key=model_key,
+        )
+    elif mode == "extension":
+        kwargs = await _extend_kwargs(
+            sdk, req, project_id, tier, aspect, model_key, inst.video_ctx
+        )
+    elif mode == "edit":
+        kwargs = _edit_kwargs(
+            req, project_id, tier, aspect, model_key, inst.video_ctx
+        )
+    else:  # t2v
+        kwargs = dict(
+            prompt=req.prompt.strip(),
+            project_id=project_id,
+            aspect_ratio=aspect,
+            paygate_tier=tier,
+            model_key=model_key,
+        )
+
+    try:
+        job = jobs.start_video_job(
+            sdk,
+            mode=mode,
+            dispatch_kwargs=kwargs,
+            job_store=inst.jobs,
+            media_store=inst.media,
+            vc_store=inst.video_ctx,
+        )
+    except Exception as exc:
+        capture_exception(request, exc, {"latency_ms": elapsed_ms(t0), "mode": mode})
+        raise
+
+    capture_engine(request, "engine_request_started", {
+        "model": req.model,
+        "mode": mode,
+        "seconds": req.seconds,
+        "dispatch_ms": elapsed_ms(t0),
+    })
+    view = _job_view(job)
+    view["model"] = req.model
+    view["model_key"] = model_key
+    view["mode"] = mode
+    return view
+
+
+@router.get("/{job_id}")
+async def get_video(
+    job_id: str, api_key: Optional[str] = Depends(require_api_key)
+) -> dict[str, Any]:
+    inst, _ = require_ready(api_key)
+    job = jobs.get(job_id, inst.jobs)
+    if job is None:
+        raise HTTPException(status_code=404, detail="job not found")
+    return _job_view(job)
+
+
+@router.get("/{job_id}/debug")
+async def get_video_debug(
+    job_id: str, api_key: Optional[str] = Depends(require_api_key)
+) -> dict[str, Any]:
+    """Full Flow dispatch + poll payload for a job — for diagnosing failures."""
+    inst, _ = require_ready(api_key)
+    job = jobs.get(job_id, inst.jobs)
+    if job is None:
+        raise HTTPException(status_code=404, detail="job not found")
+    return {"id": job.id, "status": job.status, "error": job.error, "raw": job.raw}

flow_engine/pool.py

@@ -0,0 +1,138 @@
+"""Multi-tenant pool of FlowClient instances, keyed by account_id.
+
+Each PoolInstance owns:
+  - a FlowClient (one WS connection to a Chrome extension profile)
+  - an asyncio.Semaphore(POOL_MAX_CONCURRENCY) — back-pressure, not rejection
+  - an in_flight counter (avoids relying on the private Semaphore._value)
+  - isolated JobStore / MediaStore / VideoContextStore
+
+PoolRegistry routes requests by api_key → account_id. When no explicit
+mapping is present it returns the least-busy instance (lowest in_flight).
+New instances are registered on first WS connect (ws_server._handler calls
+pool.register). The pool is a module-level singleton; import `pool` from here.
+"""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Optional
+
+from flow_engine.bridge.flow_client import FlowClient
+from flow_engine.config import POOL_MAX_CONCURRENCY
+from flow_engine.job_store import JobStore
+from flow_engine.media_store import MediaStore
+from flow_engine.video_context_store import VideoContextStore
+
+
+def _inc(inst: "PoolInstance") -> None:
+    inst.in_flight += 1
+
+
+def _dec(inst: "PoolInstance") -> None:
+    if inst.in_flight > 0:
+        inst.in_flight -= 1
+
+
+@dataclass
+class PoolInstance:
+    account_id: str
+    client: FlowClient
+    # Semaphore limits concurrent in-flight WS calls; callers queue, not 503.
+    semaphore: asyncio.Semaphore = field(
+        default_factory=lambda: asyncio.Semaphore(POOL_MAX_CONCURRENCY)
+    )
+    # Explicit counter — we never read the private Semaphore._value.
+    in_flight: int = 0
+    jobs: JobStore = field(default_factory=JobStore)
+    # MediaStore needs account_id for per-account disk subdir; set post-init.
+    media: MediaStore = field(default_factory=MediaStore)
+    video_ctx: VideoContextStore = field(default_factory=VideoContextStore)
+    # Per-instance default project (replaces the session.py module globals).
+    _default_project_id: Optional[str] = None
+    _project_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    def __post_init__(self) -> None:
+        # Re-create the MediaStore with the real account_id subdir now that
+        # we have it (the default_factory above can't capture self.account_id).
+        self.media = MediaStore(self.account_id)
+
+
+class PoolRegistry:
+    """Thread-safe registry of PoolInstance objects.
+
+    Supports:
+      - register(account_id)  — called by ws_server on WS connect
+      - get_by_api_key(key)   — called by session helpers on every request
+      - resolve_callback(data) — routes HTTP callbacks to the right instance
+      - status()              — admin view of all instances
+    """
+
+    def __init__(self) -> None:
+        self._instances: dict[str, PoolInstance] = {}  # account_id → instance
+        self._api_key_map: dict[str, str] = {}         # api_key → account_id
+
+    def register(self, account_id: str) -> PoolInstance:
+        """Return (or create) the PoolInstance for account_id.
+
+        Idempotent: reconnecting the same account reuses the existing instance
+        so its JobStore / MediaStore / VideoContextStore survive reconnects.
+        The FlowClient's in_flight callbacks are (re-)attached each time so
+        the counter always points at the live PoolInstance.
+        """
+        if account_id not in self._instances:
+            client = FlowClient()
+            inst = PoolInstance(
+                account_id=account_id,
+                client=client,
+                semaphore=asyncio.Semaphore(POOL_MAX_CONCURRENCY),
+            )
+            # Wire in_flight callbacks — increment before send, decrement after.
+            inst.client._on_send_start = lambda: _inc(inst)
+            inst.client._on_send_end = lambda: _dec(inst)
+            self._instances[account_id] = inst
+        return self._instances[account_id]
+
+    def bind_api_key(self, api_key: str, account_id: str) -> None:
+        """Associate a caller api_key with an account_id for routing."""
+        self._api_key_map[api_key] = account_id
+
+    def get_by_api_key(self, api_key: Optional[str]) -> Optional[PoolInstance]:
+        """Return the instance for the given api_key, or the least-busy
+        instance as a fallback when only one account is registered."""
+        if api_key and api_key in self._api_key_map:
+            acct = self._api_key_map[api_key]
+            return self._instances.get(acct)
+        # Single-account fallback: pick the instance with the fewest in-flight
+        # calls so single-tenant setups work without any key→account mapping.
+        if self._instances:
+            return min(self._instances.values(), key=lambda i: i.in_flight)
+        return None
+
+    def resolve_callback(self, data: dict) -> bool:
+        """Route an HTTP callback to the correct instance's pending dict.
+
+        Linear scan over all instances — bounded by N accounts (expected 1–5).
+        Returns True when a pending future matched.
+        """
+        for inst in self._instances.values():
+            if inst.client.resolve_callback(data):
+                return True
+        return False
+
+    def status(self) -> list[dict]:
+        return [
+            {
+                "account_id": inst.account_id,
+                "connected": inst.client.connected,
+                "paygate_tier": inst.client.paygate_tier,
+                "sku": inst.client.sku,
+                "credits": inst.client.credits,
+                "in_flight": inst.in_flight,
+                "capacity": POOL_MAX_CONCURRENCY,
+            }
+            for inst in self._instances.values()
+        ]
+
+
+# Module-level singleton — import this in ws_server, session, main.
+pool = PoolRegistry()

flow_engine/posthog_client.py

@@ -0,0 +1,117 @@
+"""
+PostHog client for the flow-engine (image + video generation).
+
+Initialised from POSTHOG_API_KEY + POSTHOG_HOST env vars.
+When the key is unset (dev / tests), all methods are silent no-ops so engine
+behaviour is never affected.
+
+Usage:
+    from flow_engine.posthog_client import capture_engine, flush_posthog
+
+    capture_engine(request, "engine_request_started", {"n": 1})
+    capture_engine(request, "engine_request_completed", {"latency_ms": 4500})
+    capture_exception(request, exc)
+"""
+from __future__ import annotations
+
+import logging
+import os
+import time
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+# Header names must match the Node-side constants in posthog-correlation.ts
+_DISTINCT_ID_HEADER = "x-posthog-distinct-id"
+_RUN_ID_HEADER = "x-posthog-run-id"
+
+_ENGINE_NAME = "flow"
+
+# ---------------------------------------------------------------------------
+# Client init — no-op shim when key is absent
+# ---------------------------------------------------------------------------
+
+_api_key = os.environ.get("POSTHOG_API_KEY", "")
+_host = os.environ.get("POSTHOG_HOST", "https://eu.i.posthog.com")
+
+_ph: Any = None  # real PostHog client or None (no-op path)
+
+if _api_key:
+    try:
+        from posthog import Posthog  # type: ignore[import-untyped]
+        _ph = Posthog(api_key=_api_key, host=_host, sync_mode=False)
+        logger.info("PostHog client initialised (host=%s)", _host)
+    except Exception as exc:  # noqa: BLE001
+        logger.warning("PostHog init failed (no-op): %s", exc)
+
+
+# ---------------------------------------------------------------------------
+# Public helpers
+# ---------------------------------------------------------------------------
+
+
+def capture_engine(
+    request: Any,
+    event: str,
+    properties: Optional[dict[str, Any]] = None,
+) -> None:
+    """Capture an engine lifecycle event, tagged with correlation ids from headers.
+
+    Reads X-PostHog-Distinct-Id / X-PostHog-Run-Id from ``request.headers``.
+    Silent no-op when PostHog is disabled.
+    """
+    if _ph is None:
+        return
+    try:
+        distinct_id: str = request.headers.get(_DISTINCT_ID_HEADER, "anonymous") or "anonymous"
+        run_id: str = request.headers.get(_RUN_ID_HEADER, "") or ""
+
+        props: dict[str, Any] = {"engine": _ENGINE_NAME, **(properties or {})}
+        if run_id:
+            props["$run_id"] = run_id
+
+        kwargs: dict[str, Any] = {
+            "distinct_id": distinct_id,
+            "event": event,
+            "properties": props,
+        }
+        if run_id:
+            kwargs["groups"] = {"run": run_id}
+
+        _ph.capture(**kwargs)
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("PostHog capture error (non-fatal): %s", exc)
+
+
+def capture_exception(
+    request: Any,
+    exc: Exception,
+    extra: Optional[dict[str, Any]] = None,
+) -> None:
+    """Capture an exception as an engine_request_error event."""
+    capture_engine(
+        request,
+        "engine_request_error",
+        {"error": str(exc), "error_type": type(exc).__name__, **(extra or {})},
+    )
+
+
+def flush_posthog() -> None:
+    """Flush pending events — call on app shutdown so short container lifetimes
+    don't drop queued events."""
+    if _ph is None:
+        return
+    try:
+        _ph.flush()
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("PostHog flush error (non-fatal): %s", exc)
+
+
+def start_time() -> float:
+    """Return a monotonic timestamp for latency measurement."""
+    return time.monotonic()
+
+
+def elapsed_ms(t0: float) -> int:
+    """Milliseconds elapsed since ``t0`` (from start_time())."""
+    return int((time.monotonic() - t0) * 1000)

flow_engine/session.py

@@ -0,0 +1,89 @@
+"""Session resolution + readiness gating — routes via pool.
+
+Every request resolves to a PoolInstance via pool.get_by_api_key(). Per-instance
+state (default_project_id, job_store, media_store, video_ctx) lives on the
+instance so there is no shared mutable state between accounts.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from fastapi import HTTPException
+
+from flow_engine.bridge.flow_sdk import FlowSDK
+from flow_engine.catalog import catalog
+from flow_engine.config import DEFAULT_PROJECT_TITLE
+from flow_engine.pool import PoolInstance, pool
+
+logger = logging.getLogger(__name__)
+
+
+def _get_instance(api_key: Optional[str]) -> Optional[PoolInstance]:
+    return pool.get_by_api_key(api_key)
+
+
+def get_bridge(api_key: Optional[str] = None):
+    """Resolve the FlowClient for this request via pool lookup."""
+    inst = _get_instance(api_key)
+    return inst.client if inst else None
+
+
+def get_sdk(api_key: Optional[str] = None) -> FlowSDK:
+    """A FlowSDK bound to the resolved bridge. FlowSDK is stateless, so a
+    fresh instance per call is cheap and keeps us pool-ready."""
+    return FlowSDK(get_bridge(api_key))
+
+
+def require_ready(api_key: Optional[str] = None) -> tuple[PoolInstance, str]:
+    """Gate a generation request. Raises 503 unless the extension is connected
+    AND the paygate tier has been resolved.
+
+    Returns ``(instance, paygate_tier)``.  Callers use instance.jobs /
+    instance.media / instance.video_ctx for per-tenant isolation.
+    """
+    inst = _get_instance(api_key)
+    if not inst or not inst.client.connected:
+        raise HTTPException(
+            status_code=503,
+            detail="extension_not_connected — open Chrome with the flow_engine "
+            "extension installed and signed in to labs.google",
+        )
+    tier = inst.client.paygate_tier
+    if tier is None:
+        raise HTTPException(
+            status_code=503,
+            detail="paygate_tier_unresolved — extension connected but no Flow "
+            "token captured yet; open a labs.google/fx/tools/flow tab",
+        )
+    return inst, tier
+
+
+async def get_default_project(api_key: Optional[str] = None) -> str:
+    """Return a usable Flow project id for this instance, creating one on
+    first use. Per-instance lock prevents duplicate project creation."""
+    inst = _get_instance(api_key)
+    if not inst:
+        raise HTTPException(status_code=503, detail="no_pool_instance")
+    async with inst._project_lock:
+        if inst._default_project_id:
+            return inst._default_project_id
+        sdk = FlowSDK(inst.client)
+        result = await sdk.create_project(DEFAULT_PROJECT_TITLE)
+        pid = result.get("project_id")
+        if not pid:
+            raise HTTPException(
+                status_code=502,
+                detail=f"could not create Flow project: {result.get('error') or 'unknown'}",
+            )
+        inst._default_project_id = pid
+        logger.info("flow_engine default project created for %s: %s", inst.account_id, pid)
+        return pid
+
+
+async def ensure_catalog(api_key: Optional[str] = None, *, force: bool = False) -> None:
+    """Load the dynamic model catalog (idempotent). Raises 503 on failure."""
+    project_id = await get_default_project(api_key)
+    err = await catalog.ensure_loaded(get_sdk(api_key), project_id, force=force)
+    if err:
+        raise HTTPException(status_code=502, detail=f"could not load model catalog: {err}")