From 2490ae5134331a8f83fa8764b6f95d90db6fe496 Mon Sep 17 00:00:00 2001 From: Donnie Adams Date: Wed, 18 Dec 2024 12:01:53 -0500 Subject: [PATCH] enhance: create API endpoints for setting environment variables These endpoints are added for workflows and agents. Setting the environment variables for a workflow or agent will ensure that the object's manifest has the same information. Signed-off-by: Donnie Adams --- pkg/api/handlers/env.go | 132 ++++++++++++++++++++++++++++++++++ pkg/api/handlers/workflows.go | 3 - pkg/api/router/router.go | 6 ++ pkg/gateway/server/router.go | 2 +- 4 files changed, 139 insertions(+), 4 deletions(-) create mode 100644 pkg/api/handlers/env.go diff --git a/pkg/api/handlers/env.go b/pkg/api/handlers/env.go new file mode 100644 index 00000000..4bb07946 --- /dev/null +++ b/pkg/api/handlers/env.go @@ -0,0 +1,132 @@ +package handlers + +import ( + "context" + "errors" + "fmt" + "slices" + "strings" + + "github.com/gptscript-ai/go-gptscript" + "github.com/obot-platform/obot/apiclient/types" + "github.com/obot-platform/obot/pkg/api" + v1 "github.com/obot-platform/obot/pkg/storage/apis/otto.otto8.ai/v1" + "github.com/obot-platform/obot/pkg/system" + kclient "sigs.k8s.io/controller-runtime/pkg/client" +) + +func SetEnv(req api.Context) error { + id := req.PathValue("id") + if id == "" { + return types.NewErrBadRequest("id path variable is required") + } + + var envs map[string]string + if err := req.Read(&envs); err != nil { + return err + } + + var errs []error + for key, val := range envs { + if err := req.GPTClient.DeleteCredential(req.Context(), id, key); err != nil && !strings.HasSuffix(err.Error(), "credential not found") { + errs = append(errs, fmt.Errorf("failed to remove existing credetial %q: %w", key, err)) + continue + } + + if err := req.GPTClient.CreateCredential(req.Context(), gptscript.Credential{ + Context: id, + ToolName: key, + Type: gptscript.CredentialTypeTool, + Env: map[string]string{key: val}, + }); err != nil { + errs = append(errs, fmt.Errorf("failed to create credential %q: %w", key, err)) + } + } + + if len(errs) > 0 { + return errors.Join(errs...) + } + + obj, env, err := getObjectAndEnv(req.Context(), req.Storage, req.Namespace(), id) + if err != nil { + return err + } + + for i := 0; i < len(*env); i++ { + if _, ok := envs[(*env)[i].Name]; !ok { + // Delete the credential for the store + if err := req.GPTClient.DeleteCredential(req.Context(), id, (*env)[i].Name); err != nil && !strings.HasSuffix(err.Error(), "credential not found") { + errs = append(errs, fmt.Errorf("failed to remove existing credetial %q that is not longer needed: %w", (*env)[i].Name, err)) + continue + } + // Remove the item from the slice + *env = append((*env)[:i], (*env)[i+1:]...) + i-- + } + } + + if len(errs) > 0 { + return errors.Join(errs...) + } + + for name := range envs { + if !slices.ContainsFunc(*env, func(envVar types.EnvVar) bool { + return envVar.Name == name + }) { + *env = append(*env, types.EnvVar{Name: name}) + } + } + + if err = req.Update(obj); err != nil { + return fmt.Errorf("failed to update %s: %w", obj.GetObjectKind().GroupVersionKind().Kind, err) + } + + return nil +} + +func RevealEnv(req api.Context) error { + id := req.PathValue("id") + if id == "" { + return types.NewErrBadRequest("id path variable is required") + } + + _, env, err := getObjectAndEnv(req.Context(), req.Storage, req.Namespace(), id) + if err != nil { + return err + } + + resp := make(map[string]string, len(*env)) + for _, e := range *env { + cred, err := req.GPTClient.RevealCredential(req.Context(), []string{id}, e.Name) + if err != nil && !strings.HasSuffix(err.Error(), "credential not found") { + return err + } + + resp[e.Name] = cred.Env[e.Name] + } + + return req.Write(resp) +} + +func getObjectAndEnv(ctx context.Context, client kclient.Client, namespace, id string) (kclient.Object, *[]types.EnvVar, error) { + switch { + case system.IsAgentID(id): + var agent v1.Agent + if err := client.Get(ctx, kclient.ObjectKey{Namespace: namespace, Name: id}, &agent); err != nil { + return nil, nil, err + } + + return &agent, &agent.Spec.Manifest.Env, nil + + case system.IsWorkflowID(id): + var wf v1.Workflow + if err := client.Get(ctx, kclient.ObjectKey{Namespace: namespace, Name: id}, &wf); err != nil { + return nil, nil, err + } + + return &wf, &wf.Spec.Manifest.Env, nil + + default: + return nil, nil, types.NewErrBadRequest("%s is not an agent nor workflow", id) + } +} diff --git a/pkg/api/handlers/workflows.go b/pkg/api/handlers/workflows.go index 3cf27164..b3271ef3 100644 --- a/pkg/api/handlers/workflows.go +++ b/pkg/api/handlers/workflows.go @@ -48,9 +48,6 @@ func (a *WorkflowHandler) Authenticate(req api.Context) error { } agent.Spec.Manifest.Prompt = "#!sys.echo\nDONE" - if len(agent.Spec.Credentials) == 0 { - return nil - } resp, err := a.invoker.Agent(req.Context(), req.Storage, agent, "", invoke.Options{ Synchronous: true, diff --git a/pkg/api/router/router.go b/pkg/api/router/router.go index 543bff8f..58ede7a8 100644 --- a/pkg/api/router/router.go +++ b/pkg/api/router/router.go @@ -212,6 +212,12 @@ func Router(services *services.Services) (http.Handler, error) { mux.HandleFunc("DELETE /api/workflows/{context}/credentials/{id}", handlers.DeleteCredential) mux.HandleFunc("DELETE /api/credentials/{id}", handlers.DeleteCredential) + // Environment variable credentials + mux.HandleFunc("POST /api/workflows/{id}/env", handlers.SetEnv) + mux.HandleFunc("GET /api/workflows/{id}/env", handlers.RevealEnv) + mux.HandleFunc("POST /api/agents/{id}/env", handlers.SetEnv) + mux.HandleFunc("GET /api/agents/{id}/env", handlers.RevealEnv) + // Webhooks mux.HandleFunc("POST /api/webhooks", webhooks.Create) mux.HandleFunc("GET /api/webhooks", webhooks.List) diff --git a/pkg/gateway/server/router.go b/pkg/gateway/server/router.go index cb3b4190..b8aad984 100644 --- a/pkg/gateway/server/router.go +++ b/pkg/gateway/server/router.go @@ -35,7 +35,7 @@ func (s *Server) AddRoutes(mux *server.Server) { mux.HandleFunc("GET /api/oauth/start/{id}/{service}", wrap(s.oauth)) mux.HandleFunc("/api/oauth/redirect/{service}", wrap(s.redirect)) - // CRUD routes for OAuth Apps (integrations with other service such as Microsoft 365) + // CRUD routes for OAuth Apps (integrations with other services such as Microsoft 365) mux.HandleFunc("GET /api/oauth-apps", wrap(s.listOAuthApps)) mux.HandleFunc("GET /api/oauth-apps/{id}", wrap(s.oauthAppByID)) mux.HandleFunc("POST /api/oauth-apps", wrap(s.createOAuthApp))