diff --git a/.gitignore b/.gitignore index 522b107..3edd9a2 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,5 @@ images.small/ __pycache__/ .DS_Store .vercel +dist/ +package-lock.json diff --git a/README.md b/README.md index d0139a9..337ea05 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,9 @@ This repo supports two complementary systems for image search and browsing: - Local (MLX on Apple Silicon): run a lightweight Flask app that queries your Supabase Postgres with CLIP embeddings using MLX locally. Directory: `mlx-local/`. - Cloud loader and embedding backfill (R2 + Replicate): upload images in `images/` to Cloudflare R2 and backfill missing embeddings via Replicate. Directory: `loader/`. +- Cloud browser: web interface for searching and browsing images. Directory: `browse/`. + +The `loader` and `browse` applications share common code through the `shared/` library, which provides database operations and Replicate API interactions. ## Basic workflow @@ -14,7 +17,7 @@ This repo supports two complementary systems for image search and browsing: mise run sync ``` -This scans `images/`, uploads new files to R2, and ensures a row exists in `image_embeddings`. A separate worker can backfill embeddings via Replicate. +This scans `images/`, uploads new files to R2, extracts image dimensions (width/height), and ensures a row exists in `image_embeddings`. A separate worker can backfill embeddings via Replicate. ## Local MLX browser (optional) @@ -28,15 +31,19 @@ mise setup-supabase-env mise run-mlx-local ``` +## Architecture + +The project uses a shared library (`shared/`) that contains common code for database operations and Replicate API interactions. Both `loader` and `browse` use this shared library to ensure consistency. + ## Cloud loader details See `loader/README.md` for full setup. In short, configure `loader/.env`, then: ```bash -# ensure DB schema exists +# ensure DB schema exists (includes width/height columns) (cd loader && npm run ensure-schema) -# upload images and upsert DB rows (sync is a convenience wrapper) +# upload images, extract dimensions, and upsert DB rows (sync is a convenience wrapper) mise sync # backfill missing embeddings via Replicate diff --git a/browse/package.json b/browse/package.json index 92fe00f..2627c4c 100644 --- a/browse/package.json +++ b/browse/package.json @@ -11,6 +11,7 @@ "dependencies": { "dotenv": "^16.4.5", "express": "^4.19.2", + "image-browser-shared": "file:../shared", "pg": "^8.12.0", "replicate": "^0.31.1" }, diff --git a/browse/src/replicate.ts b/browse/src/replicate.ts index b9ea396..8ee5787 100644 --- a/browse/src/replicate.ts +++ b/browse/src/replicate.ts @@ -1,53 +1,3 @@ -import Replicate from "replicate"; - -type ReplicateModelId = `${string}/${string}` | `${string}/${string}:${string}`; - -// Environment-driven configuration for Replicate text and image embeddings -const replicateToken = process.env.REPLICATE_API_TOKEN; -const DEFAULT_IMAGE_MODEL: ReplicateModelId = - "lucataco/clip-vit-base-patch32:056324d6fb78878c1016e432a3827fa76950022848c5378681dd99b7dc7dcc24"; -// Default text model to krthr/clip-embeddings which returns { embedding: number[] } -const DEFAULT_TEXT_MODEL: ReplicateModelId = - "krthr/clip-embeddings:1c0371070cb827ec3c7f2f28adcdde54b50dcd239aa6faea0bc98b174ef03fb4"; -const TEXT_MODEL: ReplicateModelId = - (process.env.REPLICATE_TEXT_MODEL as ReplicateModelId) || DEFAULT_TEXT_MODEL; -const TEXT_INPUT_KEY = process.env.REPLICATE_TEXT_INPUT_KEY || "text"; -const IMAGE_MODEL: ReplicateModelId = - (process.env.REPLICATE_IMAGE_MODEL as ReplicateModelId) || DEFAULT_IMAGE_MODEL; -const IMAGE_INPUT_KEY = process.env.REPLICATE_IMAGE_INPUT_KEY || "image"; - -function getReplicateClient(): Replicate { - if (!replicateToken) { - throw new Error("REPLICATE_API_TOKEN is not set"); - } - return new Replicate({ auth: replicateToken }); -} - -export async function getTextEmbedding(query: string): Promise { - const client = getReplicateClient(); - const input: Record = {}; - input[TEXT_INPUT_KEY] = query; - const result: unknown = await client.run(TEXT_MODEL, { input }); - // Accept either raw number[] or { embedding: number[] } - if (Array.isArray(result)) { - return result.map((x) => Number(x)); - } - if (result && typeof result === "object" && Array.isArray((result as any).embedding)) { - const { embedding } = result as { embedding: unknown[] }; - return embedding.map((x) => Number(x)); - } - throw new Error("Unexpected embedding result from Replicate (expected number[] or {embedding:number[]})"); -} - -export async function getImageEmbedding(imageUrl: string): Promise { - const client = getReplicateClient(); - const input: Record = {}; - input[IMAGE_INPUT_KEY] = imageUrl; - const result: unknown = await client.run(IMAGE_MODEL, { input }); - if (Array.isArray(result)) { - return result.map((x) => Number(x)); - } - throw new Error("Unexpected embedding result from Replicate (expected number[])"); -} - +// Re-export from shared library for backward compatibility +export { getTextEmbedding, getImageEmbedding } from "image-browser-shared"; diff --git a/browse/src/server.ts b/browse/src/server.ts index faee291..667687d 100644 --- a/browse/src/server.ts +++ b/browse/src/server.ts @@ -5,7 +5,7 @@ import express from "express"; import type { Request, Response } from "express"; import dotenv from "dotenv"; import { Pool, QueryResult } from "pg"; -import { getTextEmbedding } from "./replicate.js"; +import { getTextEmbedding, toVectorParam } from "image-browser-shared"; dotenv.config(); @@ -105,26 +105,32 @@ function listLocalImages(limit = 60): string[] { } } -async function listDbImages(limit = 60): Promise { +interface ImageData { + file_name: string; + width?: number; + height?: number; +} + +async function listDbImages(limit = 60): Promise { if (!pool) return []; - const { rows }: QueryResult<{ file_name: string }> = await pool.query( - `SELECT file_name + const { rows }: QueryResult = await pool.query( + `SELECT file_name, width, height FROM image_embeddings WHERE embedding IS NOT NULL ORDER BY created_at DESC LIMIT $1;`, [limit] ); - return rows.map((r) => r.file_name); + return rows; } -// getTextEmbedding is imported from ./replicate - -function toVectorParam(embedding: number[] | string): string { - if (typeof embedding === "string") return embedding; // expected like "[0.1,0.2,...]" - return `[${embedding.join(",")}]`; +function listLocalImagesAsData(limit = 60): ImageData[] { + const files = listLocalImages(limit); + return files.map(file_name => ({ file_name })); } +// getTextEmbedding and toVectorParam are imported from shared library + // Load external HTML template and provide a renderer const TEMPLATE_PATH = path.join(__dirname, "template.html"); let TEMPLATE_SOURCE = ""; @@ -150,13 +156,20 @@ function googleLensUrl(imageUrl: string): string { return `${base}${encodeURIComponent(imageUrl)}`; } -function renderTemplate(images: string[], query: string | null): string { +function isValidDimension(value: number | null | undefined): boolean { + return typeof value === 'number' && value > 0 && value < 100000 && Number.isInteger(value); +} + +function renderTemplate(images: ImageData[], query: string | null): string { const imagesHtml = images - .map((file) => { - const src = resolveImageUrl(file); + .map((img) => { + const src = resolveImageUrl(img.file_name); const googleLink = USING_REMOTE_IMAGES ? `G` : ""; const downloadLink = `↓`; - return `\n
\n
\n \n \"${file}\"\n \n
\n ${googleLink}${googleLink ? "\n " : ""}${downloadLink}\n
\n
\n
`; + // Add width and height attributes if available and valid for better page loading + const hasValidDimensions = isValidDimension(img.width) && isValidDimension(img.height); + const dimensionAttrs = hasValidDimensions ? ` width="${img.width}" height="${img.height}"` : ""; + return `\n
\n
\n \n \"${img.file_name}\"${dimensionAttrs}\n \n
\n ${googleLink}${googleLink ? "\n " : ""}${downloadLink}\n
\n
\n
`; }) .join(""); @@ -204,24 +217,24 @@ app.get("/", async (req: Request, res: Response) => { } const vec = toVectorParam(embedding); if (!pool) throw new Error("Database not configured"); - const { rows }: QueryResult<{ file_name: string; distance: number }> = await pool.query( - `SELECT file_name, embedding <#> $1::vector AS distance + const { rows }: QueryResult = await pool.query( + `SELECT file_name, width, height, embedding <#> $1::vector AS distance FROM image_embeddings WHERE embedding IS NOT NULL ORDER BY distance ASC LIMIT 30;`, [vec] ); - const images = rows.map((r) => r.file_name); + const images: ImageData[] = rows.map((r) => ({ file_name: r.file_name, width: r.width, height: r.height })); res.type("html").send(renderTemplate(images, q)); return; } // No query: list from DB if available, else fallback to local dir - let images: string[] = []; + let images: ImageData[] = []; if (pool) { images = await listDbImages(60); } else { - images = listLocalImages(60); + images = listLocalImagesAsData(60); } res.type("html").send(renderTemplate(images, null)); } catch (err) { @@ -235,9 +248,9 @@ app.get("/neighbors/:file_name", async (req: Request, res: Response) => { const fileName = req.params.file_name; try { if (!pool) throw new Error("Database not configured"); - // Fetch embedding for the selected image as text for reuse - const { rows: embRows }: QueryResult<{ embedding_text: string }> = await pool.query( - `SELECT embedding::text AS embedding_text + // Fetch embedding and dimensions for the selected image + const { rows: embRows }: QueryResult<{ embedding_text: string; width?: number; height?: number }> = await pool.query( + `SELECT embedding::text AS embedding_text, width, height FROM image_embeddings WHERE file_name = $1 AND embedding IS NOT NULL LIMIT 1;`, @@ -248,15 +261,21 @@ app.get("/neighbors/:file_name", async (req: Request, res: Response) => { return; } const embeddingText = embRows[0].embedding_text; - const { rows }: QueryResult<{ file_name: string; distance: number }> = await pool.query( - `SELECT file_name, embedding <#> $1::vector AS distance + const selectedImage: ImageData = { + file_name: fileName, + width: embRows[0].width, + height: embRows[0].height + }; + const { rows }: QueryResult = await pool.query( + `SELECT file_name, width, height, embedding <#> $1::vector AS distance FROM image_embeddings WHERE file_name != $2 AND embedding IS NOT NULL ORDER BY distance ASC LIMIT 30;`, [embeddingText, fileName] ); - const images = [fileName, ...rows.map((r) => r.file_name)]; + const similarImages: ImageData[] = rows.map((r) => ({ file_name: r.file_name, width: r.width, height: r.height })); + const images = [selectedImage, ...similarImages]; res.type("html").send(renderTemplate(images, null)); } catch (err) { // eslint-disable-next-line no-console diff --git a/loader/README.md b/loader/README.md index b6adbca..98dfbdb 100644 --- a/loader/README.md +++ b/loader/README.md @@ -1,6 +1,8 @@ # Loader (Cloudflare R2 + Embedding Backfill) -This project uploads local images to Cloudflare R2 and ensures there is a row in Postgres (`image_embeddings`). A second worker backfills missing image embeddings using Replicate. +This project uploads local images to Cloudflare R2, extracts image dimensions (width/height), and ensures there is a row in Postgres (`image_embeddings`). A second worker backfills missing image embeddings using Replicate. + +The loader uses shared library code from `../shared/` for database operations and Replicate API interactions, ensuring consistency with the browse application. ## Setup @@ -46,13 +48,13 @@ npm install ## Commands -- Ensure schema (creates `image_embeddings` if missing): +- Ensure schema (creates `image_embeddings` with width/height columns if missing): ```bash npm run ensure-schema ``` -- Upload images (skips ones already in R2) and upsert DB rows: +- Upload images (skips ones already in R2), extract dimensions, and upsert DB rows: ```bash npm run upload diff --git a/loader/package.json b/loader/package.json index 0cb155e..2bedbf6 100644 --- a/loader/package.json +++ b/loader/package.json @@ -6,16 +6,18 @@ "scripts": { "upload": "tsx src/upload_r2.ts", "embed": "tsx src/embed_missing.ts", - "ensure-schema": "tsx src/db.ts ensure-schema", + "ensure-schema": "tsx ../shared/src/db.ts ensure-schema", "sync": "tsx src/sync.ts" }, "dependencies": { "@aws-sdk/client-s3": "^3.614.0", "dotenv": "^16.4.5", + "image-browser-shared": "file:../shared", "mime": "^3.0.0", "p-limit": "^5.0.0", "pg": "^8.12.0", - "replicate": "^0.31.1" + "replicate": "^0.31.1", + "sharp": "^0.33.5" }, "devDependencies": { "@types/node": "^22.8.4", diff --git a/loader/src/db.ts b/loader/src/db.ts index 8c5929d..1a83476 100644 --- a/loader/src/db.ts +++ b/loader/src/db.ts @@ -1,110 +1,2 @@ -import { Pool } from "pg"; -import dotenv from "dotenv"; - -dotenv.config(); - -const DATABASE_URL = process.env.SUPABASE_DB_URL || ""; -const EXPECTED_VECTOR_DIM = process.env.EXPECTED_VECTOR_DIM ? Number(process.env.EXPECTED_VECTOR_DIM) : 768; - -if (!DATABASE_URL) { - // eslint-disable-next-line no-console - console.warn("SUPABASE_DB_URL is not set; DB operations will fail."); -} - -function sslConfigFor(url: string): any { - if (!url) return undefined; - // Enable SSL for Supabase pooler or when PGSSL=require - if (/pooler\.supabase\.com|supabase\.co/.test(url) || (process.env.PGSSL || "").toLowerCase() === "require") { - return { rejectUnauthorized: false }; - } - return undefined; -} - -let pool: Pool | null = null; -export function getPool(): Pool { - if (!pool) { - pool = new Pool({ - connectionString: DATABASE_URL, - ssl: sslConfigFor(DATABASE_URL), - keepAlive: true, - max: process.env.PG_MAX ? Number(process.env.PG_MAX) : 5, - idleTimeoutMillis: process.env.PG_IDLE ? Number(process.env.PG_IDLE) : 10000 - }); - // Avoid noisy logs on expected pool shutdown/termination signals - pool.on("error", (err: any) => { - const code = err?.code as string | undefined; - const msg = String(err?.message || ""); - if ( - code === "57P01" || // admin shutdown - code === "XX000" || // internal error (db_termination observed) - /terminat/i.test(msg) || - /db_termination/i.test(msg) - ) { - return; // ignore expected termination during/after shutdown - } - // eslint-disable-next-line no-console - console.warn("pg pool error (ignored)", err); - }); - } - return pool; -} - -export async function ensureSchema(): Promise { - const dim = EXPECTED_VECTOR_DIM || 768; - // eslint-disable-next-line no-console - console.log(`Ensuring schema image_embeddings with vector(${dim})...`); - const client = await getPool().connect(); - try { - await client.query("create extension if not exists vector;"); - await client.query( - `create table if not exists image_embeddings ( - id serial primary key, - file_name text unique, - embedding vector(${dim}), - created_at timestamp default now() - );` - ); - // Ensure unique index on file_name (in case of legacy table without constraint) - await client.query( - `do $$ begin - if not exists ( - select 1 from pg_indexes where schemaname = current_schema() and indexname = 'image_embeddings_file_name_key' - ) then - begin - alter table image_embeddings add constraint image_embeddings_file_name_key unique (file_name); - exception when duplicate_table then null; end; - end if; - end $$;` - ); - } finally { - client.release(); - } -} - -export async function getAllFileNames(): Promise { - const pool = getPool(); - const { rows } = await pool.query<{ file_name: string }>( - "select file_name from image_embeddings;" - ); - return rows.map((r) => r.file_name); -} - -export function toVectorParam(embedding: number[] | string): string { - if (typeof embedding === "string") return embedding; - return `[${embedding.join(",")}]`; -} - -// CLI support: `tsx src/db.ts ensure-schema` -if (process.argv[2] === "ensure-schema") { - ensureSchema() - .then(() => { - // eslint-disable-next-line no-console - console.log("Schema ensured."); - process.exit(0); - }) - .catch((err) => { - // eslint-disable-next-line no-console - console.error("Failed to ensure schema", err); - process.exit(1); - }); -} +// Re-export from shared library for backward compatibility +export { getPool, ensureSchema, getAllFileNames, toVectorParam } from "image-browser-shared"; diff --git a/loader/src/replicate.ts b/loader/src/replicate.ts index ef96542..f3c8c22 100644 --- a/loader/src/replicate.ts +++ b/loader/src/replicate.ts @@ -1,44 +1,2 @@ -import Replicate from "replicate"; -import dotenv from "dotenv"; - -dotenv.config(); - -type ReplicateModelId = `${string}/${string}` | `${string}/${string}:${string}`; - -const replicateToken = process.env.REPLICATE_API_TOKEN; -const DEFAULT_IMAGE_MODEL: ReplicateModelId = - (process.env.REPLICATE_IMAGE_MODEL as ReplicateModelId) || - ("krthr/clip-embeddings:1c0371070cb827ec3c7f2f28adcdde54b50dcd239aa6faea0bc98b174ef03fb4" as ReplicateModelId); -const IMAGE_INPUT_KEY = process.env.REPLICATE_IMAGE_INPUT_KEY || "image"; -const EXPECTED_VECTOR_DIM = process.env.EXPECTED_VECTOR_DIM ? Number(process.env.EXPECTED_VECTOR_DIM) : 768; - -function getReplicateClient(): Replicate { - if (!replicateToken) { - throw new Error("REPLICATE_API_TOKEN is not set"); - } - return new Replicate({ auth: replicateToken }); -} - -export async function getImageEmbedding(imageUrl: string): Promise { - const modelId = DEFAULT_IMAGE_MODEL; - const client = getReplicateClient(); - const input: Record = {}; - input[IMAGE_INPUT_KEY] = imageUrl; - const result: unknown = await client.run(modelId, { input }); - if (Array.isArray(result)) { - const embedding = result.map((x) => Number(x)); - if (EXPECTED_VECTOR_DIM && embedding.length !== EXPECTED_VECTOR_DIM) { - throw new Error(`Embedding dimension mismatch: got ${embedding.length}, expected ${EXPECTED_VECTOR_DIM}`); - } - return embedding; - } - if (result && typeof result === "object" && Array.isArray((result as any).embedding)) { - const { embedding } = result as { embedding: unknown[] }; - const vec = embedding.map((x) => Number(x)); - if (EXPECTED_VECTOR_DIM && vec.length !== EXPECTED_VECTOR_DIM) { - throw new Error(`Embedding dimension mismatch: got ${vec.length}, expected ${EXPECTED_VECTOR_DIM}`); - } - return vec; - } - throw new Error("Unexpected embedding result from Replicate (expected number[] or {embedding:number[]})"); -} +// Re-export from shared library for backward compatibility +export { getImageEmbedding } from "image-browser-shared"; diff --git a/loader/src/upload_r2.ts b/loader/src/upload_r2.ts index 0e25975..ab4367c 100644 --- a/loader/src/upload_r2.ts +++ b/loader/src/upload_r2.ts @@ -2,6 +2,7 @@ import fs from "fs"; import path from "path"; import { fileURLToPath } from "url"; import dotenv from "dotenv"; +import sharp from "sharp"; import { getR2Client, existsInBucket, uploadFile, toKey } from "./r2.js"; import { ensureSchema, getPool, getAllFileNames } from "./db.js"; @@ -29,12 +30,30 @@ function listLocalImages(): string[] { } } -async function upsertImageRow(fileName: string): Promise { +async function getImageDimensions(filePath: string): Promise<{ width: number; height: number } | null> { + try { + const metadata = await sharp(filePath).metadata(); + if (metadata.width && metadata.height) { + return { width: metadata.width, height: metadata.height }; + } + } catch (err) { + // eslint-disable-next-line no-console + console.warn(`Failed to extract dimensions for ${filePath}:`, err); + } + return null; +} + +async function upsertImageRow(fileName: string, filePath: string): Promise { const pool = getPool(); + const dims = await getImageDimensions(filePath); + + // Single query that handles both cases - with or without dimensions await pool.query( - `insert into image_embeddings (file_name) values ($1) - on conflict (file_name) do nothing;`, - [fileName] + `insert into image_embeddings (file_name, width, height) values ($1, $2, $3) + on conflict (file_name) do update set + width = COALESCE(EXCLUDED.width, image_embeddings.width), + height = COALESCE(EXCLUDED.height, image_embeddings.height);`, + [fileName, dims?.width || null, dims?.height || null] ); } @@ -90,7 +109,7 @@ async function main(): Promise { await uploadFile(s3, filePath, key); // eslint-disable-next-line no-console console.log(`Uploaded ${fileName}`); - await upsertImageRow(fileName); + await upsertImageRow(fileName, filePath); return; } @@ -105,7 +124,7 @@ async function main(): Promise { // eslint-disable-next-line no-console console.log(`Already exists in bucket, skipping upload: ${fileName}`); } - await upsertImageRow(fileName); + await upsertImageRow(fileName, filePath); }); // eslint-disable-next-line no-console console.log("Upload + DB upsert complete. Closing DB pool..."); diff --git a/shared/README.md b/shared/README.md new file mode 100644 index 0000000..bbda05b --- /dev/null +++ b/shared/README.md @@ -0,0 +1,69 @@ +# Shared Library + +This package contains shared code used by both `loader` and `browse` applications to ensure consistency across the image browser system. + +## Contents + +### Database Module (`src/db.ts`) + +Provides database connection and schema management: + +- `getPool()` - Returns a singleton PostgreSQL connection pool +- `ensureSchema()` - Creates the `image_embeddings` table with vector extension, including width/height columns +- `getAllFileNames()` - Retrieves all file names from the database +- `toVectorParam()` - Converts embedding arrays to PostgreSQL vector format + +The schema includes: +- `file_name` - Unique identifier for each image +- `embedding` - Vector embedding for similarity search +- `width` and `height` - Image dimensions for optimized page loading +- `created_at` - Timestamp + +### Replicate Module (`src/replicate.ts`) + +Provides embedding generation via Replicate API: + +- `getImageEmbedding(imageUrl)` - Generates embedding vector from an image URL +- `getTextEmbedding(query)` - Generates embedding vector from text query + +Both functions: +- Support configurable models via environment variables +- Validate embedding dimensions against `EXPECTED_VECTOR_DIM` +- Handle multiple response formats from Replicate API + +## Environment Variables + +The shared library respects the following environment variables: + +- `SUPABASE_DB_URL` - PostgreSQL connection string +- `REPLICATE_API_TOKEN` - Replicate API authentication token +- `EXPECTED_VECTOR_DIM` - Expected embedding dimension (default: 768) +- `REPLICATE_IMAGE_MODEL` - Custom image embedding model +- `REPLICATE_TEXT_MODEL` - Custom text embedding model +- `REPLICATE_IMAGE_INPUT_KEY` - Input key for image model (default: "image") +- `REPLICATE_TEXT_INPUT_KEY` - Input key for text model (default: "text") + +## Usage + +Both `loader` and `browse` reference this package via `file:../shared` in their `package.json`: + +```json +{ + "dependencies": { + "image-browser-shared": "file:../shared" + } +} +``` + +Then import the functions: + +```typescript +import { getPool, ensureSchema, getImageEmbedding, getTextEmbedding } from "image-browser-shared"; +``` + +## Benefits + +1. **Single Source of Truth**: Database schema and API interactions are defined once +2. **Consistency**: Both applications use identical embedding logic +3. **Maintainability**: Bug fixes and improvements only need to be made in one place +4. **Type Safety**: Shared TypeScript types ensure compatibility diff --git a/shared/package.json b/shared/package.json new file mode 100644 index 0000000..cd25fc9 --- /dev/null +++ b/shared/package.json @@ -0,0 +1,17 @@ +{ + "name": "image-browser-shared", + "version": "0.1.0", + "private": true, + "type": "module", + "main": "./src/index.ts", + "dependencies": { + "dotenv": "^16.4.5", + "pg": "^8.12.0", + "replicate": "^0.31.1" + }, + "devDependencies": { + "@types/node": "^22.8.4", + "@types/pg": "^8.11.6", + "typescript": "^5.6.3" + } +} diff --git a/shared/src/db.ts b/shared/src/db.ts new file mode 100644 index 0000000..28cf835 --- /dev/null +++ b/shared/src/db.ts @@ -0,0 +1,126 @@ +import { Pool } from "pg"; +import { fileURLToPath } from "url"; +import dotenv from "dotenv"; + +dotenv.config(); + +const DATABASE_URL = process.env.SUPABASE_DB_URL || ""; +const EXPECTED_VECTOR_DIM = process.env.EXPECTED_VECTOR_DIM ? Number(process.env.EXPECTED_VECTOR_DIM) : 768; + +if (!DATABASE_URL) { + // eslint-disable-next-line no-console + console.warn("SUPABASE_DB_URL is not set; DB operations will fail."); +} + +function sslConfigFor(url: string): any { + if (!url) return undefined; + // Enable SSL for Supabase pooler or when PGSSL=require + if (/pooler\.supabase\.com|supabase\.co/.test(url) || (process.env.PGSSL || "").toLowerCase() === "require") { + return { rejectUnauthorized: false }; + } + return undefined; +} + +let pool: Pool | null = null; +export function getPool(): Pool { + if (!pool) { + pool = new Pool({ + connectionString: DATABASE_URL, + ssl: sslConfigFor(DATABASE_URL), + keepAlive: true, + max: process.env.PG_MAX ? Number(process.env.PG_MAX) : 5, + idleTimeoutMillis: process.env.PG_IDLE ? Number(process.env.PG_IDLE) : 10000 + }); + // Avoid noisy logs on expected pool shutdown/termination signals + pool.on("error", (err: any) => { + const code = err?.code as string | undefined; + const msg = String(err?.message || ""); + if ( + code === "57P01" || // admin shutdown + code === "XX000" || // internal error (db_termination observed) + /terminat/i.test(msg) || + /db_termination/i.test(msg) + ) { + return; // ignore expected termination during/after shutdown + } + // eslint-disable-next-line no-console + console.warn("pg pool error (ignored)", err); + }); + } + return pool; +} + +export async function ensureSchema(): Promise { + const dim = EXPECTED_VECTOR_DIM || 768; + // eslint-disable-next-line no-console + console.log(`Ensuring schema image_embeddings with vector(${dim})...`); + const client = await getPool().connect(); + try { + await client.query("create extension if not exists vector;"); + await client.query( + `create table if not exists image_embeddings ( + id serial primary key, + file_name text unique, + embedding vector(${dim}), + width integer, + height integer, + created_at timestamp default now() + );` + ); + // Ensure unique index on file_name (in case of legacy table without constraint) + await client.query( + `do $$ begin + if not exists ( + select 1 from pg_indexes where schemaname = current_schema() and indexname = 'image_embeddings_file_name_key' + ) then + begin + alter table image_embeddings add constraint image_embeddings_file_name_key unique (file_name); + exception when duplicate_table then null; end; + end if; + end $$;` + ); + // Add width and height columns if they don't exist (for migration) + await client.query( + `do $$ begin + if not exists (select 1 from information_schema.columns where table_name = 'image_embeddings' and column_name = 'width') then + alter table image_embeddings add column width integer; + end if; + if not exists (select 1 from information_schema.columns where table_name = 'image_embeddings' and column_name = 'height') then + alter table image_embeddings add column height integer; + end if; + end $$;` + ); + } finally { + client.release(); + } +} + +export async function getAllFileNames(): Promise { + const pool = getPool(); + const { rows } = await pool.query<{ file_name: string }>( + "select file_name from image_embeddings;" + ); + return rows.map((r) => r.file_name); +} + +export function toVectorParam(embedding: number[] | string): string { + if (typeof embedding === "string") return embedding; + return `[${embedding.join(",")}]`; +} + +// CLI support: `tsx src/db.ts ensure-schema` +// Only run if this module is the main entry point +const __filename = fileURLToPath(import.meta.url); +if (process.argv[1] === __filename && process.argv[2] === "ensure-schema") { + ensureSchema() + .then(() => { + // eslint-disable-next-line no-console + console.log("Schema ensured."); + process.exit(0); + }) + .catch((err) => { + // eslint-disable-next-line no-console + console.error("Failed to ensure schema", err); + process.exit(1); + }); +} diff --git a/shared/src/index.ts b/shared/src/index.ts new file mode 100644 index 0000000..dad3bcc --- /dev/null +++ b/shared/src/index.ts @@ -0,0 +1,2 @@ +export { getPool, ensureSchema, getAllFileNames, toVectorParam } from "./db.js"; +export { getImageEmbedding, getTextEmbedding } from "./replicate.js"; diff --git a/shared/src/replicate.ts b/shared/src/replicate.ts new file mode 100644 index 0000000..dcf7142 --- /dev/null +++ b/shared/src/replicate.ts @@ -0,0 +1,61 @@ +import Replicate from "replicate"; +import dotenv from "dotenv"; + +dotenv.config(); + +type ReplicateModelId = `${string}/${string}` | `${string}/${string}:${string}`; + +const replicateToken = process.env.REPLICATE_API_TOKEN; +const DEFAULT_IMAGE_MODEL: ReplicateModelId = + (process.env.REPLICATE_IMAGE_MODEL as ReplicateModelId) || + ("krthr/clip-embeddings:1c0371070cb827ec3c7f2f28adcdde54b50dcd239aa6faea0bc98b174ef03fb4" as ReplicateModelId); +const DEFAULT_TEXT_MODEL: ReplicateModelId = + (process.env.REPLICATE_TEXT_MODEL as ReplicateModelId) || + ("krthr/clip-embeddings:1c0371070cb827ec3c7f2f28adcdde54b50dcd239aa6faea0bc98b174ef03fb4" as ReplicateModelId); +const IMAGE_INPUT_KEY = process.env.REPLICATE_IMAGE_INPUT_KEY || "image"; +const TEXT_INPUT_KEY = process.env.REPLICATE_TEXT_INPUT_KEY || "text"; +const EXPECTED_VECTOR_DIM = process.env.EXPECTED_VECTOR_DIM ? Number(process.env.EXPECTED_VECTOR_DIM) : 768; + +function getReplicateClient(): Replicate { + if (!replicateToken) { + throw new Error("REPLICATE_API_TOKEN is not set"); + } + return new Replicate({ auth: replicateToken }); +} + +function parseEmbeddingResult(result: unknown, context: string): number[] { + if (Array.isArray(result)) { + const embedding = result.map((x) => Number(x)); + if (EXPECTED_VECTOR_DIM && embedding.length !== EXPECTED_VECTOR_DIM) { + throw new Error(`${context}: Embedding dimension mismatch: got ${embedding.length}, expected ${EXPECTED_VECTOR_DIM}`); + } + return embedding; + } + if (result && typeof result === "object" && Array.isArray((result as any).embedding)) { + const { embedding } = result as { embedding: unknown[] }; + const vec = embedding.map((x) => Number(x)); + if (EXPECTED_VECTOR_DIM && vec.length !== EXPECTED_VECTOR_DIM) { + throw new Error(`${context}: Embedding dimension mismatch: got ${vec.length}, expected ${EXPECTED_VECTOR_DIM}`); + } + return vec; + } + throw new Error(`${context}: Unexpected embedding result from Replicate (expected number[] or {embedding:number[]})`); +} + +export async function getImageEmbedding(imageUrl: string): Promise { + const modelId = DEFAULT_IMAGE_MODEL; + const client = getReplicateClient(); + const input: Record = {}; + input[IMAGE_INPUT_KEY] = imageUrl; + const result: unknown = await client.run(modelId, { input }); + return parseEmbeddingResult(result, "getImageEmbedding"); +} + +export async function getTextEmbedding(query: string): Promise { + const modelId = DEFAULT_TEXT_MODEL; + const client = getReplicateClient(); + const input: Record = {}; + input[TEXT_INPUT_KEY] = query; + const result: unknown = await client.run(modelId, { input }); + return parseEmbeddingResult(result, "getTextEmbedding"); +} diff --git a/shared/tsconfig.json b/shared/tsconfig.json new file mode 100644 index 0000000..3fb41fe --- /dev/null +++ b/shared/tsconfig.json @@ -0,0 +1,16 @@ +{ + "compilerOptions": { + "target": "ES2022", + "module": "ES2022", + "lib": ["ES2022"], + "moduleResolution": "node", + "esModuleInterop": true, + "skipLibCheck": true, + "strict": true, + "resolveJsonModule": true, + "outDir": "./dist", + "rootDir": "./src" + }, + "include": ["src/**/*"], + "exclude": ["node_modules", "dist"] +} diff --git a/verify-consolidation.sh b/verify-consolidation.sh new file mode 100755 index 0000000..b21d39e --- /dev/null +++ b/verify-consolidation.sh @@ -0,0 +1,109 @@ +#!/bin/bash +set -e + +echo "=== Verifying Consolidated Image Browser ===" +echo "" + +echo "1. Checking shared library structure..." +if [ -d "shared/src" ]; then + echo " ✓ shared/src directory exists" +else + echo " ✗ shared/src directory missing" + exit 1 +fi + +if [ -f "shared/src/db.ts" ]; then + echo " ✓ shared/src/db.ts exists" +else + echo " ✗ shared/src/db.ts missing" + exit 1 +fi + +if [ -f "shared/src/replicate.ts" ]; then + echo " ✓ shared/src/replicate.ts exists" +else + echo " ✗ shared/src/replicate.ts missing" + exit 1 +fi + +if [ -f "shared/src/index.ts" ]; then + echo " ✓ shared/src/index.ts exists" +else + echo " ✗ shared/src/index.ts missing" + exit 1 +fi + +echo "" +echo "2. Checking loader integration..." +if grep -q "image-browser-shared" loader/package.json; then + echo " ✓ loader references shared library" +else + echo " ✗ loader does not reference shared library" + exit 1 +fi + +if grep -q "image-browser-shared" loader/src/db.ts; then + echo " ✓ loader/src/db.ts uses shared library" +else + echo " ✗ loader/src/db.ts does not use shared library" + exit 1 +fi + +if grep -q "sharp" loader/package.json; then + echo " ✓ loader includes sharp for image dimension extraction" +else + echo " ✗ loader missing sharp dependency" + exit 1 +fi + +echo "" +echo "3. Checking browse integration..." +if grep -q "image-browser-shared" browse/package.json; then + echo " ✓ browse references shared library" +else + echo " ✗ browse does not reference shared library" + exit 1 +fi + +if grep -q "image-browser-shared" browse/src/replicate.ts; then + echo " ✓ browse/src/replicate.ts uses shared library" +else + echo " ✗ browse/src/replicate.ts does not use shared library" + exit 1 +fi + +echo "" +echo "4. Checking database schema includes width/height..." +if grep -q "width integer" shared/src/db.ts; then + echo " ✓ Schema includes width column" +else + echo " ✗ Schema missing width column" + exit 1 +fi + +if grep -q "height integer" shared/src/db.ts; then + echo " ✓ Schema includes height column" +else + echo " ✗ Schema missing height column" + exit 1 +fi + +echo "" +echo "5. Verifying TypeScript compilation..." +cd browse +if npm run build > /dev/null 2>&1; then + echo " ✓ browse compiles successfully" +else + echo " ✗ browse compilation failed" + exit 1 +fi +cd .. + +echo "" +echo "✓ All consolidation checks passed!" +echo "" +echo "Summary:" +echo " - Shared library created with common DB and Replicate code" +echo " - Loader uses shared library and extracts image dimensions" +echo " - Browse uses shared library and displays image dimensions" +echo " - Database schema includes width and height columns"