MLflow 3.15 GenAI evaluation for image generation + editable criteria

README.md

@@ -1,6 +1,6 @@
 # Databricks Vision
 
-Production-ready image generation, editing, analysis, and semantic search on Databricks. One Python library (`image_gen.py`) drives a deployable FastAPI + React Databricks App that supports single-image generate/edit, batch jobs, gallery, semantic search, and bulk import.
+Production-ready image generation, editing, MLflow-based evaluation, and semantic search on Databricks. One Python library (`image_gen.py`) drives a deployable FastAPI + React Databricks App that supports single-image generate/edit, batch jobs, gallery, MLflow GenAI judge evaluation, semantic search, and bulk import.
 
 <p align="center">
   <img src="samples/collage.png" alt="Sample outputs" width="900" />
@@ -10,7 +10,7 @@ Production-ready image generation, editing, analysis, and semantic search on Dat
 
 - **Generate** images with `gpt-image-2` (auto-routes to `gpt-image-1.5` for transparent backgrounds). Any WxH multiple of 16 up to 3840px, total pixels 655K–8.3M, edge ratio ≤ 3:1.
 - **Edit** existing images with prompt + reference. Edits land as new gallery rows.
-- **Analyze** every generated image with `databricks-gpt-5-5`: description, tags, evaluation, 5×0-5 metric sub-scores, missing-elements, safety flags, brand conflicts, an improved-prompt suggestion, and an optional `criteria_evaluation` against user-supplied style guidelines.
+- **Evaluate** every generated image with **MLflow 3.15 GenAI judges** (`make_judge`, multimodal `{{ trace }}` judges that see the image via `get_span_image`): five 0-5 metric dimensions (quality, prompt, text legibility, safe content, + a user-defined *custom* criterion), each with a per-metric rationale, plus a generated improved-prompt suggestion for weak results. Judge criteria are **editable from the Settings page**; traces + scores land in an MLflow experiment. A structured `databricks-gpt-5-5` call still supplies enrichment (description, tags).
 - **Search** the corpus by text or by uploaded image. SigLIP-2 1152-dim embeddings, pgvector HNSW + cosine similarity, single-query join of metadata + similarity. FTS fallback available.
 - **Batch generate** via a Databricks Job: multi-image (one prompt template applied to N inputs) or variations (one source × N variation prompts).
 - **Import** local images in bulk; the app synthesizes a prompt then runs the same analyzer + embedder pipeline so imported images become searchable.
@@ -23,40 +23,48 @@ Production-ready image generation, editing, analysis, and semantic search on Dat
 
 - **Backend** — Python, FastAPI, Pydantic
 - **Frontend** — React 19, TypeScript, TanStack Router / Query, shadcn/ui, Tailwind
-- **Storage** — Unity Catalog Volumes (image bytes) + Lakebase Autoscaling Postgres with pgvector (metadata, evals, embeddings)
-- **Inference** — Databricks Model Serving wrapping the Responses API with `gpt-image-2` / `gpt-image-1.5`; analyzer backed by `databricks-gpt-5-5`; image + text embeddings from a custom SigLIP-2 SO400M/14-384 endpoint
+- **Storage** — Unity Catalog Volumes (image bytes + MLflow trace attachments) + Lakebase Autoscaling Postgres with pgvector (metadata, evals, embeddings)
+- **Evaluation** — MLflow 3.15 GenAI: `make_judge` multimodal `{{ trace }}` judges on `databricks-gpt-5-5`, traced to an MLflow experiment (UC-Volume artifact location); DB-managed, editable criteria
+- **Inference** — Databricks Model Serving wrapping the Responses API with `gpt-image-2` / `gpt-image-1.5`; judge + enrichment backed by `databricks-gpt-5-5`; image + text embeddings from a custom SigLIP-2 SO400M/14-384 endpoint
 - **Build / deploy** — Databricks Asset Bundles, [apx](https://github.com/databricks/apx) toolkit, `uv`, `bun`
 
 ## Architecture
 
 ```
 ┌────────────────────────────────────┐     ┌────────────────────────────────────┐
 │  Databricks App                    │     │  Lakebase Autoscaling Postgres     │
-│  (FastAPI + React)                 │◄───►│  (metadata + pgvector embeddings)  │
+│  (FastAPI + React)                 │◄───►│  (metadata + evals + pgvector,     │
+│                                    │     │   editable eval_criteria table)    │
 └─────────────┬──────────────────────┘     └────────────────────────────────────┘
               │
               ├──► Foundation Model serving
               │     ├── gpt-image-2 / gpt-image-1.5   (generate, edit)
-              │     └── databricks-gpt-5-5            (analyzer, prompt rewrite)
+              │     └── databricks-gpt-5-5            (enrichment + MLflow judges)
+              │
+              ├──► MLflow 3.15 GenAI evaluation
+              │     ├── make_judge {{ trace }} multimodal judges (5 criteria)
+              │     │     see the image via the get_span_image tool
+              │     └── traces + judge scores  →  MLflow experiment
+              │           (attachments stored on a UC Volume)
               │
               ├──► Custom SigLIP-2 SO400M/14-384 endpoint
               │     (text + image  →  1152-dim embeddings)
               │
               └──► Unity Catalog Volumes
-                    (PNG bytes, organised by batch / folder)
+                    (PNG bytes + MLflow trace attachments)
                                 ▲
    Batch Generation Job ────────┘
    (Databricks Job; ai_query()
     against image-generator endpoint)
 ```
 
-**Single-image flow:** UI streams partial-image events over SSE while the model generates, then the backend persists bytes to a UC Volume, writes the row to Lakebase, and runs the analyzer + embedder as background tasks that `UPDATE` the row when they finish.
+**Single-image flow:** UI streams partial-image events over SSE while the model generates, then the backend persists bytes to a UC Volume, writes the row to Lakebase, and (as background tasks) scores the image with MLflow judges + computes the SigLIP embedding, then `UPDATE`s the row. Scoring traces the image so the `{{ trace }}` judges can retrieve it via `get_span_image`; judge scores + rationales mirror back into the Lakebase eval columns the gallery reads.
 
 **Batch flow:** the app kicks off a Databricks Job that reads inputs from a UC Volume, runs `ai_query()` against the image-generator serving endpoint, writes outputs back to a Volume, and syncs metadata + embeddings to Lakebase. The gallery shows single-gen and batch images from the same table.
 
 ## What's interesting
 
-- **Image evaluation** — every generated image is scored on five 0-5 dimensions (quality, prompt adherence, purpose fit, text legibility, safe content) plus categorical safety / brand flags, with a critique paragraph and a suggested improved prompt for off-spec results. See [`ImageAnalyzer` in image_gen.py](image_gen.py).
+- **MLflow multimodal judge evaluation** — every generated image is scored by MLflow 3.15 `make_judge` `{{ trace }}` judges that *see* the image (the judge calls the `get_span_image` tool to fetch the image attachment off the trace) and return a 0-5 score + rationale per dimension. Five criteria (quality, prompt, text, safe, + a user-defined custom one) are stored in an editable `eval_criteria` table and manageable from the Settings page; scores + rationales are mirrored into Lakebase for the gallery and logged to an MLflow experiment for history/comparison. See [`app/src/databricks_vision/backend/core/_eval.py`](app/src/databricks_vision/backend/core/_eval.py).
 - **Semantic search over the corpus** — every image is embedded with SigLIP-2 (1152-dim) at ingest time and stored in pgvector with an HNSW index on cosine similarity. A single SQL query joins metadata filters with similarity ranking, so the gallery can search by text or by an uploaded image without a separate vector store. See [`ImageSearch` in image_gen.py](image_gen.py).
 - **Two-phase persistence** — generate / edit / import endpoints insert the gallery row immediately with placeholder eval fields, then run the analyzer + embedder as background tasks that `UPDATE` the row when ready. The UI polls for a short window after a generate so eval fields appear without manual refresh.
 - **Lakebase Autoscaling with OAuth-rotating connections** — the psycopg pool re-fetches a Lakebase credential token on every new connection, with `max_lifetime=2700` so connections recycle before the 1-hour token expiry. See [`VisionWorkspace` in image_gen.py](image_gen.py).
@@ -129,6 +137,7 @@ Databricks support doesn't cover this content. For questions or bugs, please [op
 | python-dotenv | .env loader | BSD-3-Clause | https://github.com/theskumar/python-dotenv |
 | httpx | HTTP client | BSD-3-Clause | https://github.com/encode/httpx |
 | OpenAI Python SDK | OpenAI client | Apache-2.0 | https://github.com/openai/openai-python |
+| MLflow | GenAI evaluation (judges + tracing) | Apache-2.0 | https://github.com/mlflow/mlflow |
 | Databricks SDK for Python | Databricks client | Apache-2.0 | https://github.com/databricks/databricks-sdk-py |
 | psycopg | PostgreSQL adapter | LGPL-3.0 | https://github.com/psycopg/psycopg |
 | pgvector-python | pgvector client | MIT | https://github.com/pgvector/pgvector-python |

app/app.yml.template

@@ -10,5 +10,7 @@ env:
     value: "${SCHEMA_NAME}"
   - name: DATABRICKS_VISION_EMBEDDING_ENDPOINT
     value: "${EMBEDDING_ENDPOINT_NAME}"
+  - name: DATABRICKS_VISION_EVAL_EXPERIMENT
+    value: "${EVAL_EXPERIMENT}"
   - name: DATABRICKS_VISION_BATCH_GEN_JOB_ID
     valueFrom: batch-gen-job

app/pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "psycopg[binary]>=3.2.0",
     "psycopg-pool>=3.2.0",
     "openai>=2.21.0",
+    "mlflow[databricks]>=3.15.0",  # GenAI eval: make_judge trace judges + tracing
     "sse-starlette>=2.2.0",
     "Pillow>=11.0.0",
     "httpx>=0.28.0",
@@ -31,7 +32,7 @@ image-gen = { path = "..", editable = true }
 [tool.apx.metadata]
 app-name = "databricks-vision"
 app-slug = "databricks_vision"
-app-entrypoint = "databricks_vision.backend.app:app"
+app-module = "databricks_vision.backend.app:app"
 api-prefix = "/api"
 metadata-path = "src/databricks_vision/_metadata.py"
 

app/src/databricks_vision/backend/core/_config.py

@@ -45,6 +45,20 @@ class AppConfig(BaseSettings):
     default_input_fidelity: str = Field(default="high")
     default_output_format: str = Field(default="png")
 
+    # MLflow GenAI evaluation. Judges score every generated image via make_judge
+    # {{ trace }} judges (see backend/core/_eval.py). eval_experiment is injected at
+    # deploy time (DATABRICKS_VISION_EVAL_EXPERIMENT); when empty, eval is skipped and
+    # the app falls back to the plain structured analyzer output.
+    eval_experiment: str = Field(default="")
+    eval_enabled: bool = Field(default=True)
+    # Judge model must be multimodal AND accept function tools; GPT-5.x require
+    # reasoning_effort='none' with tools. Defaults to databricks:/{model_name}.
+    judge_model: str = Field(default="")
+    judge_reasoning_effort: str = Field(default="none")
+    # Images are downscaled to this max edge before becoming a trace attachment,
+    # keeping them under MLflow's ~10MB attachment cap and cutting judge fetch cost.
+    eval_image_max_px: int = Field(default=1024)
+
     @property
     def output_volume(self) -> str:
         return f"/Volumes/{self.catalog}/{self.schema_name}/generated_images"
@@ -54,6 +68,8 @@ def output_volume(self) -> str:
     def model_post_init(self, __context):
         if not self.vision_volume:
             self.vision_volume = f"/Volumes/{self.catalog}/{self.schema_name}/vision_images"
+        if not self.judge_model:
+            self.judge_model = f"databricks:/{self.model_name}"
 
     @property
     def static_assets_path(self) -> Path:
@@ -63,6 +79,13 @@ def __hash__(self) -> int:
         return hash(self.app_name)
 
 
+# --- Static assets dir ---
+# The compiled frontend lives at <package>/__dist__. We derive this here rather than
+# importing a `dist_dir` symbol from the apx-generated `_metadata.py` — newer apx versions
+# no longer emit `dist_dir`, so deriving it keeps the app buildable across apx versions.
+dist_dir = Path(str(resources.files(app_slug))).joinpath("__dist__")
+
+
 # --- Logger ---
 
 logger = logging.getLogger(app_name)

app/src/databricks_vision/backend/core/_database.py

@@ -72,16 +72,75 @@ def _configure(conn: psycopg.Connection) -> None:
     return pool
 
 
+def _split_sql_statements(sql: str) -> list[str]:
+    """Split a multi-statement SQL script on top-level semicolons, ignoring semicolons and `--`
+    line comments inside single-quoted string literals (with '' escapes). _SCHEMA_SQL has no
+    dollar-quoted blocks, so this single-pass scanner is sufficient. `--` comments are handled
+    inline (only when outside a string) rather than pre-stripped by line, so a literal line that
+    begins with `--` inside a string (e.g. free-text seed instructions) is preserved intact."""
+    stmts: list[str] = []
+    buf: list[str] = []
+    in_str = False
+    i = 0
+    n = len(sql)
+    while i < n:
+        ch = sql[i]
+        if in_str:
+            if ch == "'":
+                if i + 1 < n and sql[i + 1] == "'":
+                    buf.append("''")  # escaped quote inside a string literal
+                    i += 2
+                    continue
+                in_str = False
+            buf.append(ch)
+        else:
+            if ch == "-" and i + 1 < n and sql[i + 1] == "-":
+                # line comment (outside a string): skip to end of line, keep the newline
+                nl = sql.find("\n", i)
+                i = n if nl == -1 else nl
+                continue
+            if ch == "'":
+                in_str = True
+                buf.append(ch)
+            elif ch == ";":
+                stmt = "".join(buf).strip()
+                if stmt:
+                    stmts.append(stmt)
+                buf = []
+            else:
+                buf.append(ch)
+        i += 1
+    tail = "".join(buf).strip()
+    if tail:
+        stmts.append(tail)
+    return stmts
+
+
 def _bootstrap_schema(ws, endpoint_name: str, database: str) -> None:
     """Run CREATE EXTENSION + table DDL on a one-shot autocommit connection *before* opening
     the pool, so per-connection `register_vector` callbacks find the type on every connect.
+
+    Statements are executed INDIVIDUALLY (not as one multi-statement query): a single failing
+    statement — e.g. `CREATE EXTENSION` raising "must be owner" for a non-owner service
+    principal — must not abort the transaction and skip everything after it (that previously
+    left newly-added tables like eval_criteria uncreated on existing databases). Per-statement
+    "already exists"/"must be owner" errors are expected and logged at debug; anything else is
+    a warning. All DDL is idempotent (IF NOT EXISTS / ON CONFLICT).
     """
     host, user = _resolve_pg_host_user(ws, endpoint_name)
     port = os.environ.get("PGPORT", "5432")
     conninfo = f"dbname={database} user={user} host={host} port={port} sslmode=require"
     cred = ws.postgres.generate_database_credential(endpoint=endpoint_name)
     with psycopg.connect(conninfo, password=cred.token, autocommit=True) as boot:
-        boot.execute(_SCHEMA_SQL)
+        for stmt in _split_sql_statements(_SCHEMA_SQL):
+            try:
+                boot.execute(stmt)
+            except Exception as exc:
+                msg = str(exc)
+                if "must be owner" in msg or "already exists" in msg:
+                    logger.debug("Schema stmt skipped (%s): %s", msg.splitlines()[0], stmt[:60])
+                else:
+                    logger.warning("Schema stmt failed: %s | stmt: %s", msg.splitlines()[0], stmt[:80])
 
 
 class _DatabasePoolDependency(LifespanDependency):
@@ -98,13 +157,13 @@ async def lifespan(self, app: FastAPI) -> AsyncGenerator[None, None]:
         database = os.environ.get("PGDATABASE", "postgres")
 
         try:
+            # Per-statement execution (see _bootstrap_schema): individual "already exists"/
+            # "must be owner" errors are handled inside; this guard is for a hard connection
+            # failure only. Idempotent DDL means it's safe to run on every startup.
             _bootstrap_schema(ws, endpoint_name, database)
             logger.info("Lakebase schema bootstrap complete")
         except Exception as exc:
-            if "must be owner" in str(exc) or "already exists" in str(exc):
-                logger.info("Schema bootstrap skipped (tables already exist with correct schema)")
-            else:
-                logger.warning("Schema bootstrap encountered: %s", exc)
+            logger.warning("Schema bootstrap could not run (continuing): %s", exc)
 
         pool = _make_pool(ws, endpoint_name, database)
         pool.open(wait=True, timeout=30.0)