Merge pull request #41 from bricks-cloud/v1.7.7

WIP: V1.7.7

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 1.13.0 - 2024-03-14
+### Added
+- Added `policy` APIs
+
 ## 1.12.3 - 2024-03-07
 ### Added
 - Started supporting storing streaming responses in bytes as part of `event`

cmd/bricksllm/main.go

@@ -15,6 +15,9 @@ import (
 	"github.com/bricks-cloud/bricksllm/internal/logger/zap"
 	"github.com/bricks-cloud/bricksllm/internal/manager"
 	"github.com/bricks-cloud/bricksllm/internal/message"
+	"github.com/bricks-cloud/bricksllm/internal/pii"
+	"github.com/bricks-cloud/bricksllm/internal/pii/amazon"
+	custompolicy "github.com/bricks-cloud/bricksllm/internal/policy/custom"
 	"github.com/bricks-cloud/bricksllm/internal/provider/anthropic"
 	"github.com/bricks-cloud/bricksllm/internal/provider/azure"
 	"github.com/bricks-cloud/bricksllm/internal/provider/custom"
@@ -101,6 +104,11 @@ func main() {
 		log.Sugar().Fatalf("error altering provider settings table: %v", err)
 	}
 
+	err = store.CreatePolicyTable()
+	if err != nil {
+		log.Sugar().Fatalf("error creating policies table: %v", err)
+	}
+
 	memStore, err := memdb.NewMemDb(store, log, cfg.InMemoryDbUpdateInterval)
 	if err != nil {
 		log.Sugar().Fatalf("cannot initialize memdb: %v", err)
@@ -119,7 +127,7 @@ func main() {
 	}
 	cpMemStore.Listen()
 
-	rMemStore, err := memdb.NewRoutesMemDb(store, log, cfg.InMemoryDbUpdateInterval)
+	rMemStore, err := memdb.NewRoutesMemDb(store, store, log, cfg.InMemoryDbUpdateInterval)
 	if err != nil {
 		log.Sugar().Fatalf("cannot initialize routes memdb: %v", err)
 	}
@@ -196,7 +204,9 @@ func main() {
 	cpm := manager.NewCustomProvidersManager(store, cpMemStore)
 	rm := manager.NewRouteManager(store, store, rMemStore, psMemStore)
 
-	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, cfg.AdminPass)
+	pm := manager.NewPolicyManager(store, rMemStore)
+
+	as, err := admin.NewAdminServer(log, *modePtr, m, krm, psm, cpm, rm, pm, cfg.AdminPass)
 	if err != nil {
 		log.Sugar().Fatalf("error creating admin http server: %v", err)
 	}
@@ -232,7 +242,15 @@ func main() {
 	eventConsumer := message.NewConsumer(eventMessageChan, log, 4, handler.HandleEventWithRequestAndResponse)
 	eventConsumer.StartEventMessageConsumers()
 
-	ps, err := proxy.NewProxyServer(log, *modePtr, *privacyPtr, c, m, rm, a, psm, cpm, store, memStore, ce, ace, aoe, v, rec, messageBus, rlm, cfg.ProxyTimeout, accessCache)
+	detector, err := amazon.NewClient(cfg.AmazonRequestTimeout, cfg.AmazonConnectionTimeout)
+	if err != nil {
+		log.Sugar().Infof("error when connecting to amazon: %v", err)
+	}
+
+	scanner := pii.NewScanner(detector)
+	cd := custompolicy.NewOpenAiDetector(cfg.CustomPolicyDetectionTimeout, cfg.OpenAiApiKey)
+
+	ps, err := proxy.NewProxyServer(log, *modePtr, *privacyPtr, c, m, rm, a, psm, cpm, store, memStore, ce, ace, aoe, v, rec, messageBus, rlm, cfg.ProxyTimeout, accessCache, pm, scanner, cd)
 	if err != nil {
 		log.Sugar().Fatalf("error creating proxy http server: %v", err)
 	}

docker-compose.yml

@@ -18,22 +18,22 @@ services:
       - '5432:5432'
     volumes: 
       - postgresql:/var/lib/postgresql/data
-  bricksllm:
-    depends_on: 
-      - redis
-      - postgresql
-    image: luyuanxin1995/bricksllm
-    environment:
-      - POSTGRESQL_USERNAME=postgres
-      - POSTGRESQL_PASSWORD=postgres
-      - REDIS_PASSWORD=eYVX7EwVmmxKPCDmwMtyKVge8oLd2t81
-      - POSTGRESQL_HOSTS=postgresql
-      - REDIS_HOSTS=redis
-    ports:
-      - '8001:8001'
-      - '8002:8002'
-    command:
-      - '-m=dev'
+  # bricksllm:
+  #   depends_on: 
+  #     - redis
+  #     - postgresql
+  #   image: luyuanxin1995/bricksllm
+  #   environment:
+  #     - POSTGRESQL_USERNAME=postgres
+  #     - POSTGRESQL_PASSWORD=postgres
+  #     - REDIS_PASSWORD=eYVX7EwVmmxKPCDmwMtyKVge8oLd2t81
+  #     - POSTGRESQL_HOSTS=postgresql
+  #     - REDIS_HOSTS=redis
+  #   ports:
+  #     - '8001:8001'
+  #     - '8002:8002'
+  #   command:
+  #     - '-m=dev'
 volumes:
   redis:
     driver: local

go.mod

@@ -22,6 +22,25 @@ require (
 	github.com/asticode/go-astikit v0.20.0 // indirect
 	github.com/asticode/go-astisub v0.26.2 // indirect
 	github.com/asticode/go-astits v1.8.0 // indirect
+	github.com/aws/aws-sdk-go-v2 v1.25.3 // indirect
+	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 // indirect
+	github.com/aws/aws-sdk-go-v2/config v1.27.7 // indirect
+	github.com/aws/aws-sdk-go-v2/credentials v1.17.7 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.3 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.3 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.3 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.3 // indirect
+	github.com/aws/aws-sdk-go-v2/service/comprehend v1.31.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.5 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.5 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.3 // indirect
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.52.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.20.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.2 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.28.4 // indirect
+	github.com/aws/smithy-go v1.20.1 // indirect
 	github.com/bytedance/sonic v1.9.1 // indirect
 	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect

go.sum

@@ -8,6 +8,44 @@ github.com/asticode/go-astisub v0.26.2 h1:cdEXcm+SUSmYCEPTQYbbfCECnmQoIFfH6pF8wD
 github.com/asticode/go-astisub v0.26.2/go.mod h1:WTkuSzFB+Bp7wezuSf2Oxulj5A8zu2zLRVFf6bIFQK8=
 github.com/asticode/go-astits v1.8.0 h1:rf6aiiGn/QhlFjNON1n5plqF3Fs025XLUwiQ0NB6oZg=
 github.com/asticode/go-astits v1.8.0/go.mod h1:DkOWmBNQpnr9mv24KfZjq4JawCFX1FCqjLVGvO0DygQ=
+github.com/aws/aws-sdk-go-v2 v1.25.3 h1:xYiLpZTQs1mzvz5PaI6uR0Wh57ippuEthxS4iK5v0n0=
+github.com/aws/aws-sdk-go-v2 v1.25.3/go.mod h1:35hUlJVYd+M++iLI3ALmVwMOyRYMmRqUXpTtRGW+K9I=
+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 h1:gTK2uhtAPtFcdRRJilZPx8uJLL2J85xK11nKtWL0wfU=
+github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1/go.mod h1:sxpLb+nZk7tIfCWChfd+h4QwHNUR57d8hA1cleTkjJo=
+github.com/aws/aws-sdk-go-v2/config v1.27.7 h1:JSfb5nOQF01iOgxFI5OIKWwDiEXWTyTgg1Mm1mHi0A4=
+github.com/aws/aws-sdk-go-v2/config v1.27.7/go.mod h1:PH0/cNpoMO+B04qET699o5W92Ca79fVtbUnvMIZro4I=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.7 h1:WJd+ubWKoBeRh7A5iNMnxEOs982SyVKOJD+K8HIezu4=
+github.com/aws/aws-sdk-go-v2/credentials v1.17.7/go.mod h1:UQi7LMR0Vhvs+44w5ec8Q+VS+cd10cjwgHwiVkE0YGU=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.3 h1:p+y7FvkK2dxS+FEwRIDHDe//ZX+jDhP8HHE50ppj4iI=
+github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.15.3/go.mod h1:/fYB+FZbDlwlAiynK9KDXlzZl3ANI9JkD0Uhz5FjNT4=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.3 h1:ifbIbHZyGl1alsAhPIYsHOg5MuApgqOvVeI8wIugXfs=
+github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.3/go.mod h1:oQZXg3c6SNeY6OZrDY+xHcF4VGIEoNotX2B4PrDeoJI=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.3 h1:Qvodo9gHG9F3E8SfYOspPeBt0bjSbsevK8WhRAUHcoY=
+github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.3/go.mod h1:vCKrdLXtybdf/uQd/YfVR2r5pcbNuEYKzMQpcxmeSJw=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 h1:hT8rVHwugYE2lEfdFE0QWVo81lF7jMrYJVDWI+f+VxU=
+github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0/go.mod h1:8tu/lYfQfFe6IGnaOdrpVgEL2IrrDOf6/m9RQum4NkY=
+github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.3 h1:mDnFOE2sVkyphMWtTH+stv0eW3k0OTx94K63xpxHty4=
+github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.3/go.mod h1:V8MuRVcCRt5h1S+Fwu8KbC7l/gBGo3yBAyUbJM2IJOk=
+github.com/aws/aws-sdk-go-v2/service/comprehend v1.31.2 h1:iAnydKItgi2m2rOPFfyolvjXuZimVZgRPxGlYg6Vt5U=
+github.com/aws/aws-sdk-go-v2/service/comprehend v1.31.2/go.mod h1:4jJr/hungAbvS0vQqkZQvxBqxJ4oUSEpvezYM75q2e4=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 h1:EyBZibRTVAs6ECHZOw5/wlylS9OcTzwyjeQMudmREjE=
+github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1/go.mod h1:JKpmtYhhPs7D97NL/ltqz7yCkERFW5dOlHyVl66ZYF8=
+github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.5 h1:mbWNpfRUTT6bnacmvOTKXZjR/HycibdWzNpfbrbLDIs=
+github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.5/go.mod h1:FCOPWGjsshkkICJIn9hq9xr6dLKtyaWpuUojiN3W1/8=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.5 h1:K/NXvIftOlX+oGgWGIa3jDyYLDNsdVhsjHmsBH2GLAQ=
+github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.5/go.mod h1:cl9HGLV66EnCmMNzq4sYOti+/xo8w34CsgzVtm2GgsY=
+github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.3 h1:4t+QEX7BsXz98W8W1lNvMAG+NX8qHz2CjLBxQKku40g=
+github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.3/go.mod h1:oFcjjUq5Hm09N9rpxTdeMeLeQcxS7mIkBkL8qUKng+A=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.52.0 h1:k7gL76sSR0e2pLphjfmjD/+pDDtoOHvWp8ezpTsdyes=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.52.0/go.mod h1:MGTaf3x/+z7ZGugCGvepnx2DS6+caCYYqKhzVoLNYPk=
+github.com/aws/aws-sdk-go-v2/service/sso v1.20.2 h1:XOPfar83RIRPEzfihnp+U6udOveKZJvPQ76SKWrLRHc=
+github.com/aws/aws-sdk-go-v2/service/sso v1.20.2/go.mod h1:Vv9Xyk1KMHXrR3vNQe8W5LMFdTjSeWk0gBZBzvf3Qa0=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.2 h1:pi0Skl6mNl2w8qWZXcdOyg197Zsf4G97U7Sso9JXGZE=
+github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.2/go.mod h1:JYzLoEVeLXk+L4tn1+rrkfhkxl6mLDEVaDSvGq9og90=
+github.com/aws/aws-sdk-go-v2/service/sts v1.28.4 h1:Ppup1nVNAOWbBOrcoOxaxPeEnSFB2RnnQdguhXpmeQk=
+github.com/aws/aws-sdk-go-v2/service/sts v1.28.4/go.mod h1:+K1rNPVyGxkRuv9NNiaZ4YhBFuyw2MMA9SlIJ1Zlpz8=
+github.com/aws/smithy-go v1.20.1 h1:4SZlSlMr36UEqC7XOyRVb27XMeZubNcBNN+9IgEPIQw=
+github.com/aws/smithy-go v1.20.1/go.mod h1:krry+ya/rV9RDcV/Q16kpu6ypI4K2czasz0NC3qS14E=
 github.com/benbjohnson/clock v1.1.0 h1:Q92kusRqC1XV2MjkWETPvjJVqKetz1OzxZB7mHJLju8=
 github.com/bsm/ginkgo/v2 v2.7.0 h1:ItPMPH90RbmZJt5GtkcNvIRuGEdwlBItdNVoyzaNQao=
 github.com/bsm/gomega v1.26.0 h1:LhQm+AFcgV2M0WyKroMASzAzCAJVpAxQXv4SaI9a69Y=

internal/config/config.go

@@ -27,6 +27,10 @@ type Config struct {
 	AdminPass                     string        `env:"ADMIN_PASS"`
 	ProxyTimeout                  time.Duration `env:"PROXY_TIMEOUT" envDefault:"600s"`
 	NumberOfEventMessageConsumers int           `env:"NUMBER_OF_EVENT_MESSAGE_CONSUMERS" envDefault:"3"`
+	OpenAiApiKey                  string        `env:"OPENAI_API_KEY"`
+	CustomPolicyDetectionTimeout  time.Duration `env:"CUSTOM_POLICY_DETECTION_TIMEOUT" envDefault:"10m"`
+	AmazonRequestTimeout          time.Duration `env:"AMAZON_REQUEST_TIMEOUT" envDefault:"5s"`
+	AmazonConnectionTimeout       time.Duration `env:"AMAZON_CONNECTION_TIMEOUT" envDefault:"10s"`
 }
 
 func ParseEnvVariables() (*Config, error) {

internal/errors/blocked_err.go

@@ -0,0 +1,17 @@
+package errors
+
+type BlockedError struct {
+	message string
+}
+
+func NewBlockedError(msg string) *BlockedError {
+	return &BlockedError{
+		message: msg,
+	}
+}
+
+func (be *BlockedError) Error() string {
+	return be.message
+}
+
+func (be *BlockedError) Blocked() {}

internal/errors/warning_err.go

@@ -0,0 +1,17 @@
+package errors
+
+type WarningError struct {
+	message string
+}
+
+func NewWarningError(msg string) *WarningError {
+	return &WarningError{
+		message: msg,
+	}
+}
+
+func (we *WarningError) Error() string {
+	return we.message
+}
+
+func (we *WarningError) Warnings() {}

internal/key/key.go

@@ -28,6 +28,7 @@ type UpdateKey struct {
 	ShouldLogRequest       *bool         `json:"shouldLogRequest"`
 	ShouldLogResponse      *bool         `json:"shouldLogResponse"`
 	RotationEnabled        *bool         `json:"rotationEnabled"`
+	PolicyId               *string       `json:"policyId"`
 }
 
 func (uk *UpdateKey) Validate() error {
@@ -55,7 +56,7 @@ func (uk *UpdateKey) Validate() error {
 	if len(uk.SettingIds) != 0 {
 		for index, id := range uk.SettingIds {
 			if len(id) == 0 {
-				invalid = append(invalid, fmt.Sprintf("settingIds.[%d]", index))
+				invalid = append(invalid, fmt.Sprintf("settingIds[%d]", index))
 			}
 		}
 	}
@@ -74,6 +75,12 @@ func (uk *UpdateKey) Validate() error {
 		}
 	}
 
+	if uk.PolicyId != nil {
+		if len(*uk.PolicyId) == 0 {
+			invalid = append(invalid, "policyId")
+		}
+	}
+
 	if len(invalid) > 0 {
 		return internal_errors.NewValidationError(fmt.Sprintf("fields [%s] are invalid", strings.Join(invalid, ", ")))
 	}
@@ -161,6 +168,7 @@ type RequestKey struct {
 	ShouldLogRequest       bool         `json:"shouldLogRequest"`
 	ShouldLogResponse      bool         `json:"shouldLogResponse"`
 	RotationEnabled        bool         `json:"rotationEnabled"`
+	PolicyId               string       `json:"policyId"`
 }
 
 func (rk *RequestKey) Validate() error {
@@ -204,7 +212,7 @@ func (rk *RequestKey) Validate() error {
 	if len(rk.SettingIds) != 0 {
 		for index, id := range rk.SettingIds {
 			if len(id) == 0 {
-				invalid = append(invalid, fmt.Sprintf("settingIds.[%d]", index))
+				invalid = append(invalid, fmt.Sprintf("settingIds[%d]", index))
 			}
 		}
 	}
@@ -308,6 +316,7 @@ type ResponseKey struct {
 	ShouldLogRequest       bool         `json:"shouldLogRequest"`
 	ShouldLogResponse      bool         `json:"shouldLogResponse"`
 	RotationEnabled        bool         `json:"rotationEnabled"`
+	PolicyId               string       `json:"policyId"`
 }
 
 func (rk *ResponseKey) GetSettingIds() []string {

internal/manager/key.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/bricks-cloud/bricksllm/internal/encrypter"
 	"github.com/bricks-cloud/bricksllm/internal/key"
+	"github.com/bricks-cloud/bricksllm/internal/policy"
 	"github.com/bricks-cloud/bricksllm/internal/provider"
 	"github.com/bricks-cloud/bricksllm/internal/util"
 )
@@ -16,6 +17,7 @@ type Storage interface {
 	CreateKey(key *key.RequestKey) (*key.ResponseKey, error)
 	DeleteKey(id string) error
 	GetProviderSetting(id string) (*provider.Setting, error)
+	GetPolicyById(id string) (*policy.Policy, error)
 	GetProviderSettings(withSecret bool, ids []string) ([]*provider.Setting, error)
 	GetKey(keyId string) (*key.ResponseKey, error)
 }
@@ -83,6 +85,13 @@ func (m *Manager) CreateKey(rk *key.RequestKey) (*key.ResponseKey, error) {
 		}
 	}
 
+	if len(rk.PolicyId) != 0 {
+		_, err := m.s.GetPolicyById(rk.PolicyId)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return m.s.CreateKey(rk)
 }
 
@@ -131,6 +140,15 @@ func (m *Manager) UpdateKey(id string, uk *key.UpdateKey) (*key.ResponseKey, err
 		}
 	}
 
+	if uk.PolicyId != nil {
+		if len(*uk.PolicyId) != 0 {
+			_, err := m.s.GetPolicyById(*uk.PolicyId)
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
+
 	return m.s.UpdateKey(id, uk)
 }