Skip to content

Commit

Permalink
Merge pull request #999 from Abirdcfly/newgpts
Browse files Browse the repository at this point in the history 
fix: common user can only access public app
  • Loading branch information
@bjwswang
bjwswang authored Apr 9, 2024
2 parents 8e2af87 + 81c92be commit ff9f9d1
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 17 deletions.
2 changes: 1 addition & 1 deletion apiserver/pkg/chat/chat_docs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ func (cs *ChatServer) BuildConversationKnowledgeBase(ctx context.Context, req Co
return err
}
// create or update the conversation knowledgebase
_, err = controllerutil.CreateOrUpdate(ctx, cs.cli, kb, func() error {
_, err = controllerutil.CreateOrUpdate(ctx, cs.systemCli, kb, func() error {
if err := controllerutil.SetControllerReference(app, kb, pkgclient.Scheme); err != nil {
return err
}
Expand Down
52 changes: 36 additions & 16 deletions apiserver/pkg/chat/chat_server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,16 +54,16 @@ import (
)

type ChatServer struct {
cli runtimeclient.Client
storage storage.Storage
once sync.Once
isGpts bool
systemCli runtimeclient.Client
storage storage.Storage
once sync.Once
isGpts bool
}

func NewChatServer(cli runtimeclient.Client, isGpts bool) *ChatServer {
return &ChatServer{
cli: cli,
isGpts: isGpts,
systemCli: cli,
isGpts: isGpts,
}
}

Expand All @@ -81,7 +81,7 @@ func (cs *ChatServer) Storage() storage.Storage {
cs.storage = storage.NewMemoryStorage()
return
}
pg, err := datasource.GetPostgreSQLPool(ctx, cs.cli, ds)
pg, err := datasource.GetPostgreSQLPool(ctx, cs.systemCli, ds)
if err != nil {
klog.Errorf("get postgresql pool failed : %s", err.Error())
cs.storage = storage.NewMemoryStorage()
Expand Down Expand Up @@ -154,12 +154,12 @@ func (cs *ChatServer) AppRun(ctx context.Context, req ChatReqBody, respStream ch
Answer: "",
})
// since authenticattion already passed by http handler,we should use chatserver's client which is also the system client to new/ini appruntime
appRun, err := appruntime.NewAppOrGetFromCache(ctx, cs.cli, app)
appRun, err := appruntime.NewAppOrGetFromCache(ctx, cs.systemCli, app)
if err != nil {
return nil, err
}
klog.FromContext(ctx).Info("begin to run application", "appName", req.APPName, "appNamespace", req.AppNamespace)
out, err := appRun.Run(ctx, cs.cli, respStream, appruntime.Input{Question: req.Query, Files: req.Files, NeedStream: req.ResponseMode.IsStreaming(), History: history, ConversationID: req.ConversationID})
out, err := appRun.Run(ctx, cs.systemCli, respStream, appruntime.Input{Question: req.Query, Files: req.Files, NeedStream: req.ResponseMode.IsStreaming(), History: history, ConversationID: req.ConversationID})
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -254,19 +254,19 @@ func (cs *ChatServer) ListPromptStarters(ctx context.Context, req APPMetadata, l
switch baseNode.Kind() {
case "llmchain":
ch := appruntimechain.NewLLMChain(baseNode)
if err := ch.Init(ctx, cs.cli, nil); err != nil {
if err := ch.Init(ctx, cs.systemCli, nil); err != nil {
klog.Infof("init llmchain err:%s, will use empty chain config", err)
}
chainOptions = appruntimechain.GetChainOptions(ch.Instance.Spec.CommonChainConfig)
case "retrievalqachain":
ch := appruntimechain.NewRetrievalQAChain(baseNode)
if err := ch.Init(ctx, cs.cli, nil); err != nil {
if err := ch.Init(ctx, cs.systemCli, nil); err != nil {
klog.Infof("init retrievalqachain err:%s, will use empty chain config", err)
}
chainOptions = appruntimechain.GetChainOptions(ch.Instance.Spec.CommonChainConfig)
case "apichain":
ch := appruntimechain.NewAPIChain(baseNode)
if err := ch.Init(ctx, cs.cli, nil); err != nil {
if err := ch.Init(ctx, cs.systemCli, nil); err != nil {
klog.Infof("init apichain err:%s, will use empty chain config", err)
}
chainOptions = appruntimechain.GetChainOptions(ch.Instance.Spec.CommonChainConfig)
Expand All @@ -277,14 +277,14 @@ func (cs *ChatServer) ListPromptStarters(ctx context.Context, req APPMetadata, l
switch baseNode.Kind() {
case "llm":
l := llm.NewLLM(baseNode)
if err := l.Init(ctx, cs.cli, nil); err != nil {
if err := l.Init(ctx, cs.systemCli, nil); err != nil {
klog.Infof("init llm err:%s, abort", err)
return nil, err
}
model = l.Model
case "knowledgebase":
k := knowledgebase.NewKnowledgebase(baseNode)
if err := k.Init(ctx, cs.cli, nil); err != nil {
if err := k.Init(ctx, cs.systemCli, nil); err != nil {
klog.Infof("init knowledgebase err:%s, abort", err)
return nil, err
}
Expand All @@ -296,7 +296,7 @@ func (cs *ChatServer) ListPromptStarters(ctx context.Context, req APPMetadata, l
content := bytes.Buffer{}
// if there is a knowledgebase, use it to generate prompt starter
if kb != nil {
outArg, finish, err := retriever.GenerateKnowledgebaseRetriever(ctx, cs.cli, kb.Name, kb.Namespace, apiretriever.CommonRetrieverConfig{NumDocuments: limit * 2}, map[string]any{"question": "开始"})
outArg, finish, err := retriever.GenerateKnowledgebaseRetriever(ctx, cs.systemCli, kb.Name, kb.Namespace, apiretriever.CommonRetrieverConfig{NumDocuments: limit * 2}, map[string]any{"question": "开始"})
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -394,12 +394,15 @@ The question you asked is:`

func (cs *ChatServer) GetApp(ctx context.Context, appName, appNamespace string) (*v1alpha1.Application, error) {
app := &v1alpha1.Application{}
if err := cs.cli.Get(ctx, types.NamespacedName{Namespace: appNamespace, Name: appName}, app); err != nil {
if err := cs.systemCli.Get(ctx, types.NamespacedName{Namespace: appNamespace, Name: appName}, app); err != nil {
return nil, fmt.Errorf("failed to get application: %w", err)
}
if !app.Status.IsReady() {
return nil, fmt.Errorf("application not ready: %s", app.Status.GetCondition(v1alpha1.TypeReady).Message)
}
if !cs.IsGPTUserHasPermissionForApp(ctx, app) {
return nil, fmt.Errorf("user don't have permission for app: %s", app.Name)
}
return app, nil
}

Expand Down Expand Up @@ -446,3 +449,20 @@ func (cs *ChatServer) FillAppIconToConversations(ctx context.Context, conversati
}
return nil
}

func (cs *ChatServer) IsGPTUserHasPermissionForApp(ctx context.Context, app *v1alpha1.Application) (ok bool) {
if !cs.isGpts {
return true
}
// currentUser, _ := ctx.Value(auth.UserNameContextKey).(string)
if app.Spec.IsPublic {
return true
}
gptCofig, err := pkgconfig.GetGPTsConfig(ctx, cs.systemCli)
if err != nil {
klog.FromContext(ctx).Error(err, "failed to get gpt config")
return false
}
publicNs := gptCofig.PublicNamespace
return app.Namespace == publicNs
}

0 comments on commit ff9f9d1

Please sign in to comment.