100 lines
2.6 KiB
Go
100 lines
2.6 KiB
Go
|
package memoir
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/dify"
|
||
|
|
"data-terminal-cloud/internal/extraction"
|
||
|
|
"sync"
|
||
|
|
)
|
||
|
|
|
||
|
|
type BookWorkflow interface {
|
||
|
|
Generate(context.Context, []Source, func(Document) error) (Document, [][]byte, error)
|
||
|
|
}
|
||
|
|
type DifyWorkflow struct {
|
||
|
|
Client *dify.Client
|
||
|
|
Bridge *dify.Bridge
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewDifyWorkflow(base, key, bridgeURL, bridgeSecret string) (*DifyWorkflow, error) {
|
||
|
|
client, err := dify.NewClient(base, key)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return &DifyWorkflow{client, dify.NewBridge(bridgeURL, bridgeSecret)}, nil
|
||
|
|
}
|
||
|
|
func (g *DifyWorkflow) Generate(ctx context.Context, sources []Source, publish func(Document) error) (Document, [][]byte, error) {
|
||
|
|
ticket, err := g.Bridge.Register(ctx, "memoir", map[string]any{"stories": sources})
|
||
|
|
if err != nil {
|
||
|
|
return Document{}, nil, err
|
||
|
|
}
|
||
|
|
defer g.Bridge.Forget(ticket)
|
||
|
|
var response struct {
|
||
|
|
Document Document `json:"document"`
|
||
|
|
ErrorCode string `json:"error_code"`
|
||
|
|
}
|
||
|
|
published := false
|
||
|
|
run, runErr := g.Client.Run(ctx, map[string]any{"context_id": ticket.ID, "step_token": ticket.Token}, func(e dify.Event) error {
|
||
|
|
if e.Data.NodeID != "validate-book" || published {
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
if err := g.Bridge.Result(ctx, ticket, &response); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
if err := ValidateDocument(response.Document, sources); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
response.Document.WorkflowRunID = e.WorkflowRunID
|
||
|
|
if err := publish(response.Document); err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
published = true
|
||
|
|
return nil
|
||
|
|
})
|
||
|
|
if runErr != nil {
|
||
|
|
_ = g.Bridge.Result(ctx, ticket, &response)
|
||
|
|
if response.ErrorCode == "MODEL_ACCOUNT_OVERDUE" {
|
||
|
|
return Document{}, nil, extraction.ErrModelAccountOverdue
|
||
|
|
}
|
||
|
|
return Document{}, nil, runErr
|
||
|
|
}
|
||
|
|
if err = g.Bridge.Result(ctx, ticket, &response); err != nil {
|
||
|
|
return Document{}, nil, err
|
||
|
|
}
|
||
|
|
if err = ValidateDocument(response.Document, sources); err != nil {
|
||
|
|
return Document{}, nil, err
|
||
|
|
}
|
||
|
|
response.Document.WorkflowRunID = run.ID
|
||
|
|
if !published {
|
||
|
|
if err = publish(response.Document); err != nil {
|
||
|
|
return Document{}, nil, err
|
||
|
|
}
|
||
|
|
}
|
||
|
|
pictures := make([][]byte, len(response.Document.Chapters))
|
||
|
|
sem := make(chan struct{}, 3)
|
||
|
|
var wg sync.WaitGroup
|
||
|
|
var mu sync.Mutex
|
||
|
|
var imageErr error
|
||
|
|
for i := range pictures {
|
||
|
|
wg.Add(1)
|
||
|
|
go func(i int) {
|
||
|
|
defer wg.Done()
|
||
|
|
select {
|
||
|
|
case sem <- struct{}{}:
|
||
|
|
case <-ctx.Done():
|
||
|
|
return
|
||
|
|
}
|
||
|
|
defer func() { <-sem }()
|
||
|
|
b, err := g.Bridge.Image(ctx, ticket, i)
|
||
|
|
mu.Lock()
|
||
|
|
defer mu.Unlock()
|
||
|
|
if err != nil {
|
||
|
|
imageErr = err
|
||
|
|
} else {
|
||
|
|
pictures[i] = b
|
||
|
|
}
|
||
|
|
}(i)
|
||
|
|
}
|
||
|
|
wg.Wait()
|
||
|
|
return response.Document, pictures, imageErr
|
||
|
|
}
|