story/internal/dify/client.go

package dify

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"strings"
	"time"
)

var ErrUnavailable = errors.New("DIFY_WORKFLOW_UNAVAILABLE")
var ErrFailed = errors.New("DIFY_WORKFLOW_FAILED")

type Client struct {
	BaseURL, APIKey string
	HTTP            *http.Client
}
type Event struct {
	Event         string `json:"event"`
	WorkflowRunID string `json:"workflow_run_id"`
	TaskID        string `json:"task_id"`
	Data          struct {
		ID      string                     `json:"id"`
		NodeID  string                     `json:"node_id"`
		Status  string                     `json:"status"`
		Error   string                     `json:"error"`
		Outputs map[string]json.RawMessage `json:"outputs"`
	} `json:"data"`
}
type Run struct {
	ID      string
	Outputs map[string]json.RawMessage
}

func NewClient(base, key string) (*Client, error) {
	u, err := url.Parse(base)
	if err != nil || u.Host == "" || (u.Scheme != "https" && u.Scheme != "http") || strings.TrimSpace(key) == "" {
		return nil, errors.New("invalid Dify configuration")
	}
	return &Client{BaseURL: strings.TrimRight(base, "/"), APIKey: key, HTTP: &http.Client{Timeout: 13 * time.Minute}}, nil
}

// Streaming avoids the reverse proxy timeout for a book with many pictures,
// and gives the caller a chance to publish text before illustration completes.
func (c *Client) Run(ctx context.Context, inputs map[string]any, onNode func(Event) error) (Run, error) {
	raw, _ := json.Marshal(map[string]any{"inputs": inputs, "response_mode": "streaming", "user": "story-platform"})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/workflows/run", bytes.NewReader(raw))
	if err != nil {
		return Run{}, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	resp, err := c.HTTP.Do(req)
	if err != nil {
		return Run{}, ErrUnavailable
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Run{}, ErrUnavailable
	}
	scanner := bufio.NewScanner(io.LimitReader(resp.Body, 16<<20))
	scanner.Buffer(make([]byte, 4096), 2<<20)
	var run Run
	finished := false
	var taskID string
	defer func() {
		if !finished && taskID != "" {
			c.stop(taskID)
		}
	}()
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "[DONE]" {
			break
		}
		var event Event
		if json.Unmarshal([]byte(payload), &event) != nil {
			return run, ErrFailed
		}
		if event.WorkflowRunID != "" {
			run.ID = event.WorkflowRunID
		}
		if event.TaskID != "" {
			taskID = event.TaskID
		}
		if event.Event == "node_finished" && event.Data.Status == "succeeded" && onNode != nil {
			if err = onNode(event); err != nil {
				return run, err
			}
		}
		if event.Event == "workflow_finished" {
			finished = true
			if event.Data.Status != "succeeded" {
				return run, ErrFailed
			}
			run.Outputs = event.Data.Outputs
			if run.ID == "" {
				run.ID = event.Data.ID
			}
		}
		if event.Event == "error" {
			return run, ErrFailed
		}
	}
	if scanner.Err() != nil || !finished {
		return run, ErrUnavailable
	}
	return run, nil
}

func (c *Client) stop(taskID string) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	body := strings.NewReader(`{"user":"story-platform"}`)
	req, err := http.NewRequestWithContext(ctx, "POST", c.BaseURL+"/workflows/tasks/"+taskID+"/stop", body)
	if err != nil {
		return
	}
	req.Header.Set("Authorization", "Bearer "+c.APIKey)
	req.Header.Set("Content-Type", "application/json")
	resp, err := c.HTTP.Do(req)
	if err == nil {
		resp.Body.Close()
	}
}