47 lines
1.4 KiB
Go
47 lines
1.4 KiB
Go
|
package extraction
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"data-terminal-cloud/internal/dify"
|
||
|
|
"errors"
|
||
|
|
)
|
||
|
|
|
||
|
|
type DifyGenerator struct {
|
||
|
|
Client *dify.Client
|
||
|
|
Bridge *dify.Bridge
|
||
|
|
}
|
||
|
|
|
||
|
|
func (*DifyGenerator) SupportsEventStories() bool { return true }
|
||
|
|
func NewDifyGenerator(base, key, bridgeURL, bridgeSecret string) (*DifyGenerator, error) {
|
||
|
|
client, err := dify.NewClient(base, key)
|
||
|
|
if err != nil {
|
||
|
|
return nil, err
|
||
|
|
}
|
||
|
|
return &DifyGenerator{client, dify.NewBridge(bridgeURL, bridgeSecret)}, nil
|
||
|
|
}
|
||
|
|
func (g *DifyGenerator) Generate(ctx context.Context, in GenerationInput) (GenerationResult, error) {
|
||
|
|
ticket, err := g.Bridge.Register(ctx, "extract", in)
|
||
|
|
if err != nil {
|
||
|
|
return GenerationResult{}, err
|
||
|
|
}
|
||
|
|
defer g.Bridge.Forget(ticket)
|
||
|
|
run, runErr := g.Client.Run(ctx, map[string]any{"context_id": ticket.ID, "step_token": ticket.Token}, nil)
|
||
|
|
var response struct {
|
||
|
|
Result GenerationResult `json:"result"`
|
||
|
|
ErrorCode string `json:"error_code"`
|
||
|
|
}
|
||
|
|
if err = g.Bridge.Result(ctx, ticket, &response); err != nil {
|
||
|
|
return GenerationResult{}, err
|
||
|
|
}
|
||
|
|
if runErr != nil {
|
||
|
|
if response.ErrorCode == "MODEL_ACCOUNT_OVERDUE" {
|
||
|
|
return GenerationResult{}, ErrModelAccountOverdue
|
||
|
|
}
|
||
|
|
return GenerationResult{}, runErr
|
||
|
|
}
|
||
|
|
if len(run.Outputs) == 0 {
|
||
|
|
return GenerationResult{}, errors.New("empty Dify extraction output")
|
||
|
|
}
|
||
|
|
return ValidateWorkflowResult(in, response.Result, "dify:"+run.ID+":deepseek-v4-1-flash-260910")
|
||
|
|
}
|