@chainfuse/ai-tools 0.1.2 → 0.2.1

package/dist/base.mjs

@@ -6,6 +6,6 @@ export class AiBase {
         this._config = config;
     }
     get config() {
-        return this._config;
+        return Object.freeze(this._config);
     }
 }

package/dist/providers/customProviders.d.mts

@@ -1,10 +1,14 @@
+import type { GoogleGenerativeAIProvider } from '@ai-sdk/google';
+import type { OpenAICompatibleProvider } from '@ai-sdk/openai-compatible';
 import { AiBase } from '../base.mjs';
 import type { AiRequestConfig } from '../types.mjs';
-import type { AzureOpenAIProvider, CloudflareOpenAIProvider } from './types.mjs';
+import type { AzureOpenAIProvider } from './types.mjs';
 export declare class AiCustomProviders extends AiBase {
     oaiOpenai(args: AiRequestConfig): Promise<import("@ai-sdk/openai").OpenAIProvider>;
     azOpenai(args: AiRequestConfig, [server, ...servers]?: import("../serverSelector/types.mjs").Server[]): Promise<AzureOpenAIProvider>;
     anthropic(args: AiRequestConfig): Promise<import("@ai-sdk/anthropic").AnthropicProvider>;
     private static workersAiIsRest;
-    cfWorkersAi(args: AiRequestConfig): Promise<CloudflareOpenAIProvider>;
+    cfWorkersAi(args: AiRequestConfig): Promise<OpenAICompatibleProvider<"@cf/meta/llama-2-7b-chat-int8" | "@cf/mistral/mistral-7b-instruct-v0.1" | "@cf/meta/llama-2-7b-chat-fp16" | "@hf/thebloke/llama-2-13b-chat-awq" | "@hf/thebloke/mistral-7b-instruct-v0.1-awq" | "@hf/thebloke/zephyr-7b-beta-awq" | "@hf/thebloke/openhermes-2.5-mistral-7b-awq" | "@hf/thebloke/neural-chat-7b-v3-1-awq" | "@hf/thebloke/llamaguard-7b-awq" | "@hf/thebloke/deepseek-coder-6.7b-base-awq" | "@hf/thebloke/deepseek-coder-6.7b-instruct-awq" | "@cf/deepseek-ai/deepseek-math-7b-instruct" | "@cf/defog/sqlcoder-7b-2" | "@cf/openchat/openchat-3.5-0106" | "@cf/tiiuae/falcon-7b-instruct" | "@cf/thebloke/discolm-german-7b-v1-awq" | "@cf/qwen/qwen1.5-0.5b-chat" | "@cf/qwen/qwen1.5-7b-chat-awq" | "@cf/qwen/qwen1.5-14b-chat-awq" | "@cf/tinyllama/tinyllama-1.1b-chat-v1.0" | "@cf/microsoft/phi-2" | "@cf/qwen/qwen1.5-1.8b-chat" | "@cf/mistral/mistral-7b-instruct-v0.2-lora" | "@hf/nousresearch/hermes-2-pro-mistral-7b" | "@hf/nexusflow/starling-lm-7b-beta" | "@hf/google/gemma-7b-it" | "@cf/meta-llama/llama-2-7b-chat-hf-lora" | "@cf/google/gemma-2b-it-lora" | "@cf/google/gemma-7b-it-lora" | "@hf/mistral/mistral-7b-instruct-v0.2" | "@cf/meta/llama-3-8b-instruct" | "@cf/fblgit/una-cybertron-7b-v2-bf16" | "@cf/meta/llama-3-8b-instruct-awq" | "@hf/meta-llama/meta-llama-3-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct-fp8" | "@cf/meta/llama-3.1-8b-instruct-awq" | "@cf/meta/llama-3.2-3b-instruct" | "@cf/meta/llama-3.2-1b-instruct" | "@cf/meta/llama-3.3-70b-instruct-fp8-fast" | "@cf/meta/llama-3.2-11b-vision-instruct", "@cf/meta/llama-2-7b-chat-int8" | "@cf/mistral/mistral-7b-instruct-v0.1" | "@cf/meta/llama-2-7b-chat-fp16" | "@hf/thebloke/llama-2-13b-chat-awq" | "@hf/thebloke/mistral-7b-instruct-v0.1-awq" | "@hf/thebloke/zephyr-7b-beta-awq" | "@hf/thebloke/openhermes-2.5-mistral-7b-awq" | "@hf/thebloke/neural-chat-7b-v3-1-awq" | "@hf/thebloke/llamaguard-7b-awq" | "@hf/thebloke/deepseek-coder-6.7b-base-awq" | "@hf/thebloke/deepseek-coder-6.7b-instruct-awq" | "@cf/deepseek-ai/deepseek-math-7b-instruct" | "@cf/defog/sqlcoder-7b-2" | "@cf/openchat/openchat-3.5-0106" | "@cf/tiiuae/falcon-7b-instruct" | "@cf/thebloke/discolm-german-7b-v1-awq" | "@cf/qwen/qwen1.5-0.5b-chat" | "@cf/qwen/qwen1.5-7b-chat-awq" | "@cf/qwen/qwen1.5-14b-chat-awq" | "@cf/tinyllama/tinyllama-1.1b-chat-v1.0" | "@cf/microsoft/phi-2" | "@cf/qwen/qwen1.5-1.8b-chat" | "@cf/mistral/mistral-7b-instruct-v0.2-lora" | "@hf/nousresearch/hermes-2-pro-mistral-7b" | "@hf/nexusflow/starling-lm-7b-beta" | "@hf/google/gemma-7b-it" | "@cf/meta-llama/llama-2-7b-chat-hf-lora" | "@cf/google/gemma-2b-it-lora" | "@cf/google/gemma-7b-it-lora" | "@hf/mistral/mistral-7b-instruct-v0.2" | "@cf/meta/llama-3-8b-instruct" | "@cf/fblgit/una-cybertron-7b-v2-bf16" | "@cf/meta/llama-3-8b-instruct-awq" | "@hf/meta-llama/meta-llama-3-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct-fp8" | "@cf/meta/llama-3.1-8b-instruct-awq" | "@cf/meta/llama-3.2-3b-instruct" | "@cf/meta/llama-3.2-1b-instruct" | "@cf/meta/llama-3.3-70b-instruct-fp8-fast" | "@cf/meta/llama-3.2-11b-vision-instruct", "@cf/baai/bge-base-en-v1.5" | "@cf/baai/bge-small-en-v1.5" | "@cf/baai/bge-large-en-v1.5">>;
+    custom(args: AiRequestConfig): Promise<OpenAICompatibleProvider<string, string, string>>;
+    googleAi(args: AiRequestConfig): Promise<GoogleGenerativeAIProvider>;
 }

package/dist/providers/customProviders.mjs

@@ -1,5 +1,5 @@
 import { Helpers } from '@chainfuse/helpers';
-import { enabledCloudflareLlmEmbeddingProviders, enabledCloudflareLlmProviders } from '@chainfuse/types';
+import { AiModels, enabledCloudflareLlmProviders } from '@chainfuse/types';
 import { APICallError, experimental_customProvider as customProvider, TypeValidationError, experimental_wrapLanguageModel as wrapLanguageModel } from 'ai';
 import { ZodError } from 'zod';
 import { AiBase } from '../base.mjs';
@@ -17,7 +17,7 @@ export class AiCustomProviders extends AiBase {
                 const acc = await accPromise;
                 // @ts-expect-error override for types
                 acc[model] = wrapLanguageModel({
-                    model: (await raw.azOpenai(args, server.id))(model),
+                    model: (await raw.azOpenai(args, server))(model),
                     middleware: {
                         wrapGenerate: async ({ doGenerate, model, params }) => {
                             try {
@@ -40,7 +40,7 @@ export class AiCustomProviders extends AiBase {
                                             if (args.logging ?? this.config.environment !== 'production')
                                                 console.error('ai', 'custom provider', this.chalk.rgb(...Helpers.uniqueIdColor(idempotencyId))(`[${idempotencyId}]`), this.chalk.blue('FALLBACK'), nextServer.id, 'REMAINING', JSON.stringify(leftOverServers.slice(leftOverServers.indexOf(nextServer) + 1).map((s) => s.id)));
                                             // Must be double awaited to prevent a promise from being returned
-                                            return await (await raw.azOpenai({ ...args, idempotencyId }, nextServer.id))(model.modelId).doGenerate(params);
+                                            return await (await raw.azOpenai({ ...args, idempotencyId }, nextServer))(model.modelId).doGenerate(params);
                                         }
                                         catch (nextServerError) {
                                             if (APICallError.isInstance(nextServerError)) {
@@ -70,7 +70,7 @@ export class AiCustomProviders extends AiBase {
             textEmbeddingModels: await server.textEmbeddingModelAvailability.reduce(async (accPromise, model) => {
                 const acc = await accPromise;
                 // @ts-expect-error override for types
-                acc[model] = (await raw.azOpenai(args, server.id)).textEmbeddingModel(model);
+                acc[model] = (await raw.azOpenai(args, server)).textEmbeddingModel(model);
                 return acc;
             }, Promise.resolve({})),
             // An optional fallback provider to use when a requested model is not found in the custom provider.
@@ -139,13 +139,7 @@ export class AiCustomProviders extends AiBase {
                     });
                     return acc;
                 }, Promise.resolve({})),
-                // @ts-expect-error override for types
-                textEmbeddingModels: await enabledCloudflareLlmEmbeddingProviders.reduce(async (accPromise, model) => {
-                    const acc = await accPromise;
-                    // @ts-expect-error override for types
-                    acc[model] = (await raw.restWorkersAi(args)).textEmbeddingModel(model);
-                    return acc;
-                }, Promise.resolve({})),
+                fallbackProvider: await raw.restWorkersAi(args),
             });
         }
         else {
@@ -164,4 +158,20 @@ export class AiCustomProviders extends AiBase {
             }) as CloudflareOpenAIProvider;*/
         }
     }
+    custom(args) {
+        return new AiRawProviders(this.config).custom(args);
+    }
+    async googleAi(args) {
+        const fallbackProvider = await new AiRawProviders(this.config).googleAi(args);
+        return customProvider({
+            languageModels: {
+                // provider:actual model name:extra
+                [AiModels.LanguageModels.GoogleGenerativeAi.gemini_flash_beta_search.split(':').slice(1).join(':')]: fallbackProvider(AiModels.LanguageModels.GoogleGenerativeAi.gemini_flash_beta_search.split(':')[1], { useSearchGrounding: true }),
+                [AiModels.LanguageModels.GoogleGenerativeAi.gemini_flash_search.split(':').slice(1).join(':')]: fallbackProvider(AiModels.LanguageModels.GoogleGenerativeAi.gemini_flash_search.split(':')[1], { useSearchGrounding: true }),
+                [AiModels.LanguageModels.GoogleGenerativeAi.gemini_pro_search.split(':').slice(1).join(':')]: fallbackProvider(AiModels.LanguageModels.GoogleGenerativeAi.gemini_pro_search.split(':')[1], { useSearchGrounding: true }),
+            },
+            fallbackProvider,
+            // GoogleGenerativeAi
+        });
+    }
 }

package/dist/providers/rawProviders.d.mts

@@ -1,9 +1,15 @@
+import type { OpenAICompatibleProvider } from '@ai-sdk/openai-compatible';
+import type { cloudflareModelPossibilities } from '@chainfuse/types';
 import { AiBase } from '../base.mjs';
+import type { Server } from '../serverSelector/types.mjs';
 import type { AiRequestConfig } from '../types.mjs';
 export declare class AiRawProviders extends AiBase {
     private readonly cacheTtl;
+    private updateGatewayLog;
     oaiOpenai(args: AiRequestConfig): Promise<import("@ai-sdk/openai").OpenAIProvider>;
-    azOpenai(args: AiRequestConfig, server: string): Promise<import("@ai-sdk/azure").AzureOpenAIProvider>;
+    azOpenai(args: AiRequestConfig, server: Server): Promise<import("@ai-sdk/azure").AzureOpenAIProvider>;
     anthropic(args: AiRequestConfig): Promise<import("@ai-sdk/anthropic").AnthropicProvider>;
-    restWorkersAi(args: AiRequestConfig): Promise<import("@ai-sdk/openai").OpenAIProvider>;
+    custom(args: AiRequestConfig): Promise<OpenAICompatibleProvider<string, string, string>>;
+    googleAi(args: AiRequestConfig): Promise<import("@ai-sdk/google").GoogleGenerativeAIProvider>;
+    restWorkersAi(args: AiRequestConfig): Promise<OpenAICompatibleProvider<cloudflareModelPossibilities<'Text Generation'>, cloudflareModelPossibilities<'Text Generation'>, cloudflareModelPossibilities<'Text Embeddings'>>>;
 }

package/dist/providers/rawProviders.mjs

@@ -1,8 +1,38 @@
-import { BufferHelpers, CryptoHelpers, Helpers } from '@chainfuse/helpers';
+import { BufferHelpers, CryptoHelpers, DnsHelpers, Helpers } from '@chainfuse/helpers';
+import haversine from 'haversine-distance';
+import { z } from 'zod';
 import { AiBase } from '../base.mjs';
 export class AiRawProviders extends AiBase {
     // 2628288 seconds is what cf defines as 1 month in their cache rules
     cacheTtl = 2628288;
+    async updateGatewayLog(response, metadataHeader, startRoundTrip, modelTime) {
+        /**
+         * @todo `cloudflare` rest package not updated to this endpoint yet
+         */
+        const updateMetadata = fetch(new URL(['client', 'v4', 'accounts', this.config.gateway.accountId, 'ai-gateway', 'gateways', this.config.environment, 'logs', response.headers.get('cf-aig-log-id')].join('/'), 'https://api.cloudflare.com'), {
+            method: 'PATCH',
+            headers: {
+                Authorization: `Bearer ${this.config.gateway.apiToken}`,
+                'Content-Type': 'application/json',
+            },
+            body: JSON.stringify({
+                metadata: {
+                    ...metadataHeader,
+                    timing: JSON.stringify({
+                        fromCache: response.headers.get('cf-aig-cache-status')?.toLowerCase() === 'hit',
+                        totalRoundtripTime: performance.now() - startRoundTrip,
+                        modelTime,
+                    }),
+                },
+            }),
+        });
+        if (this.config.backgroundContext) {
+            this.config.backgroundContext.waitUntil(updateMetadata);
+        }
+        else {
+            await updateMetadata;
+        }
+    }
     oaiOpenai(args) {
         return import('@ai-sdk/openai').then(async ({ createOpenAI }) => createOpenAI({
             baseURL: new URL(['v1', this.config.gateway.accountId, this.config.environment, 'openai'].join('/'), 'https://gateway.ai.cloudflare.com').toString(),
@@ -29,6 +59,7 @@ export class AiRawProviders extends AiBase {
             },
             compatibility: 'strict',
             fetch: async (input, rawInit) => {
+                const startRoundTrip = performance.now();
                 const headers = new Headers(rawInit?.headers);
                 const metadataHeader = JSON.parse(headers.get('cf-aig-metadata'));
                 if (metadataHeader.idempotencyId.split('-').length === 4) {
@@ -40,6 +71,7 @@ export class AiRawProviders extends AiBase {
                 return fetch(input, { ...rawInit, headers }).then(async (response) => {
                     if (args.logging ?? this.config.environment !== 'production')
                         console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                    await this.updateGatewayLog(response, metadataHeader, startRoundTrip, response.headers.has('openai-processing-ms') ? parseInt(response.headers.get('openai-processing-ms')) : undefined);
                     // Inject it to have it available for retries
                     const mutableHeaders = new Headers(response.headers);
                     mutableHeaders.set('X-Idempotency-Id', metadataHeader.idempotencyId);
@@ -57,13 +89,13 @@ export class AiRawProviders extends AiBase {
     }
     azOpenai(args, server) {
         return import('@ai-sdk/azure').then(async ({ createAzure }) => createAzure({
-            apiKey: this.config.providers.azureOpenAi.apiTokens[`AZURE_API_KEY_${server.toUpperCase().replaceAll('-', '_')}`],
+            apiKey: this.config.providers.azureOpenAi.apiTokens[`AZURE_API_KEY_${server.id.toUpperCase().replaceAll('-', '_')}`],
             /**
              * @link https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#api-specs
              * From the table, pick the `Latest GA release` for `Data plane - inference`
              */
             apiVersion: '2024-10-21',
-            baseURL: new URL(['v1', this.config.gateway.accountId, this.config.environment, 'azure-openai', server.toLowerCase()].join('/'), 'https://gateway.ai.cloudflare.com').toString(),
+            baseURL: new URL(['v1', this.config.gateway.accountId, this.config.environment, 'azure-openai', server.id.toLowerCase()].join('/'), 'https://gateway.ai.cloudflare.com').toString(),
             headers: {
                 'cf-aig-authorization': `Bearer ${this.config.gateway.apiToken}`,
                 'cf-aig-metadata': JSON.stringify({
@@ -72,7 +104,11 @@ export class AiRawProviders extends AiBase {
                     // Generate incomplete id because we don't have the body to hash yet. Fill it in in the `fetch()`
                     idempotencyId: args.idempotencyId ?? (await BufferHelpers.generateUuid).utf8.slice(0, 23),
                     serverInfo: JSON.stringify({
-                        name: 'openai',
+                        name: `azure-${server.id}`,
+                        distance: haversine({
+                            lat: Helpers.precisionFloat(this.config.geoRouting?.userCoordinate?.lat ?? '0'),
+                            lon: Helpers.precisionFloat(this.config.geoRouting?.userCoordinate?.lon ?? '0'),
+                        }, server.coordinate),
                     }),
                     /**
                      * Blank at first, add after request finishes
@@ -84,6 +120,7 @@ export class AiRawProviders extends AiBase {
                 ...(args.skipCache && { 'cf-aig-skip-cache': 'true' }),
             },
             fetch: async (input, rawInit) => {
+                const startRoundTrip = performance.now();
                 const headers = new Headers(rawInit?.headers);
                 const metadataHeader = JSON.parse(headers.get('cf-aig-metadata'));
                 if (metadataHeader.idempotencyId.split('-').length === 4) {
@@ -95,6 +132,7 @@ export class AiRawProviders extends AiBase {
                 return fetch(input, { ...rawInit, headers }).then(async (response) => {
                     if (args.logging ?? this.config.environment !== 'production')
                         console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                    await this.updateGatewayLog(response, metadataHeader, startRoundTrip, response.headers.has('x-envoy-upstream-service-time') ? parseInt(response.headers.get('x-envoy-upstream-service-time')) : undefined);
                     // Inject it to have it available for retries
                     const mutableHeaders = new Headers(response.headers);
                     mutableHeaders.set('X-Idempotency-Id', metadataHeader.idempotencyId);
@@ -122,7 +160,168 @@ export class AiRawProviders extends AiBase {
                     // Generate incomplete id because we don't have the body to hash yet. Fill it in in the `fetch()`
                     idempotencyId: args.idempotencyId ?? (await BufferHelpers.generateUuid).utf8.slice(0, 23),
                     serverInfo: JSON.stringify({
-                        name: 'openai',
+                        name: 'anthropic',
+                    }),
+                    /**
+                     * Blank at first, add after request finishes
+                     * CF AI Gateway allows only editing existing metadata not creating new ones after the request is made
+                     */
+                    timing: JSON.stringify({}),
+                }),
+                ...(args.cache && { 'cf-aig-cache-ttl': (typeof args.cache === 'boolean' ? (args.cache ? this.cacheTtl : 0) : args.cache).toString() }),
+                ...(args.skipCache && { 'cf-aig-skip-cache': 'true' }),
+            },
+            fetch: async (input, rawInit) => {
+                const startRoundTrip = performance.now();
+                const headers = new Headers(rawInit?.headers);
+                const metadataHeader = JSON.parse(headers.get('cf-aig-metadata'));
+                if (metadataHeader.idempotencyId.split('-').length === 4) {
+                    metadataHeader.idempotencyId = `${metadataHeader.idempotencyId}-${(await CryptoHelpers.getHash('SHA-256', await new Request(input, rawInit).arrayBuffer())).slice(0, 12)}`;
+                    headers.set('cf-aig-metadata', JSON.stringify(metadataHeader));
+                }
+                if (args.logging ?? this.config.environment !== 'production')
+                    console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), this.chalk.magenta(rawInit?.method), this.chalk.magenta(new URL(new Request(input).url).pathname));
+                return fetch(input, { ...rawInit, headers }).then(async (response) => {
+                    if (args.logging ?? this.config.environment !== 'production')
+                        console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                    await this.updateGatewayLog(response, metadataHeader, startRoundTrip);
+                    // Inject it to have it available for retries
+                    const mutableHeaders = new Headers(response.headers);
+                    mutableHeaders.set('X-Idempotency-Id', metadataHeader.idempotencyId);
+                    if (response.ok) {
+                        return new Response(response.body, { ...response, headers: mutableHeaders });
+                    }
+                    else {
+                        const [body1, body2] = response.body.tee();
+                        console.error('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), this.chalk.red(JSON.stringify(await new Response(body1, response).json())));
+                        return new Response(body2, { ...response, headers: mutableHeaders });
+                    }
+                });
+            },
+        }));
+    }
+    custom(args) {
+        if (this.config.providers.custom?.url) {
+            // Verify that the custom provider url is a valid URL
+            return z
+                .string()
+                .trim()
+                .url()
+                .transform((url) => new URL(url))
+                .parseAsync(this.config.providers.custom.url)
+                .then((customProviderUrl) => 
+            // Verify that the custom provider url is not an IP address
+            z
+                .string()
+                .trim()
+                .ip()
+                .safeParseAsync(customProviderUrl.hostname)
+                .then(async ({ success }) => {
+                if (success) {
+                    throw new Error('IP custom providers not allowed');
+                }
+                else {
+                    // Run domain through ZT policies
+                    const doh = new DnsHelpers(new URL('dns-query', `https://${this.config.providers.custom?.dohId}.cloudflare-gateway.com`));
+                    const aCheck = doh.query(customProviderUrl.hostname, 'A', undefined, undefined, 2 * 1000);
+                    const aaaaCheck = doh.query(customProviderUrl.hostname, 'AAAA', undefined, undefined, 2 * 1000);
+                    return Promise.allSettled([aCheck, aaaaCheck]).then((checks) => {
+                        const fulfulledChecks = checks.filter((check) => check.status === 'fulfilled');
+                        /**
+                         * Blocked domains return 0.0.0.0 or :: as the answer
+                         * @link https://developers.cloudflare.com/cloudflare-one/policies/gateway/block-page/
+                         */
+                        if (fulfulledChecks.length > 0 && fulfulledChecks.some((obj) => 'Answer' in obj.value && Array.isArray(obj.value.Answer) && obj.value.Answer.some((answer) => answer.data !== '0.0.0.0' && answer.data !== '::'))) {
+                            // ZT Pass, perform the calls
+                            return import('@ai-sdk/openai-compatible').then(async ({ createOpenAICompatible }) => createOpenAICompatible({
+                                baseURL: customProviderUrl.toString(),
+                                ...(this.config.providers.custom?.apiToken && { apiKey: this.config.providers.custom.apiToken }),
+                                headers: {
+                                    // ZT Auth if present
+                                    ...(this.config.providers.custom &&
+                                        'clientId' in this.config.providers.custom &&
+                                        this.config.providers.custom.clientId &&
+                                        'clientSecret' in this.config.providers.custom &&
+                                        this.config.providers.custom.clientSecret && {
+                                        'CF-Access-Client-Id': this.config.providers.custom.clientId,
+                                        'CF-Access-Client-Secret': this.config.providers.custom.clientSecret,
+                                    }),
+                                    'X-Dataspace-Id': (await BufferHelpers.uuidConvert(args.dataspaceId)).utf8,
+                                    'X-Executor': JSON.stringify(args.executor),
+                                    // Generate incomplete id because we don't have the body to hash yet. Fill it in in the `fetch()`
+                                    'X-Idempotency-Id': args.idempotencyId ?? (await BufferHelpers.generateUuid).utf8.slice(0, 23),
+                                    // Request to skip or custom cache duration (no guarantee that upstream server will respect it)
+                                    // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
+                                    ...((args.skipCache || args.cache) && { 'Cache-Control': [args.skipCache && 'no-cache', args.cache && `max-age=${typeof args.cache === 'boolean' ? (args.cache ? this.cacheTtl : 0) : args.cache}`].join(', ') }),
+                                },
+                                name: 'custom',
+                                fetch: async (input, rawInit) => {
+                                    const headers = new Headers(rawInit?.headers);
+                                    let idempotencyId = headers.get('X-Idempotency-Id');
+                                    if (idempotencyId.split('-').length === 4) {
+                                        idempotencyId = `${idempotencyId}-${(await CryptoHelpers.getHash('SHA-256', await new Request(input, rawInit).arrayBuffer())).slice(0, 12)}`;
+                                        headers.set('X-Idempotency-Id', idempotencyId);
+                                    }
+                                    if (args.logging ?? this.config.environment !== 'production')
+                                        console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(idempotencyId))(`[${idempotencyId}]`), this.chalk.magenta(rawInit?.method), this.chalk.magenta(new URL(new Request(input).url).pathname));
+                                    return fetch(input, { ...rawInit, headers }).then(async (response) => {
+                                        if (args.logging ?? this.config.environment !== 'production')
+                                            console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(idempotencyId))(`[${idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                                        // Inject it to have it available for retries
+                                        const mutableHeaders = new Headers(response.headers);
+                                        mutableHeaders.set('X-Idempotency-Id', idempotencyId);
+                                        if (response.ok) {
+                                            return new Response(response.body, { ...response, headers: mutableHeaders });
+                                        }
+                                        else {
+                                            const [body1, body2] = response.body.tee();
+                                            console.error('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(idempotencyId))(`[${idempotencyId}]`), this.chalk.red(JSON.stringify(await new Response(body1, response).json())));
+                                            return new Response(body2, { ...response, headers: mutableHeaders });
+                                        }
+                                    });
+                                },
+                            }));
+                        }
+                        else {
+                            throw new Error('Failed ZT check on custom provider url');
+                        }
+                    });
+                }
+            }))
+                .catch(() => {
+                throw new Error('Invalid custom provider url');
+            });
+        }
+        else {
+            // This always gets called, only throw error if actually being used
+            return import('@ai-sdk/openai-compatible').then(async ({ createOpenAICompatible }) => createOpenAICompatible({
+                // Dummy url that'll never be hit
+                baseURL: 'https://sushidata.com',
+                name: 'custom',
+                // eslint-disable-next-line @typescript-eslint/require-await
+                fetch: async () => {
+                    throw new Error('Custom provider not configured');
+                },
+            }));
+        }
+    }
+    googleAi(args) {
+        return import('@ai-sdk/google').then(async ({ createGoogleGenerativeAI }) => createGoogleGenerativeAI({
+            /**
+             * `v1beta` is the only one that supports function calls as of now
+             * @link https://ai.google.dev/gemini-api/docs/api-versions
+             */
+            baseURL: new URL(['v1', this.config.gateway.accountId, this.config.environment, 'google-ai-studio', 'v1beta'].join('/'), 'https://gateway.ai.cloudflare.com').toString(),
+            apiKey: this.config.providers.googleAi.apiToken,
+            headers: {
+                'cf-aig-authorization': `Bearer ${this.config.gateway.apiToken}`,
+                'cf-aig-metadata': JSON.stringify({
+                    dataspaceId: (await BufferHelpers.uuidConvert(args.dataspaceId)).utf8,
+                    executor: JSON.stringify(args.executor),
+                    // Generate incomplete id because we don't have the body to hash yet. Fill it in in the `fetch()`
+                    idempotencyId: args.idempotencyId ?? (await BufferHelpers.generateUuid).utf8.slice(0, 23),
+                    serverInfo: JSON.stringify({
+                        name: 'googleai',
                     }),
                     /**
                      * Blank at first, add after request finishes
@@ -134,6 +333,7 @@ export class AiRawProviders extends AiBase {
                 ...(args.skipCache && { 'cf-aig-skip-cache': 'true' }),
             },
             fetch: async (input, rawInit) => {
+                const startRoundTrip = performance.now();
                 const headers = new Headers(rawInit?.headers);
                 const metadataHeader = JSON.parse(headers.get('cf-aig-metadata'));
                 if (metadataHeader.idempotencyId.split('-').length === 4) {
@@ -145,6 +345,7 @@ export class AiRawProviders extends AiBase {
                 return fetch(input, { ...rawInit, headers }).then(async (response) => {
                     if (args.logging ?? this.config.environment !== 'production')
                         console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                    await this.updateGatewayLog(response, metadataHeader, startRoundTrip);
                     // Inject it to have it available for retries
                     const mutableHeaders = new Headers(response.headers);
                     mutableHeaders.set('X-Idempotency-Id', metadataHeader.idempotencyId);
@@ -161,7 +362,7 @@ export class AiRawProviders extends AiBase {
         }));
     }
     restWorkersAi(args) {
-        return import('@ai-sdk/openai').then(async ({ createOpenAI }) => createOpenAI({
+        return import('@ai-sdk/openai-compatible').then(async ({ createOpenAICompatible }) => createOpenAICompatible({
             baseURL: new URL(['v1', this.config.gateway.accountId, this.config.environment, 'workers-ai', 'v1'].join('/'), 'https://gateway.ai.cloudflare.com').toString(),
             apiKey: this.config.providers.workersAi.apiToken,
             headers: {
@@ -172,7 +373,7 @@ export class AiRawProviders extends AiBase {
                     // Generate incomplete id because we don't have the body to hash yet. Fill it in in the `fetch()`
                     idempotencyId: args.idempotencyId ?? (await BufferHelpers.generateUuid).utf8.slice(0, 23),
                     serverInfo: JSON.stringify({
-                        name: 'openai',
+                        name: 'cloudflare',
                     }),
                     /**
                      * Blank at first, add after request finishes
@@ -183,9 +384,9 @@ export class AiRawProviders extends AiBase {
                 ...(args.cache && { 'cf-aig-cache-ttl': (typeof args.cache === 'boolean' ? (args.cache ? this.cacheTtl : 0) : args.cache).toString() }),
                 ...(args.skipCache && { 'cf-aig-skip-cache': 'true' }),
             },
-            compatibility: 'compatible',
             name: 'workersai',
             fetch: async (input, rawInit) => {
+                const startRoundTrip = performance.now();
                 const headers = new Headers(rawInit?.headers);
                 const metadataHeader = JSON.parse(headers.get('cf-aig-metadata'));
                 if (metadataHeader.idempotencyId.split('-').length === 4) {
@@ -197,6 +398,7 @@ export class AiRawProviders extends AiBase {
                 return fetch(input, { ...rawInit, headers }).then(async (response) => {
                     if (args.logging ?? this.config.environment !== 'production')
                         console.info('ai', 'raw provider', this.chalk.rgb(...Helpers.uniqueIdColor(metadataHeader.idempotencyId))(`[${metadataHeader.idempotencyId}]`), response.ok ? this.chalk.green(response.status) : this.chalk.red(response.status), response.ok ? this.chalk.green(new URL(response.url).pathname) : this.chalk.red(new URL(response.url).pathname));
+                    await this.updateGatewayLog(response, metadataHeader, startRoundTrip);
                     // Inject it to have it available for retries
                     const mutableHeaders = new Headers(response.headers);
                     mutableHeaders.set('X-Idempotency-Id', metadataHeader.idempotencyId);

package/dist/providers/types.d.mts

@@ -1,6 +1,6 @@
 import type { OpenAIChatSettings, OpenAIEmbeddingSettings } from '@ai-sdk/openai/internal';
 import type { EmbeddingModelV1, LanguageModelV1 } from '@ai-sdk/provider';
-import type { AzureChatModels, AzureEmbeddingModels, cloudflareModelPossibilities } from '@chainfuse/types';
+import type { AzureChatModels, AzureEmbeddingModels } from '@chainfuse/types';
 import type { Provider } from 'ai';
 export interface AzureOpenAIProvider extends Provider {
     (deploymentId: AzureChatModels, settings?: OpenAIChatSettings): LanguageModelV1;
@@ -13,14 +13,3 @@ export interface AzureOpenAIProvider extends Provider {
      */
     textEmbeddingModel(deploymentId: AzureEmbeddingModels, settings?: OpenAIEmbeddingSettings): EmbeddingModelV1<string>;
 }
-export interface CloudflareOpenAIProvider extends Provider {
-    (deploymentId: cloudflareModelPossibilities<'Text Generation'>, settings?: OpenAIChatSettings): LanguageModelV1;
-    /**
-  Creates an Azure OpenAI chat model for text generation.
-     */
-    languageModel(deploymentId: cloudflareModelPossibilities<'Text Generation'>, settings?: OpenAIChatSettings): LanguageModelV1;
-    /**
-  Creates an Azure OpenAI model for text embeddings.
-     */
-    textEmbeddingModel(deploymentId: cloudflareModelPossibilities<'Text Embeddings'>, settings?: OpenAIEmbeddingSettings): EmbeddingModelV1<string>;
-}

package/dist/registry.d.mts

@@ -5,7 +5,9 @@ export declare class AiRegistry extends AiBase {
         openai: import("@ai-sdk/openai").OpenAIProvider;
         azure: import("./providers/types.mjs").AzureOpenAIProvider;
         anthropic: import("@ai-sdk/anthropic").AnthropicProvider;
-        workersai: import("./providers/types.mjs").CloudflareOpenAIProvider;
+        custom: import("@ai-sdk/openai-compatible").OpenAICompatibleProvider<string, string, string>;
+        'google.generative-ai': import("@ai-sdk/google").GoogleGenerativeAIProvider;
+        workersai: import("@ai-sdk/openai-compatible").OpenAICompatibleProvider<"@cf/meta/llama-2-7b-chat-int8" | "@cf/mistral/mistral-7b-instruct-v0.1" | "@cf/meta/llama-2-7b-chat-fp16" | "@hf/thebloke/llama-2-13b-chat-awq" | "@hf/thebloke/mistral-7b-instruct-v0.1-awq" | "@hf/thebloke/zephyr-7b-beta-awq" | "@hf/thebloke/openhermes-2.5-mistral-7b-awq" | "@hf/thebloke/neural-chat-7b-v3-1-awq" | "@hf/thebloke/llamaguard-7b-awq" | "@hf/thebloke/deepseek-coder-6.7b-base-awq" | "@hf/thebloke/deepseek-coder-6.7b-instruct-awq" | "@cf/deepseek-ai/deepseek-math-7b-instruct" | "@cf/defog/sqlcoder-7b-2" | "@cf/openchat/openchat-3.5-0106" | "@cf/tiiuae/falcon-7b-instruct" | "@cf/thebloke/discolm-german-7b-v1-awq" | "@cf/qwen/qwen1.5-0.5b-chat" | "@cf/qwen/qwen1.5-7b-chat-awq" | "@cf/qwen/qwen1.5-14b-chat-awq" | "@cf/tinyllama/tinyllama-1.1b-chat-v1.0" | "@cf/microsoft/phi-2" | "@cf/qwen/qwen1.5-1.8b-chat" | "@cf/mistral/mistral-7b-instruct-v0.2-lora" | "@hf/nousresearch/hermes-2-pro-mistral-7b" | "@hf/nexusflow/starling-lm-7b-beta" | "@hf/google/gemma-7b-it" | "@cf/meta-llama/llama-2-7b-chat-hf-lora" | "@cf/google/gemma-2b-it-lora" | "@cf/google/gemma-7b-it-lora" | "@hf/mistral/mistral-7b-instruct-v0.2" | "@cf/meta/llama-3-8b-instruct" | "@cf/fblgit/una-cybertron-7b-v2-bf16" | "@cf/meta/llama-3-8b-instruct-awq" | "@hf/meta-llama/meta-llama-3-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct-fp8" | "@cf/meta/llama-3.1-8b-instruct-awq" | "@cf/meta/llama-3.2-3b-instruct" | "@cf/meta/llama-3.2-1b-instruct" | "@cf/meta/llama-3.3-70b-instruct-fp8-fast" | "@cf/meta/llama-3.2-11b-vision-instruct", "@cf/meta/llama-2-7b-chat-int8" | "@cf/mistral/mistral-7b-instruct-v0.1" | "@cf/meta/llama-2-7b-chat-fp16" | "@hf/thebloke/llama-2-13b-chat-awq" | "@hf/thebloke/mistral-7b-instruct-v0.1-awq" | "@hf/thebloke/zephyr-7b-beta-awq" | "@hf/thebloke/openhermes-2.5-mistral-7b-awq" | "@hf/thebloke/neural-chat-7b-v3-1-awq" | "@hf/thebloke/llamaguard-7b-awq" | "@hf/thebloke/deepseek-coder-6.7b-base-awq" | "@hf/thebloke/deepseek-coder-6.7b-instruct-awq" | "@cf/deepseek-ai/deepseek-math-7b-instruct" | "@cf/defog/sqlcoder-7b-2" | "@cf/openchat/openchat-3.5-0106" | "@cf/tiiuae/falcon-7b-instruct" | "@cf/thebloke/discolm-german-7b-v1-awq" | "@cf/qwen/qwen1.5-0.5b-chat" | "@cf/qwen/qwen1.5-7b-chat-awq" | "@cf/qwen/qwen1.5-14b-chat-awq" | "@cf/tinyllama/tinyllama-1.1b-chat-v1.0" | "@cf/microsoft/phi-2" | "@cf/qwen/qwen1.5-1.8b-chat" | "@cf/mistral/mistral-7b-instruct-v0.2-lora" | "@hf/nousresearch/hermes-2-pro-mistral-7b" | "@hf/nexusflow/starling-lm-7b-beta" | "@hf/google/gemma-7b-it" | "@cf/meta-llama/llama-2-7b-chat-hf-lora" | "@cf/google/gemma-2b-it-lora" | "@cf/google/gemma-7b-it-lora" | "@hf/mistral/mistral-7b-instruct-v0.2" | "@cf/meta/llama-3-8b-instruct" | "@cf/fblgit/una-cybertron-7b-v2-bf16" | "@cf/meta/llama-3-8b-instruct-awq" | "@hf/meta-llama/meta-llama-3-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct" | "@cf/meta/llama-3.1-8b-instruct-fp8" | "@cf/meta/llama-3.1-8b-instruct-awq" | "@cf/meta/llama-3.2-3b-instruct" | "@cf/meta/llama-3.2-1b-instruct" | "@cf/meta/llama-3.3-70b-instruct-fp8-fast" | "@cf/meta/llama-3.2-11b-vision-instruct", "@cf/baai/bge-base-en-v1.5" | "@cf/baai/bge-small-en-v1.5" | "@cf/baai/bge-large-en-v1.5">;
     }>>;
     registry(args: AiRequestConfig): Promise<import("ai").Provider>;
 }

package/dist/registry.mjs

@@ -7,6 +7,8 @@ export class AiRegistry extends AiBase {
             openai: await new AiCustomProviders(this.config).oaiOpenai(args),
             azure: await new AiCustomProviders(this.config).azOpenai(args),
             anthropic: await new AiCustomProviders(this.config).anthropic(args),
+            custom: await new AiCustomProviders(this.config).custom(args),
+            'google.generative-ai': await new AiCustomProviders(this.config).googleAi(args),
             workersai: await new AiCustomProviders(this.config).cfWorkersAi(args),
         });
     }

package/dist/types.d.mts

@@ -1,5 +1,5 @@
 import type { PrefixedUuid, RawCoordinate, UuidExport } from '@chainfuse/types';
-import type { Ai, IncomingRequestCfProperties } from '@cloudflare/workers-types/experimental';
+import type { Ai, ExecutionContext, IncomingRequestCfProperties } from '@cloudflare/workers-types/experimental';
 import type haversine from 'haversine-distance';
 export interface AiConfig {
     gateway: {
@@ -13,10 +13,13 @@ export interface AiConfig {
     };
     environment: 'production' | 'preview';
     providers: AiConfigProviders;
+    backgroundContext?: ExecutionContext;
 }
 export interface AiConfigProviders {
     anthropic: AiConfigAnthropic;
     azureOpenAi: AiConfigAzOpenai;
+    custom?: AiConfigCustom;
+    googleAi: AiConfigGoogleAi;
     openAi: AiConfigOaiOpenai;
     workersAi: AiConfigWorkersai;
 }
@@ -26,10 +29,27 @@ export interface AiConfigAnthropic {
 export interface AiConfigAzOpenai {
     apiTokens: Record<`AZURE_API_KEY_${string}`, string>;
 }
+export interface AiConfigCustomBase {
+    /**
+     * For safety resons, no direct IPs allowed.
+     */
+    url: string;
+    dohId: string;
+    apiToken?: string;
+}
+export type AiConfigCustom = AiConfigCustomBase | (AiConfigCustomBase & AiConfigCustomZT);
+export interface AiConfigCustomZT {
+    clientId: string;
+    clientSecret: string;
+}
+export interface AiConfigGoogleAi {
+    apiToken: string;
+}
 export interface AiConfigOaiOpenai {
     apiToken: `sk-${string}`;
     organization: `org-${string}`;
 }
+export type AiConfigWorkersai = AiConfigWorkersaiRest | AiConfigWorkersaiBinding;
 export interface AiConfigWorkersaiRest {
     apiToken: string;
 }
@@ -37,7 +57,6 @@ export interface AiConfigWorkersaiRest {
  * @deprecated Not functional. Use REST instead
  */
 export type AiConfigWorkersaiBinding = Ai;
-export type AiConfigWorkersai = AiConfigWorkersaiRest | AiConfigWorkersaiBinding;
 /**
  * It's a UUID, but the last block is SHA256 of the request body
  */
@@ -74,7 +93,7 @@ export interface AiRequestConfig {
     skipCache?: boolean;
 }
 export interface AiRequestMetadataServerInfo {
-    name: 'anthropic' | 'cloudflare' | 'openai';
+    name: 'anthropic' | 'cloudflare' | 'googleai' | 'openai';
 }
 export interface AiRequestMetadataServerInfoWithLocation {
     name: `${'azure' | 'google'}-${string}`;

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@chainfuse/ai-tools",
-	"version": "0.1.2",
+	"version": "0.2.1",
 	"description": "",
 	"author": "ChainFuse",
 	"homepage": "https://github.com/ChainFuse/packages/tree/main/packages/ai-tools#readme",
@@ -48,19 +48,21 @@
 	},
 	"prettier": "@demosjarco/prettier-config",
 	"dependencies": {
-		"@ai-sdk/anthropic": "^1.0.6",
-		"@ai-sdk/azure": "^1.0.15",
+		"@ai-sdk/anthropic": "^1.0.8",
+		"@ai-sdk/azure": "^1.0.18",
+		"@ai-sdk/google": "^1.0.14",
 		"@ai-sdk/openai": "^1.0.5",
-		"@chainfuse/helpers": "^0.5.0",
-		"@chainfuse/types": "^1.4.0",
-		"ai": "^4.0.25",
+		"@ai-sdk/openai-compatible": "^0.0.15",
+		"@chainfuse/helpers": "^0.6.0",
+		"@chainfuse/types": "^1.4.1",
+		"ai": "^4.0.31",
 		"chalk": "^5.4.1",
 		"haversine-distance": "^1.2.3",
 		"workers-ai-provider": "^0.0.10"
 	},
 	"devDependencies": {
-		"@cloudflare/workers-types": "^4.20241230.0",
-		"openai": "^4.77.3"
+		"@cloudflare/workers-types": "^4.20250109.0",
+		"openai": "^4.77.4"
 	},
-	"gitHead": "54f874d4afcfc72b7405dca76313fb7c8294b738"
+	"gitHead": "875629e9d3ae2cb62f053f78551c3a5e4d4443f5"
 }