mirror of
https://github.com/songquanpeng/one-api.git
synced 2025-09-29 22:56:39 +08:00
55 lines
1.6 KiB
Go
55 lines
1.6 KiB
Go
package category
|
|
|
|
import (
|
|
"errors"
|
|
"net/http"
|
|
"one-api/common/requester"
|
|
"one-api/providers/base"
|
|
"one-api/types"
|
|
"strings"
|
|
)
|
|
|
|
var CategoryMap = map[string]Category{}
|
|
|
|
type Category struct {
|
|
ModelName string
|
|
ChatComplete ChatCompletionConvert
|
|
ResponseChatComplete ChatCompletionResponse
|
|
ResponseChatCompleteStrem ChatCompletionStreamResponse
|
|
}
|
|
|
|
func GetCategory(modelName string) (*Category, error) {
|
|
modelName = GetModelName(modelName)
|
|
|
|
// 点分割
|
|
provider := strings.Split(modelName, ".")[0]
|
|
|
|
if category, exists := CategoryMap[provider]; exists {
|
|
category.ModelName = modelName
|
|
return &category, nil
|
|
}
|
|
|
|
return nil, errors.New("category_not_found")
|
|
}
|
|
|
|
func GetModelName(modelName string) string {
|
|
bedrockMap := map[string]string{
|
|
"claude-3-sonnet-20240229": "anthropic.claude-3-sonnet-20240229-v1:0",
|
|
"claude-3-haiku-20240307": "anthropic.claude-3-haiku-20240307-v1:0",
|
|
"claude-2.1": "anthropic.claude-v2:1",
|
|
"claude-2.0": "anthropic.claude-v2",
|
|
"claude-instant-1.2": "anthropic.claude-instant-v1",
|
|
}
|
|
|
|
if value, exists := bedrockMap[modelName]; exists {
|
|
modelName = value
|
|
}
|
|
|
|
return modelName
|
|
}
|
|
|
|
type ChatCompletionConvert func(*types.ChatCompletionRequest) (any, *types.OpenAIErrorWithStatusCode)
|
|
type ChatCompletionResponse func(base.ProviderInterface, *http.Response, *types.ChatCompletionRequest) (*types.ChatCompletionResponse, *types.OpenAIErrorWithStatusCode)
|
|
|
|
type ChatCompletionStreamResponse func(base.ProviderInterface, *types.ChatCompletionRequest) requester.HandlerPrefix[string]
|