Skip to content

Commit

Permalink
✨ feat: support open router auto model list
Browse files Browse the repository at this point in the history
  • Loading branch information
arvinxx committed Apr 10, 2024
1 parent 1583011 commit 1ba90d3
Show file tree
Hide file tree
Showing 5 changed files with 70 additions and 8 deletions.
9 changes: 7 additions & 2 deletions src/app/api/chat/models/[provider]/route.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import { NextResponse } from 'next/server';

import { getPreferredRegion } from '@/app/api/config';
import { createErrorResponse } from '@/app/api/errorResponse';
import { ChatCompletionErrorPayload } from '@/libs/agent-runtime';
import { ChatCompletionErrorPayload, ModelProvider } from '@/libs/agent-runtime';
import { ChatErrorType } from '@/types/fetch';

import AgentRuntime from '../../agentRuntime';
Expand All @@ -16,7 +16,12 @@ export const GET = checkAuth(async (req, { params, jwtPayload }) => {
const { provider } = params;

try {
const agentRuntime = await AgentRuntime.initializeWithUserPayload(provider, jwtPayload);
const openRouterKey = jwtPayload.apiKey || 'openrouter-dont-need-api-key-for-model-list';

const agentRuntime = await AgentRuntime.initializeWithUserPayload(provider, {
...jwtPayload,
apiKey: provider === ModelProvider.OpenRouter ? openRouterKey : jwtPayload.apiKey,
});

const list = await agentRuntime.models();

Expand Down
23 changes: 23 additions & 0 deletions src/libs/agent-runtime/openrouter/index.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
import { LOBE_DEFAULT_MODEL_LIST } from '@/config/modelProviders';

import { AgentRuntimeErrorType } from '../error';
import { ModelProvider } from '../types';
import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
import { OpenRouterModelCard } from './type';

export const LobeOpenRouterAI = LobeOpenAICompatibleFactory({
baseURL: 'https://openrouter.ai/api/v1',
Expand All @@ -13,9 +16,29 @@ export const LobeOpenRouterAI = LobeOpenAICompatibleFactory({
debug: {
chatCompletion: () => process.env.DEBUG_OPENROUTER_CHAT_COMPLETION === '1',
},

errorType: {
bizError: AgentRuntimeErrorType.OpenRouterBizError,
invalidAPIKey: AgentRuntimeErrorType.InvalidOpenRouterAPIKey,
},
models: {
transformModel: (m) => {
const model = m as unknown as OpenRouterModelCard;

return {
description: model.description,
displayName: model.name,
enabled: LOBE_DEFAULT_MODEL_LIST.find((m) => model.id.endsWith(m.id))?.enabled || false,
functionCall: model.description.includes('function calling'),
id: model.id,
maxTokens:
typeof model.top_provider.max_completion_tokens === 'number'
? model.top_provider.max_completion_tokens
: undefined,
tokens: model.context_length,
vision: model.description.includes('vision'),
};
},
},
provider: ModelProvider.OpenRouter,
});
28 changes: 28 additions & 0 deletions src/libs/agent-runtime/openrouter/type.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
interface ModelPricing {
completion: string;
image: string;
prompt: string;
request: string;
}

interface ModelArchitecture {
instruct_type: string | null;
modality: string;
tokenizer: string;
}

interface ModelTopProvider {
is_moderated: boolean;
max_completion_tokens: number | null;
}

export interface OpenRouterModelCard {
architecture: ModelArchitecture;
context_length: number;
description: string;
id: string;
name: string;
per_request_limits: any | null;
pricing: ModelPricing;
top_provider: ModelTopProvider;
}
16 changes: 11 additions & 5 deletions src/libs/agent-runtime/utils/openaiCompatibleFactory/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,9 @@ interface OpenAICompatibleFactoryOptions {
bizError: ILobeAgentRuntimeErrorType;
invalidAPIKey: ILobeAgentRuntimeErrorType;
};
models?: {
transformModel?: (model: OpenAI.Model) => ChatModelCard;
};
provider: string;
}

Expand All @@ -35,6 +38,7 @@ export const LobeOpenAICompatibleFactory = ({
debug,
constructorOptions,
chatCompletion,
models,
}: OpenAICompatibleFactoryOptions) =>
class LobeOpenAICompatibleAI implements LobeRuntimeAI {
client: OpenAI;
Expand Down Expand Up @@ -107,15 +111,17 @@ export const LobeOpenAICompatibleFactory = ({
async models() {
const list = await this.client.models.list();

const modelIds = list.data.map((i) => i.id);
return list.data
.map((item) => {
if (models?.transformModel) {
return models.transformModel(item);
}

return modelIds
.map((id) => {
const knownModel = LOBE_DEFAULT_MODEL_LIST.find((model) => model.id === id);
const knownModel = LOBE_DEFAULT_MODEL_LIST.find((model) => model.id === item.id);

if (knownModel) return knownModel;

return { id: id };
return { id: item.id };
})
.filter(Boolean) as ChatModelCard[];
}
Expand Down
2 changes: 1 addition & 1 deletion src/services/models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,9 @@ class ModelsService {
headers: { 'Content-Type': 'application/json' },
provider,
});

try {
const res = await fetch(API_ENDPOINTS.chatModels(provider), { headers });
if (!res.ok) return;

return res.json();
} catch {
Expand Down

0 comments on commit 1ba90d3

Please sign in to comment.