[pull] trunk from spiceai:trunk

Cargo.toml

@@ -120,6 +120,7 @@ paste = "1.0.15"
 pem = "3.0.4"
 percent-encoding = "2.3.1"
 pin-project = "1.1"
+postcard = { version = "1.1.1", features = ["use-std"] }
 prometheus = "0.13"
 r2d2 = "0.8.10"
 regex = "1.10.3"

README.md

@@ -14,7 +14,7 @@
   <a href="https://github.com/spiceai/spiceai/actions/workflows/pr.yml"><img alt="GitHub Actions Workflow Status - unit tests" src="https://img.shields.io/github/actions/workflow/status/spiceai/spiceai/pr.yml?event=merge_group&label=unit%20tests" /></a>
   <a href="https://github.com/spiceai/spiceai/actions/workflows/integration.yml"><img alt="GitHub Actions Workflow Status - integration tests" src="https://img.shields.io/github/actions/workflow/status/spiceai/spiceai/integration.yml?event=pull_request&label=integration%20tests" /></a>
   <a href="https://github.com/spiceai/spiceai/actions/workflows/integration_models.yml"><img alt="GitHub Actions Workflow Status - integration tests (models)" src="https://img.shields.io/github/actions/workflow/status/spiceai/spiceai/integration_models.yml?event=pull_request&label=integration%20tests%20(models)" /></a>
-  <a href="https://github.com/spiceai/spiceai/actions/workflows/benchmarks.yml"><img alt="GitHub Actions Workflow Status - benchmark tests" src="https://img.shields.io/github/actions/workflow/status/spiceai/spiceai/benchmarks.yml?branch=trunk&label=benchmark%20tests" /></a>
+  <a href="https://github.com/spiceai/spiceai/actions/workflows/benchmarks.yml"><img alt="GitHub Actions Workflow Status - benchmark tests" src="https://img.shields.io/github/actions/workflow/status/spiceai/spiceai/testoperator_run_bench.yml?branch=trunk&label=benchmark%20tests" /></a>
 </p>
 
 <p align="center">

bin/spice/cmd/add.go

@@ -50,10 +50,10 @@ func init() {
 
 func getAddOrConnectCmdHandler(connect bool) func(cmd *cobra.Command, args []string) {
 	return func(cmd *cobra.Command, args []string) {
-		ctx := context.NewContext()
-		err := ctx.Init()
+		rtcontext := context.NewContext()
+		err := rtcontext.Init(cmd.Flags())
 		if err != nil {
-			slog.Error("could not initialize runtime context", "error", err)
+			slog.Error("failed to initialize runtime context", "error", err)
 			os.Exit(1)
 		}
 
@@ -62,20 +62,21 @@ func getAddOrConnectCmdHandler(connect bool) func(cmd *cobra.Command, args []str
 		slog.Info(fmt.Sprintf("Getting Spicepod %s ...\n", podPath))
 
 		if connect {
-			if ctx.GetApiKey() == "" {
-				slog.Error("A valid Spice.ai Cloud Platform API key was not provided. Run `spice login` to authenticate before proceeding.")
+			apiKey, err := rtcontext.GetApiKey()
+			if err != nil || apiKey == "" {
+				slog.Error("Missing or invalid Spice.ai Cloud API key. Run `spice login` to authenticate and continue.")
 				os.Exit(1)
 			}
 
 			headers := map[string]string{
 				"Spice-Target-Source": "spice.ai",
-				"X-API-Key":           ctx.GetApiKey(),
+				"X-API-Key":           apiKey,
 			}
-			ctx.AddHeaders(headers)
+			rtcontext.AddHeaders(headers)
 		}
 
 		r := registry.GetRegistry(podPath)
-		downloadPath, err := r.GetPod(ctx, podPath)
+		downloadPath, err := r.GetPod(rtcontext, podPath)
 		if err != nil {
 			var itemNotFound *registry.RegistryItemNotFound
 			if errors.As(err, &itemNotFound) {
@@ -146,7 +147,7 @@ func getAddOrConnectCmdHandler(connect bool) func(cmd *cobra.Command, args []str
 
 		slog.Info(fmt.Sprintf("added %s\n", relativePath))
 
-		err = checkLatestCliReleaseVersion()
+		err = checkLatestCliReleaseVersion(rtcontext)
 		if err != nil && util.IsDebug() {
 			slog.Error("failed to check for latest CLI release version", "error", err)
 		}

bin/spice/cmd/catalogs.go

@@ -32,14 +32,16 @@ var catalogsCmd = &cobra.Command{
 spice catalogs
 `,
 	Run: func(cmd *cobra.Command, args []string) {
-		rtcontext := context.NewContext()
+		var rtcontext *context.RuntimeContext
 		if rootCertPath, err := cmd.Flags().GetString("tls-root-certificate-file"); err == nil && rootCertPath != "" {
 			rtcontext = context.NewHttpsContext(rootCertPath)
+		} else {
+			rtcontext = context.NewContext()
 		}
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			return
 		}
 
 		catalogs, err := api.GetData[api.Catalog](rtcontext, "/v1/catalogs")

bin/spice/cmd/chat.go

@@ -149,17 +149,12 @@ spice chat --model <model> "What is Spice.ai?"
 	Run: func(cmd *cobra.Command, args []string) {
 		cloud, _ := cmd.Flags().GetBool(cloudKeyFlag)
 		rtcontext := context.NewContext().WithCloud(cloud)
-		err := rtcontext.Init()
+		err := rtcontext.Init(cmd.Flags())
 		if err != nil {
-			slog.Error("could not initialize runtime context", "error", err)
+			slog.Error("failed to initialize runtime context", "error", err)
 			os.Exit(1)
 		}
 
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
-		}
-
 		temperature, err := cmd.Flags().GetFloat32("temperature")
 		if err != nil {
 			slog.Error("could not get temperature flag", "error", err)
@@ -181,15 +176,6 @@ spice chat --model <model> "What is Spice.ai?"
 			rtcontext.RequireModelsFlavor(cmd)
 		}
 
-		httpEndpoint, err := cmd.Flags().GetString("http-endpoint")
-		if err != nil {
-			slog.Error("could not get http-endpoint flag", "error", err)
-			os.Exit(1)
-		}
-		if httpEndpoint != "" {
-			rtcontext.SetHttpEndpoint(httpEndpoint)
-		}
-
 		model, err := cmd.Flags().GetString(modelKeyFlag)
 		if err != nil {
 			slog.Error("could not get model flag", "error", err)

bin/spice/cmd/datasets.go

@@ -36,10 +36,10 @@ spice datasets
 		if rootCertPath, err := cmd.Flags().GetString("tls-root-certificate-file"); err == nil && rootCertPath != "" {
 			rtcontext = context.NewHttpsContext(rootCertPath)
 		}
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			return
 		}
 
 		datasets, err := api.GetDatasetsWithStatus(rtcontext)

bin/spice/cmd/install.go

@@ -23,7 +23,6 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/spiceai/spiceai/bin/spice/pkg/constants"
 	"github.com/spiceai/spiceai/bin/spice/pkg/context"
-	"github.com/spiceai/spiceai/bin/spice/pkg/runtime"
 	"github.com/spiceai/spiceai/bin/spice/pkg/util"
 )
 
@@ -39,9 +38,16 @@ spice install ai
 # See more at: https://spiceai.org/docs/
 `,
 	Run: func(cmd *cobra.Command, args []string) {
+		rtcontext := context.NewContext()
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			os.Exit(1)
+		}
+
 		slog.Info("Checking for latest Spice runtime release...")
 
-		err := checkLatestCliReleaseVersion()
+		err = checkLatestCliReleaseVersion(rtcontext)
 		if err != nil && util.IsDebug() {
 			slog.Error("failed to check for latest CLI release version", "error", err)
 		}
@@ -75,20 +81,14 @@ spice install ai
 		}
 
 		if force {
-			rtcontext := context.NewContext()
-			err := rtcontext.Init()
-			if err != nil {
-				slog.Error("initializing runtime context", "error", err)
-				os.Exit(1)
-			}
 			err = rtcontext.InstallMatchingRuntime(flavor, !cpu)
 			if err != nil {
 				slog.Error("installing runtime", "error", err)
 				os.Exit(1)
 			}
 			installed = true
 		} else {
-			installed, err = runtime.EnsureInstalled(flavor, true, !cpu)
+			installed, err = rtcontext.EnsureInstalled(flavor, true, !cpu)
 			if err != nil {
 				slog.Error("verifying runtime install", "error", err)
 				os.Exit(1)

bin/spice/cmd/models.go

@@ -36,10 +36,10 @@ spice models
 		if rootCertPath, err := cmd.Flags().GetString("tls-root-certificate-file"); err == nil && rootCertPath != "" {
 			rtcontext = context.NewHttpsContext(rootCertPath)
 		}
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			return
 		}
 
 		model_statuses, _, err := api.GetComponentStatuses(rtcontext)

bin/spice/cmd/nsql.go

@@ -72,10 +72,9 @@ nsql> How much money have I made in each country?
 
 		cloud, _ := cmd.Flags().GetBool(cloudKeyFlag)
 		rtcontext := context.NewContext().WithCloud(cloud)
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		if err := rtcontext.Init(cmd.Flags()); err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			os.Exit(1)
 		}
 
 		userAgent, _ := cmd.Flags().GetString("user-agent")
@@ -127,15 +126,6 @@ nsql> How much money have I made in each country?
 			model = selectedModel
 		}
 
-		httpEndpoint, err := cmd.Flags().GetString(httpEndpointKeyFlag)
-		if err != nil {
-			slog.Error("getting http-endpoint flag", "error", err)
-			os.Exit(1)
-		}
-		if httpEndpoint != "" {
-			rtcontext.SetHttpEndpoint(httpEndpoint)
-		}
-
 		cmd.Println("")
 		cmd.Println("Enter a query in natural language.")
 

bin/spice/cmd/pods.go

@@ -36,10 +36,10 @@ spice pods
 		if rootCertPath, err := cmd.Flags().GetString("tls-root-certificate-file"); err == nil && rootCertPath != "" {
 			rtcontext = context.NewHttpsContext(rootCertPath)
 		}
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			return
 		}
 
 		spicepods, err := api.GetData[api.Spicepod](rtcontext, "/v1/spicepods")

bin/spice/cmd/refresh.go

@@ -91,10 +91,10 @@ spice refresh taxi_trips
 		if rootCertPath, err := cmd.Flags().GetString("tls-root-certificate-file"); err == nil && rootCertPath != "" {
 			rtcontext = context.NewHttpsContext(rootCertPath)
 		}
-
-		apiKey, _ := cmd.Flags().GetString("api-key")
-		if apiKey != "" {
-			rtcontext.SetApiKey(apiKey)
+		err := rtcontext.Init(cmd.Flags())
+		if err != nil {
+			slog.Error("failed to initialize runtime context", "error", err)
+			return
 		}
 
 		url := fmt.Sprintf("/v1/datasets/%s/acceleration/refresh", dataset)