70 lines
1.4 KiB
Go
70 lines
1.4 KiB
Go
package store
|
|
|
|
import (
|
|
"database/sql/driver"
|
|
"encoding/json"
|
|
"fmt"
|
|
"maps"
|
|
|
|
"github.com/klauspost/compress/zstd"
|
|
"modernc.org/sqlite"
|
|
|
|
"g.fc.run/ai/byre/internal/api"
|
|
)
|
|
|
|
var (
|
|
zstdEncoder = must(zstd.NewWriter(nil, zstd.WithEncoderLevel(zstd.SpeedBetterCompression)))
|
|
zstdDecoder = must(zstd.NewReader(nil))
|
|
)
|
|
|
|
func must[T any](v T, err error) T {
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return v
|
|
}
|
|
|
|
func init() {
|
|
sqlite.MustRegisterDeterministicScalarFunction("zstd", 1, func(ctx *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
|
|
switch v := args[0].(type) {
|
|
case nil:
|
|
return nil, nil
|
|
case string:
|
|
return zstdEncoder.EncodeAll([]byte(v), nil), nil
|
|
default:
|
|
return nil, fmt.Errorf("zstd: want text, got %T", v)
|
|
}
|
|
})
|
|
}
|
|
|
|
func splitContent(ev api.Event) (map[string]any, []byte, error) {
|
|
c, ok := ev.Payload["content"]
|
|
if ev.Kind != "agent.tool_result" || !ok {
|
|
return ev.Payload, nil, nil
|
|
}
|
|
raw, err := json.Marshal(c)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
payload := maps.Clone(ev.Payload)
|
|
delete(payload, "content")
|
|
return payload, zstdEncoder.EncodeAll(raw, nil), nil
|
|
}
|
|
|
|
func joinContent(payload map[string]any, content []byte) error {
|
|
if content == nil {
|
|
return nil
|
|
}
|
|
raw, err := zstdDecoder.DecodeAll(content, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var c any
|
|
err = json.Unmarshal(raw, &c)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
payload["content"] = c
|
|
return nil
|
|
}
|