1package kronk
2
3import (
4 "encoding/json"
5
6 "charm.land/fantasy"
7)
8
9// Global type identifiers for Kronk-specific provider data.
10const (
11 TypeProviderOptions = Name + ".options"
12 TypeProviderMetadata = Name + ".metadata"
13)
14
15// Register Kronk provider-specific types with the global registry.
16func init() {
17 fantasy.RegisterProviderType(TypeProviderOptions, func(data []byte) (fantasy.ProviderOptionsData, error) {
18 var v ProviderOptions
19 if err := json.Unmarshal(data, &v); err != nil {
20 return nil, err
21 }
22 return &v, nil
23 })
24
25 fantasy.RegisterProviderType(TypeProviderMetadata, func(data []byte) (fantasy.ProviderOptionsData, error) {
26 var v ProviderMetadata
27 if err := json.Unmarshal(data, &v); err != nil {
28 return nil, err
29 }
30 return &v, nil
31 })
32}
33
34// ProviderMetadata represents additional metadata from Kronk provider.
35type ProviderMetadata struct {
36 TokensPerSecond float64 `json:"tokens_per_second"`
37 OutputTokens int64 `json:"output_tokens"`
38}
39
40// Options implements the ProviderOptionsData interface.
41func (*ProviderMetadata) Options() {}
42
43// MarshalJSON implements custom JSON marshaling with type info for ProviderMetadata.
44func (m ProviderMetadata) MarshalJSON() ([]byte, error) {
45 type plain ProviderMetadata
46 return fantasy.MarshalProviderType(TypeProviderMetadata, plain(m))
47}
48
49// UnmarshalJSON implements custom JSON unmarshaling with type info for ProviderMetadata.
50func (m *ProviderMetadata) UnmarshalJSON(data []byte) error {
51 type plain ProviderMetadata
52 var p plain
53 if err := fantasy.UnmarshalProviderType(data, &p); err != nil {
54 return err
55 }
56 *m = ProviderMetadata(p)
57 return nil
58}
59
60// ProviderOptions represents additional options for Kronk provider.
61type ProviderOptions struct {
62 TopK *int64 `json:"top_k"`
63 RepeatPenalty *float64 `json:"repeat_penalty"`
64 Seed *int64 `json:"seed"`
65 MinP *float64 `json:"min_p"`
66 NumPredict *int64 `json:"num_predict"`
67 Stop []string `json:"stop"`
68}
69
70// Options implements the ProviderOptionsData interface.
71func (*ProviderOptions) Options() {}
72
73// MarshalJSON implements custom JSON marshaling with type info for ProviderOptions.
74func (o ProviderOptions) MarshalJSON() ([]byte, error) {
75 type plain ProviderOptions
76 return fantasy.MarshalProviderType(TypeProviderOptions, plain(o))
77}
78
79// UnmarshalJSON implements custom JSON unmarshaling with type info for ProviderOptions.
80func (o *ProviderOptions) UnmarshalJSON(data []byte) error {
81 type plain ProviderOptions
82 var p plain
83 if err := fantasy.UnmarshalProviderType(data, &p); err != nil {
84 return err
85 }
86 *o = ProviderOptions(p)
87 return nil
88}
89
90// NewProviderOptions creates new provider options for Kronk.
91func NewProviderOptions(opts *ProviderOptions) fantasy.ProviderOptions {
92 return fantasy.ProviderOptions{
93 Name: opts,
94 }
95}
96
97// ParseOptions parses provider options from a map.
98func ParseOptions(data map[string]any) (*ProviderOptions, error) {
99 var options ProviderOptions
100 if err := fantasy.ParseOptions(data, &options); err != nil {
101 return nil, err
102 }
103 return &options, nil
104}