Make genai prices auto-update in the worker

Author: petyosi

gateway/package.json

@@ -11,7 +11,7 @@
     "@opentelemetry/api": "^1.9.0",
     "@opentelemetry/exporter-trace-otlp-http": "^0.203.0",
     "@opentelemetry/resources": "^2.0.1",
-    "@pydantic/genai-prices": "~0.0.36",
+    "@pydantic/genai-prices": "~0.0.38",
     "@pydantic/logfire-api": "^0.9.0",
     "eventsource-parser": "^3.0.6",
     "mime-types": "^3.0.1",

gateway/src/index.ts

@@ -19,6 +19,7 @@ import type { KeysDb, LimitDb } from './db'
 import { gateway } from './gateway'
 import type { DefaultProviderProxy, Middleware, Next } from './providers/default'
 import type { RateLimiter } from './rateLimiter'
+import { setupPriceAutoUpdate } from './setupPriceAutoUpdate'
 import type { SubFetch } from './types'
 import { ctHeader, ResponseError, response405, textResponse } from './utils'
 
@@ -48,6 +49,7 @@ export async function gatewayFetch(
   ctx: ExecutionContext,
   options: GatewayOptions,
 ): Promise<Response> {
+  await setupPriceAutoUpdate()
   let { pathname: proxyPath, search: queryString } = url
   if (options.proxyPrefixLength) {
     proxyPath = proxyPath.slice(options.proxyPrefixLength)

gateway/src/setupPriceAutoUpdate.ts

@@ -0,0 +1,71 @@
+import { type Provider, updatePrices, waitForUpdate } from '@pydantic/genai-prices'
+
+// data will be refetched every 30 minutes
+const PRICE_TTL = 1000 * 60 * 30
+let genaiData: Provider[] | null = null
+let genaiDataTimestamp: number | null = null
+let isFetching = false
+let isSetup = false
+
+export async function setupPriceAutoUpdate() {
+  if (!isSetup) {
+    doSetupPriceAutoUpdate()
+    isSetup = true
+  }
+
+  // this will await the eventual fresh genai price fetching
+  // Note: if the data is fresh, the updatePrices will immediately resolve, so there's no performance cost to always await this.
+  await waitForUpdate()
+}
+
+function doSetupPriceAutoUpdate() {
+  updatePrices(({ setProviderData, remoteDataUrl }) => {
+    if (genaiDataTimestamp !== null) {
+      console.debug('genai prices local storage data found')
+
+      if (genaiData !== null) {
+        setProviderData(genaiData)
+      }
+
+      if (Date.now() - genaiDataTimestamp < PRICE_TTL) {
+        // this will be the most frequent, cheap path
+        console.debug('genai prices local storage data is fresh')
+        return
+      } else {
+        console.debug('genai prices local storage data is stale, attempting to fetch remote data')
+      }
+    }
+
+    if (isFetching) {
+      console.debug('genai-prices data fetch already in progress, skipping')
+      return
+    }
+
+    console.debug('genai-prices data is stale')
+    isFetching = true
+
+    // It's important **not** to await this promise
+    const freshDataPromise = fetch(remoteDataUrl)
+      .then(async (response) => {
+        if (!response.ok) {
+          console.error('Failed fetching provider data, response status %d', response.status)
+          return null
+        }
+
+        const freshData = (await response.json()) as Provider[]
+        console.debug('Updated genai prices data, %d providers', freshData.length)
+        genaiDataTimestamp = Date.now()
+        genaiData = freshData
+        return freshData
+      })
+      .catch((error: unknown) => {
+        console.error('Failed fetching provider data err: %o', error)
+        return null
+      })
+      .finally(() => {
+        isFetching = false
+      })
+
+    setProviderData(freshDataPromise)
+  })
+}