From 5fda8b936b492268ff0ad9e299ce22565b60034d Mon Sep 17 00:00:00 2001 From: Kelvin Chiu Date: Sun, 10 Sep 2023 10:51:31 +0800 Subject: [PATCH] feat: add support of OpenRouter (#619) --- go.mod | 1 + go.sum | 2 + model/openrouter.go | 130 ++++++++++++++++++++++++++++++++++++++++++++ model/provider.go | 2 + web/src/Setting.js | 26 +++++++++ 5 files changed, 161 insertions(+) create mode 100644 model/openrouter.go diff --git a/go.mod b/go.mod index d25aff4..6bb95ff 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/casbin/casibase go 1.18 require ( + github.com/Lok-Lu/go-openrouter v0.0.0-20230807015935-ab5cee433ad3 github.com/aliyun/alibaba-cloud-sdk-go v1.61.1585 github.com/aliyun/aliyun-oss-go-sdk v2.2.2+incompatible github.com/anhao/go-ernie v1.0.4 diff --git a/go.sum b/go.sum index 20a1b83..fd30677 100644 --- a/go.sum +++ b/go.sum @@ -38,6 +38,8 @@ github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym github.com/DATA-DOG/go-sqlmock v1.3.3/go.mod h1:f/Ixk793poVmq4qj/V1dPUg2JEAKC73Q5eFN3EC/SaM= github.com/Knetic/govaluate v3.0.0+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0= github.com/Knetic/govaluate v3.0.1-0.20171022003610-9aa49832a739+incompatible/go.mod h1:r7JcOSlj0wfOMncg0iLm8Leh48TZaKVeNIfJntJ2wa0= +github.com/Lok-Lu/go-openrouter v0.0.0-20230807015935-ab5cee433ad3 h1:E15Zr1fu3QEIleqINqXpl4SlYuWkwWPJ4cx6UVaJITE= +github.com/Lok-Lu/go-openrouter v0.0.0-20230807015935-ab5cee433ad3/go.mod h1:nM0kITDAJEkwn9x2DQPdfJynThaLe0vGNpPcqJ8vjAI= github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs= github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo= github.com/Shopify/toxiproxy v2.1.4+incompatible/go.mod h1:OXgGpZ6Cli1/URJOF1DMxUHB2q5Ap20/P/eIdh4G0pI= diff --git a/model/openrouter.go b/model/openrouter.go new file mode 100644 index 0000000..211dbe6 --- /dev/null +++ b/model/openrouter.go @@ -0,0 +1,130 @@ +// Copyright 2023 The casbin Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package model + +import ( + "context" + "fmt" + "io" + "net/http" + "strings" + + "github.com/Lok-Lu/go-openrouter" + "github.com/casbin/casibase/proxy" +) + +type OpenRouterModelProvider struct { + subType string + secretKey string + siteName string + siteUrl string +} + +func NewOpenRouterModelProvider(subType string, secretKey string) (*OpenRouterModelProvider, error) { + p := &OpenRouterModelProvider{ + subType: subType, + secretKey: secretKey, + siteName: "Casibase", + siteUrl: "https://casibase.org", + } + return p, nil +} + +func (p *OpenRouterModelProvider) getProxyClientFromToken() *openrouter.Client { + config, err := openrouter.DefaultConfig(p.secretKey, p.siteName, p.siteUrl) + if err != nil { + panic(err) + } + + config.HTTPClient = proxy.ProxyHttpClient + + c := openrouter.NewClientWithConfig(config) + return c +} + +func (p *OpenRouterModelProvider) QueryText(question string, writer io.Writer, builder *strings.Builder) error { + client := p.getProxyClientFromToken() + + ctx := context.Background() + flusher, ok := writer.(http.Flusher) + if !ok { + return fmt.Errorf("writer does not implement http.Flusher") + } + + model := p.subType + if model == "" { + model = openrouter.Gpt35Turbo + } + + promptTokens, err := GetTokenSize(model, question) + if err != nil { + return err + } + + maxTokens := 4097 - promptTokens + + respStream, err := client.CreateChatCompletionStream( + ctx, + &openrouter.ChatCompletionRequest{ + Model: p.subType, + Messages: []openrouter.ChatCompletionMessage{ + { + Role: openrouter.ChatMessageRoleSystem, + Content: "You are a helpful assistant.", + }, + { + Role: openrouter.ChatMessageRoleUser, + Content: question, + }, + }, + Stream: false, + Temperature: nil, + TopP: nil, + MaxTokens: maxTokens, + }, + ) + if err != nil { + return err + } + defer respStream.Close() + + isLeadingReturn := true + for { + completion, streamErr := respStream.Recv() + if streamErr != nil { + if streamErr == io.EOF { + break + } + return streamErr + } + + data := completion.Choices[0].Message.Content + if isLeadingReturn && len(data) != 0 { + if strings.Count(data, "\n") == len(data) { + continue + } else { + isLeadingReturn = false + } + } + + if _, err = fmt.Fprintf(writer, "event: message\ndata: %s\n\n", data); err != nil { + return err + } + flusher.Flush() + builder.WriteString(data) + } + + return nil +} diff --git a/model/provider.go b/model/provider.go index 5e0fe1a..c5a3615 100644 --- a/model/provider.go +++ b/model/provider.go @@ -30,6 +30,8 @@ func GetModelProvider(typ string, subType string, clientId string, clientSecret p, err = NewOpenAiModelProvider(subType, clientSecret) } else if typ == "Hugging Face" { p, err = NewHuggingFaceModelProvider(subType, clientSecret) + } else if typ == "OpenRouter" { + p, err = NewOpenRouterModelProvider(subType, clientSecret) } else if typ == "Ernie" { p, err = NewErnieModelProvider(subType, clientId, clientSecret) } diff --git a/web/src/Setting.js b/web/src/Setting.js index ff19e09..f82de6f 100644 --- a/web/src/Setting.js +++ b/web/src/Setting.js @@ -659,6 +659,7 @@ export function getProviderTypeOptions(category) { [ {id: "OpenAI", name: "OpenAI"}, {id: "Hugging Face", name: "Hugging Face"}, + {id: "OpenRouter", name: "OpenRouter"}, {id: "Ernie", name: "Ernie"}, ] ); @@ -739,6 +740,31 @@ export function getProviderSubTypeOptions(category, type) { {id: "THUDM/chatglm2-6b", name: "THUDM/chatglm2-6b"}, ] ); + } else if (type === "OpenRouter") { + return ( + [ + {id: "google/palm-2-codechat-bison", name: "google/palm-2-codechat-bison"}, + {id: "google/palm-2-chat-bison", name: "google/palm-2-chat-bison"}, + {id: "openai/gpt-3.5-turbo", name: "openai/gpt-3.5-turbo"}, + {id: "openai/gpt-3.5-turbo-16k", name: "openai/gpt-3.5-turbo-16k"}, + {id: "openai/gpt-4", name: "openai/gpt-4"}, + {id: "openai/gpt-4-32k", name: "openai/gpt-4-32k"}, + {id: "anthropic/claude-2", name: "anthropic/claude-2"}, + {id: "anthropic/claude-instant-v1", name: "anthropic/claude-instant-v1"}, + {id: "meta-llama/llama-2-13b-chat", name: "meta-llama/llama-2-13b-chat"}, + {id: "meta-llama/llama-2-70b-chat", name: "meta-llama/llama-2-70b-chat"}, + {id: "palm-2-codechat-bison", name: "palm-2-codechat-bison"}, + {id: "palm-2-chat-bison", name: "palm-2-chat-bison"}, + {id: "gpt-3.5-turbo", name: "gpt-3.5-turbo"}, + {id: "gpt-3.5-turbo-16k", name: "gpt-3.5-turbo-16k"}, + {id: "gpt-4", name: "gpt-4"}, + {id: "gpt-4-32k", name: "gpt-4-32k"}, + {id: "claude-2", name: "claude-2"}, + {id: "claude-instant-v1", name: "claude-instant-v1"}, + {id: "llama-2-13b-chat", name: "llama-2-13b-chat"}, + {id: "llama-2-70b-chat", name: "llama-2-70b-chat"}, + ] + ); } else if (type === "Ernie") { return ( [