Chore/return model in response

src/handlers/responseHandlers.ts

@@ -102,7 +102,8 @@ export async function responseHandler(
         provider,
         responseTransformerFunction,
         requestURL,
-        strictOpenAiCompliance
+        strictOpenAiCompliance,
+        gatewayRequest
       ),
       responseJson: null,
     };
@@ -148,7 +149,8 @@ export async function responseHandler(
     response,
     responseTransformerFunction,
     strictOpenAiCompliance,
-    gatewayRequestUrl
+    gatewayRequestUrl,
+    gatewayRequest
   );
 
   return {

src/handlers/streamHandler.ts

@@ -11,6 +11,7 @@ import {
 import { VertexLlamaChatCompleteStreamChunkTransform } from '../providers/google-vertex-ai/chatComplete';
 import { OpenAIChatCompleteResponse } from '../providers/openai/chatComplete';
 import { OpenAICompleteResponse } from '../providers/openai/complete';
+import { Params } from '../types/requestBody';
 import { getStreamModeSplitPattern, type SplitPatternType } from '../utils';
 
 function readUInt32BE(buffer: Uint8Array, offset: number) {
@@ -49,7 +50,9 @@ function concatenateUint8Arrays(a: Uint8Array, b: Uint8Array): Uint8Array {
 export async function* readAWSStream(
   reader: ReadableStreamDefaultReader,
   transformFunction: Function | undefined,
-  fallbackChunkId: string
+  fallbackChunkId: string,
+  strictOpenAiCompliance: boolean,
+  gatewayRequest: Params
 ) {
   let buffer = new Uint8Array();
   let expectedLength = 0;
@@ -68,7 +71,9 @@ export async function* readAWSStream(
             const transformedChunk = transformFunction(
               payload,
               fallbackChunkId,
-              streamState
+              streamState,
+              strictOpenAiCompliance,
+              gatewayRequest
             );
             if (Array.isArray(transformedChunk)) {
               for (const item of transformedChunk) {
@@ -102,7 +107,9 @@ export async function* readAWSStream(
         const transformedChunk = transformFunction(
           payload,
           fallbackChunkId,
-          streamState
+          streamState,
+          strictOpenAiCompliance,
+          gatewayRequest
         );
         if (Array.isArray(transformedChunk)) {
           for (const item of transformedChunk) {
@@ -124,7 +131,8 @@ export async function* readStream(
   transformFunction: Function | undefined,
   isSleepTimeRequired: boolean,
   fallbackChunkId: string,
-  strictOpenAiCompliance: boolean
+  strictOpenAiCompliance: boolean,
+  gatewayRequest: Params
 ) {
   let buffer = '';
   const decoder = new TextDecoder();
@@ -140,7 +148,8 @@ export async function* readStream(
             buffer,
             fallbackChunkId,
             streamState,
-            strictOpenAiCompliance
+            strictOpenAiCompliance,
+            gatewayRequest
           );
         } else {
           yield buffer;
@@ -171,7 +180,8 @@ export async function* readStream(
               part,
               fallbackChunkId,
               streamState,
-              strictOpenAiCompliance
+              strictOpenAiCompliance,
+              gatewayRequest
             );
             if (transformedChunk !== undefined) {
               yield transformedChunk;
@@ -215,7 +225,8 @@ export async function handleNonStreamingMode(
   response: Response,
   responseTransformer: Function | undefined,
   strictOpenAiCompliance: boolean,
-  gatewayRequestUrl: string
+  gatewayRequestUrl: string,
+  gatewayRequest: Params
 ): Promise<{
   response: Response;
   json: Record<string, any>;
@@ -241,7 +252,8 @@ export async function handleNonStreamingMode(
       response.status,
       response.headers,
       strictOpenAiCompliance,
-      gatewayRequestUrl
+      gatewayRequestUrl,
+      gatewayRequest
     );
   }
 
@@ -270,7 +282,8 @@ export function handleStreamingMode(
   proxyProvider: string,
   responseTransformer: Function | undefined,
   requestURL: string,
-  strictOpenAiCompliance: boolean
+  strictOpenAiCompliance: boolean,
+  gatewayRequest: Params
 ): Response {
   const splitPattern = getStreamModeSplitPattern(proxyProvider, requestURL);
   // If the provider doesn't supply completion id,
@@ -291,7 +304,9 @@ export function handleStreamingMode(
       for await (const chunk of readAWSStream(
         reader,
         responseTransformer,
-        fallbackChunkId
+        fallbackChunkId,
+        strictOpenAiCompliance,
+        gatewayRequest
       )) {
         await writer.write(encoder.encode(chunk));
       }
@@ -305,7 +320,8 @@ export function handleStreamingMode(
         responseTransformer,
         isSleepTimeRequired,
         fallbackChunkId,
-        strictOpenAiCompliance
+        strictOpenAiCompliance,
+        gatewayRequest
       )) {
         await writer.write(encoder.encode(chunk));
       }

src/providers/ai21/embed.ts

@@ -39,7 +39,6 @@ export const AI21EmbedResponseTransform: (
     );
     if (errorResposne) return errorResposne;
   }
-
   if ('results' in response) {
     return {
       object: 'list',

src/providers/anthropic/chatComplete.ts

@@ -410,6 +410,7 @@ export interface AnthropicChatCompleteStreamResponse {
       cache_creation_input_tokens?: number;
       cache_read_input_tokens?: number;
     };
+    model?: string;
   };
   error?: AnthropicErrorObject;
 }
@@ -529,7 +530,7 @@ export const AnthropicChatCompleteStreamChunkTransform: (
   response: string,
   fallbackId: string,
   streamState: AnthropicStreamState,
-  strictOpenAiCompliance: boolean
+  _strictOpenAiCompliance: boolean
 ) => string | undefined = (
   responseChunk,
   fallbackId,
@@ -585,6 +586,7 @@ export const AnthropicChatCompleteStreamChunkTransform: (
     parsedChunk.message?.usage?.cache_creation_input_tokens;
 
   if (parsedChunk.type === 'message_start' && parsedChunk.message?.usage) {
+    streamState.model = parsedChunk?.message?.model ?? '';
     streamState.usage = {
       prompt_tokens: parsedChunk.message?.usage?.input_tokens,
       ...(shouldSendCacheUsage && {
@@ -599,7 +601,7 @@ export const AnthropicChatCompleteStreamChunkTransform: (
         id: fallbackId,
         object: 'chat.completion.chunk',
         created: Math.floor(Date.now() / 1000),
-        model: '',
+        model: streamState.model,
         provider: ANTHROPIC,
         choices: [
           {
@@ -626,7 +628,7 @@ export const AnthropicChatCompleteStreamChunkTransform: (
         id: fallbackId,
         object: 'chat.completion.chunk',
         created: Math.floor(Date.now() / 1000),
-        model: '',
+        model: streamState.model,
         provider: ANTHROPIC,
         choices: [
           {
@@ -689,7 +691,7 @@ export const AnthropicChatCompleteStreamChunkTransform: (
       id: fallbackId,
       object: 'chat.completion.chunk',
       created: Math.floor(Date.now() / 1000),
-      model: '',
+      model: streamState.model,
       provider: ANTHROPIC,
       choices: [
         {

src/providers/anthropic/types.ts

@@ -6,4 +6,5 @@ export type AnthropicStreamState = {
     cache_read_input_tokens?: number;
     cache_creation_input_tokens?: number;
   };
+  model?: string;
 };

src/providers/bedrock/chatComplete.ts

@@ -402,12 +402,16 @@ export const BedrockChatCompleteResponseTransform: (
   response: BedrockChatCompletionResponse | BedrockErrorResponse,
   responseStatus: number,
   responseHeaders: Headers,
-  strictOpenAiCompliance: boolean
+  strictOpenAiCompliance: boolean,
+  _gatewayRequestUrl: string,
+  gatewayRequest: Params
 ) => ChatCompletionResponse | ErrorResponse = (
   response,
   responseStatus,
-  _responseHeaders,
-  strictOpenAiCompliance
+  responseHeaders,
+  strictOpenAiCompliance,
+  _gatewayRequestUrl,
+  gatewayRequest
 ) => {
   if (responseStatus !== 200) {
     const errorResponse = BedrockErrorResponseTransform(
@@ -430,7 +434,7 @@ export const BedrockChatCompleteResponseTransform: (
       id: Date.now().toString(),
       object: 'chat.completion',
       created: Math.floor(Date.now() / 1000),
-      model: '',
+      model: gatewayRequest.model || '',
       provider: BEDROCK,
       choices: [
         {
@@ -512,12 +516,14 @@ export const BedrockChatCompleteStreamChunkTransform: (
   response: string,
   fallbackId: string,
   streamState: BedrockStreamState,
-  strictOpenAiCompliance: boolean
+  strictOpenAiCompliance: boolean,
+  gatewayRequest: Params
 ) => string | string[] = (
   responseChunk,
   fallbackId,
   streamState,
-  strictOpenAiCompliance
+  strictOpenAiCompliance,
+  gatewayRequest
 ) => {
   const parsedChunk: BedrockChatCompleteStreamChunk = JSON.parse(responseChunk);
   if (parsedChunk.stopReason) {
@@ -533,7 +539,7 @@ export const BedrockChatCompleteStreamChunkTransform: (
         id: fallbackId,
         object: 'chat.completion.chunk',
         created: Math.floor(Date.now() / 1000),
-        model: '',
+        model: gatewayRequest.model || '',
         provider: BEDROCK,
         choices: [
           {
@@ -597,7 +603,7 @@ export const BedrockChatCompleteStreamChunkTransform: (
     id: fallbackId,
     object: 'chat.completion.chunk',
     created: Math.floor(Date.now() / 1000),
-    model: '',
+    model: gatewayRequest.model || '',
     provider: BEDROCK,
     choices: [
       {
@@ -802,11 +808,17 @@ export const BedrockCohereChatCompleteConfig: ProviderConfig = {
 export const BedrockCohereChatCompleteResponseTransform: (
   response: BedrockCohereCompleteResponse | BedrockErrorResponse,
   responseStatus: number,
-  responseHeaders: Headers
+  responseHeaders: Headers,
+  strictOpenAiCompliance: boolean,
+  gatewayRequestUrl: string,
+  gatewayRequest: Params
 ) => ChatCompletionResponse | ErrorResponse = (
   response,
   responseStatus,
-  responseHeaders
+  responseHeaders,
+  _strictOpenAiCompliance,
+  _gatewayRequestUrl,
+  gatewayRequest
 ) => {
   if (responseStatus !== 200) {
     const errorResposne = BedrockErrorResponseTransform(
@@ -815,6 +827,8 @@ export const BedrockCohereChatCompleteResponseTransform: (
     if (errorResposne) return errorResposne;
   }
 
+  const model = gatewayRequest.model || '';
+
   if ('generations' in response) {
     const prompt_tokens =
       Number(responseHeaders.get('X-Amzn-Bedrock-Input-Token-Count')) || 0;
@@ -824,7 +838,7 @@ export const BedrockCohereChatCompleteResponseTransform: (
       id: Date.now().toString(),
       object: 'chat.completion',
       created: Math.floor(Date.now() / 1000),
-      model: '',
+      model,
       provider: BEDROCK,
       choices: response.generations.map((generation, index) => ({
         index: index,
@@ -847,21 +861,31 @@ export const BedrockCohereChatCompleteResponseTransform: (
 
 export const BedrockCohereChatCompleteStreamChunkTransform: (
   response: string,
-  fallbackId: string
-) => string | string[] = (responseChunk, fallbackId) => {
+  fallbackId: string,
+  _streamState: Record<string, any>,
+  _strictOpenAiCompliance: boolean,
+  gatewayRequest: Params
+) => string | string[] = (
+  responseChunk,
+  fallbackId,
+  _streamState,
+  _strictOpenAiCompliance,
+  gatewayRequest
+) => {
   let chunk = responseChunk.trim();
   chunk = chunk.replace(/^data: /, '');
   chunk = chunk.trim();
   const parsedChunk: BedrockCohereStreamChunk = JSON.parse(chunk);
 
+  const model = gatewayRequest.model || '';
   // discard the last cohere chunk as it sends the whole response combined.
   if (parsedChunk.is_finished) {
     return [
       `data: ${JSON.stringify({
         id: fallbackId,
         object: 'chat.completion.chunk',
         created: Math.floor(Date.now() / 1000),
-        model: '',
+        model,
         provider: BEDROCK,
         choices: [
           {
@@ -888,7 +912,7 @@ export const BedrockCohereChatCompleteStreamChunkTransform: (
     id: fallbackId,
     object: 'chat.completion.chunk',
     created: Math.floor(Date.now() / 1000),
-    model: '',
+    model,
     provider: BEDROCK,
     choices: [
       {
@@ -978,11 +1002,17 @@ export const BedrockAI21ChatCompleteConfig: ProviderConfig = {
 export const BedrockAI21ChatCompleteResponseTransform: (
   response: BedrockAI21CompleteResponse | BedrockErrorResponse,
   responseStatus: number,
-  responseHeaders: Headers
+  responseHeaders: Headers,
+  strictOpenAiCompliance: boolean,
+  _gatewayRequestUrl: string,
+  gatewayRequest: Params
 ) => ChatCompletionResponse | ErrorResponse = (
   response,
   responseStatus,
-  responseHeaders
+  responseHeaders,
+  _strictOpenAiCompliance,
+  _gatewayRequestUrl,
+  gatewayRequest
 ) => {
   if (responseStatus !== 200) {
     const errorResposne = BedrockErrorResponseTransform(
@@ -1000,7 +1030,7 @@ export const BedrockAI21ChatCompleteResponseTransform: (
       id: response.id.toString(),
       object: 'chat.completion',
       created: Math.floor(Date.now() / 1000),
-      model: '',
+      model: gatewayRequest.model ?? '',
       provider: BEDROCK,
       choices: response.completions.map((completion, index) => ({
         index: index,