feat: add support of OpenRouter (#619)

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

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=

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
+}

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)
 	}

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 (
       [