Merge pull request #42 from warjiang/feat-langchain

Feat langchain

Author: stulzq

.gitignore

@@ -19,3 +19,4 @@ bin/
 release/
 docker-compose.yml
 dist/
+config/config.yaml

README.md

@@ -19,6 +19,7 @@ Verified support projects：
 | -------------------------------------------------------- | ------ |
 | [chatgpt-web](https://github.com/Chanzhaoyu/chatgpt-web) | √   |
 | [chatbox](https://github.com/Bin-Huang/chatbox)          | √    |
+| [langchain](https://python.langchain.com/en/latest/)     | √    |
 
 ## Get Start
 
@@ -56,11 +57,17 @@ API Key: This value can be found in the **Keys & Endpoint** section when examini
 ### Use Docker
 
 ````shell
+# config by environment 
 docker run -d -p 8080:8080 --name=azure-openai-proxy \
   --env AZURE_OPENAI_ENDPOINT=your_azure_endpoint \
   --env AZURE_OPENAI_API_VER=your_azure_api_ver \
   --env AZURE_OPENAI_MODEL_MAPPER=your_azure_deploy_mapper \
   stulzq/azure-openai-proxy:latest
+
+# config by file
+docker run -d -p 8080:8080 --name=azure-openai-proxy \
+  -v /path/to/config-file.yaml:/app/config/config.yaml \
+  stulzq/azure-openai-proxy:latest
 ````
 
 Call API:

azure/init.go

@@ -1,11 +1,12 @@
 package azure
 
 import (
+	"fmt"
+	"github.com/spf13/viper"
 	"github.com/stulzq/azure-openai-proxy/constant"
+	"github.com/stulzq/azure-openai-proxy/util"
 	"log"
 	"net/url"
-	"os"
-	"regexp"
 	"strings"
 )
 
@@ -14,43 +15,92 @@ const (
 )
 
 var (
-	AzureOpenAIEndpoint      = ""
-	AzureOpenAIEndpointParse *url.URL
-
-	AzureOpenAIAPIVer = ""
-
-	AzureOpenAIModelMapper = map[string]string{
-		"gpt-3.5-turbo": "gpt-35-turbo",
-	}
-	fallbackModelMapper = regexp.MustCompile(`[.:]`)
+	C                     Config
+	ModelDeploymentConfig = map[string]DeploymentConfig{}
 )
 
-func Init() {
-	AzureOpenAIAPIVer = os.Getenv(constant.ENV_AZURE_OPENAI_API_VER)
-	AzureOpenAIEndpoint = os.Getenv(constant.ENV_AZURE_OPENAI_ENDPOINT)
+func Init() error {
+	var (
+		apiVersion        string
+		endpoint          string
+		openaiModelMapper string
+		err               error
+	)
 
-	if AzureOpenAIAPIVer == "" {
-		AzureOpenAIAPIVer = "2023-03-15-preview"
+	apiVersion = viper.GetString(constant.ENV_AZURE_OPENAI_API_VER)
+	endpoint = viper.GetString(constant.ENV_AZURE_OPENAI_ENDPOINT)
+	openaiModelMapper = viper.GetString(constant.ENV_AZURE_OPENAI_MODEL_MAPPER)
+	if endpoint != "" && openaiModelMapper != "" {
+		if apiVersion == "" {
+			apiVersion = "2023-03-15-preview"
+		}
+		InitFromEnvironmentVariables(apiVersion, endpoint, openaiModelMapper)
+	} else {
+		if err = InitFromConfigFile(); err != nil {
+			return err
+		}
 	}
 
-	var err error
-	AzureOpenAIEndpointParse, err = url.Parse(AzureOpenAIEndpoint)
-	if err != nil {
-		log.Fatal("parse AzureOpenAIEndpoint error: ", err)
+	// ensure apiBase likes /v1
+	apiBase := viper.GetString("api_base")
+	if !strings.HasPrefix(apiBase, "/") {
+		apiBase = "/" + apiBase
+	}
+	if strings.HasSuffix(apiBase, "/") {
+		apiBase = apiBase[:len(apiBase)-1]
+	}
+	viper.Set("api_base", apiBase)
+	log.Printf("apiBase is: %s", apiBase)
+	for _, itemConfig := range C.DeploymentConfig {
+		u, err := url.Parse(itemConfig.Endpoint)
+		if err != nil {
+			return fmt.Errorf("parse endpoint error: %w", err)
+		}
+		itemConfig.EndpointUrl = u
+		ModelDeploymentConfig[itemConfig.ModelName] = itemConfig
 	}
+	return err
+}
 
-	if v := os.Getenv(constant.ENV_AZURE_OPENAI_MODEL_MAPPER); v != "" {
-		for _, pair := range strings.Split(v, ",") {
+func InitFromEnvironmentVariables(apiVersion, endpoint, openaiModelMapper string) {
+	log.Println("Init from environment variables")
+	if openaiModelMapper != "" {
+		// openaiModelMapper example:
+		// gpt-3.5-turbo=deployment_name_for_gpt_model,text-davinci-003=deployment_name_for_davinci_model
+		for _, pair := range strings.Split(openaiModelMapper, ",") {
 			info := strings.Split(pair, "=")
 			if len(info) != 2 {
 				log.Fatalf("error parsing %s, invalid value %s", constant.ENV_AZURE_OPENAI_MODEL_MAPPER, pair)
 			}
-
-			AzureOpenAIModelMapper[info[0]] = info[1]
+			modelName, deploymentName := info[0], info[1]
+			ModelDeploymentConfig[modelName] = DeploymentConfig{
+				DeploymentName: deploymentName,
+				ModelName:      modelName,
+				Endpoint:       endpoint,
+				ApiKey:         "",
+				ApiVersion:     apiVersion,
+			}
 		}
 	}
+}
+
+func InitFromConfigFile() error {
+	log.Println("Init from config file")
+	workDir := util.GetWorkdir()
+	viper.SetConfigName("config")
+	viper.SetConfigType("yaml")
+	viper.AddConfigPath(fmt.Sprintf("%s/config", workDir))
+	if err := viper.ReadInConfig(); err != nil {
+		log.Printf("read config file error: %+v\n", err)
+		return err
+	}
 
-	log.Println("AzureOpenAIAPIVer: ", AzureOpenAIAPIVer)
-	log.Println("AzureOpenAIEndpoint: ", AzureOpenAIEndpoint)
-	log.Println("AzureOpenAIModelMapper: ", AzureOpenAIModelMapper)
+	if err := viper.Unmarshal(&C); err != nil {
+		log.Printf("unmarshal config file error: %+v\n", err)
+		return err
+	}
+	for _, configItem := range C.DeploymentConfig {
+		ModelDeploymentConfig[configItem.ModelName] = configItem
+	}
+	return nil
 }

azure/model.go

@@ -0,0 +1,100 @@
+package azure
+
+import (
+	"bytes"
+	"fmt"
+	"github.com/pkg/errors"
+	"log"
+	"net/http"
+	"net/url"
+	"path"
+	"strings"
+	"text/template"
+)
+
+type DeploymentConfig struct {
+	DeploymentName string   `yaml:"deployment_name" json:"deployment_name" mapstructure:"deployment_name"` // azure openai deployment name
+	ModelName      string   `yaml:"model_name" json:"model_name" mapstructure:"model_name"`                // corresponding model name in openai
+	Endpoint       string   `yaml:"endpoint" json:"endpoint" mapstructure:"endpoint"`                      // deployment endpoint
+	ApiKey         string   `yaml:"api_key" json:"api_key" mapstructure:"api_key"`                         // secrect key1 or 2
+	ApiVersion     string   `yaml:"api_version" json:"api_version" mapstructure:"api_version"`             // deployment version, not required
+	EndpointUrl    *url.URL // url.URL form deployment endpoint
+}
+
+type Config struct {
+	ApiBase          string             `yaml:"api_base" mapstructure:"api_base"`                   // if you use openai、langchain as sdk, it will be useful
+	DeploymentConfig []DeploymentConfig `yaml:"deployment_config" mapstructure:"deployment_config"` // deployment config
+}
+
+type RequestConverter interface {
+	Name() string
+	Convert(req *http.Request, config *DeploymentConfig) (*http.Request, error)
+}
+
+type StripPrefixConverter struct {
+	Prefix string
+}
+
+func (c *StripPrefixConverter) Name() string {
+	return "StripPrefix"
+}
+func (c *StripPrefixConverter) Convert(req *http.Request, config *DeploymentConfig) (*http.Request, error) {
+	req.Host = config.EndpointUrl.Host
+	req.URL.Scheme = config.EndpointUrl.Scheme
+	req.URL.Host = config.EndpointUrl.Host
+	req.URL.Path = path.Join(fmt.Sprintf("/openai/deployments/%s", config.DeploymentName), strings.Replace(req.URL.Path, c.Prefix+"/", "/", 1))
+	req.URL.RawPath = req.URL.EscapedPath()
+
+	query := req.URL.Query()
+	query.Add("api-version", config.ApiVersion)
+	req.URL.RawQuery = query.Encode()
+	return req, nil
+}
+func NewStripPrefixConverter(prefix string) *StripPrefixConverter {
+	return &StripPrefixConverter{
+		Prefix: prefix,
+	}
+}
+
+type TemplateConverter struct {
+	Tpl      string
+	Tempalte *template.Template
+}
+
+func (c *TemplateConverter) Name() string {
+	return "Template"
+}
+func (c *TemplateConverter) Convert(req *http.Request, config *DeploymentConfig) (*http.Request, error) {
+	data := map[string]interface{}{
+		"DeploymentName": config.DeploymentName,
+		"ModelName":      config.ModelName,
+		"Endpoint":       config.Endpoint,
+		"ApiKey":         config.ApiKey,
+		"ApiVersion":     config.ApiVersion,
+	}
+	buff := new(bytes.Buffer)
+	if err := c.Tempalte.Execute(buff, data); err != nil {
+		return req, errors.Wrap(err, "template execute error")
+	}
+
+	req.Host = config.EndpointUrl.Host
+	req.URL.Scheme = config.EndpointUrl.Scheme
+	req.URL.Host = config.EndpointUrl.Host
+	req.URL.Path = buff.String()
+	req.URL.RawPath = req.URL.EscapedPath()
+
+	query := req.URL.Query()
+	query.Add("api-version", config.ApiVersion)
+	req.URL.RawQuery = query.Encode()
+	return req, nil
+}
+func NewTemplateConverter(tpl string) *TemplateConverter {
+	_template, err := template.New("template").Parse(tpl)
+	if err != nil {
+		log.Fatalf("template parse error: %s", err.Error())
+	}
+	return &TemplateConverter{
+		Tpl:      tpl,
+		Tempalte: _template,
+	}
+}

azure/proxy.go

@@ -8,16 +8,21 @@ import (
 	"log"
 	"net/http"
 	"net/http/httputil"
-	"path"
 	"strings"
 
 	"github.com/bytedance/sonic"
 	"github.com/gin-gonic/gin"
 	"github.com/pkg/errors"
 )
 
+func ProxyWithConverter(requestConverter RequestConverter) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		Proxy(c, requestConverter)
+	}
+}
+
 // Proxy Azure OpenAI
-func Proxy(c *gin.Context) {
+func Proxy(c *gin.Context, requestConverter RequestConverter) {
 	if c.Request.Method == http.MethodOptions {
 		c.Header("Access-Control-Allow-Origin", "*")
 		c.Header("Access-Control-Allow-Methods", "GET, OPTIONS, POST")
@@ -34,38 +39,48 @@ func Proxy(c *gin.Context) {
 		body, _ := io.ReadAll(req.Body)
 		req.Body = io.NopCloser(bytes.NewBuffer(body))
 
-		// get model from body
-		model, err := sonic.Get(body, "model")
-		if err != nil {
-			util.SendError(c, errors.Wrap(err, "get model error"))
-			return
+		// get model from url params or body
+		model := c.Param("model")
+		if model == "" {
+			_model, err := sonic.Get(body, "model")
+			if err != nil {
+				util.SendError(c, errors.Wrap(err, "get model error"))
+				return
+			}
+			_modelStr, err := _model.String()
+			if err != nil {
+				util.SendError(c, errors.Wrap(err, "get model name error"))
+				return
+			}
+			model = _modelStr
 		}
 
 		// get deployment from request
-		deployment, err := model.String()
+		deployment, err := GetDeploymentByModel(model)
 		if err != nil {
-			util.SendError(c, errors.Wrap(err, "get deployment error"))
+			util.SendError(c, err)
 			return
 		}
-		deployment = GetDeploymentByModel(deployment)
 
-		// get auth token from header
-		rawToken := req.Header.Get("Authorization")
-		token := strings.TrimPrefix(rawToken, "Bearer ")
+		// get auth token from header or deployemnt config
+		token := deployment.ApiKey
+		if token == "" {
+			rawToken := req.Header.Get("Authorization")
+			token = strings.TrimPrefix(rawToken, "Bearer ")
+		}
+		if token == "" {
+			util.SendError(c, errors.New("token is empty"))
+			return
+		}
 		req.Header.Set(AuthHeaderKey, token)
 		req.Header.Del("Authorization")
 
 		originURL := req.URL.String()
-		req.Host = AzureOpenAIEndpointParse.Host
-		req.URL.Scheme = AzureOpenAIEndpointParse.Scheme
-		req.URL.Host = AzureOpenAIEndpointParse.Host
-		req.URL.Path = path.Join(fmt.Sprintf("/openai/deployments/%s", deployment), strings.Replace(req.URL.Path, "/v1/", "/", 1))
-		req.URL.RawPath = req.URL.EscapedPath()
-
-		query := req.URL.Query()
-		query.Add("api-version", AzureOpenAIAPIVer)
-		req.URL.RawQuery = query.Encode()
-
+		req, err = requestConverter.Convert(req, deployment)
+		if err != nil {
+			util.SendError(c, errors.Wrap(err, "convert request error"))
+			return
+		}
 		log.Printf("proxying request [%s] %s -> %s", model, originURL, req.URL.String())
 	}
 
@@ -80,10 +95,10 @@ func Proxy(c *gin.Context) {
 	}
 }
 
-func GetDeploymentByModel(model string) string {
-	if v, ok := AzureOpenAIModelMapper[model]; ok {
-		return v
+func GetDeploymentByModel(model string) (*DeploymentConfig, error) {
+	deploymentConfig, exist := ModelDeploymentConfig[model]
+	if !exist {
+		return nil, errors.New(fmt.Sprintf("deployment config for %s not found", model))
 	}
-
-	return fallbackModelMapper.ReplaceAllString(model, "")
+	return &deploymentConfig, nil
 }

cmd/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"github.com/spf13/viper"
 	"github.com/stulzq/azure-openai-proxy/azure"
 	"log"
 	"net/http"
@@ -22,6 +23,7 @@ var (
 )
 
 func main() {
+	viper.AutomaticEnv()
 	parseFlag()
 
 	azure.Init()