diff --git a/src/adapters/gemini.ts b/src/adapters/gemini.ts index e2bd28c..54a8020 100644 --- a/src/adapters/gemini.ts +++ b/src/adapters/gemini.ts @@ -38,14 +38,28 @@ export const geminiAdapter: ProviderAdapter = { throw new Error("TokenFirewall: Invalid Gemini response format"); } - // Extract model from request or use default + // Extract model from request context or use default let model = "gemini-1.5-flash"; if (request && typeof request === "object") { - const req = request as { model?: string }; - if (req.model) { - model = req.model; + const req = request as { url?: string; model?: string }; + + // Try to extract model from the Gemini URL: /models/MODEL_NAME:generateContent + if (req.url) { + const match = req.url.match(/\/models\/([^/:]+)[:/]/); + if (match) { + model = match[1]; + } + } + + // Fallback: explicit model field in request body + if (!model || model === "gemini-1.5-flash") { + if (req.model) { + model = req.model; + } } } + + // modelVersion in the response always wins if present (most accurate) if (resp.modelVersion) { model = resp.modelVersion; } diff --git a/src/interceptors/fetchInterceptor.ts b/src/interceptors/fetchInterceptor.ts index 3048ec4..05c8d25 100644 --- a/src/interceptors/fetchInterceptor.ts +++ b/src/interceptors/fetchInterceptor.ts @@ -75,9 +75,15 @@ async function standardFetch( // Process response and track budget BEFORE returning try { const responseData = await clonedResponse.json(); - + + // Build a request context so adapters can read the model/URL + const url = typeof input === 'string' + ? input + : (input instanceof Request ? input.url : String(input)); + const requestContext = { url, body: init?.body }; + // Try to process with adapter registry - const normalizedUsage = adapterRegistry.process(responseData); + const normalizedUsage = adapterRegistry.process(responseData, requestContext); if (normalizedUsage) { // Calculate cost diff --git a/tsconfig.json b/tsconfig.json index 759c90a..6b808c9 100644 --- a/tsconfig.json +++ b/tsconfig.json @@ -2,7 +2,7 @@ "compilerOptions": { "target": "ES2020", "module": "commonjs", - "lib": ["ES2020"], + "lib": ["ES2020", "DOM"], "declaration": true, "outDir": "./dist", "rootDir": "./src",