88 lines
2.4 KiB
Go
88 lines
2.4 KiB
Go
|
package dify
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"io"
|
||
|
|
"net/http"
|
||
|
|
"strconv"
|
||
|
|
"strings"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Only the backend has the bridge credential. Dify receives a short-lived step
|
||
|
|
// token, and records counts and IDs rather than elder speech or image bytes.
|
||
|
|
type Bridge struct {
|
||
|
|
BaseURL, Secret string
|
||
|
|
HTTP *http.Client
|
||
|
|
}
|
||
|
|
type Ticket struct {
|
||
|
|
ID string `json:"context_id"`
|
||
|
|
Token string `json:"step_token"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewBridge(base, secret string) *Bridge {
|
||
|
|
return &Bridge{strings.TrimRight(base, "/"), secret, &http.Client{Timeout: 30 * time.Second}}
|
||
|
|
}
|
||
|
|
func (b *Bridge) request(ctx context.Context, method, path string, payload any, out any) error {
|
||
|
|
var body io.Reader
|
||
|
|
if payload != nil {
|
||
|
|
raw, _ := json.Marshal(payload)
|
||
|
|
body = bytes.NewReader(raw)
|
||
|
|
}
|
||
|
|
req, err := http.NewRequestWithContext(ctx, method, b.BaseURL+path, body)
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
req.Header.Set("Authorization", "Bearer "+b.Secret)
|
||
|
|
req.Header.Set("Content-Type", "application/json")
|
||
|
|
resp, err := b.HTTP.Do(req)
|
||
|
|
if err != nil {
|
||
|
|
return ErrUnavailable
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
if resp.StatusCode != 200 {
|
||
|
|
return ErrUnavailable
|
||
|
|
}
|
||
|
|
if out != nil {
|
||
|
|
return json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(out)
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
func (b *Bridge) Register(ctx context.Context, kind string, input any) (Ticket, error) {
|
||
|
|
var t Ticket
|
||
|
|
err := b.request(ctx, "POST", "/contexts", map[string]any{"kind": kind, "input": input}, &t)
|
||
|
|
return t, err
|
||
|
|
}
|
||
|
|
func (b *Bridge) Result(ctx context.Context, t Ticket, out any) error {
|
||
|
|
return b.request(ctx, "GET", "/contexts/"+t.ID, nil, out)
|
||
|
|
}
|
||
|
|
func (b *Bridge) Forget(t Ticket) {
|
||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||
|
|
defer cancel()
|
||
|
|
_ = b.request(ctx, "DELETE", "/contexts/"+t.ID, nil, nil)
|
||
|
|
}
|
||
|
|
func (b *Bridge) Image(ctx context.Context, t Ticket, index int) ([]byte, error) {
|
||
|
|
path := b.BaseURL + "/contexts/" + t.ID + "/images/" + strconv.Itoa(index)
|
||
|
|
req, _ := http.NewRequestWithContext(ctx, "GET", path, nil)
|
||
|
|
req.Header.Set("Authorization", "Bearer "+b.Secret)
|
||
|
|
resp, err := b.HTTP.Do(req)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
if resp.StatusCode == 404 {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
if resp.StatusCode != 200 {
|
||
|
|
return nil, ErrUnavailable
|
||
|
|
}
|
||
|
|
raw, err := io.ReadAll(io.LimitReader(resp.Body, (10<<20)+1))
|
||
|
|
if err != nil || len(raw) > 10<<20 || http.DetectContentType(raw) != "image/jpeg" {
|
||
|
|
return nil, errors.New("invalid workflow illustration")
|
||
|
|
}
|
||
|
|
return raw, nil
|
||
|
|
}
|