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 { refreshGenaiPrices } from './refreshGenaiPrices'
 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> {
+  ctx.waitUntil(refreshGenaiPrices())
   let { pathname: proxyPath, search: queryString } = url
   if (options.proxyPrefixLength) {
     proxyPath = proxyPath.slice(options.proxyPrefixLength)

gateway/src/refreshGenaiPrices.ts

@@ -0,0 +1,60 @@
+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
+
+export function refreshGenaiPrices() {
+  updatePrices(({ setProviderData, remoteDataUrl }) => {
+    if (genaiDataTimestamp !== null) {
+      console.debug('genai prices in-memory cache found')
+
+      if (genaiData !== null) {
+        setProviderData(genaiData)
+      }
+
+      if (Date.now() - genaiDataTimestamp < PRICE_TTL) {
+        // this will be the most frequent, cheap path
+        console.debug('genai prices in-memory data is fresh')
+        return
+      } else {
+        console.debug('genai prices in-memory cache is stale, attempting to fetch remote data')
+      }
+    }
+
+    if (isFetching) {
+      console.debug('genai-prices data fetch already in progress, skipping')
+      return
+    }
+
+    console.debug('Fetching genai-prices data')
+    isFetching = true
+
+    // Note: **DO NOT** 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)
+  })
+  return waitForUpdate()
+}