@sixb/pg 0.1.21 → 0.1.22

package/src/ontology-storage/unique-values.ts

@@ -0,0 +1,274 @@
+import {
+  uniquePropertyNotPrepared,
+  uniqueValueViolation,
+} from "@sixb/core/internal/ontology-storage-provider"
+import type {
+  ExactUniqueValueWrite,
+  OntologyUniqueValueStorage,
+  PreparedUniqueProperty,
+  UniqueValueEntry,
+  UniqueValueRef,
+} from "@sixb/core/storage"
+import type postgres from "postgres"
+import type { SQLClient } from "../pg-client"
+import { jsonParameter, type PgRootOperation } from "./shared"
+
+type Query = postgres.PendingQuery<postgres.Row[]>
+
+interface PgClaimCounts {
+  readonly staged: number | string
+  readonly written: number | string
+  readonly unprepared_type: string | null
+  readonly unprepared_property: string | null
+}
+
+/**
+ * Gives each candidate value to its object, set-based. `candidates` selects `object_type_id`,
+ * `property_id`, `value_key` and `primary_id`, and with `requirePrepared` also the claim's
+ * `key_kind` and `key_version`, which must match how the property was prepared.
+ *
+ * The insert goes first and conflicts do nothing, as the plan writer's statements do: only when
+ * it writes fewer rows than it was given does a second query look for the holder to name. A value
+ * another object holds, or that two candidates give to different objects, throws
+ * `UniqueValueViolationError`; a value already held by its own object is no conflict. A value a
+ * concurrent transaction claimed after this one's snapshot fails serialization instead, so the
+ * commit retries and then finds it held.
+ */
+export async function claimPgUniqueValues(
+  sql: SQLClient,
+  projectId: string,
+  candidates: Query,
+  options: { readonly requirePrepared: boolean }
+): Promise<number> {
+  const unprepared = options.requirePrepared
+    ? sql`
+        SELECT claimed.object_type_id, claimed.property_id
+        FROM (
+          SELECT DISTINCT object_type_id, property_id, key_kind, key_version FROM candidates
+        ) AS claimed
+        WHERE NOT EXISTS (
+          SELECT 1 FROM object_unique_properties AS prepared
+          WHERE prepared.project_id = ${projectId}
+            AND prepared.object_type_id = claimed.object_type_id
+            AND prepared.property_id = claimed.property_id
+            AND prepared.key_kind = claimed.key_kind
+            AND prepared.key_version = claimed.key_version
+        )
+        LIMIT 1
+      `
+    : sql`SELECT NULL::text AS object_type_id, NULL::text AS property_id WHERE false`
+
+  const [counts] = await sql<PgClaimCounts[]>`
+    WITH candidates AS MATERIALIZED (${candidates}),
+    unprepared AS (${unprepared}),
+    written AS (
+      INSERT INTO object_unique_values (
+        project_id, object_type_id, property_id, value_hash, value_key, primary_id
+      )
+      SELECT ${projectId}, object_type_id, property_id,
+        sha256(convert_to(value_key, 'UTF8')), value_key, primary_id
+      FROM candidates
+      WHERE NOT EXISTS (SELECT 1 FROM unprepared)
+      ON CONFLICT DO NOTHING
+      RETURNING 1
+    )
+    SELECT (SELECT COUNT(*) FROM candidates) AS staged, (SELECT COUNT(*) FROM written) AS written,
+      (SELECT object_type_id FROM unprepared) AS unprepared_type,
+      (SELECT property_id FROM unprepared) AS unprepared_property
+  `
+
+  if (counts?.unprepared_type && counts.unprepared_property) {
+    throw uniquePropertyNotPrepared({
+      objectTypeId: counts.unprepared_type,
+      propertyId: counts.unprepared_property,
+    })
+  }
+
+  const written = Number(counts?.written ?? 0)
+  if (written !== Number(counts?.staged ?? 0)) await assertNoOtherHolder(sql, projectId, candidates)
+  return written
+}
+
+/** Names a candidate whose value another object holds, now that every candidate was inserted. */
+async function assertNoOtherHolder(
+  sql: SQLClient,
+  projectId: string,
+  candidates: Query
+): Promise<void> {
+  const [conflict] = await sql<
+    {
+      readonly object_type_id: string
+      readonly property_id: string
+      readonly holder: string
+      readonly claimant: string
+    }[]
+  >`
+    WITH candidates AS MATERIALIZED (${candidates})
+    SELECT candidates.object_type_id, candidates.property_id,
+      held.primary_id AS holder, candidates.primary_id AS claimant
+    FROM candidates
+    CROSS JOIN LATERAL (${heldValue(sql, projectId)} LIMIT 1) AS held
+    WHERE held.primary_id <> candidates.primary_id
+    LIMIT 1
+  `
+  if (!conflict) return
+
+  throw uniqueValueViolation(
+    { objectTypeId: conflict.object_type_id, propertyId: conflict.property_id },
+    conflict.holder,
+    conflict.claimant
+  )
+}
+
+/**
+ * The stored row of `candidates`' value, by its whole primary key: a lateral lookup per candidate
+ * keeps the planner on the index whatever it guesses of the table.
+ */
+export function heldValue(sql: SQLClient, projectId: string) {
+  return sql`
+    SELECT stored.primary_id, stored.value_hash
+    FROM object_unique_values AS stored
+    WHERE stored.project_id = ${projectId}
+      AND stored.object_type_id = candidates.object_type_id
+      AND stored.property_id = candidates.property_id
+      AND stored.value_hash = sha256(convert_to(candidates.value_key, 'UTF8'))
+      AND stored.value_key = candidates.value_key
+  `
+}
+
+/** The objects holding these unique values; a value no object holds is left out. */
+export async function pgUniqueHolders(
+  sql: SQLClient,
+  projectId: string,
+  values: readonly UniqueValueRef[]
+): Promise<ExactUniqueValueWrite[]> {
+  if (values.length === 0) return []
+  const requested = values.map((value) => ({
+    object_type_id: value.objectTypeId,
+    property_id: value.propertyId,
+    value_key: value.valueKey,
+  }))
+  const rows = await sql<
+    {
+      readonly object_type_id: string
+      readonly property_id: string
+      readonly value_key: string
+      readonly primary_id: string
+    }[]
+  >`
+    WITH candidates AS MATERIALIZED (
+      SELECT * FROM jsonb_to_recordset(${jsonParameter(sql, requested)}::jsonb)
+        AS requested(object_type_id text, property_id text, value_key text)
+    )
+    SELECT candidates.object_type_id, candidates.property_id, candidates.value_key,
+      held.primary_id
+    FROM candidates
+    CROSS JOIN LATERAL (${heldValue(sql, projectId)} LIMIT 1) AS held
+  `
+  return rows.map((row) => ({
+    objectTypeId: row.object_type_id,
+    propertyId: row.property_id,
+    valueKey: row.value_key,
+    primaryId: row.primary_id,
+  }))
+}
+
+export class PgOntologyUniqueValueStorage implements OntologyUniqueValueStorage {
+  constructor(private readonly runRootOperation: PgRootOperation) {}
+
+  /**
+   * `LOCK TABLE` takes no snapshot, so as the transaction's first statement it waits before the
+   * serializable snapshot is taken, and a preparer that waited sees the values the previous one
+   * built instead of building them again. An advisory lock would be taken by a `SELECT`, after
+   * the snapshot. Only preparation writes this table, so claims, which read it, never wait.
+   */
+  async lockPreparation(_input: { readonly projectId: string }): Promise<void> {
+    await this.runRootOperation(async (sql) => {
+      await sql`LOCK TABLE object_unique_properties IN SHARE ROW EXCLUSIVE MODE`
+    })
+  }
+
+  async listPrepared(input: {
+    readonly projectId: string
+  }): Promise<readonly PreparedUniqueProperty[]> {
+    return this.runRootOperation(async (sql) => {
+      const rows = await sql<
+        {
+          readonly object_type_id: string
+          readonly property_id: string
+          readonly key_kind: PreparedUniqueProperty["keyKind"]
+          readonly key_version: number
+        }[]
+      >`
+        SELECT object_type_id, property_id, key_kind, key_version FROM object_unique_properties
+        WHERE project_id = ${input.projectId}
+        ORDER BY object_type_id, property_id
+      `
+      return rows.map((row) => ({
+        objectTypeId: row.object_type_id,
+        propertyId: row.property_id,
+        keyKind: row.key_kind,
+        keyVersion: row.key_version,
+      }))
+    })
+  }
+
+  async clear(input: {
+    readonly projectId: string
+    readonly objectTypeId: string
+    readonly propertyId: string
+  }): Promise<void> {
+    await this.runRootOperation(async (sql) => {
+      await sql`
+        DELETE FROM object_unique_values
+        WHERE project_id = ${input.projectId} AND object_type_id = ${input.objectTypeId}
+          AND property_id = ${input.propertyId}
+      `
+      await sql`
+        DELETE FROM object_unique_properties
+        WHERE project_id = ${input.projectId} AND object_type_id = ${input.objectTypeId}
+          AND property_id = ${input.propertyId}
+      `
+    })
+  }
+
+  async insert(input: {
+    readonly projectId: string
+    readonly objectTypeId: string
+    readonly propertyId: string
+    readonly entries: readonly UniqueValueEntry[]
+  }): Promise<void> {
+    await this.runRootOperation((sql) =>
+      claimPgUniqueValues(
+        sql,
+        input.projectId,
+        sql`
+          SELECT ${input.objectTypeId}::text AS object_type_id,
+            ${input.propertyId}::text AS property_id,
+            entry->>'valueKey' AS value_key, entry->>'primaryId' AS primary_id
+          FROM jsonb_array_elements(${jsonParameter(sql, input.entries)}::jsonb) AS entry
+        `,
+        // The values are being prepared: nothing to require yet.
+        { requirePrepared: false }
+      )
+    )
+  }
+
+  async markPrepared(
+    input: PreparedUniqueProperty & { readonly projectId: string; readonly preparedAt: string }
+  ): Promise<void> {
+    await this.runRootOperation(async (sql) => {
+      await sql`
+        INSERT INTO object_unique_properties (
+          project_id, object_type_id, property_id, key_kind, key_version, prepared_at
+        ) VALUES (
+          ${input.projectId}, ${input.objectTypeId}, ${input.propertyId}, ${input.keyKind},
+          ${input.keyVersion}, ${input.preparedAt}::timestamptz
+        )
+        ON CONFLICT (project_id, object_type_id, property_id)
+        DO UPDATE SET key_kind = EXCLUDED.key_kind, key_version = EXCLUDED.key_version,
+          prepared_at = EXCLUDED.prepared_at
+      `
+    })
+  }
+}

package/src/pg-ai-cost-storage.ts

@@ -1,6 +1,7 @@
 import {
   type AiAccountingAggregateFragment,
   type AiAccountingRecordSetItem,
+  aiAccountingModelKind,
   aiModelCallCostDetails,
   aiModelCallCostMatchesUsage,
   buildAiAccountingOverviewFromFragments,
@@ -103,6 +104,7 @@ export class PgAiCostStorage implements AiCostStorage {
       WITH filtered AS (
         SELECT usage.provider_id,
           usage.requested_model_id,
+          usage.model_kind,
           usage.occurred_at,
           usage.input_tokens,
           usage.output_tokens,
@@ -112,6 +114,10 @@ export class PgAiCostStorage implements AiCostStorage {
           usage.text_output_tokens,
           usage.reasoning_output_tokens,
           usage.audio_duration_ms,
+          usage.image_output_tokens,
+          usage.audio_output_tokens,
+          usage.output_images,
+          usage.input_characters,
           usage.reporting_status,
           CASE
             WHEN direct_agent.kind = 'conversation' THEN 'agent'
@@ -174,6 +180,8 @@ export class PgAiCostStorage implements AiCostStorage {
         UNION ALL
         SELECT 'model', provider_id, requested_model_id, NULL, NULL, classified.* FROM classified
         UNION ALL
+        SELECT 'kind', model_kind, NULL, NULL, NULL, classified.* FROM classified
+        UNION ALL
         SELECT 'agent', attribution_agent_kind, attribution_workflow_id, attribution_agent_step_id, NULL, classified.*
         FROM classified WHERE attribution_agent_kind IS NOT NULL
         UNION ALL
@@ -206,7 +214,15 @@ export class PgAiCostStorage implements AiCostStorage {
         COUNT(reasoning_output_tokens)::text AS reasoning_output_tokens_count,
         SUM(reasoning_output_tokens)::text AS reasoning_output_tokens_sum,
         COUNT(audio_duration_ms)::text AS audio_duration_ms_count,
+        COUNT(image_output_tokens)::text AS image_output_tokens_count,
+        COUNT(audio_output_tokens)::text AS audio_output_tokens_count,
+        COUNT(output_images)::text AS output_images_count,
+        COUNT(input_characters)::text AS input_characters_count,
         SUM(audio_duration_ms)::text AS audio_duration_ms_sum,
+        SUM(image_output_tokens)::text AS image_output_tokens_sum,
+        SUM(audio_output_tokens)::text AS audio_output_tokens_sum,
+        SUM(output_images)::text AS output_images_sum,
+        SUM(input_characters)::text AS input_characters_sum,
         COUNT(*) FILTER (WHERE valuation_status = 'rated')::text AS rated_call_count,
         COUNT(*) FILTER (WHERE valuation_status = 'unpriceable')::text
           AS unpriceable_call_count,
@@ -347,6 +363,10 @@ interface UsageRow {
   readonly text_output_tokens: number | string | null
   readonly reasoning_output_tokens: number | string | null
   readonly audio_duration_ms: number | string | null
+  readonly image_output_tokens: number | string | null
+  readonly audio_output_tokens: number | string | null
+  readonly output_images: number | string | null
+  readonly input_characters: number | string | null
   readonly reporting_status: AiModelCallUsageRecord["usage"]["reportingStatus"]
   readonly raw_usage: ReadonlyJsonObject | string | null
   readonly provider_ids: Record<string, unknown> | string | null
@@ -402,7 +422,15 @@ interface AggregateRow {
   readonly reasoning_output_tokens_count: string
   readonly reasoning_output_tokens_sum: string | null
   readonly audio_duration_ms_count: string
+  readonly image_output_tokens_count: string
+  readonly audio_output_tokens_count: string
+  readonly output_images_count: string
+  readonly input_characters_count: string
   readonly audio_duration_ms_sum: string | null
+  readonly image_output_tokens_sum: string | null
+  readonly audio_output_tokens_sum: string | null
+  readonly output_images_sum: string | null
+  readonly input_characters_sum: string | null
   readonly rated_call_count: string
   readonly unpriceable_call_count: string
   readonly unvalued_call_count: string
@@ -447,18 +475,20 @@ function aggregateFragmentFromRow(row: AggregateRow): AiAccountingAggregateFragm
   const dimensionKeys =
     row.dimension === "model"
       ? { providerId: required(row.key_1), modelId: required(row.key_2) }
-      : row.dimension === "agent"
-        ? {
-            agentKind: required(row.key_1) as "agent" | "workflowAgent",
-            ...(row.key_1 === "workflowAgent"
-              ? { workflowId: required(row.key_2), agentStepId: required(row.key_3) }
-              : {}),
-          }
-        : row.dimension === "workflow"
-          ? { workflowId: required(row.key_1) }
-          : row.dimension === "series"
-            ? { start: new Date(required(row.bucket_start)) }
-            : {}
+      : row.dimension === "kind"
+        ? withModelKind(aiAccountingModelKind(row.key_1))
+        : row.dimension === "agent"
+          ? {
+              agentKind: required(row.key_1) as "agent" | "workflowAgent",
+              ...(row.key_1 === "workflowAgent"
+                ? { workflowId: required(row.key_2), agentStepId: required(row.key_3) }
+                : {}),
+            }
+          : row.dimension === "workflow"
+            ? { workflowId: required(row.key_1) }
+            : row.dimension === "series"
+              ? { start: new Date(required(row.bucket_start)) }
+              : {}
   return {
     dimension: row.dimension,
     ...dimensionKeys,
@@ -486,6 +516,10 @@ function aggregateFragmentFromRow(row: AggregateRow): AiAccountingAggregateFragm
         row.reasoning_output_tokens_sum
       ),
       audioDurationMs: meterFragment(row.audio_duration_ms_count, row.audio_duration_ms_sum),
+      imageOutputTokens: meterFragment(row.image_output_tokens_count, row.image_output_tokens_sum),
+      audioOutputTokens: meterFragment(row.audio_output_tokens_count, row.audio_output_tokens_sum),
+      outputImages: meterFragment(row.output_images_count, row.output_images_sum),
+      inputCharacters: meterFragment(row.input_characters_count, row.input_characters_sum),
     },
     costs: {
       ...(row.amount_currency === null || row.amount_nanos === null
@@ -587,6 +621,18 @@ function usageFromRow(row: UsageRow, requesterGroupIds: readonly string[]): AiMo
       ...(number(row.audio_duration_ms) === undefined
         ? {}
         : { audioDurationMs: number(row.audio_duration_ms) }),
+      ...(number(row.image_output_tokens) === undefined
+        ? {}
+        : { imageOutputTokens: number(row.image_output_tokens) }),
+      ...(number(row.audio_output_tokens) === undefined
+        ? {}
+        : { audioOutputTokens: number(row.audio_output_tokens) }),
+      ...(number(row.output_images) === undefined
+        ? {}
+        : { outputImages: number(row.output_images) }),
+      ...(number(row.input_characters) === undefined
+        ? {}
+        : { inputCharacters: number(row.input_characters) }),
       reportingStatus: row.reporting_status,
     },
     ...(row.raw_usage === null ? {} : { rawUsage: parseJsonObject(row.raw_usage) }),
@@ -777,3 +823,7 @@ function required<T>(value: T | null | undefined): T {
   }
   return value
 }
+
+function withModelKind(modelKind: AiAccountingAggregateFragment["modelKind"]) {
+  return modelKind === undefined ? {} : { modelKind }
+}

package/src/pg-ai-usage-storage.ts

@@ -1,14 +1,23 @@
-import { assertAiUsageExecutionId, normalizeAiModelCallRecord } from "@sixb/core/internal/storage"
+import type { FileRef } from "@sixb/core/blob-storage/server"
+import {
+  assertAiUsageExecutionId,
+  assertSameCallExecution,
+  normalizeAiModelCallFilesRecord,
+  normalizeAiModelCallRecord,
+} from "@sixb/core/internal/storage"
 import { isModelReasoning, normalizeModelProviderIds } from "@sixb/core/models"
 import type { ReadonlyJsonObject } from "@sixb/core/storage"
 import {
+  type AiModelCallFilesRecord,
   type AiModelCallUsageInput,
   type AiModelCallUsageRecord,
   type AiUsageExecutionSummary,
   type AiUsageStorage,
   AiUsageStorageError,
   aggregateAiModelCallUsage,
+  type GetAiModelCallFilesInput,
   type GetLatestAiModelCallForExecutionInput,
+  type RecordAiModelCallFilesInput,
   type RecordAiModelCallInput,
   type RecordAiModelCallResult,
   type SummarizeAiUsageExecutionInput,
@@ -51,6 +60,10 @@ export class PgAiUsageStorage implements AiUsageStorage {
           text_output_tokens,
           reasoning_output_tokens,
           audio_duration_ms,
+          image_output_tokens,
+          audio_output_tokens,
+          output_images,
+          input_characters,
           reporting_status,
           raw_usage,
           occurred_at,
@@ -81,6 +94,10 @@ export class PgAiUsageStorage implements AiUsageStorage {
           ${record.usage.textOutputTokens ?? null},
           ${record.usage.reasoningOutputTokens ?? null},
           ${record.usage.audioDurationMs ?? null},
+          ${record.usage.imageOutputTokens ?? null},
+          ${record.usage.audioOutputTokens ?? null},
+          ${record.usage.outputImages ?? null},
+          ${record.usage.inputCharacters ?? null},
           ${record.usage.reportingStatus},
           ${record.rawUsage === undefined ? null : JSON.stringify(record.rawUsage)}::text::jsonb,
           ${record.occurredAt},
@@ -235,6 +252,48 @@ export class PgAiUsageStorage implements AiUsageStorage {
     `
   }
 
+  async recordModelCallFiles(input: RecordAiModelCallFilesInput): Promise<AiModelCallFilesRecord> {
+    const record = normalizeAiModelCallFilesRecord(input)
+    return runPgTransaction(this.sql, async (tx) => {
+      const existing = await this.findCallFiles(tx, record.projectId, record.callId)
+      if (existing) return assertSameCallExecution(existing, record)
+      await this.assertExecutionExists(tx, record.projectId, record.executionId)
+      const inserted = await tx<CallFilesRow[]>`
+        INSERT INTO ai_model_call_files (project_id, call_id, execution_id, files, recorded_at)
+        VALUES (
+          ${record.projectId},
+          ${record.callId},
+          ${record.executionId},
+          ${JSON.stringify(record.files)}::text::jsonb,
+          ${record.recordedAt}
+        )
+        ON CONFLICT (project_id, call_id) DO NOTHING
+        RETURNING *
+      `
+      if (inserted[0]) return callFilesFromRow(inserted[0])
+      // Another writer recorded this call between the read and the insert.
+      const concurrent = await this.findCallFiles(tx, record.projectId, record.callId)
+      if (!concurrent) throw new Error("[SixbPg] AI model-call files disappeared after a conflict.")
+      return assertSameCallExecution(concurrent, record)
+    })
+  }
+
+  async getModelCallFiles(input: GetAiModelCallFilesInput): Promise<AiModelCallFilesRecord | null> {
+    assertNonBlankProjectId(input.projectId)
+    return this.findCallFiles(this.sql, input.projectId, input.callId)
+  }
+
+  private async findCallFiles(
+    sql: PgStoreClient,
+    projectId: string,
+    callId: string
+  ): Promise<AiModelCallFilesRecord | null> {
+    const [row] = await sql<CallFilesRow[]>`
+      SELECT * FROM ai_model_call_files WHERE project_id = ${projectId} AND call_id = ${callId}
+    `
+    return row ? callFilesFromRow(row) : null
+  }
+
   private async assertExecutionExists(
     sql: PgStoreClient,
     projectId: string,
@@ -253,6 +312,24 @@ export class PgAiUsageStorage implements AiUsageStorage {
   }
 }
 
+interface CallFilesRow {
+  readonly project_id: string
+  readonly call_id: string
+  readonly execution_id: string
+  readonly files: FileRef[] | string
+  readonly recorded_at: Date | string
+}
+
+function callFilesFromRow(row: CallFilesRow): AiModelCallFilesRecord {
+  return {
+    projectId: row.project_id,
+    executionId: row.execution_id,
+    callId: row.call_id,
+    files: typeof row.files === "string" ? (JSON.parse(row.files) as FileRef[]) : row.files,
+    recordedAt: new Date(row.recorded_at),
+  }
+}
+
 interface AiUsageRow {
   readonly project_id: string
   readonly id: string
@@ -275,6 +352,10 @@ interface AiUsageRow {
   readonly text_output_tokens: number | string | null
   readonly reasoning_output_tokens: number | string | null
   readonly audio_duration_ms: number | string | null
+  readonly image_output_tokens: number | string | null
+  readonly audio_output_tokens: number | string | null
+  readonly output_images: number | string | null
+  readonly input_characters: number | string | null
   readonly reporting_status: AiModelCallUsageRecord["usage"]["reportingStatus"]
   readonly raw_usage: ReadonlyJsonObject | string | null
   readonly occurred_at: Date | string
@@ -319,6 +400,14 @@ function usageFromRow(row: AiUsageRow): AiModelCallUsageInput {
       ? {}
       : { reasoningOutputTokens: Number(row.reasoning_output_tokens) }),
     ...(row.audio_duration_ms === null ? {} : { audioDurationMs: Number(row.audio_duration_ms) }),
+    ...(row.image_output_tokens === null
+      ? {}
+      : { imageOutputTokens: Number(row.image_output_tokens) }),
+    ...(row.audio_output_tokens === null
+      ? {}
+      : { audioOutputTokens: Number(row.audio_output_tokens) }),
+    ...(row.output_images === null ? {} : { outputImages: Number(row.output_images) }),
+    ...(row.input_characters === null ? {} : { inputCharacters: Number(row.input_characters) }),
   }
 }