From 27ea9aea462c64a84d4fcb3aaa4e3557f1d745e8 Mon Sep 17 00:00:00 2001
From: mudler <mudler@localai.io>
Date: Mon, 16 Dec 2024 18:29:24 +0100
Subject: [PATCH] chore: shortcut client when result is available with errors

Signed-off-by: mudler <mudler@localai.io>
---
 cmd/tee-worker/config.go | 33 ++++++++++++++++++---------------
 pkg/client/http.go       | 16 ++++++++--------
 pkg/client/result.go     |  7 ++++---
 3 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/cmd/tee-worker/config.go b/cmd/tee-worker/config.go
index de00f86..141b2cc 100644
--- a/cmd/tee-worker/config.go
+++ b/cmd/tee-worker/config.go
@@ -11,12 +11,18 @@ import (
 )
 
 func readConfig() types.JobConfiguration {
+	// The jobs will then unmarshal from this configuration to the specific configuration
+	// that is needed for the job
+	jc := types.JobConfiguration{}
+
 	dataDir := os.Getenv("DATA_DIR")
 	if dataDir == "" {
 		dataDir = "/home/masa"
 		os.Setenv("DATA_DIR", dataDir)
 	}
 
+	jc["data_dir"] = dataDir
+
 	// Read the env file
 	if err := godotenv.Load(filepath.Join(dataDir, ".env")); err != nil {
 		fmt.Println("Failed reading env file!")
@@ -24,27 +30,24 @@ func readConfig() types.JobConfiguration {
 	}
 
 	webScraperBlacklist := os.Getenv("WEBSCRAPER_BLACKLIST")
-
-	blacklistURLs := strings.Split(webScraperBlacklist, ",")
-	for i, u := range blacklistURLs {
-		blacklistURLs[i] = strings.TrimSpace(u)
+	if webScraperBlacklist != "" {
+		blacklistURLs := strings.Split(webScraperBlacklist, ",")
+		for i, u := range blacklistURLs {
+			blacklistURLs[i] = strings.TrimSpace(u)
+		}
+		jc["webscraper_blacklist"] = blacklistURLs
 	}
 
 	twitterAccount := os.Getenv("TWITTER_ACCOUNTS")
+	if twitterAccount != "" {
+		twitterAccounts := strings.Split(twitterAccount, ",")
+		for i, u := range twitterAccounts {
+			twitterAccounts[i] = strings.TrimSpace(u)
+		}
 
-	twitterAccounts := strings.Split(twitterAccount, ",")
-	for i, u := range twitterAccounts {
-		twitterAccounts[i] = strings.TrimSpace(u)
+		jc["twitter_accounts"] = twitterAccounts
 	}
 
-	// Read the .env file and set the global configuration for all the jobs
-	// The jobs will then unmarshal from this configuration to the specific configuration
-	// that is needed for the job
-	jc := types.JobConfiguration{}
-	jc["webscraper_blacklist"] = blacklistURLs
-	jc["twitter_accounts"] = twitterAccounts
-	jc["data_dir"] = dataDir
-
 	return jc
 }
 
diff --git a/pkg/client/http.go b/pkg/client/http.go
index c71652a..c91c7bd 100644
--- a/pkg/client/http.go
+++ b/pkg/client/http.go
@@ -86,27 +86,27 @@ func (c *Client) Decrypt(encryptedResult string) (string, error) {
 }
 
 // GetJobResult retrieves the encrypted result of a job.
-func (c *Client) GetResult(jobUUID string) (string, error) {
+func (c *Client) GetResult(jobUUID string) (string, bool, error) {
 	resp, err := c.HTTPClient.Get(c.BaseURL + "/job/" + jobUUID)
 	if err != nil {
-		return "", fmt.Errorf("error sending GET request: %w", err)
+		return "", false, fmt.Errorf("error sending GET request: %w", err)
 	}
 	defer resp.Body.Close()
 
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return "", fmt.Errorf("error reading response body: %w", err)
+		return "", false, fmt.Errorf("error reading response body: %w", err)
 	}
 
 	if resp.StatusCode == http.StatusNotFound {
-		err = fmt.Errorf("job not found or not ready")
+		return "", false, fmt.Errorf("job not found")
 	}
 
-	if resp.StatusCode != http.StatusOK {
-		respErr := types.JobError{}
-		json.Unmarshal(body, &respErr)
+	respErr := types.JobError{}
+	json.Unmarshal(body, &respErr)
+	if respErr.Error != "" {
 		err = fmt.Errorf("error: %s", respErr.Error)
 	}
 
-	return string(body), err
+	return string(body), true, err
 }
diff --git a/pkg/client/result.go b/pkg/client/result.go
index 40f6d45..91d0822 100644
--- a/pkg/client/result.go
+++ b/pkg/client/result.go
@@ -21,13 +21,14 @@ func (jr *JobResult) SetDelay(delay time.Duration) {
 }
 
 // GetJobResult retrieves the encrypted result of a job.
-func (jr *JobResult) getResult() (string, error) {
+func (jr *JobResult) getResult() (string, bool, error) {
 	return jr.client.GetResult(jr.UUID)
 }
 
 // Get polls the server until the job result is ready or a timeout occurs.
 func (jr *JobResult) Get() (result string, err error) {
 	retries := 0
+	var resultIsAvailable bool
 
 	for {
 		if retries >= jr.maxRetries {
@@ -35,8 +36,8 @@ func (jr *JobResult) Get() (result string, err error) {
 		}
 		retries++
 
-		result, err = jr.getResult()
-		if err == nil {
+		result, resultIsAvailable, err = jr.getResult()
+		if err == nil || resultIsAvailable {
 			break
 		}
 		time.Sleep(jr.delay)