137 lines
3.3 KiB
Go
137 lines
3.3 KiB
Go
package appendorder
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
|
)
|
|
|
|
const Key = "appendorder"
|
|
|
|
var _ contracts.Merger = (*Merger)(nil)
|
|
|
|
type Merger struct{}
|
|
|
|
func New() *Merger {
|
|
return &Merger{}
|
|
}
|
|
|
|
func (m *Merger) Key() string {
|
|
return Key
|
|
}
|
|
|
|
func (m *Merger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
|
|
if m == nil {
|
|
return contracts.MergeResult{}, mergerErrorf("merger must not be nil")
|
|
}
|
|
if ctx == nil {
|
|
return contracts.MergeResult{}, mergerErrorf("context must not be nil")
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return contracts.MergeResult{}, mergerErrorf("context error before merge: %w", err)
|
|
}
|
|
|
|
if len(req.ExtractOutputs) == 1 {
|
|
payload := cloneRawPayload(req.ExtractOutputs[0].Payload)
|
|
return contracts.MergeResult{
|
|
Output: contracts.MergeOutput{
|
|
LaneID: req.LaneID,
|
|
MergerKey: Key,
|
|
SourceID: req.ExtractOutputs[0].SourceID,
|
|
Schema: req.ExtractOutputs[0].Schema,
|
|
Payload: payload,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
content, err := orderedContent(req.ExtractOutputs)
|
|
if err != nil {
|
|
return contracts.MergeResult{}, err
|
|
}
|
|
return contracts.MergeResult{
|
|
Output: contracts.MergeOutput{
|
|
LaneID: req.LaneID,
|
|
MergerKey: Key,
|
|
SourceID: sourceID(req.ExtractOutputs),
|
|
Payload: contracts.RawPayload{
|
|
Content: content,
|
|
MediaType: "application/json",
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
func ModuleSpec() pipeline.ModuleSpec {
|
|
return pipeline.ModuleSpec{
|
|
Key: Key,
|
|
Stage: pipeline.StageMerge,
|
|
Provides: []string{"merged"},
|
|
}
|
|
}
|
|
|
|
func Register(registry *pipeline.MergerRegistry) error {
|
|
return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Merger, error) {
|
|
return New(), nil
|
|
})
|
|
}
|
|
|
|
func orderedContent(outputs []contracts.ExtractOutput) ([]byte, error) {
|
|
items := make([]map[string]any, 0, len(outputs))
|
|
for _, output := range outputs {
|
|
item := map[string]any{
|
|
"chunk_id": output.ChunkID,
|
|
"chunk_index": output.ChunkIndex,
|
|
"media_type": output.Payload.MediaType,
|
|
}
|
|
if json.Valid(output.Payload.Content) {
|
|
item["content"] = json.RawMessage(append([]byte(nil), output.Payload.Content...))
|
|
} else {
|
|
item["content"] = string(output.Payload.Content)
|
|
}
|
|
items = append(items, item)
|
|
}
|
|
content, err := json.Marshal(struct {
|
|
Outputs []map[string]any `json:"outputs"`
|
|
}{Outputs: items})
|
|
if err != nil {
|
|
return nil, mergerErrorf("encode merged output: %w", err)
|
|
}
|
|
return content, nil
|
|
}
|
|
|
|
func sourceID(outputs []contracts.ExtractOutput) string {
|
|
for _, output := range outputs {
|
|
if output.SourceID != "" {
|
|
return output.SourceID
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func cloneRawPayload(payload contracts.RawPayload) contracts.RawPayload {
|
|
return contracts.RawPayload{
|
|
Content: append([]byte(nil), payload.Content...),
|
|
MediaType: payload.MediaType,
|
|
Metadata: cloneMetadata(payload.Metadata),
|
|
Warnings: append([]contracts.Warning(nil), payload.Warnings...),
|
|
}
|
|
}
|
|
|
|
func cloneMetadata(metadata map[string]any) map[string]any {
|
|
if len(metadata) == 0 {
|
|
return nil
|
|
}
|
|
out := make(map[string]any, len(metadata))
|
|
for key, value := range metadata {
|
|
out[key] = value
|
|
}
|
|
return out
|
|
}
|
|
|
|
func mergerErrorf(format string, args ...any) error {
|
|
return fmt.Errorf("appendorder merger: "+format, args...)
|
|
}
|