Add production default pipeline modules
This commit is contained in:
241
internal/modules/chunk/generic/chunker.go
Normal file
241
internal/modules/chunk/generic/chunker.go
Normal file
@@ -0,0 +1,241 @@
|
||||
package generic
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"strconv"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||
)
|
||||
|
||||
const Key = "generic"
|
||||
|
||||
const (
|
||||
defaultMaxUnits = 50
|
||||
defaultOverlapUnits = 0
|
||||
)
|
||||
|
||||
var _ contracts.Chunker = (*Chunker)(nil)
|
||||
|
||||
type Chunker struct{}
|
||||
|
||||
func New() *Chunker {
|
||||
return &Chunker{}
|
||||
}
|
||||
|
||||
func (c *Chunker) Key() string {
|
||||
return Key
|
||||
}
|
||||
|
||||
func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
if c == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
|
||||
}
|
||||
if ctx == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
}
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
|
||||
}
|
||||
if len(req.Source.Units) == 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
|
||||
}
|
||||
if err := source.ValidateDocument(req.Source); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
}
|
||||
|
||||
opts, err := chunkOptionsFrom(req.Options)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, err
|
||||
}
|
||||
|
||||
step := opts.maxUnits - opts.overlapUnits
|
||||
chunks := make([]contracts.SourceChunk, 0, (len(req.Source.Units)+step-1)/step)
|
||||
for start := 0; start < len(req.Source.Units); start += step {
|
||||
end := start + opts.maxUnits
|
||||
if end > len(req.Source.Units) {
|
||||
end = len(req.Source.Units)
|
||||
}
|
||||
units := cloneUnits(req.Source.Units[start:end])
|
||||
chunks = append(chunks, contracts.SourceChunk{
|
||||
ID: fmt.Sprintf("chunk-%06d", len(chunks)+1),
|
||||
SourceID: req.Source.ID,
|
||||
Index: len(chunks),
|
||||
Units: units,
|
||||
Metadata: map[string]any{
|
||||
"start_unit_id": units[0].ID,
|
||||
"end_unit_id": units[len(units)-1].ID,
|
||||
"unit_count": len(units),
|
||||
},
|
||||
})
|
||||
if end == len(req.Source.Units) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return contracts.ChunkResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func ModuleSpec() pipeline.ModuleSpec {
|
||||
return pipeline.ModuleSpec{
|
||||
Key: Key,
|
||||
Stage: pipeline.StageChunk,
|
||||
Provides: []string{"chunks"},
|
||||
}
|
||||
}
|
||||
|
||||
func Register(registry *pipeline.ChunkerRegistry) error {
|
||||
return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Chunker, error) {
|
||||
return New(), nil
|
||||
})
|
||||
}
|
||||
|
||||
type chunkOptions struct {
|
||||
maxUnits int
|
||||
overlapUnits int
|
||||
}
|
||||
|
||||
func chunkOptionsFrom(options map[string]any) (chunkOptions, error) {
|
||||
opts := chunkOptions{
|
||||
maxUnits: defaultMaxUnits,
|
||||
overlapUnits: defaultOverlapUnits,
|
||||
}
|
||||
var err error
|
||||
if value, ok := options["max_units"]; ok {
|
||||
opts.maxUnits, err = positiveIntOption("max_units", value)
|
||||
if err != nil {
|
||||
return chunkOptions{}, err
|
||||
}
|
||||
}
|
||||
if value, ok := options["overlap_units"]; ok {
|
||||
opts.overlapUnits, err = nonNegativeIntOption("overlap_units", value)
|
||||
if err != nil {
|
||||
return chunkOptions{}, err
|
||||
}
|
||||
}
|
||||
if opts.overlapUnits >= opts.maxUnits {
|
||||
return chunkOptions{}, chunkerErrorf("overlap_units must be less than max_units")
|
||||
}
|
||||
return opts, nil
|
||||
}
|
||||
|
||||
func positiveIntOption(name string, value any) (int, error) {
|
||||
got, err := intOption(name, value)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if got <= 0 {
|
||||
return 0, chunkerErrorf("%s must be positive", name)
|
||||
}
|
||||
return got, nil
|
||||
}
|
||||
|
||||
func nonNegativeIntOption(name string, value any) (int, error) {
|
||||
got, err := intOption(name, value)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if got < 0 {
|
||||
return 0, chunkerErrorf("%s must be non-negative", name)
|
||||
}
|
||||
return got, nil
|
||||
}
|
||||
|
||||
func intOption(name string, value any) (int, error) {
|
||||
switch typed := value.(type) {
|
||||
case int:
|
||||
return typed, nil
|
||||
case int8:
|
||||
return int(typed), nil
|
||||
case int16:
|
||||
return int(typed), nil
|
||||
case int32:
|
||||
return int(typed), nil
|
||||
case int64:
|
||||
if typed > maxInt() || typed < minInt() {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(typed), nil
|
||||
case uint:
|
||||
if uint64(typed) > uint64(maxInt()) {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(typed), nil
|
||||
case uint8:
|
||||
return int(typed), nil
|
||||
case uint16:
|
||||
return int(typed), nil
|
||||
case uint32:
|
||||
if uint64(typed) > uint64(maxInt()) {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(typed), nil
|
||||
case uint64:
|
||||
if typed > uint64(maxInt()) {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(typed), nil
|
||||
case float64:
|
||||
if typed != math.Trunc(typed) {
|
||||
return 0, chunkerErrorf("%s must be an integer", name)
|
||||
}
|
||||
if typed > float64(maxInt()) || typed < float64(minInt()) {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(typed), nil
|
||||
case json.Number:
|
||||
parsed, err := typed.Int64()
|
||||
if err != nil {
|
||||
return 0, chunkerErrorf("%s must be an integer", name)
|
||||
}
|
||||
if parsed > maxInt() || parsed < minInt() {
|
||||
return 0, chunkerErrorf("%s is outside supported integer range", name)
|
||||
}
|
||||
return int(parsed), nil
|
||||
default:
|
||||
return 0, chunkerErrorf("%s must be an integer", name)
|
||||
}
|
||||
}
|
||||
|
||||
func maxInt() int64 {
|
||||
return int64(1<<(strconv.IntSize-1) - 1)
|
||||
}
|
||||
|
||||
func minInt() int64 {
|
||||
return -maxInt() - 1
|
||||
}
|
||||
|
||||
func cloneUnits(units []source.SourceUnit) []source.SourceUnit {
|
||||
out := make([]source.SourceUnit, 0, len(units))
|
||||
for _, unit := range units {
|
||||
out = append(out, source.SourceUnit{
|
||||
ID: unit.ID,
|
||||
Kind: unit.Kind,
|
||||
Text: unit.Text,
|
||||
Metadata: cloneMetadata(unit.Metadata),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
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 chunkerErrorf(format string, args ...any) error {
|
||||
return fmt.Errorf("generic chunker: "+format, args...)
|
||||
}
|
||||
Reference in New Issue
Block a user