149 lines
4.7 KiB
TypeScript
149 lines
4.7 KiB
TypeScript
import {
|
|
pgTable,
|
|
uuid,
|
|
varchar,
|
|
text,
|
|
integer,
|
|
jsonb,
|
|
timestamp,
|
|
pgEnum,
|
|
index,
|
|
check,
|
|
type AnyPgColumn,
|
|
} from 'drizzle-orm/pg-core';
|
|
import { sql } from 'drizzle-orm';
|
|
import { projects } from './projects';
|
|
import { flows } from './flows';
|
|
import { apiKeys } from './apiKeys';
|
|
|
|
// Enum for generation status
|
|
export const generationStatusEnum = pgEnum('generation_status', [
|
|
'pending',
|
|
'processing',
|
|
'success',
|
|
'failed',
|
|
]);
|
|
|
|
// Type for referenced images JSONB
|
|
export type ReferencedImage = {
|
|
imageId: string;
|
|
alias: string;
|
|
};
|
|
|
|
export const generations = pgTable(
|
|
'generations',
|
|
{
|
|
id: uuid('id').primaryKey().defaultRandom(),
|
|
|
|
// Relations
|
|
projectId: uuid('project_id')
|
|
.notNull()
|
|
.references(() => projects.id, { onDelete: 'cascade' }),
|
|
flowId: uuid('flow_id').references(() => flows.id, { onDelete: 'set null' }),
|
|
pendingFlowId: text('pending_flow_id'), // Temporary UUID for lazy flow pattern
|
|
apiKeyId: uuid('api_key_id').references(() => apiKeys.id, { onDelete: 'set null' }),
|
|
|
|
// Status
|
|
status: generationStatusEnum('status').notNull().default('pending'),
|
|
|
|
// Prompts (Section 2.1: Reversed semantics)
|
|
// prompt: The prompt that was ACTUALLY USED for generation (enhanced OR original)
|
|
// originalPrompt: User's ORIGINAL input, only stored if autoEnhance was used
|
|
prompt: text('prompt').notNull(), // Prompt used for generation
|
|
originalPrompt: text('original_prompt'), // User's original (nullable, only if enhanced)
|
|
|
|
// Generation parameters
|
|
aspectRatio: varchar('aspect_ratio', { length: 10 }),
|
|
width: integer('width'),
|
|
height: integer('height'),
|
|
|
|
// AI Model
|
|
modelName: varchar('model_name', { length: 100 }).notNull().default('gemini-flash-image-001'),
|
|
modelVersion: varchar('model_version', { length: 50 }),
|
|
|
|
// Result
|
|
outputImageId: uuid('output_image_id').references(
|
|
(): AnyPgColumn => {
|
|
const { images } = require('./images');
|
|
return images.id;
|
|
},
|
|
{ onDelete: 'set null' },
|
|
),
|
|
|
|
// Referenced images used in generation
|
|
// Format: [{ "imageId": "uuid", "alias": "@product" }, ...]
|
|
referencedImages: jsonb('referenced_images').$type<ReferencedImage[]>(),
|
|
|
|
// Error handling
|
|
errorMessage: text('error_message'),
|
|
errorCode: varchar('error_code', { length: 50 }),
|
|
retryCount: integer('retry_count').notNull().default(0),
|
|
|
|
// Metrics
|
|
processingTimeMs: integer('processing_time_ms'),
|
|
cost: integer('cost'), // In cents (USD)
|
|
|
|
// Request context
|
|
requestId: uuid('request_id'),
|
|
userAgent: text('user_agent'),
|
|
ipAddress: text('ip_address'),
|
|
|
|
// Metadata
|
|
meta: jsonb('meta').$type<Record<string, unknown>>().notNull().default({}),
|
|
|
|
// Audit
|
|
createdAt: timestamp('created_at').notNull().defaultNow(),
|
|
updatedAt: timestamp('updated_at')
|
|
.notNull()
|
|
.defaultNow()
|
|
.$onUpdate(() => new Date()),
|
|
},
|
|
(table) => ({
|
|
// CHECK constraints
|
|
statusSuccessCheck: check(
|
|
'status_success_check',
|
|
sql`(${table.status} = 'success' AND ${table.outputImageId} IS NOT NULL) OR (${table.status} != 'success')`,
|
|
),
|
|
statusFailedCheck: check(
|
|
'status_failed_check',
|
|
sql`(${table.status} = 'failed' AND ${table.errorMessage} IS NOT NULL) OR (${table.status} != 'failed')`,
|
|
),
|
|
retryCountCheck: check('retry_count_check', sql`${table.retryCount} >= 0`),
|
|
processingTimeCheck: check(
|
|
'processing_time_check',
|
|
sql`${table.processingTimeMs} IS NULL OR ${table.processingTimeMs} >= 0`,
|
|
),
|
|
costCheck: check('cost_check', sql`${table.cost} IS NULL OR ${table.cost} >= 0`),
|
|
|
|
// Indexes
|
|
// Index for querying generations by project and status
|
|
projectStatusIdx: index('idx_generations_project_status').on(
|
|
table.projectId,
|
|
table.status,
|
|
table.createdAt.desc(),
|
|
),
|
|
|
|
// Index for flow-scoped generations (partial index)
|
|
flowIdx: index('idx_generations_flow')
|
|
.on(table.flowId, table.createdAt.desc())
|
|
.where(sql`${table.flowId} IS NOT NULL`),
|
|
|
|
// Index for pending flow-scoped generations (partial index)
|
|
pendingFlowIdx: index('idx_generations_pending_flow')
|
|
.on(table.pendingFlowId, table.createdAt.desc())
|
|
.where(sql`${table.pendingFlowId} IS NOT NULL`),
|
|
|
|
// Index for output image lookup
|
|
outputIdx: index('idx_generations_output').on(table.outputImageId),
|
|
|
|
// Index for request correlation
|
|
requestIdx: index('idx_generations_request').on(table.requestId),
|
|
|
|
// Index for API key audit trail
|
|
apiKeyIdx: index('idx_generations_api_key').on(table.apiKeyId),
|
|
}),
|
|
);
|
|
|
|
export type Generation = typeof generations.$inferSelect;
|
|
export type NewGeneration = typeof generations.$inferInsert;
|