diff --git a/app/api/chat/messages/route.ts b/app/api/chat/messages/route.ts index 0090f6b4..d83f3b06 100644 --- a/app/api/chat/messages/route.ts +++ b/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 }); } diff --git a/config/provider/index.ts b/config/provider/index.ts index a72dabd0..ba1c2df7 100644 --- a/config/provider/index.ts +++ b/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, diff --git a/config/provider/litellm.ts b/config/provider/litellm.ts new file mode 100644 index 00000000..9115d64d --- /dev/null +++ b/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, + }, +]; diff --git a/lib/provider/LiteLLM.ts b/lib/provider/LiteLLM.ts new file mode 100644 index 00000000..8ca56a76 --- /dev/null +++ b/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', +});