feat: add support for s3 workspace provider

Signed-off-by: Donnie Adams <[email protected]>
obot-platform · Oct 18, 2024 · 5dbb7c6 · 5dbb7c6
1 parent b093f91
commit 5dbb7c6
Show file tree

Hide file tree

Showing 8 changed files with 57 additions and 49 deletions.
diff --git a/pkg/api/handlers/threads.go b/pkg/api/handlers/threads.go
@@ -244,7 +244,7 @@ func (a *ThreadHandler) DeleteFile(req api.Context) error {
 
 	for _, workspace := range workspaces.Items {
 		if !workspace.Spec.IsKnowledge {
-			return deleteFileFromWorkspaceID(req.Context(), req, a.gptscript, workspace.Spec.WorkspaceID)
+			return deleteFileFromWorkspaceID(req.Context(), req, a.gptscript, workspace.Status.WorkspaceID)
 		}
 	}
 

diff --git a/pkg/controller/handlers/knowledge/knowledge.go b/pkg/controller/handlers/knowledge/knowledge.go
@@ -28,18 +28,16 @@ import (
 )
 
 type Handler struct {
-	gptscript         *gptscript.GPTScript
-	ingester          *knowledge.Ingester
-	events            *events.Emitter
-	workspaceProvider string
+	gptscript *gptscript.GPTScript
+	ingester  *knowledge.Ingester
+	events    *events.Emitter
 }
 
-func New(gClient *gptscript.GPTScript, ingester *knowledge.Ingester, wp string, events *events.Emitter) *Handler {
+func New(gClient *gptscript.GPTScript, ingester *knowledge.Ingester, events *events.Emitter) *Handler {
 	return &Handler{
-		gptscript:         gClient,
-		ingester:          ingester,
-		workspaceProvider: wp,
-		events:            events,
+		gptscript: gClient,
+		ingester:  ingester,
+		events:    events,
 	}
 }
 

diff --git a/pkg/controller/handlers/knowledgeset/knowledgeset.go b/pkg/controller/handlers/knowledgeset/knowledgeset.go
@@ -90,7 +90,6 @@ func (h *Handler) CreateWorkspace(req router.Request, resp router.Response) erro
 	}
 
 	ws := &v1.Workspace{
-		TypeMeta: metav1.TypeMeta{},
 		ObjectMeta: metav1.ObjectMeta{
 			GenerateName: system.WorkspacePrefix,
 			Namespace:    ks.Namespace,
@@ -101,7 +100,6 @@ func (h *Handler) CreateWorkspace(req router.Request, resp router.Response) erro
 			KnowledgeSetName: ks.Name,
 			IsKnowledge:      true,
 		},
-		Status: v1.WorkspaceStatus{},
 	}
 
 	if err := req.Client.Create(req.Ctx, ws); err != nil {

diff --git a/pkg/controller/handlers/workspace/workspace.go b/pkg/controller/handlers/workspace/workspace.go
@@ -29,7 +29,13 @@ func (a *Handler) CreateWorkspace(req router.Request, _ router.Response) error {
 		return nil
 	}
 
-	workspaceID, err := a.gptscript.CreateWorkspace(req.Ctx, a.workspaceProvider, ws.Spec.FromWorkspaces...)
+	providerType := a.workspaceProvider
+	if ws.Spec.IsKnowledge {
+		// Knowledge files should be stored locally.
+		providerType = "directory"
+	}
+
+	workspaceID, err := a.gptscript.CreateWorkspace(req.Ctx, providerType, ws.Spec.FromWorkspaces...)
 	if err != nil {
 		return err
 	}

diff --git a/pkg/controller/routes.go b/pkg/controller/routes.go
@@ -28,8 +28,8 @@ func (c *Controller) setupRoutes() error {
 	workflowStep := workflowstep.New(c.services.Invoker)
 	ingester := knowledge.NewIngester(c.services.Invoker)
 	toolRef := toolreference.New(c.services.GPTClient, c.services.ToolRegistryURL)
-	workspace := workspace.New(c.services.GPTClient, "directory")
-	knowledge := knowledgehandler.New(c.services.GPTClient, ingester, "directory", c.services.Events)
+	workspace := workspace.New(c.services.GPTClient, c.services.WorkspaceProviderType)
+	knowledge := knowledgehandler.New(c.services.GPTClient, ingester, c.services.Events)
 	knowledgeset := knowledgeset.New(c.services.AIHelper)
 	uploads := uploads.New(c.services.Invoker, c.services.GPTClient, "directory")
 	runs := runs.New(c.services.Invoker)

diff --git a/pkg/invoke/invoker.go b/pkg/invoke/invoker.go
@@ -39,13 +39,13 @@ type Invoker struct {
 	threadWorkspaceProvider string
 }
 
-func NewInvoker(c kclient.Client, gptClient *gptscript.GPTScript, tokenService *jwt.TokenService, events *events.Emitter) *Invoker {
+func NewInvoker(c kclient.Client, gptClient *gptscript.GPTScript, workspaceProviderType string, tokenService *jwt.TokenService, events *events.Emitter) *Invoker {
 	return &Invoker{
 		uncached:                c,
 		gptClient:               gptClient,
 		tokenService:            tokenService,
 		events:                  events,
-		threadWorkspaceProvider: "directory",
+		threadWorkspaceProvider: workspaceProviderType,
 	}
 }
 

diff --git a/pkg/services/config.go b/pkg/services/config.go
@@ -42,31 +42,33 @@ type (
 )
 
 type Config struct {
-	HTTPListenPort int    `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
-	DevMode        bool   `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OTTO_DEV_MODE"`
-	DevUIPort      int    `usage:"The port on localhost running the dev instance of the UI" default:"5173"`
-	AllowedOrigin  string `usage:"Allowed origin for CORS"`
-	ToolRegistry   string `usage:"The tool reference for the tool registry" default:"github.com/gptscript-ai/tools"`
+	HTTPListenPort        int    `usage:"HTTP port to listen on" default:"8080" name:"http-listen-port"`
+	DevMode               bool   `usage:"Enable development mode" default:"false" name:"dev-mode" env:"OTTO_DEV_MODE"`
+	DevUIPort             int    `usage:"The port on localhost running the dev instance of the UI" default:"5173"`
+	AllowedOrigin         string `usage:"Allowed origin for CORS"`
+	ToolRegistry          string `usage:"The tool reference for the tool registry" default:"github.com/gptscript-ai/tools"`
+	WorkspaceProviderType string `usage:"The type of workspace provider to use for non-knowledge workspaces" default:"directory" env:"OTTO_WORKSPACE_PROVIDER_TYPE"`
 
 	AuthConfig
 	GatewayConfig
 	services.Config
 }
 
 type Services struct {
-	ToolRegistryURL string
-	DevUIPort       int
-	Events          *events.Emitter
-	StorageClient   storage.Client
-	Router          *router.Router
-	GPTClient       *gptscript.GPTScript
-	Invoker         *invoke.Invoker
-	TokenServer     *jwt.TokenService
-	APIServer       *server.Server
-	AIHelper        *aihelper.AIHelper
-	Started         chan struct{}
-	ProxyServer     *proxy.Proxy
-	GatewayServer   *gserver.Server
+	ToolRegistryURL       string
+	WorkspaceProviderType string
+	DevUIPort             int
+	Events                *events.Emitter
+	StorageClient         storage.Client
+	Router                *router.Router
+	GPTClient             *gptscript.GPTScript
+	Invoker               *invoke.Invoker
+	TokenServer           *jwt.TokenService
+	APIServer             *server.Server
+	AIHelper              *aihelper.AIHelper
+	Started               chan struct{}
+	ProxyServer           *proxy.Proxy
+	GatewayServer         *gserver.Server
 }
 
 func newGPTScript(ctx context.Context) (*gptscript.GPTScript, error) {
@@ -181,18 +183,19 @@ func New(ctx context.Context, config Config) (*Services, error) {
 
 	// For now, always auto-migrate the gateway database
 	return &Services{
-		DevUIPort:       devPort,
-		ToolRegistryURL: config.ToolRegistry,
-		Events:          events,
-		StorageClient:   storageClient,
-		Router:          r,
-		GPTClient:       c,
-		APIServer:       server.NewServer(storageClient, c, authn.NewAuthenticator(authenticators), authz.NewAuthorizer()),
-		TokenServer:     tokenServer,
-		Invoker:         invoke.NewInvoker(storageClient, c, tokenServer, events),
-		AIHelper:        aihelper.New(c, config.HelperModel),
-		GatewayServer:   gatewayServer,
-		ProxyServer:     proxyServer,
+		WorkspaceProviderType: config.WorkspaceProviderType,
+		DevUIPort:             devPort,
+		ToolRegistryURL:       config.ToolRegistry,
+		Events:                events,
+		StorageClient:         storageClient,
+		Router:                r,
+		GPTClient:             c,
+		APIServer:             server.NewServer(storageClient, c, authn.NewAuthenticator(authenticators), authz.NewAuthorizer()),
+		TokenServer:           tokenServer,
+		Invoker:               invoke.NewInvoker(storageClient, c, config.WorkspaceProviderType, tokenServer, events),
+		AIHelper:              aihelper.New(c, config.HelperModel),
+		GatewayServer:         gatewayServer,
+		ProxyServer:           proxyServer,
 	}, nil
 }
 

diff --git a/pkg/workspace/workspace.go b/pkg/workspace/workspace.go
@@ -5,6 +5,9 @@ import (
 )
 
 func GetDir(workspaceID string) string {
-	_, path, _ := strings.Cut(workspaceID, "://")
-	return path
+	provider, path, _ := strings.Cut(workspaceID, "://")
+	if provider == "directory" {
+		return path
+	}
+	return ""
 }