@truefoundry/ai-gateway-sdk 0.0.0-stage → 0.0.2

package/dist/providers/google-vertex-ai/getBatchOutput.js

@@ -0,0 +1,160 @@
+/**
+ * ================================
+ * Inherited upstream file
+ * But needs some custom changes for TFY
+ * 1. Line Splitter will be coming from streamHandler.ts
+ * 2. Custom id will be row['custom_id'] fallback to the upstream logic
+ * 3. Use const instead of let for variables that are not reassigned (for linting)
+ * 4. Remove unused vars for linting
+ * ================================
+ */
+import { BatchEndpoints, GOOGLE_VERTEX_AI } from '../../globals.js';
+import { createLineSplitter } from '../../handlers/streamHandler.js';
+import { safeFetch } from '../../http/index.js';
+import { sdkRequestUrl } from '../../utils/sdkRequestUrl.js';
+import { responseTransformers } from '../open-ai-base/index.js';
+import { providerApiArgsFromHandler } from '../requestHandlerApiArgs.js';
+import GoogleApiConfig from './api.js';
+import { GoogleChatCompleteResponseTransform, VertexAnthropicChatCompleteResponseTransform, VertexLlamaChatCompleteResponseTransform, } from './chatComplete.js';
+import { GoogleEmbedResponseTransform } from './embed.js';
+import { getModelAndProvider, isEmbeddingModel } from './utils.js';
+const responseTransforms = {
+    google: {
+        [BatchEndpoints.CHAT_COMPLETIONS]: GoogleChatCompleteResponseTransform,
+        [BatchEndpoints.EMBEDDINGS]: GoogleEmbedResponseTransform,
+    },
+    anthropic: {
+        [BatchEndpoints.CHAT_COMPLETIONS]: VertexAnthropicChatCompleteResponseTransform,
+        [BatchEndpoints.EMBEDDINGS]: null,
+    },
+    meta: {
+        [BatchEndpoints.CHAT_COMPLETIONS]: VertexLlamaChatCompleteResponseTransform,
+        [BatchEndpoints.EMBEDDINGS]: null,
+    },
+    endpoints: {
+        [BatchEndpoints.CHAT_COMPLETIONS]: responseTransformers(GOOGLE_VERTEX_AI, {
+            chatComplete: true,
+        }).chatComplete,
+        [BatchEndpoints.EMBEDDINGS]: responseTransformers(GOOGLE_VERTEX_AI, {
+            embed: true,
+        }).embed,
+    },
+};
+const getOpenAIBatchRow = ({ row, batchId, transform, endpoint, modelName, }) => {
+    const response = endpoint === BatchEndpoints.EMBEDDINGS
+        ? (row ?? {})
+        : (row['response'] ?? {});
+    const id = `batch-${batchId}-${response.responseId ? `-${response.responseId}` : ''}`;
+    let error = null;
+    try {
+        error = JSON.parse(row.status);
+    }
+    catch {
+        error = row.status;
+    }
+    return {
+        id,
+        custom_id: row['custom_id'] || row.requestId || row?.instance?.requestId || response.responseId,
+        response: {
+            ...(!error && { status_code: 200 }),
+            request_id: id,
+            body: !error && transform(response, 200, {}, false, '', { model: modelName }),
+        },
+        error: error,
+    };
+};
+export const BatchOutputRequestHandler = async (args) => {
+    const requestBody = args.request.body;
+    const gatewayUrl = sdkRequestUrl(args.request);
+    const retrieveArgs = providerApiArgsFromHandler(args, { endpoint: 'retrieveBatch', method: 'GET' });
+    const headers = await GoogleApiConfig.headers({
+        ...retrieveArgs,
+        transformedRequestBody: requestBody,
+        transformedRequestUrl: gatewayUrl,
+    });
+    const options = {
+        method: 'GET',
+        headers,
+        signal: args.signal,
+    };
+    // URL: <gateway>/v1/batches/<batchId>/output
+    const batchId = args.request.path.split('/').at(-2);
+    const batchPath = args.request.path.replace(/\/output$/, '');
+    const detailsArgs = providerApiArgsFromHandler(args, {
+        endpoint: 'retrieveBatch',
+        method: 'GET',
+        path: batchPath,
+    });
+    const baseURL = await GoogleApiConfig.getBaseURL(detailsArgs);
+    const endpoint = GoogleApiConfig.getEndpoint(detailsArgs);
+    const batchesURL = `${baseURL}${endpoint}`;
+    let modelName = '';
+    let outputURL;
+    try {
+        const response = await safeFetch(batchesURL, { ssrf: 'guard', ...options });
+        if (!response.ok) {
+            const error = await response.text();
+            throw new Error(error);
+        }
+        const data = (await response.json());
+        outputURL = data.outputInfo?.gcsOutputDirectory ?? '';
+        modelName = data.model;
+    }
+    catch (error) {
+        const errorMessage = error.message || 'Failed to retrieve batch output';
+        throw new Error(errorMessage);
+    }
+    if (!outputURL) {
+        throw new Error('Failed to retrieve batch details');
+    }
+    const { provider } = getModelAndProvider(modelName ?? '');
+    const batchEndpoint = isEmbeddingModel(modelName) ? BatchEndpoints.EMBEDDINGS : BatchEndpoints.CHAT_COMPLETIONS;
+    const providerConfigMap = responseTransforms[provider];
+    const providerConfig = providerConfigMap?.[batchEndpoint] ?? responseTransforms['endpoints'][batchEndpoint];
+    if (!providerConfig) {
+        throw new Error(`Endpoint ${endpoint} not supported for provider ${provider}`);
+    }
+    outputURL = outputURL.replace('gs://', 'https://storage.googleapis.com/');
+    const predictionFileId = endpoint === BatchEndpoints.EMBEDDINGS ? '000000000000.jsonl' : 'predictions.jsonl';
+    const outputResponse = await safeFetch(`${outputURL}/${predictionFileId}`, { ssrf: 'guard', ...options });
+    const reader = outputResponse.body;
+    if (!reader) {
+        throw new Error('Failed to retrieve batch output');
+    }
+    const encoder = new TextEncoder();
+    // Prepare a transform stream to process complete lines.
+    const responseStream = new TransformStream({
+        transform(chunk, controller) {
+            let buffer;
+            try {
+                const json = JSON.parse(chunk.toString());
+                const row = getOpenAIBatchRow({
+                    row: json,
+                    batchId: batchId ?? '',
+                    transform: providerConfig,
+                    endpoint: batchEndpoint,
+                    modelName,
+                });
+                buffer = JSON.stringify(row);
+            }
+            catch ( /* eslint-disable @typescript-eslint/no-unused-vars */error) {
+                return;
+            }
+            controller.enqueue(encoder.encode(buffer + '\n'));
+        },
+        flush(controller) {
+            controller.terminate();
+        },
+    });
+    const [safeStream] = responseStream.readable.tee();
+    // Pipe the node stream through the line splitter and then to the response stream.
+    const lineSplitter = createLineSplitter();
+    reader.pipeThrough(lineSplitter).pipeTo(responseStream.writable);
+    return new Response(safeStream, {
+        headers: { 'Content-Type': 'application/octet-stream' },
+    });
+};
+export const BatchOutputResponseTransform = async (response) => {
+    return response;
+};
+//# sourceMappingURL=getBatchOutput.js.map

package/dist/providers/google-vertex-ai/getBatchOutput.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"getBatchOutput.js","sourceRoot":"","sources":["../../../src/providers/google-vertex-ai/getBatchOutput.ts"],"names":[],"mappings":"AAAA;;;;;;;;;GASG;AACH,OAAO,EAAE,cAAc,EAAE,gBAAgB,EAAE,MAAM,eAAe,CAAC;AACjE,OAAO,EAAE,kBAAkB,EAAE,MAAM,8BAA8B,CAAC;AAClE,OAAO,EAAE,SAAS,EAAE,MAAM,YAAY,CAAC;AACvC,OAAO,EAAE,aAAa,EAAE,MAAM,2BAA2B,CAAC;AAC1D,OAAO,EAAE,oBAAoB,EAAE,MAAM,iBAAiB,CAAC;AACvD,OAAO,EAAE,0BAA0B,EAAE,MAAM,0BAA0B,CAAC;AAEtE,OAAO,eAAe,MAAM,OAAO,CAAC;AACpC,OAAO,EACL,mCAAmC,EACnC,4CAA4C,EAC5C,wCAAwC,GACzC,MAAM,gBAAgB,CAAC;AACxB,OAAO,EAAE,4BAA4B,EAAE,MAAM,SAAS,CAAC;AAEvD,OAAO,EAAE,mBAAmB,EAAE,gBAAgB,EAAE,MAAM,SAAS,CAAC;AAEhE,MAAM,kBAAkB,GAAG;IACzB,MAAM,EAAE;QACN,CAAC,cAAc,CAAC,gBAAgB,CAAC,EAAE,mCAAmC;QACtE,CAAC,cAAc,CAAC,UAAU,CAAC,EAAE,4BAA4B;KAC1D;IACD,SAAS,EAAE;QACT,CAAC,cAAc,CAAC,gBAAgB,CAAC,EAAE,4CAA4C;QAC/E,CAAC,cAAc,CAAC,UAAU,CAAC,EAAE,IAAI;KAClC;IACD,IAAI,EAAE;QACJ,CAAC,cAAc,CAAC,gBAAgB,CAAC,EAAE,wCAAwC;QAC3E,CAAC,cAAc,CAAC,UAAU,CAAC,EAAE,IAAI;KAClC;IACD,SAAS,EAAE;QACT,CAAC,cAAc,CAAC,gBAAgB,CAAC,EAAE,oBAAoB,CAAC,gBAAgB,EAAE;YACxE,YAAY,EAAE,IAAI;SACnB,CAAC,CAAC,YAAY;QACf,CAAC,cAAc,CAAC,UAAU,CAAC,EAAE,oBAAoB,CAAC,gBAAgB,EAAE;YAClE,KAAK,EAAE,IAAI;SACZ,CAAC,CAAC,KAAK;KACT;CACF,CAAC;AAWF,MAAM,iBAAiB,GAAG,CAAC,EACzB,GAAG,EACH,OAAO,EACP,SAAS,EACT,QAAQ,EACR,SAAS,GAOV,EAAE,EAAE;IACH,MAAM,QAAQ,GACZ,QAAQ,KAAK,cAAc,CAAC,UAAU;QACpC,CAAC,CAAE,CAAC,GAAG,IAAI,EAAE,CAA6B;QAC1C,CAAC,CAAE,CAAC,GAAG,CAAC,UAAU,CAAC,IAAI,EAAE,CAA6B,CAAC;IAC3D,MAAM,EAAE,GAAG,SAAS,OAAO,IAAI,QAAQ,CAAC,UAAU,CAAC,CAAC,CAAC,IAAI,QAAQ,CAAC,UAAU,EAAE,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC;IAEtF,IAAI,KAAK,GAAG,IAAI,CAAC;IACjB,IAAI,CAAC;QACH,KAAK,GAAG,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,MAAgB,CAAC,CAAC;IAC3C,CAAC;IAAC,MAAM,CAAC;QACP,KAAK,GAAG,GAAG,CAAC,MAAM,CAAC;IACrB,CAAC;IAED,OAAO;QACL,EAAE;QACF,SAAS,EAAE,GAAG,CAAC,WAAW,CAAC,IAAI,GAAG,CAAC,SAAS,IAAK,GAAG,EAAE,QAAgB,EAAE,SAAS,IAAI,QAAQ,CAAC,UAAU;QACxG,QAAQ,EAAE;YACR,GAAG,CAAC,CAAC,KAAK,IAAI,EAAE,WAAW,EAAE,GAAG,EAAE,CAAC;YACnC,UAAU,EAAE,EAAE;YACd,IAAI,EAAE,CAAC,KAAK,IAAI,SAAS,CAAC,QAAQ,EAAE,GAAG,EAAE,EAAE,EAAE,KAAK,EAAE,EAAE,EAAE,EAAE,KAAK,EAAE,SAAS,EAAE,CAAC;SAC9E;QACD,KAAK,EAAE,KAAK;KACb,CAAC;AACJ,CAAC,CAAC;AAEF,MAAM,CAAC,MAAM,yBAAyB,GAAmB,KAAK,EAAC,IAAI,EAAC,EAAE;IACpE,MAAM,WAAW,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC;IACtC,MAAM,UAAU,GAAG,aAAa,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;IAC/C,MAAM,YAAY,GAAG,0BAA0B,CAAC,IAAI,EAAE,EAAE,QAAQ,EAAE,eAAe,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC,CAAC;IACpG,MAAM,OAAO,GAAG,MAAM,eAAe,CAAC,OAAO,CAAC;QAC5C,GAAG,YAAY;QACf,sBAAsB,EAAE,WAAW;QACnC,qBAAqB,EAAE,UAAU;KAClC,CAAC,CAAC;IAEH,MAAM,OAAO,GAAG;QACd,MAAM,EAAE,KAAK;QACb,OAAO;QACP,MAAM,EAAE,IAAI,CAAC,MAAM;KACpB,CAAC;IAEF,6CAA6C;IAC7C,MAAM,OAAO,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;IAEpD,MAAM,SAAS,GAAG,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,OAAO,CAAC,WAAW,EAAE,EAAE,CAAC,CAAC;IAC7D,MAAM,WAAW,GAAG,0BAA0B,CAAC,IAAI,EAAE;QACnD,QAAQ,EAAE,eAAe;QACzB,MAAM,EAAE,KAAK;QACb,IAAI,EAAE,SAAS;KAChB,CAAC,CAAC;IAEH,MAAM,OAAO,GAAG,MAAM,eAAe,CAAC,UAAU,CAAC,WAAW,CAAC,CAAC;IAC9D,MAAM,QAAQ,GAAG,eAAe,CAAC,WAAW,CAAC,WAAW,CAAC,CAAC;IAE1D,MAAM,UAAU,GAAG,GAAG,OAAO,GAAG,QAAQ,EAAE,CAAC;IAC3C,IAAI,SAAS,GAAG,EAAE,CAAC;IACnB,IAAI,SAAS,CAAC;IACd,IAAI,CAAC;QACH,MAAM,QAAQ,GAAG,MAAM,SAAS,CAAC,UAAU,EAAE,EAAE,IAAI,EAAE,OAAO,EAAE,GAAG,OAAO,EAAE,CAAC,CAAC;QAC5E,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;YACjB,MAAM,KAAK,GAAG,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAC;YACpC,MAAM,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC;QACzB,CAAC;QACD,MAAM,IAAI,GAAG,CAAC,MAAM,QAAQ,CAAC,IAAI,EAAE,CAAsB,CAAC;QAC1D,SAAS,GAAG,IAAI,CAAC,UAAU,EAAE,kBAAkB,IAAI,EAAE,CAAC;QACtD,SAAS,GAAG,IAAI,CAAC,KAAK,CAAC;IACzB,CAAC;IAAC,OAAO,KAAK,EAAE,CAAC;QACf,MAAM,YAAY,GAAI,KAAe,CAAC,OAAO,IAAI,iCAAiC,CAAC;QACnF,MAAM,IAAI,KAAK,CAAC,YAAY,CAAC,CAAC;IAChC,CAAC;IAED,IAAI,CAAC,SAAS,EAAE,CAAC;QACf,MAAM,IAAI,KAAK,CAAC,kCAAkC,CAAC,CAAC;IACtD,CAAC;IAED,MAAM,EAAE,QAAQ,EAAE,GAAG,mBAAmB,CAAC,SAAS,IAAI,EAAE,CAAC,CAAC;IAE1D,MAAM,aAAa,GAAG,gBAAgB,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,cAAc,CAAC,UAAU,CAAC,CAAC,CAAC,cAAc,CAAC,gBAAgB,CAAC;IAEhH,MAAM,iBAAiB,GAAG,kBAAkB,CAAC,QAA2C,CAAC,CAAC;IAC1F,MAAM,cAAc,GAAG,iBAAiB,EAAE,CAAC,aAAa,CAAC,IAAI,kBAAkB,CAAC,WAAW,CAAC,CAAC,aAAa,CAAC,CAAC;IAE5G,IAAI,CAAC,cAAc,EAAE,CAAC;QACpB,MAAM,IAAI,KAAK,CAAC,YAAY,QAAQ,+BAA+B,QAAQ,EAAE,CAAC,CAAC;IACjF,CAAC;IAED,SAAS,GAAG,SAAS,CAAC,OAAO,CAAC,OAAO,EAAE,iCAAiC,CAAC,CAAC;IAE1E,MAAM,gBAAgB,GAAG,QAAQ,KAAK,cAAc,CAAC,UAAU,CAAC,CAAC,CAAC,oBAAoB,CAAC,CAAC,CAAC,mBAAmB,CAAC;IAE7G,MAAM,cAAc,GAAG,MAAM,SAAS,CAAC,GAAG,SAAS,IAAI,gBAAgB,EAAE,EAAE,EAAE,IAAI,EAAE,OAAO,EAAE,GAAG,OAAO,EAAE,CAAC,CAAC;IAE1G,MAAM,MAAM,GAAG,cAAc,CAAC,IAAI,CAAC;IACnC,IAAI,CAAC,MAAM,EAAE,CAAC;QACZ,MAAM,IAAI,KAAK,CAAC,iCAAiC,CAAC,CAAC;IACrD,CAAC;IAED,MAAM,OAAO,GAAG,IAAI,WAAW,EAAE,CAAC;IAElC,wDAAwD;IACxD,MAAM,cAAc,GAAG,IAAI,eAAe,CAAC;QACzC,SAAS,CAAC,KAAK,EAAE,UAAU;YACzB,IAAI,MAAM,CAAC;YACX,IAAI,CAAC;gBACH,MAAM,IAAI,GAAG,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,QAAQ,EAAE,CAAC,CAAC;gBAC1C,MAAM,GAAG,GAAG,iBAAiB,CAAC;oBAC5B,GAAG,EAAE,IAAI;oBACT,OAAO,EAAE,OAAO,IAAI,EAAE;oBACtB,SAAS,EAAE,cAAmC;oBAC9C,QAAQ,EAAE,aAAa;oBACvB,SAAS;iBACV,CAAC,CAAC;gBACH,MAAM,GAAG,IAAI,CAAC,SAAS,CAAC,GAAG,CAAC,CAAC;YAC/B,CAAC;YAAC,QAAO,sDAAuD,KAAK,EAAE,CAAC;gBACtE,OAAO;YACT,CAAC;YACD,UAAU,CAAC,OAAO,CAAC,OAAO,CAAC,MAAM,CAAC,MAAM,GAAG,IAAI,CAAC,CAAC,CAAC;QACpD,CAAC;QACD,KAAK,CAAC,UAAU;YACd,UAAU,CAAC,SAAS,EAAE,CAAC;QACzB,CAAC;KACF,CAAC,CAAC;IAEH,MAAM,CAAC,UAAU,CAAC,GAAG,cAAc,CAAC,QAAQ,CAAC,GAAG,EAAE,CAAC;IAEnD,kFAAkF;IAClF,MAAM,YAAY,GAAG,kBAAkB,EAAE,CAAC;IAC1C,MAAM,CAAC,WAAW,CAAC,YAAY,CAAC,CAAC,MAAM,CAAC,cAAc,CAAC,QAAQ,CAAC,CAAC;IAEjE,OAAO,IAAI,QAAQ,CAAC,UAAU,EAAE;QAC9B,OAAO,EAAE,EAAE,cAAc,EAAE,0BAA0B,EAAE;KACxD,CAAC,CAAC;AACL,CAAC,CAAC;AAEF,MAAM,CAAC,MAAM,4BAA4B,GAAG,KAAK,EAAE,QAAkB,EAAE,EAAE;IACvE,OAAO,QAAQ,CAAC;AAClB,CAAC,CAAC"}

package/dist/providers/google-vertex-ai/imageEdit.d.ts

@@ -0,0 +1,56 @@
+import { ErrorResponse, ImageGenerateResponse, ProviderConfig, RequestHandler } from '../types.js';
+import { GoogleErrorResponse } from './types.js';
+/**
+ * Derive MIME type from a File using its name (extension) only.
+ * Exported for use by Google AI Studio image edit. We do not use file.type because in Node/non-browser
+ * runtimes File may come from a polyfill where type is missing or inconsistent.
+ */
+export declare function getMimeFromFile(file: File): string;
+/**
+ * Gemini image edit config. Uses generateContent (contents + generationConfig); unsupported params stripped.
+ * @see https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/image-generation#edit-an-image
+ * @see https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/GenerationConfig#ImageConfig
+ */
+export declare const GoogleGeminiImageEditConfig: ProviderConfig;
+export declare const GoogleImageEditConfig: ProviderConfig;
+/** Multi-turn (and thought signatures for gemini-3-pro-image-preview): parse previous_messages; pass-through. */
+export declare function parsePreviousMessages(raw: unknown): unknown[] | undefined;
+export declare const GoogleImageEditRequestHandler: RequestHandler<FormData>;
+interface GoogleImageEditResponse {
+    predictions: Array<{
+        bytesBase64Encoded?: string;
+        mimeType: string;
+        raiFilteredReason?: string;
+        safetyAttributes: {
+            categories: string | string[];
+            scores: any;
+        };
+    }>;
+}
+/** Gemini generateContent image edit response (candidates + usageMetadata). */
+interface GoogleGeminiImageEditResponse {
+    candidates?: Array<{
+        content?: {
+            parts?: Array<{
+                inlineData?: {
+                    data?: string;
+                    mimeType?: string;
+                };
+            }>;
+        };
+    }>;
+    usageMetadata?: {
+        promptTokenCount?: number;
+        candidatesTokenCount?: number;
+        totalTokenCount?: number;
+        promptTokensDetails?: Array<{
+            modality?: string;
+            tokenCount?: number;
+        }>;
+    };
+}
+/** Transforms Gemini image edit response (candidates) to gateway format; includes usage when present. */
+export declare const GoogleGeminiImageEditResponseTransform: (response: GoogleGeminiImageEditResponse | GoogleErrorResponse, responseStatus: number) => ImageGenerateResponse | ErrorResponse;
+/** Routes Vertex image edit response to Gemini or Imagen transformer by response shape (candidates vs predictions). */
+export declare const GoogleImageEditResponseTransform: (response: GoogleGeminiImageEditResponse | GoogleImageEditResponse | GoogleErrorResponse, responseStatus: number) => ImageGenerateResponse | ErrorResponse;
+export {};