improvement: updated ai21 integration

src/providers/ai21/api.ts

@@ -8,17 +8,10 @@ const AI21APIConfig: ProviderAPIConfig = {
     };
     return headers;
   },
-  getEndpoint: ({ fn, gatewayRequestBodyJSON }) => {
-    const { model } = gatewayRequestBodyJSON;
+  getEndpoint: ({ fn }) => {
     switch (fn) {
-      case 'complete': {
-        return `/${model}/complete`;
-      }
       case 'chatComplete': {
-        return `/${model}/chat`;
-      }
-      case 'embed': {
-        return `/embed`;
+        return `/chat/completions`;
       }
       default:
         return '';

src/providers/ai21/chatComplete.ts

@@ -1,5 +1,5 @@
 import { AI21 } from '../../globals';
-import { Params, SYSTEM_MESSAGE_ROLES } from '../../types/requestBody';
+import { Params } from '../../types/requestBody';
 import {
   ChatCompletionResponse,
   ErrorResponse,
@@ -12,112 +12,115 @@ import {
 import { AI21ErrorResponse } from './complete';
 
 export const AI21ChatCompleteConfig: ProviderConfig = {
+  model: {
+    param: 'model',
+    required: true,
+  },
   messages: [
     {
       param: 'messages',
       required: true,
       transform: (params: Params) => {
-        let inputMessages: any = [];
-
-        if (
-          params.messages?.[0]?.role &&
-          SYSTEM_MESSAGE_ROLES.includes(params.messages?.[0]?.role)
-        ) {
-          inputMessages = params.messages.slice(1);
-        } else if (params.messages) {
-          inputMessages = params.messages;
-        }
+        return params.messages?.map((msg: any) => {
+          let textContent: string = '';
+          if (Array.isArray(msg.content)) {
+            msg.content.map((c: any) => {
+              if (c.type === 'text') {
+                textContent = c.text;
+              }
+            });
+          } else {
+            textContent = msg.content;
+          }
 
-        return inputMessages.map((msg: any) => ({
-          text: msg.content,
-          role: msg.role,
-        }));
+          return {
+            role: msg.role,
+            content: textContent,
+          };
+        });
       },
     },
     {
-      param: 'system',
-      required: false,
+      param: 'documents',
       transform: (params: Params) => {
-        if (
-          params.messages?.[0]?.role &&
-          SYSTEM_MESSAGE_ROLES.includes(params.messages?.[0]?.role)
-        ) {
-          return params.messages?.[0].content;
-        }
+        const documents: any[] = [];
+        params.messages?.forEach((msg: any) => {
+          if (Array.isArray(msg.content)) {
+            msg.content.forEach((c: any) => {
+              if (c.type !== 'text') {
+                documents.push({
+                  content: c.content,
+                  value: c.type,
+                });
+              }
+            });
+          }
+        });
+
+        return documents.length == 0 ? documents : undefined;
       },
     },
   ],
+  tools: {
+    param: 'tools',
+  },
   n: {
     param: 'numResults',
     default: 1,
+    min: 1,
+    max: 16,
   },
-  max_tokens: {
-    param: 'maxTokens',
-    default: 16,
+  response_format: {
+    param: 'responseFormat',
   },
-  max_completion_tokens: {
+  max_tokens: {
     param: 'maxTokens',
     default: 16,
   },
-  minTokens: {
-    param: 'minTokens',
-    default: 0,
-  },
   temperature: {
     param: 'temperature',
-    default: 0.7,
+    default: 0.4,
     min: 0,
-    max: 1,
+    max: 2,
   },
   top_p: {
     param: 'topP',
     default: 1,
-  },
-  top_k: {
-    param: 'topKReturn',
-    default: 0,
+    min: 0,
+    max: 1,
   },
   stop: {
     param: 'stopSequences',
   },
-  presence_penalty: {
-    param: 'presencePenalty',
-    transform: (params: Params) => {
-      return {
-        scale: params.presence_penalty,
-      };
-    },
-  },
-  frequency_penalty: {
-    param: 'frequencyPenalty',
-    transform: (params: Params) => {
-      return {
-        scale: params.frequency_penalty,
-      };
-    },
-  },
-  countPenalty: {
-    param: 'countPenalty',
-  },
-  frequencyPenalty: {
-    param: 'frequencyPenalty',
-  },
-  presencePenalty: {
-    param: 'presencePenalty',
+  stream: {
+    param: 'stream',
   },
 };
 
 interface AI21ChatCompleteResponse {
   id: string;
-  outputs: {
-    text: string;
-    role: string;
-    finishReason: {
-      reason: string;
-      length: number | null;
-      sequence: string | null;
+  model: string;
+  choices: {
+    index: {};
+    message: {
+      role: string;
+      content: string;
     };
+    tool_calls?: {
+      id: string;
+      type: 'function';
+      function: {
+        name: string;
+        arguments: Record<string, any>;
+      };
+    }[];
   }[];
+  finish_reason: 'stop' | 'length';
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
 }
 
 export const AI21ErrorResponseTransform: (
@@ -138,28 +141,37 @@ export const AI21ChatCompleteResponseTransform: (
   responseStatus: number
 ) => ChatCompletionResponse | ErrorResponse = (response, responseStatus) => {
   if (responseStatus !== 200) {
-    const errorResposne = AI21ErrorResponseTransform(
+    const errorResponse = AI21ErrorResponseTransform(
       response as AI21ErrorResponse
     );
-    if (errorResposne) return errorResposne;
+    if (errorResponse) return errorResponse;
   }
 
-  if ('outputs' in response) {
+  if ('choices' in response) {
     return {
       id: response.id,
       object: 'chat.completion',
-      created: Math.floor(Date.now() / 1000),
-      model: '',
-      provider: AI21,
-      choices: response.outputs.map((o, index) => ({
+      model: response.model,
+      choices: response.choices?.map((c, index) => ({
+        index: index,
         message: {
           role: 'assistant',
-          content: o.text,
+          content: c.message.content,
         },
-        index: index,
-        logprobs: null,
-        finish_reason: o.finishReason?.reason,
+        tool_calls: c.tool_calls?.map((toolCall: any) => {
+          return {
+            id: toolCall.id,
+            type: toolCall.type,
+            function: {
+              name: toolCall.function.name,
+              arguments: toolCall.function.arguments,
+            },
+          };
+        }),
       })),
+      finish_reason: response.finish_reason,
+      usage: response.usage,
+      created: Math.floor(Date.now() / 1000),
     };
   }
 

src/providers/types.ts

@@ -200,7 +200,7 @@ export interface GroundingMetadata {
 export interface ChatChoice {
   index: number;
   message: Message;
-  finish_reason: string;
+  finish_reason?: string;
   logprobs?: object | null;
   groundingMetadata?: GroundingMetadata;
 }