Merge pull request #1235 from shiwo6324/feat/302AI-integration

feat: add 302AI provider integration

Author: VisargD

src/globals.ts

@@ -99,6 +99,7 @@ export const BYTEZ: string = 'bytez';
 export const FEATHERLESS_AI: string = 'featherless-ai';
 export const KRUTRIM: string = 'krutrim';
 export const QDRANT: string = 'qdrant';
+export const THREE_ZERO_TWO_AI: string = '302ai';
 
 export const VALID_PROVIDERS = [
   ANTHROPIC,
@@ -163,6 +164,7 @@ export const VALID_PROVIDERS = [
   FEATHERLESS_AI,
   KRUTRIM,
   QDRANT,
+  THREE_ZERO_TWO_AI,
 ];
 
 export const CONTENT_TYPES = {

src/providers/302ai/api.ts

@@ -0,0 +1,18 @@
+import { ProviderAPIConfig } from '../types';
+
+const AI302APIConfig: ProviderAPIConfig = {
+  getBaseURL: () => 'https://api.302.ai',
+  headers: ({ providerOptions }) => {
+    return { Authorization: `Bearer ${providerOptions.apiKey}` };
+  },
+  getEndpoint: ({ fn }) => {
+    switch (fn) {
+      case 'chatComplete':
+        return '/v1/chat/completions';
+      default:
+        return '';
+    }
+  },
+};
+
+export default AI302APIConfig;

src/providers/302ai/chatComplete.ts

@@ -0,0 +1,154 @@
+import { THREE_ZERO_TWO_AI } from '../../globals';
+import { OpenAIErrorResponseTransform } from '../openai/utils';
+import {
+  ChatCompletionResponse,
+  ErrorResponse,
+  ProviderConfig,
+} from '../types';
+import { generateInvalidProviderResponseError } from '../utils';
+
+export const AI302ChatCompleteConfig: ProviderConfig = {
+  model: {
+    param: 'model',
+    required: true,
+    default: 'gpt-3.5-turbo',
+  },
+  messages: {
+    param: 'messages',
+    default: '',
+  },
+  max_tokens: {
+    param: 'max_tokens',
+    default: 100,
+    min: 0,
+  },
+  temperature: {
+    param: 'temperature',
+    default: 1,
+    min: 0,
+    max: 2,
+  },
+  top_p: {
+    param: 'top_p',
+    default: 1,
+    min: 0,
+    max: 1,
+  },
+  stream: {
+    param: 'stream',
+    default: false,
+  },
+  frequency_penalty: {
+    param: 'frequency_penalty',
+    default: 0,
+    min: -2,
+    max: 2,
+  },
+  presence_penalty: {
+    param: 'presence_penalty',
+    default: 0,
+    min: -2,
+    max: 2,
+  },
+  stop: {
+    param: 'stop',
+    default: null,
+  },
+};
+
+interface AI302ChatCompleteResponse extends ChatCompletionResponse {
+  id: string;
+  object: string;
+  created: number;
+  model: string;
+  usage?: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
+}
+
+interface AI302StreamChunk {
+  id: string;
+  object: string;
+  created: number;
+  model: string;
+  choices: {
+    delta: {
+      role?: string | null;
+      content?: string;
+    };
+    index: number;
+    finish_reason: string | null;
+  }[];
+}
+
+export const AI302ChatCompleteResponseTransform: (
+  response: AI302ChatCompleteResponse | ErrorResponse,
+  responseStatus: number
+) => ChatCompletionResponse | ErrorResponse = (response, responseStatus) => {
+  if ('error' in response && responseStatus !== 200) {
+    return OpenAIErrorResponseTransform(response, THREE_ZERO_TWO_AI);
+  }
+
+  if ('choices' in response) {
+    return {
+      id: response.id,
+      object: response.object,
+      created: response.created,
+      model: response.model,
+      provider: THREE_ZERO_TWO_AI,
+      choices: response.choices.map((c) => ({
+        index: c.index,
+        message: {
+          role: c.message.role,
+          content: c.message.content,
+        },
+        finish_reason: c.finish_reason,
+      })),
+      usage: {
+        prompt_tokens: response.usage?.prompt_tokens || 0,
+        completion_tokens: response.usage?.completion_tokens || 0,
+        total_tokens: response.usage?.total_tokens || 0,
+      },
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, THREE_ZERO_TWO_AI);
+};
+
+export const AI302ChatCompleteStreamChunkTransform: (
+  response: string
+) => string = (responseChunk) => {
+  let chunk = responseChunk.trim();
+  chunk = chunk.replace(/^data: /, '');
+  chunk = chunk.trim();
+
+  if (chunk === '[DONE]') {
+    return `data: ${chunk}\n\n`;
+  }
+
+  try {
+    const parsedChunk: AI302StreamChunk = JSON.parse(chunk);
+
+    return (
+      `data: ${JSON.stringify({
+        id: parsedChunk.id,
+        object: parsedChunk.object,
+        created: parsedChunk.created,
+        model: parsedChunk.model,
+        provider: THREE_ZERO_TWO_AI,
+        choices: [
+          {
+            index: parsedChunk.choices[0]?.index ?? 0,
+            delta: parsedChunk.choices[0]?.delta ?? {},
+            finish_reason: parsedChunk.choices[0]?.finish_reason ?? null,
+          },
+        ],
+      })}` + '\n\n'
+    );
+  } catch (error) {
+    console.error('Error parsing 302AI stream chunk:', error);
+    return `data: ${chunk}\n\n`;
+  }
+};

src/providers/302ai/index.ts

@@ -0,0 +1,18 @@
+import { ProviderConfigs } from '../types';
+import AI302APIConfig from './api';
+import {
+  AI302ChatCompleteConfig,
+  AI302ChatCompleteResponseTransform,
+  AI302ChatCompleteStreamChunkTransform,
+} from './chatComplete';
+
+const AI302Config: ProviderConfigs = {
+  chatComplete: AI302ChatCompleteConfig,
+  api: AI302APIConfig,
+  responseTransforms: {
+    chatComplete: AI302ChatCompleteResponseTransform,
+    'stream-chatComplete': AI302ChatCompleteStreamChunkTransform,
+  },
+};
+
+export default AI302Config;

src/providers/index.ts

@@ -62,6 +62,7 @@ import NscaleConfig from './nscale';
 import HyperbolicConfig from './hyperbolic';
 import { FeatherlessAIConfig } from './featherless-ai';
 import KrutrimConfig from './krutrim';
+import AI302Config from './302ai';
 
 const Providers: { [key: string]: ProviderConfigs } = {
   openai: OpenAIConfig,
@@ -124,6 +125,7 @@ const Providers: { [key: string]: ProviderConfigs } = {
   bytez: BytezConfig,
   'featherless-ai': FeatherlessAIConfig,
   krutrim: KrutrimConfig,
+  '302ai': AI302Config,
 };
 
 export default Providers;