feat: add LiteLLM as AI gateway provider

app/api/chat/messages/route.ts

@@ -9,6 +9,7 @@ import { CohereClient } from 'cohere-ai';
 import Groq from 'groq-sdk';
 import OpenAI from 'openai';
 
+import { litellm } from '@/lib/provider/LiteLLM';
 import { Provider } from '@/config/provider';
 import { ApiConfig } from '@/types/app';
 import { toCohereRole } from '@/utils/provider/cohere';
@@ -185,6 +186,15 @@ export async function POST(req: Request) {
             const output = OpenAIStream(response);
             return new StreamingTextResponse(output);
         }
+        case Provider.LiteLLM: {
+            const response = await litellm.chat.completions.create({
+                model: config.model.model_id,
+                stream: true,
+                messages,
+            });
+            const output = OpenAIStream(response);
+            return new StreamingTextResponse(output);
+        }
         default:
             return new Response('Invalid Provider', { status: 400 });
     }

config/provider/index.ts

@@ -9,6 +9,7 @@ import { HuggingFaceModelId, HuggingFaceModelName } from '@/config/provider/hugg
 import { MistralModelId, MistralModelName } from '@/config/provider/mistral';
 import { OpenAIModelId, OpenAIModelName } from '@/config/provider/openai';
 import { PerplexityModelId, PerplexityModelName } from '@/config/provider/perplexity';
+import { LiteLLMModelId, LiteLLMModelName } from '@/config/provider/litellm';
 
 export type AllModelId =
     | AmazonModelId
@@ -21,7 +22,8 @@ export type AllModelId =
     | MistralModelId
     | GoogleModelId
     | OpenAIModelId
-    | PerplexityModelId;
+    | PerplexityModelId
+    | LiteLLMModelId;
 
 export type AllModelName =
     | AmazonModelName
@@ -34,7 +36,8 @@ export type AllModelName =
     | MistralModelName
     | GoogleModelName
     | OpenAIModelName
-    | PerplexityModelName;
+    | PerplexityModelName
+    | LiteLLMModelName;
 
 export enum Provider {
     Amazon = 'Amazon',
@@ -48,6 +51,7 @@ export enum Provider {
     HuggingFace = 'HuggingFace',
     Mistral = 'Mistral',
     Perplexity = 'Perplexity',
+    LiteLLM = 'LiteLLM',
 
     Custom = 'Custom',
 }
@@ -100,6 +104,10 @@ export const Providers: {
         id: 'perplexity',
         name: Provider.Perplexity,
     },
+    {
+        id: 'litellm',
+        name: Provider.LiteLLM,
+    },
     {
         id: 'custom',
         name: Provider.Custom,

config/provider/litellm.ts

@@ -0,0 +1,31 @@
+import { Model } from '@/types/model';
+
+export type LiteLLMModelId = string;
+export type LiteLLMModelName = string;
+
+export const model: Model[] = [
+    {
+        id: 'openai/gpt-4o-mini',
+        name: 'GPT-4o Mini (via LiteLLM)',
+        maxInputTokens: null,
+        maxOutputTokens: 16384,
+        maxTokens: 128000,
+        price: 0.6,
+    },
+    {
+        id: 'anthropic/claude-3-5-sonnet-20240620',
+        name: 'Claude 3.5 Sonnet (via LiteLLM)',
+        maxInputTokens: null,
+        maxOutputTokens: 8192,
+        maxTokens: 200000,
+        price: 18.0,
+    },
+    {
+        id: 'openai/gpt-4o',
+        name: 'GPT-4o (via LiteLLM)',
+        maxInputTokens: null,
+        maxOutputTokens: 16384,
+        maxTokens: 128000,
+        price: 15.0,
+    },
+];

lib/provider/LiteLLM.ts

@@ -0,0 +1,6 @@
+import OpenAI from 'openai';
+
+export const litellm = new OpenAI({
+    apiKey: process.env.LITELLM_API_KEY ?? '',
+    baseURL: process.env.LITELLM_BASE_URL ?? 'http://localhost:4000/v1',
+});