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

flow_engine/media_store.py

@@ -0,0 +1,202 @@
+"""Per-instance media store.
+
+Extracted from the module-level ``_registry`` dict in media.py so each pool
+instance owns an isolated in-memory registry. On-disk bytes are written to a
+per-account subdir (``storage/media/<account_id>/``) to prevent cross-tenant
+file collisions at the OS path layer.
+"""
+from __future__ import annotations
+
+import logging
+import threading
+from pathlib import Path
+from typing import Any, Optional
+
+import httpx
+
+from flow_engine.config import STORAGE_DIR
+from flow_engine.media import (
+    Asset,
+    _ALLOWED_URL_PREFIXES,
+    _EXT_BY_MIME,
+    _MEDIA_ID_RE,
+    _mime_from_ext,
+    _url_allowed,
+    is_valid_media_id,
+    normalize_media_id,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class MediaStore:
+    """Thread-safe media registry + disk cache, one per PoolInstance.
+
+    Each instance writes cached bytes under ``storage/media/<account_id>/``
+    so accounts never share or accidentally serve each other's files.
+    """
+
+    def __init__(self, account_id: str = "default") -> None:
+        self._registry: dict[str, Asset] = {}
+        self._lock = threading.Lock()
+        # Per-account subdir prevents cross-tenant file collisions.
+        self._cache_dir: Path = STORAGE_DIR / "media" / account_id
+        self._cache_dir.mkdir(parents=True, exist_ok=True)
+
+    # ── helpers ────────────────────────────────────────────────────────────
+
+    def _cache_glob(self, media_id: str) -> Optional[Path]:
+        if not is_valid_media_id(media_id):
+            return None
+        for p in self._cache_dir.glob(f"{media_id}.*"):
+            if p.is_file():
+                return p
+        return None
+
+    def _get(self, media_id: str) -> Optional[Asset]:
+        with self._lock:
+            return self._registry.get(media_id)
+
+    # ── write ──────────────────────────────────────────────────────────────
+
+    def ingest_urls(self, urls: list[dict[str, Any]]) -> int:
+        """Register Asset rows keyed by uuid media_id. Returns count touched."""
+        touched = 0
+        with self._lock:
+            for entry in urls or []:
+                if not isinstance(entry, dict):
+                    continue
+                media_id = entry.get("media_id") or entry.get("mediaId")
+                url = entry.get("url")
+                kind = entry.get("mediaType") or entry.get("kind") or "image"
+                if not isinstance(media_id, str) or not is_valid_media_id(media_id):
+                    continue
+                if not _url_allowed(url):
+                    logger.warning(
+                        "media_store: skip non-allowed url for %s: %r",
+                        media_id, (url or "")[:80],
+                    )
+                    continue
+                row = self._registry.get(media_id)
+                if row is None:
+                    self._registry[media_id] = Asset(media_id=media_id, url=url, kind=kind)
+                else:
+                    row.url = url
+                    if not row.kind:
+                        row.kind = kind
+                touched += 1
+        if touched:
+            logger.info("media_store: ingested %d url(s)", touched)
+        return touched
+
+    def ingest_inline_bytes(
+        self,
+        media_id: str,
+        data: bytes,
+        *,
+        kind: str = "video",
+        mime: str = "video/mp4",
+    ) -> bool:
+        """Cache pre-fetched media bytes (workflow-mode base64 MP4) and mark
+        the Asset as locally available."""
+        if not is_valid_media_id(media_id) or not data:
+            return False
+        ext = _EXT_BY_MIME.get(mime, ".mp4")
+        path = self._cache_dir / f"{media_id}{ext}"
+        try:
+            path.write_bytes(data)
+        except OSError as exc:
+            logger.error("media_store: failed to write inline cache %s: %s", path, exc)
+            return False
+        with self._lock:
+            row = self._registry.get(media_id) or Asset(media_id=media_id, kind=kind)
+            row.local_path = str(path)
+            row.mime = mime
+            if not row.kind:
+                row.kind = kind
+            self._registry[media_id] = row
+        logger.info("media_store: ingested %d inline bytes for %s", len(data), media_id)
+        return True
+
+    # ── read ───────────────────────────────────────────────────────────────
+
+    async def fetch_and_cache(
+        self, media_id: str
+    ) -> Optional[tuple[bytes, str, Path]]:
+        """Fetch bytes from the signed URL, cache to disk, return
+        ``(bytes, mime, path)``; else ``None``."""
+        if not is_valid_media_id(media_id):
+            return None
+        row = self._get(media_id)
+        if row is None or not row.url:
+            return None
+        url = row.url
+        if not _url_allowed(url):
+            logger.warning("media_store: refusing non-allowed URL: %s", url[:60])
+            return None
+        try:
+            async with httpx.AsyncClient(timeout=30.0) as client:
+                resp = await client.get(url)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("media_store: GCS fetch failed for %s: %s", media_id, exc)
+            return None
+        if resp.status_code != 200:
+            logger.warning(
+                "media_store: GCS fetch %s returned %d", media_id, resp.status_code
+            )
+            return None
+        mime = (
+            resp.headers.get("content-type", "application/octet-stream")
+            .split(";")[0]
+            .strip()
+        )
+        if not (mime.startswith("image/") or mime.startswith("video/")):
+            logger.warning(
+                "media_store: refusing non-media mime from GCS: %s (media_id=%s)",
+                mime, media_id,
+            )
+            return None
+        ext = _EXT_BY_MIME.get(mime, ".bin")
+        path = self._cache_dir / f"{media_id}{ext}"
+        try:
+            path.write_bytes(resp.content)
+        except OSError as exc:
+            logger.error("media_store: failed to write cache %s: %s", path, exc)
+            return None
+        with self._lock:
+            r = self._registry.get(media_id)
+            if r is not None:
+                r.local_path = str(path)
+                r.mime = mime
+        return resp.content, mime, path
+
+    async def ensure_cached(self, media_id: str) -> Optional[Path]:
+        """Return a local file path for the media, fetching on first access."""
+        media_id = normalize_media_id(media_id)
+        cached = self._cache_glob(media_id)
+        if cached is not None:
+            return cached
+        got = await self.fetch_and_cache(media_id)
+        return got[2] if got else None
+
+    def status(self, media_id: str) -> dict:
+        media_id = normalize_media_id(media_id)
+        if not is_valid_media_id(media_id):
+            return {"available": False, "has_url": False, "reason": "invalid_id"}
+        cached = self._cache_glob(media_id)
+        if cached is not None:
+            return {"available": True, "has_url": True, "mime": _mime_from_ext(cached.suffix)}
+        row = self._get(media_id)
+        if row is None:
+            return {"available": False, "has_url": False, "reason": "unknown_media"}
+        if row.url:
+            return {"available": False, "has_url": True, "reason": "not_cached_yet"}
+        return {"available": False, "has_url": False, "reason": "no_url_yet"}
+
+    def media_mime(self, media_id: str) -> str:
+        media_id = normalize_media_id(media_id)
+        cached = self._cache_glob(media_id)
+        if cached is not None:
+            return _mime_from_ext(cached.suffix)
+        row = self._get(media_id)
+        return row.mime if (row and row.mime) else "application/octet-stream"

flow_engine/openai/__init__.py

@@ -0,0 +1 @@
+"""OpenAI-compatible HTTP surface mapped onto Google Flow generation."""

flow_engine/openai/_util.py

@@ -0,0 +1,59 @@
+"""Shared helpers for the OpenAI-compatible routes."""
+from __future__ import annotations
+
+import base64
+from typing import Optional
+
+from fastapi import Header, HTTPException
+
+from flow_engine import media as media_service
+from flow_engine.config import API_KEY, PUBLIC_BASE_URL
+
+
+def require_api_key(authorization: Optional[str] = Header(default=None)) -> Optional[str]:
+    """FastAPI dependency. When FLOWPROXY_API_KEY is set, enforce a bearer
+    token and return it (used later as the pool routing key). No-op otherwise."""
+    if not API_KEY:
+        return None
+    token = ""
+    if authorization and authorization.lower().startswith("bearer "):
+        token = authorization[7:].strip()
+    if token != API_KEY:
+        raise HTTPException(status_code=401, detail="invalid_api_key")
+    return token
+
+
+def media_url(media_id: str) -> str:
+    """Absolute, client-fetchable URL for a Flow media id."""
+    return f"{PUBLIC_BASE_URL.rstrip('/')}/media/{media_service.normalize_media_id(media_id)}"
+
+
+async def media_b64(media_id: str) -> Optional[str]:
+    """Base64 of the cached media bytes (for response_format='b64_json')."""
+    path = await media_service.ensure_cached(media_id)
+    if path is None:
+        return None
+    return base64.b64encode(path.read_bytes()).decode("ascii")
+
+
+# OpenAI `size` ("WxH") → Flow image aspect-ratio enum.
+def image_aspect_from_size(size: Optional[str]) -> str:
+    if not isinstance(size, str) or "x" not in size:
+        return "IMAGE_ASPECT_RATIO_LANDSCAPE"
+    try:
+        w, h = (int(x) for x in size.lower().split("x", 1))
+    except ValueError:
+        return "IMAGE_ASPECT_RATIO_LANDSCAPE"
+    if w == h:
+        return "IMAGE_ASPECT_RATIO_SQUARE"
+    return "IMAGE_ASPECT_RATIO_PORTRAIT" if h > w else "IMAGE_ASPECT_RATIO_LANDSCAPE"
+
+
+def video_aspect_from_size(size: Optional[str]) -> str:
+    if not isinstance(size, str) or "x" not in size:
+        return "VIDEO_ASPECT_RATIO_LANDSCAPE"
+    try:
+        w, h = (int(x) for x in size.lower().split("x", 1))
+    except ValueError:
+        return "VIDEO_ASPECT_RATIO_LANDSCAPE"
+    return "VIDEO_ASPECT_RATIO_PORTRAIT" if h > w else "VIDEO_ASPECT_RATIO_LANDSCAPE"

flow_engine/openai/images.py

@@ -0,0 +1,134 @@
+"""POST /v1/images/generations — OpenAI-compatible image generation.
+
+Maps onto Flow's synchronous ``batchGenerateImages`` (gen_image).
+
+Only ``NANO_BANANA_2`` is accepted as the image model. Passing any other
+non-null model value returns HTTP 400. Omitting ``model`` (or passing null)
+silently forces ``NANO_BANANA_2``.
+
+Non-standard extras accepted in the body:
+  - ``model``: must be "NANO_BANANA_2" or omitted
+  - ``project``: existing Flow project id (else flow_engine's default project)
+  - ``ref_media_ids``: Flow media ids to condition on (reference images)
+  - ``ref_images``: raw reference images (data URL / http(s) URL / base64) that
+    flow_engine uploads for you and conditions on — merged with ``ref_media_ids``
+"""
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from pydantic import BaseModel
+
+from flow_engine import ingest
+from flow_engine.openai._util import (
+    image_aspect_from_size,
+    media_b64,
+    media_url,
+    require_api_key,
+)
+from flow_engine.posthog_client import capture_engine, capture_exception, elapsed_ms, start_time
+from flow_engine.session import get_default_project, get_sdk, require_ready
+
+logger = logging.getLogger(__name__)
+router = APIRouter(prefix="/v1/images", tags=["images"])
+
+# The only accepted image model. Any other non-null value → HTTP 400.
+ALLOWED_IMAGE_MODEL = "NANO_BANANA_2"
+
+
+class ImageRequest(BaseModel):
+    prompt: str
+    model: Optional[str] = None
+    n: int = 1
+    size: Optional[str] = None
+    response_format: str = "url"  # "url" | "b64_json"
+    # flow_engine extensions
+    project: Optional[str] = None
+    ref_media_ids: Optional[list[str]] = None
+    ref_images: Optional[list[str]] = None  # raw refs — uploaded for you
+
+
+@router.post("/generations")
+async def generate_images(
+    req: ImageRequest,
+    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 explicitly-passed non-allowed value.
+    if req.model is not None and req.model != ALLOWED_IMAGE_MODEL:
+        raise HTTPException(
+            status_code=400,
+            detail=f"only model={ALLOWED_IMAGE_MODEL!r} is supported",
+        )
+    # Always use the pinned model regardless of what was passed.
+    image_model = ALLOWED_IMAGE_MODEL
+
+    t0 = start_time()
+    capture_engine(request, "engine_request_started", {"model": image_model, "n": req.n})
+
+    try:
+        inst, tier = require_ready(api_key)
+        project_id = req.project or await get_default_project(api_key)
+
+        sdk = get_sdk(api_key)
+        ref_ids = list(req.ref_media_ids or [])
+        ref_ids += await ingest.resolve_image_inputs(sdk, project_id, req.ref_images)
+
+        result = await sdk.gen_image(
+            prompt=req.prompt.strip(),
+            project_id=project_id,
+            aspect_ratio=image_aspect_from_size(req.size),
+            paygate_tier=tier,
+            ref_media_ids=ref_ids or None,
+            variant_count=max(1, min(int(req.n or 1), 4)),
+            image_model=image_model,
+        )
+        if result.get("error"):
+            capture_engine(request, "engine_request_error", {
+                "error": str(result["error"])[:200],
+                "latency_ms": elapsed_ms(t0),
+            })
+            raise HTTPException(status_code=502, detail=str(result["error"])[:200])
+
+        entries = [e for e in (result.get("media_entries") or []) if isinstance(e, dict)]
+        url_entries = [e for e in entries if e.get("url")]
+        if url_entries:
+            # Ingest into the per-instance media store.
+            inst.media.ingest_urls(url_entries)
+
+        data: list[dict[str, Any]] = []
+        for e in entries:
+            mid = e.get("media_id")
+            if not isinstance(mid, str):
+                continue
+            if req.response_format == "b64_json":
+                b64 = await media_b64(mid)
+                data.append({"b64_json": b64})
+            else:
+                data.append({"url": media_url(mid)})
+
+        if not data:
+            capture_engine(request, "engine_request_error", {
+                "error": "no media returned",
+                "latency_ms": elapsed_ms(t0),
+            })
+            raise HTTPException(status_code=502, detail="no media returned by Flow")
+
+    except HTTPException:
+        raise
+    except Exception as exc:
+        capture_exception(request, exc, {"latency_ms": elapsed_ms(t0)})
+        raise
+
+    capture_engine(request, "engine_request_completed", {
+        "latency_ms": elapsed_ms(t0),
+        "returned": len(data),
+        "model": image_model,
+    })
+    return {"created": int(time.time()), "data": data}

flow_engine/openai/models.py

@@ -0,0 +1,38 @@
+"""GET /v1/models — dynamic catalog sourced from Flow's projectInitialData.
+
+Lists the real video/image model families available to the signed-in account
+(filtered to the user's service tier), plus the preset audio voices. Pass
+``?refresh=1`` to force a re-fetch of the upstream catalog.
+"""
+from __future__ import annotations
+
+import time
+from typing import Any, Optional
+
+from fastapi import APIRouter, Depends
+
+from flow_engine.catalog import catalog
+from flow_engine.openai._util import require_api_key
+from flow_engine.session import ensure_catalog, require_ready
+
+router = APIRouter(prefix="/v1", tags=["models"])
+
+
+@router.get("/models")
+async def list_models(
+    refresh: bool = False, _: Optional[str] = Depends(require_api_key)
+) -> dict[str, Any]:
+    require_ready()
+    await ensure_catalog(force=refresh)
+    created = int(time.time())
+    data = catalog.image_models_api() + catalog.video_models_api()
+    for m in data:
+        m["created"] = created
+    return {"object": "list", "data": data}
+
+
+@router.get("/audio/voices")
+async def list_voices(_: Optional[str] = Depends(require_api_key)) -> dict[str, Any]:
+    require_ready()
+    await ensure_catalog()
+    return {"object": "list", "data": catalog.audio_voices()}

flow_engine/openai/uploads.py

@@ -0,0 +1,70 @@
+"""POST /v1/flow/uploads — turn raw media bytes into a Flow media id.
+
+Generation endpoints need Flow media ids (reference images for veo-3.1, a start
+frame for i2v, or a reference video for v2v/edit/extend). This wraps
+``flow_sdk.upload_image`` / ``flow_sdk.upload_video`` so a client can upload a
+file once and reuse the returned ``media_id`` across requests.
+
+Routing is by mime: ``video/*`` goes through the resumable video-upload flow,
+everything else through ``uploadImage``. (Most callers don't need this endpoint
+at all — the image/video generation routes accept raw ``ref_images`` /
+``start_images`` and upload inline.)
+
+Accepts EITHER a JSON body ``{image_base64, mime_type, file_name}`` OR a
+multipart file upload.
+"""
+from __future__ import annotations
+
+import base64
+from typing import Any, Optional
+
+from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile
+
+from flow_engine.openai._util import require_api_key
+from flow_engine.session import get_default_project, get_sdk, require_ready
+
+router = APIRouter(prefix="/v1/flow", tags=["uploads"])
+
+
+@router.post("/uploads")
+async def upload(
+    file: Optional[UploadFile] = File(default=None),
+    image_base64: Optional[str] = Form(default=None),
+    mime_type: Optional[str] = Form(default=None),
+    file_name: Optional[str] = Form(default=None),
+    project: Optional[str] = Form(default=None),
+    api_key: Optional[str] = Depends(require_api_key),
+) -> dict[str, Any]:
+    _, _tier = require_ready(api_key)
+    if file is not None:
+        raw = await file.read()
+        b64 = base64.b64encode(raw).decode("ascii")
+        mime = mime_type or file.content_type or "image/png"
+        name = file_name or file.filename
+    elif image_base64:
+        b64 = image_base64.split(",", 1)[-1]  # tolerate data: prefix
+        mime = mime_type or "image/png"
+        name = file_name
+    else:
+        raise HTTPException(status_code=400, detail="provide a file or image_base64")
+
+    is_video = mime.lower().startswith("video/")
+    name = name or ("upload.mp4" if is_video else "upload.png")
+
+    project_id = project or await get_default_project(api_key)
+    sdk = get_sdk(api_key)
+    if is_video:
+        result = await sdk.upload_video(
+            video_base64=b64, mime_type=mime, project_id=project_id, file_name=name
+        )
+    else:
+        result = await sdk.upload_image(
+            image_base64=b64, mime_type=mime, project_id=project_id, file_name=name
+        )
+    if result.get("error"):
+        raise HTTPException(status_code=502, detail=str(result["error"])[:200])
+    return {
+        "media_id": result.get("media_id"),
+        "project": project_id,
+        "kind": "video" if is_video else "image",
+    }