Procházet zdrojové kódy

feat(db): add fn_upsert_usage_ledger trigger function

ding113 před 2 dny
rodič
revize
3bc82b420f
1 změnil soubory, kde provedl 83 přidání a 0 odebrání
  1. 83 0
      src/lib/ledger-backfill/trigger.sql

+ 83 - 0
src/lib/ledger-backfill/trigger.sql

@@ -0,0 +1,83 @@
+CREATE OR REPLACE FUNCTION fn_upsert_usage_ledger()
+RETURNS TRIGGER AS $$
+DECLARE
+  v_final_provider_id integer;
+  v_is_success boolean;
+BEGIN
+  IF NEW.blocked_by = 'warmup' THEN
+    RETURN NEW;
+  END IF;
+
+  IF NEW.provider_chain IS NOT NULL
+     AND jsonb_typeof(NEW.provider_chain) = 'array'
+     AND jsonb_array_length(NEW.provider_chain) > 0
+     AND jsonb_typeof(NEW.provider_chain -> -1) = 'object'
+     AND (NEW.provider_chain -> -1 ? 'id')
+     AND (NEW.provider_chain -> -1 ->> 'id') ~ '^[0-9]+$' THEN
+    v_final_provider_id := (NEW.provider_chain -> -1 ->> 'id')::integer;
+  ELSE
+    v_final_provider_id := NEW.provider_id;
+  END IF;
+
+  v_is_success := (NEW.error_message IS NULL OR NEW.error_message = '');
+
+  INSERT INTO usage_ledger (
+    request_id, user_id, key, provider_id, final_provider_id,
+    model, original_model, endpoint, api_type, session_id,
+    status_code, is_success, blocked_by,
+    cost_usd, cost_multiplier,
+    input_tokens, output_tokens,
+    cache_creation_input_tokens, cache_read_input_tokens,
+    cache_creation_5m_input_tokens, cache_creation_1h_input_tokens,
+    cache_ttl_applied, context_1m_applied, swap_cache_ttl_applied,
+    duration_ms, ttfb_ms, created_at
+  ) VALUES (
+    NEW.id, NEW.user_id, NEW.key, NEW.provider_id, v_final_provider_id,
+    NEW.model, NEW.original_model, NEW.endpoint, NEW.api_type, NEW.session_id,
+    NEW.status_code, v_is_success, NEW.blocked_by,
+    NEW.cost_usd, NEW.cost_multiplier,
+    NEW.input_tokens, NEW.output_tokens,
+    NEW.cache_creation_input_tokens, NEW.cache_read_input_tokens,
+    NEW.cache_creation_5m_input_tokens, NEW.cache_creation_1h_input_tokens,
+    NEW.cache_ttl_applied, NEW.context_1m_applied, NEW.swap_cache_ttl_applied,
+    NEW.duration_ms, NEW.ttfb_ms, NEW.created_at
+  )
+  ON CONFLICT (request_id) DO UPDATE SET
+    user_id = EXCLUDED.user_id,
+    key = EXCLUDED.key,
+    provider_id = EXCLUDED.provider_id,
+    final_provider_id = EXCLUDED.final_provider_id,
+    model = EXCLUDED.model,
+    original_model = EXCLUDED.original_model,
+    endpoint = EXCLUDED.endpoint,
+    api_type = EXCLUDED.api_type,
+    session_id = EXCLUDED.session_id,
+    status_code = EXCLUDED.status_code,
+    is_success = EXCLUDED.is_success,
+    blocked_by = EXCLUDED.blocked_by,
+    cost_usd = EXCLUDED.cost_usd,
+    cost_multiplier = EXCLUDED.cost_multiplier,
+    input_tokens = EXCLUDED.input_tokens,
+    output_tokens = EXCLUDED.output_tokens,
+    cache_creation_input_tokens = EXCLUDED.cache_creation_input_tokens,
+    cache_read_input_tokens = EXCLUDED.cache_read_input_tokens,
+    cache_creation_5m_input_tokens = EXCLUDED.cache_creation_5m_input_tokens,
+    cache_creation_1h_input_tokens = EXCLUDED.cache_creation_1h_input_tokens,
+    cache_ttl_applied = EXCLUDED.cache_ttl_applied,
+    context_1m_applied = EXCLUDED.context_1m_applied,
+    swap_cache_ttl_applied = EXCLUDED.swap_cache_ttl_applied,
+    duration_ms = EXCLUDED.duration_ms,
+    ttfb_ms = EXCLUDED.ttfb_ms,
+    created_at = EXCLUDED.created_at;
+
+  RETURN NEW;
+EXCEPTION WHEN OTHERS THEN
+  RAISE WARNING 'fn_upsert_usage_ledger failed for request_id=%: %', NEW.id, SQLERRM;
+  RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE TRIGGER trg_upsert_usage_ledger
+AFTER INSERT OR UPDATE ON message_request
+FOR EACH ROW
+EXECUTE FUNCTION fn_upsert_usage_ledger();