kling

Author: xushiwei

spec/geno/geno.go

@@ -19,6 +19,7 @@ package geno
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"net/http"
 
 	xai "github.com/goplus/xai/spec"
@@ -244,7 +245,7 @@ type responseAdapter interface {
 	Done(action xai.Action, body map[string]any) bool
 	Sleep(action xai.Action, body map[string]any)
 	Results(action xai.Action, body map[string]any) xai.Results
-	QueryOpInfo(action xai.Action, body map[string]any) QueryOpInfo
+	QueryOpInfo(action xai.Action, body map[string]any) (QueryOpInfo, error)
 }
 
 type OperationResponse[T responseAdapter] struct {
@@ -266,12 +267,21 @@ func (p *OperationResponse[T]) Sleep() {
 	p.adapter.Sleep(p.action, p.body)
 }
 
+var (
+	ErrMissingOperationID = errors.New("missing operation ID in response body")
+)
+
 func (p *OperationResponse[T]) Retry(ctx context.Context, svc xai.Service, opts xai.OptionBuilder) (resp xai.OperationResponse, err error) {
-	qoi := p.adapter.QueryOpInfo(p.action, p.body)
+	qoi, err := p.adapter.QueryOpInfo(p.action, p.body)
+	if err != nil {
+		return
+	}
 	req, err := p.c.NewRequest(http.MethodGet, qoi.Path)
 	if err != nil {
 		return
 	}
+	// query operation has no body,
+	// so we can directly call it without setting body
 	return call(ctx, req, qoi.NewResponse, opts)
 }
 

spec/kling/kling.go

@@ -29,6 +29,12 @@ import (
 
 // -----------------------------------------------------------------------------
 
+func newGenImageResponse(c *geno.Client, body map[string]any) (xai.OperationResponse, error) {
+	return geno.NewOperationResponse[adapter](c, xai.GenImage, body), nil
+}
+
+// -----------------------------------------------------------------------------
+
 type adapter struct{}
 
 func (adapter) Actions(model xai.Model) []xai.Action {
@@ -49,50 +55,52 @@ func (adapter) ActionInfo(action xai.Action) geno.ActionInfo {
 	}
 }
 
-func (adapter) QueryOpInfo(action xai.Action, body map[string]any) geno.QueryOpInfo {
+func (adapter) QueryOpInfo(action xai.Action, body map[string]any) (ret geno.QueryOpInfo, err error) {
 	switch action {
 	case xai.GenImage:
 		data, _ := body["data"].(map[string]any)
-		id, _ := data["task_id"].(string)
-		return geno.QueryOpInfo{
-			Path:        "/v1/images/generations/" + id,
-			NewResponse: newGenImageResponse,
+		if id, ok := data["task_id"].(string); ok {
+			ret.Path = "/v1/images/generations/" + id
+			ret.NewResponse = newGenImageResponse
+		} else {
+			err = geno.ErrMissingOperationID
 		}
 	default:
 		panic("unexpected action: " + action)
 	}
+	return
 }
 
 func (adapter) Results(action xai.Action, body map[string]any) xai.Results {
-	data, _ := body["task_result"].(map[string]any)
+	result, _ := body["task_result"].(map[string]any)
 	switch action {
 	case xai.GenImage:
-		return geno.NewImageResults[adapter](data, "images")
+		return geno.NewImageResults[adapter](result, "images")
 	default:
 		panic("unexpected action: " + action)
 	}
 }
 
+func (adapter) Sleep(action xai.Action, body map[string]any) {
+	switch action {
+	case xai.GenVideo:
+		// sleep 10s for video operations
+		time.Sleep(10 * time.Second)
+	default:
+		// sleep 0.5s for image operations
+		time.Sleep(time.Second / 2)
+	}
+}
+
 func (adapter) Done(action xai.Action, body map[string]any) bool {
 	data, _ := body["data"].(map[string]any)
 	return data["task_status"] == "succeed" // submitted, processing, succeed, failed
 }
 
-func (adapter) Sleep(action xai.Action, body map[string]any) {
-	if action == xai.GenVideo {
-		time.Sleep(10 * time.Second) // sleep 10s for video operations
-	}
-	time.Sleep(time.Second / 2) // sleep 0.5s for image operations
-}
-
 func (adapter) OutputImage(item any) *xai.OutputImage {
 	panic("todo")
 }
 
-func newGenImageResponse(c *geno.Client, body map[string]any) (xai.OperationResponse, error) {
-	return geno.NewOperationResponse[adapter](c, xai.GenImage, body), nil
-}
-
 // -----------------------------------------------------------------------------
 
 const (