1// Package kronk provides an implementation of the fantasy AI SDK for local
2// models using the Kronk SDK.
3package kronk
4
5import (
6 "context"
7 "fmt"
8 "sync"
9
10 "charm.land/fantasy"
11 "github.com/ardanlabs/kronk/sdk/kronk"
12 "github.com/ardanlabs/kronk/sdk/tools/catalog"
13 "github.com/ardanlabs/kronk/sdk/tools/libs"
14 "github.com/ardanlabs/kronk/sdk/tools/models"
15)
16
17const (
18 // Name is the name of the Kronk provider.
19 Name = "kronk"
20)
21
22type provider struct {
23 options options
24 mu sync.Mutex
25 kronks map[string]*kronk.Kronk
26}
27
28// New creates a new Kronk provider with the given options.
29func New(opts ...Option) (fantasy.Provider, error) {
30 providerOptions := options{
31 languageModelOptions: make([]LanguageModelOption, 0),
32 }
33
34 for _, o := range opts {
35 o(&providerOptions)
36 }
37
38 if providerOptions.name == "" {
39 providerOptions.name = Name
40 }
41
42 p := provider{
43 options: providerOptions,
44 kronks: make(map[string]*kronk.Kronk),
45 }
46
47 return &p, nil
48}
49
50// Name implements fantasy.Provider.
51func (p *provider) Name() string {
52 return p.options.name
53}
54
55// LanguageModel implements fantasy.Provider.
56// The modelURL parameter should be a URL to a GGUF model file (e.g., from Hugging Face).
57func (p *provider) LanguageModel(ctx context.Context, modelURL string) (fantasy.LanguageModel, error) {
58 p.mu.Lock()
59 defer p.mu.Unlock()
60
61 if krn, ok := p.kronks[modelURL]; ok {
62 opts := append(p.options.languageModelOptions, WithLanguageModelObjectMode(p.options.objectMode))
63 return newLanguageModel(modelURL, p.options.name, krn, opts...), nil
64 }
65
66 mp, err := p.installSystem(ctx, modelURL)
67 if err != nil {
68 return nil, fmt.Errorf("failed to install system: %w", err)
69 }
70
71 krn, err := p.newKronk(mp)
72 if err != nil {
73 return nil, fmt.Errorf("failed to create kronk instance: %w", err)
74 }
75
76 p.kronks[modelURL] = krn
77
78 opts := append(p.options.languageModelOptions, WithLanguageModelObjectMode(p.options.objectMode))
79
80 return newLanguageModel(modelURL, p.options.name, krn, opts...), nil
81}
82
83// Close unloads all Kronk instances. Call this when done with the provider.
84func (p *provider) Close(ctx context.Context) error {
85 p.mu.Lock()
86 defer p.mu.Unlock()
87
88 var errs []error
89
90 for url, krn := range p.kronks {
91 if err := krn.Unload(ctx); err != nil {
92 errs = append(errs, fmt.Errorf("failed to unload model %s: %w", url, err))
93 }
94
95 delete(p.kronks, url)
96 }
97
98 if len(errs) > 0 {
99 return errs[0]
100 }
101
102 return nil
103}
104
105func (p *provider) installSystem(ctx context.Context, modelURL string) (models.Path, error) {
106 logger := p.options.logger
107 if logger == nil {
108 logger = func(context.Context, string, ...any) {}
109 }
110
111 lbs, err := libs.New()
112 if err != nil {
113 return models.Path{}, fmt.Errorf("unable to create libs: %w", err)
114 }
115
116 if _, err := lbs.Download(ctx, libs.Logger(logger)); err != nil {
117 return models.Path{}, fmt.Errorf("unable to install llama.cpp: %w", err)
118 }
119
120 ctlg, err := catalog.New()
121 if err != nil {
122 return models.Path{}, fmt.Errorf("unable to create catalog system: %w", err)
123 }
124
125 if err := ctlg.Download(ctx); err != nil {
126 return models.Path{}, fmt.Errorf("unable to download catalog: %w", err)
127 }
128
129 mdls, err := models.New()
130 if err != nil {
131 return models.Path{}, fmt.Errorf("unable to create models: %w", err)
132 }
133
134 mp, err := mdls.Download(ctx, models.Logger(logger), modelURL, "")
135 if err != nil {
136 return models.Path{}, fmt.Errorf("unable to install model: %w", err)
137 }
138
139 return mp, nil
140}
141
142func (p *provider) newKronk(mp models.Path) (*kronk.Kronk, error) {
143 if err := kronk.Init(); err != nil {
144 return nil, fmt.Errorf("unable to init kronk: %w", err)
145 }
146
147 cfg := p.options.modelConfig
148 cfg.ModelFiles = mp.ModelFiles
149
150 krn, err := kronk.New(cfg)
151 if err != nil {
152 return nil, fmt.Errorf("unable to create inference model: %w", err)
153 }
154
155 return krn, nil
156}