@sixb/pg 0.1.6 → 0.1.8

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

@@ -10,6 +10,7 @@ import {
   parseAiModelCallCostDetails,
   toAccountingItem,
 } from "@sixb/core/internal/ai-cost-storage-provider"
+import { normalizeModelProviderIds } from "@sixb/core/models"
 import type { ReadonlyJsonObject } from "@sixb/core/storage"
 import {
   type AiAccountingOverview,
@@ -20,9 +21,13 @@ import {
   type AiModelCallUsageRecord,
   type ListAiModelCallAccountingInput,
   type ListAiModelCallAccountingResult,
+  type ListAiModelCallGroupsInput,
+  type ListAiModelCallGroupsResult,
   type QueryAiAccountingOverviewInput,
   type SummarizeAiCostExecutionsInput,
 } from "@sixb/core/storage"
+import { listAiModelCallGroups } from "./ai-cost-groups"
+import { requestedReasoningFromRow } from "./pg-ai-usage-storage"
 import type { PgStoreClient } from "./transactions"
 import { runPgTransaction } from "./transactions"
 
@@ -30,6 +35,12 @@ import { runPgTransaction } from "./transactions"
 export class PgAiCostStorage implements AiCostStorage {
   constructor(private readonly sql: PgStoreClient) {}
 
+  async listModelCallGroups(
+    input: ListAiModelCallGroupsInput
+  ): Promise<ListAiModelCallGroupsResult> {
+    return listAiModelCallGroups(this.sql, input)
+  }
+
   async recordModelCallCost(input: AiModelCallCostRecord): Promise<void> {
     const record = normalizeAiModelCallCostRecord(input)
     await runPgTransaction(this.sql, async (tx) => {
@@ -100,7 +111,12 @@ export class PgAiCostStorage implements AiCostStorage {
           usage.text_output_tokens,
           usage.reasoning_output_tokens,
           usage.reporting_status,
-          COALESCE(direct_agent.agent_id, workflow_agent.agent_id) AS attribution_agent_id,
+          CASE
+            WHEN direct_agent.kind = 'conversation' THEN 'agent'
+            WHEN workflow_agent.node_run_id IS NOT NULL THEN 'workflowAgent'
+            ELSE NULL
+          END AS attribution_agent_kind,
+          workflow_node.node_id AS attribution_agent_step_id,
           workflow_node.workflow_id AS attribution_workflow_id,
           cost.status AS cost_status,
           cost.currency AS cost_currency,
@@ -135,11 +151,12 @@ export class PgAiCostStorage implements AiCostStorage {
         SELECT 'totals'::text AS dimension,
           NULL::text AS key_1,
           NULL::text AS key_2,
+          NULL::text AS key_3,
           NULL::timestamptz AS bucket_start,
           classified.*
         FROM classified
         UNION ALL
-        SELECT 'series', NULL, NULL,
+        SELECT 'series', NULL, NULL, NULL,
           CASE ${input.bucket}::text
             WHEN 'hour' THEN
               date_trunc('hour', occurred_at AT TIME ZONE 'UTC') AT TIME ZONE 'UTC'
@@ -150,17 +167,18 @@ export class PgAiCostStorage implements AiCostStorage {
           classified.*
         FROM classified
         UNION ALL
-        SELECT 'model', provider_id, requested_model_id, NULL, classified.* FROM classified
+        SELECT 'model', provider_id, requested_model_id, NULL, NULL, classified.* FROM classified
         UNION ALL
-        SELECT 'agent', attribution_agent_id, NULL, NULL, classified.*
-        FROM classified WHERE attribution_agent_id IS NOT NULL
+        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
-        SELECT 'workflow', attribution_workflow_id, NULL, NULL, classified.*
+        SELECT 'workflow', attribution_workflow_id, NULL, NULL, NULL, classified.*
         FROM classified WHERE attribution_workflow_id IS NOT NULL
       )
       SELECT dimension,
         key_1,
         key_2,
+        key_3,
         bucket_start,
         amount_currency,
         COUNT(*)::text AS model_call_count,
@@ -187,7 +205,7 @@ export class PgAiCostStorage implements AiCostStorage {
           AS unvalued_call_count,
         SUM(amount_nanos)::text AS amount_nanos
       FROM expanded
-      GROUP BY dimension, key_1, key_2, bucket_start, amount_currency
+      GROUP BY dimension, key_1, key_2, key_3, bucket_start, amount_currency
     `
     return buildAiAccountingOverviewFromFragments(input, rows.map(aggregateFragmentFromRow))
   }
@@ -221,13 +239,15 @@ export class PgAiCostStorage implements AiCostStorage {
     const rows = await this.sql<JoinedRow[]>`
       SELECT usage.*,
         CASE
-          WHEN direct_agent.id IS NOT NULL THEN 'agent'
+          WHEN direct_agent.kind = 'conversation' THEN 'agent'
+          WHEN direct_agent.kind = 'subagent' THEN 'subagent'
           WHEN workflow_agent.node_run_id IS NOT NULL THEN 'workflowAgent'
           ELSE NULL
         END AS attribution_kind,
-        COALESCE(direct_agent.agent_id, workflow_agent.agent_id) AS attribution_agent_id,
+        workflow_node.node_id AS attribution_agent_step_id,
         direct_agent.id AS attribution_agent_run_id,
         direct_agent.thread_id AS attribution_thread_id,
+        direct_agent.parent_run_id AS attribution_parent_run_id,
         workflow_agent.node_run_id AS attribution_node_run_id,
         workflow_node.workflow_id AS attribution_workflow_id,
         workflow_node.workflow_run_id AS attribution_workflow_run_id,
@@ -303,6 +323,7 @@ interface UsageRow {
   readonly call_id: string
   readonly provider_id: string
   readonly requested_model_id: string
+  readonly requested_reasoning: string | Record<string, unknown> | null
   readonly response_model_id: string | null
   readonly response_id: string
   readonly input_tokens: number | string | null
@@ -315,15 +336,17 @@ interface UsageRow {
   readonly reasoning_output_tokens: number | string | null
   readonly reporting_status: AiModelCallUsageRecord["usage"]["reportingStatus"]
   readonly raw_usage: ReadonlyJsonObject | string | null
+  readonly provider_ids: Record<string, unknown> | string | null
   readonly occurred_at: Date | string
   readonly recorded_at: Date | string
 }
 
 interface JoinedRow extends UsageRow {
-  readonly attribution_kind: "agent" | "workflowAgent" | null
-  readonly attribution_agent_id: string | null
+  readonly attribution_kind: "agent" | "subagent" | "workflowAgent" | null
+  readonly attribution_agent_step_id: string | null
   readonly attribution_agent_run_id: string | null
   readonly attribution_thread_id: string | null
+  readonly attribution_parent_run_id: string | null
   readonly attribution_node_run_id: string | null
   readonly attribution_workflow_id: string | null
   readonly attribution_workflow_run_id: string | null
@@ -342,6 +365,7 @@ interface AggregateRow {
   readonly dimension: AiAccountingAggregateFragment["dimension"]
   readonly key_1: string | null
   readonly key_2: string | null
+  readonly key_3: string | null
   readonly bucket_start: Date | string | null
   readonly amount_currency: string | null
   readonly model_call_count: string
@@ -405,7 +429,12 @@ function aggregateFragmentFromRow(row: AggregateRow): AiAccountingAggregateFragm
     row.dimension === "model"
       ? { providerId: required(row.key_1), modelId: required(row.key_2) }
       : row.dimension === "agent"
-        ? { agentId: required(row.key_1) }
+        ? {
+            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"
@@ -466,17 +495,22 @@ function accountingItemFromRow(row: JoinedRow): AiAccountingRecordSetItem {
       : row.attribution_kind === "agent"
         ? {
             kind: "agent" as const,
-            agentId: required(row.attribution_agent_id),
             agentRunId: required(row.attribution_agent_run_id),
             threadId: required(row.attribution_thread_id),
           }
-        : {
-            kind: "workflowAgent" as const,
-            agentId: required(row.attribution_agent_id),
-            nodeRunId: required(row.attribution_node_run_id),
-            workflowId: required(row.attribution_workflow_id),
-            workflowRunId: required(row.attribution_workflow_run_id),
-          }
+        : row.attribution_kind === "subagent"
+          ? {
+              kind: "subagent" as const,
+              subagentRunId: required(row.attribution_agent_run_id),
+              parentRunId: required(row.attribution_parent_run_id),
+            }
+          : {
+              kind: "workflowAgent" as const,
+              agentStepId: required(row.attribution_agent_step_id),
+              nodeRunId: required(row.attribution_node_run_id),
+              workflowId: required(row.attribution_workflow_id),
+              workflowRunId: required(row.attribution_workflow_run_id),
+            }
   return {
     usage,
     ...(attribution ? { attribution } : {}),
@@ -494,7 +528,17 @@ function usageFromRow(row: UsageRow, requesterGroupIds: readonly string[]): AiMo
     callId: row.call_id,
     requesterGroupIds,
     providerId: row.provider_id,
+    ...(row.provider_ids === null
+      ? {}
+      : {
+          providerIds: normalizeModelProviderIds(
+            typeof row.provider_ids === "string" ? JSON.parse(row.provider_ids) : row.provider_ids
+          ),
+        }),
     requestedModelId: row.requested_model_id,
+    ...(row.requested_reasoning === null
+      ? {}
+      : { requestedReasoning: requestedReasoningFromRow(row.requested_reasoning) }),
     ...(row.response_model_id === null ? {} : { responseModelId: row.response_model_id }),
     responseId: row.response_id,
     usage: {
@@ -569,14 +613,15 @@ function costRecord(input: {
   ratedAt: Date | string
 }): AiModelCallCostRecord {
   const details = parseAiModelCallCostDetails(input.details)
+  const priceSource =
+    details.priceSource === undefined
+      ? undefined
+      : { ...details.priceSource, observedAt: new Date(details.priceSource.observedAt) }
   const base = {
+    ...(details.estimate === undefined ? {} : { estimate: details.estimate }),
     projectId: input.projectId,
     usageRecordId: input.usageId,
     pricingContext: details.pricingContext,
-    priceSource: {
-      ...details.priceSource,
-      observedAt: new Date(details.priceSource.observedAt),
-    },
     ratedAt: new Date(input.ratedAt),
   }
   if (input.status === "rated") {
@@ -584,6 +629,7 @@ function costRecord(input: {
       ...base,
       status: "rated",
       billingIdentity: { providerId: required(input.providerId), modelId: required(input.modelId) },
+      priceSource: required(priceSource),
       money: { currency: required(input.currency), amountNanos: required(input.amountNanos) },
       components: required(details.components),
     })
@@ -594,6 +640,7 @@ function costRecord(input: {
     ...(input.providerId && input.modelId
       ? { billingIdentity: { providerId: input.providerId, modelId: input.modelId } }
       : {}),
+    ...(priceSource === undefined ? {} : { priceSource }),
     reason: required(input.reason) as Extract<
       AiModelCallCostRecord,
       { status: "unpriceable" }