fix: use correct URL + timeouts

pilksoc · Mar 3, 2024 · 3737022 · 3737022
1 parent 215c6f9
commit 3737022
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 10 deletions.
diff --git a/kube_cache/aiStuff/ai.go b/kube_cache/aiStuff/ai.go
@@ -8,16 +8,20 @@ import (
 	"io"
 	"log"
 	"net/http"
+	"time"
 )
 
 type KubeAi struct {
 	Endpoint, Apikey, ModelId string
 	Metrics                   *metrics.Metrics
+	LastAccess                time.Time
 }
 
+const apiRestTime = time.Second * 5
+
 func New(metrics *metrics.Metrics, endpoint, apiKey, modelId string) *KubeAi {
 	return &KubeAi{
-    Metrics: metrics,
+		Metrics:  metrics,
 		Endpoint: endpoint,
 		Apikey:   apiKey,
 		ModelId:  modelId,
@@ -82,7 +86,7 @@ type aiReq struct {
 }
 
 func (ai *KubeAi) generateKubeRecipe(kubeName1, kubeName2 string) (string, error) {
-	url := fmt.Sprintf("%s/openai/deployments/Dalle3/images/generations?api-version=2024-02-15-preview", ai.Endpoint)
+	url := fmt.Sprintf("%s/openai/deployments/%s/chats/completions/?api-version=2024-02-15-preview", ai.Endpoint, ai.ModelId)
 
 	postReq := aiReq{
 		Messages: []aiMessage{
@@ -135,10 +139,10 @@ func (ai *KubeAi) generateKubeRecipe(kubeName1, kubeName2 string) (string, error
 		return string(body), nil
 	}
 
-  if len(aiResponse.Choices) == 0 {
-    log.Printf("The silly server sent %s, this is very bad", body)
-    return string(body), nil
-  }
+	if len(aiResponse.Choices) == 0 {
+		log.Printf("The silly server sent %s, this is very bad", body)
+		return string(body), nil
+	}
 
 	actualLegitMessage := aiResponse.Choices[0].Message.Content
 
@@ -159,6 +163,11 @@ func (ai *KubeAi) generateKubeRecipe(kubeName1, kubeName2 string) (string, error
 }
 
 func (ai *KubeAi) GenerateKubeRecipe(kubeName1, kubeName2 string) (string, error) {
+	if time.Since(ai.LastAccess) < apiRestTime {
+		time.Sleep(apiRestTime - time.Since(ai.LastAccess))
+	}
+	ai.LastAccess = time.Now()
+
 	res, err := ai.generateKubeRecipe(kubeName1, kubeName2)
 	if err != nil {
 		ai.Metrics.IncrementDalleErrors()

diff --git a/kube_cache/aiStuff/dalle.go b/kube_cache/aiStuff/dalle.go
@@ -9,6 +9,7 @@ import (
 	"log"
 	"net/http"
 	"os"
+	"time"
 
 	imgresize "github.com/CosmicKube/kube_cache/imgResize"
 )
@@ -81,8 +82,8 @@ func (ai *KubeAi) generateDalleForKube(kubeName string) ([]byte, error) {
 	}
 
 	if len(dalleResp.Data) == 0 {
-    log.Printf("The silly server sent %s, this is very bad", body)
-    return nil, errors.New("No data in response")
+		log.Printf("The silly server sent %s, this is very bad", body)
+		return nil, errors.New("No data in response")
 	}
 
 	log.Println("Downloading dalle response")
@@ -101,6 +102,11 @@ func (ai *KubeAi) generateDalleForKube(kubeName string) ([]byte, error) {
 }
 
 func (ai *KubeAi) GenerateDalleForKube(kubeName string) ([]byte, error) {
+	if time.Since(ai.LastAccess) < apiRestTime {
+		time.Sleep(apiRestTime - time.Since(ai.LastAccess))
+	}
+
+	ai.LastAccess = time.Now()
 	img, err := ai.generateDalleForKube(kubeName)
 	if err != nil {
 		ai.Metrics.IncrementDalleErrors()

diff --git a/kube_cache/main.go b/kube_cache/main.go
@@ -24,7 +24,7 @@ func main() {
 		log.Println(err)
 	}
 
-  metrics := metrics.New()
+	metrics := metrics.New()
 
 	log.Println("Creating AI client...")
 	ai := aiStuff.New(metrics, os.Getenv("OPENAI_ENDPOINT"),
@@ -54,7 +54,7 @@ func main() {
 		}
 		return url
 	}
-  
+
 	p.Use(router)
 
 	router.Use(cors.New(cors.Config{