| 💾 Saved. d722711 k33g 7h ago | 1 | package engine |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "net/http" |
| 8 | "os" |
| 9 | "strings" |
| 10 | "time" |
| 11 | |
| 12 | "mm/internal/config" |
| 13 | |
| 14 | "github.com/firebase/genkit/go/ai" |
| 15 | "github.com/firebase/genkit/go/genkit" |
| 16 | "github.com/firebase/genkit/go/plugins/compat_oai" |
| 17 | ) |
| 18 | |
| 19 | // openaiCompat is the Provider for every server that speaks the OpenAI |
| 20 | // chat-completions protocol. Docker Model Runner and llama-server are two |
| 21 | // parametrisations of it; Ollama's /v1, vLLM and the remote APIs would be more |
| 22 | // entries in init() below — the point of keeping the struct data-only. |
| 23 | type openaiCompat struct { |
| 24 | name string |
| 25 | |
| 26 | baseURL string // used when the config gives none |
| 27 | fallback string // tried when baseURL does not answer; "" = no fallback |
| 28 | |
| 29 | // legacyBaseURLEnv keeps DMR_BASE_URL working for the `dmr` provider: the |
| 30 | // variable is in every README and in the sandbox start scripts. The |
| 31 | // provider-neutral AGENT_BASE_URL wins over it. |
| 32 | legacyBaseURLEnv string |
| 33 | |
| 34 | // apiKeyEnv names the variable holding the key; keyRequired says whether an |
| 35 | // empty one is an error (a paid API) or just means "this server does not |
| 36 | // check" (llama-server without --api-key). dummyKey is sent in that case |
| 37 | // because the OpenAI client wants SOMETHING in the header. |
| 38 | apiKeyEnv string |
| 39 | keyRequired bool |
| 40 | dummyKey string |
| 41 | |
| 42 | // Words for Explain: how to start the server, how to get the model. |
| 43 | startHint string |
| 44 | notFoundHint string // %s = model |
| 45 | |
| 46 | // probe is the server-specific part of Probe; nil = reachability only. |
| 47 | probe func(ctx context.Context, b Backend) Info |
| 48 | } |
| 49 | |
| 50 | func init() { |
| 51 | register(&openaiCompat{ |
| 52 | name: "dmr", |
| 53 | baseURL: "http://localhost:12434/engines/v1", |
| 54 | fallback: "http://host.docker.internal:12434/engines/v1", |
| 55 | legacyBaseURLEnv: "DMR_BASE_URL", |
| 56 | dummyKey: "not-needed", // DMR ignores the Authorization header |
| 57 | startHint: "is Docker Model Runner running? (docker model status)", |
| 58 | notFoundHint: "docker model pull %s", |
| 59 | probe: probeModels, |
| 60 | }) |
| 61 | register(&openaiCompat{ |
| 62 | name: "llamacpp", |
| 63 | baseURL: "http://127.0.0.1:8080/v1", // llama-server's default host:port, plus the OpenAI prefix |
| 64 | // llama-server only checks a key when started with --api-key; most |
| 65 | // local set-ups are not, so an unset LLAMA_API_KEY is normal. |
| 66 | apiKeyEnv: "LLAMA_API_KEY", |
| 67 | dummyKey: "not-needed", |
| 68 | startHint: "start it: llama-server -m <model.gguf> --jinja --port 8080", |
| 69 | notFoundHint: "llama-server serves one model; check baseUrl ends with /v1 (model: %s)", |
| 70 | probe: probeLlamaCpp, |
| 71 | }) |
| 72 | } |
| 73 | |
| 74 | func (p *openaiCompat) Name() string { return p.name } |
| 75 | |
| 76 | // Resolve applies the same precedence the `dmr` package used for its URL — |
| 77 | // environment, then config, then fallback if the first does not answer — and |
| 78 | // extends it to the key and the model. |
| 79 | func (p *openaiCompat) Resolve(cfg config.Config) (Backend, error) { |
| 80 | b := Backend{Provider: p.name, Model: cfg.Model, ContextWindow: cfg.ContextWindow} |
| 81 | |
| 82 | // An explicit environment override is taken as-is, without the fallback |
| 83 | // probe: whoever set it knows where the server is. |
| 84 | env := os.Getenv("AGENT_BASE_URL") |
| 85 | if env == "" && p.legacyBaseURLEnv != "" { |
| 86 | env = os.Getenv(p.legacyBaseURLEnv) |
| 87 | } |
| 88 | switch { |
| 89 | case env != "": |
| 90 | b.BaseURL = env |
| 91 | default: |
| 92 | base := cfg.BaseURL |
| 93 | if base == "" { |
| 94 | base = p.baseURL |
| 95 | } |
| 96 | // nil = the key is absent from the file, use the provider's fallback; |
| 97 | // an explicit "" disables it (probe.yaml does that to pin the fake |
| 98 | // engine). The distinction needs the pointer. |
| 99 | fallback := p.fallback |
| 100 | if cfg.Fallback != nil { |
| 101 | fallback = *cfg.Fallback |
| 102 | } |
| 103 | b.BaseURL = base |
| 104 | if fallback != "" && !reachable(base) { |
| 105 | b.BaseURL = fallback |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | keyEnv := cfg.APIKeyEnv |
| 110 | if keyEnv == "" { |
| 111 | keyEnv = p.apiKeyEnv |
| 112 | } |
| 113 | b.APIKeyEnv = keyEnv |
| 114 | if keyEnv != "" { |
| 115 | b.APIKey = os.Getenv(keyEnv) |
| 116 | } |
| 117 | if b.APIKey == "" { |
| 118 | if p.keyRequired { |
| 119 | return b, fmt.Errorf("%s: %s is not set (the key never goes in the YAML: export it)", p.name, keyEnv) |
| 120 | } |
| 121 | b.APIKey = p.dummyKey |
| 122 | } |
| 123 | return b, nil |
| 124 | } |
| 125 | |
| 126 | // Open is the former dmr.Init with the constants replaced by the Backend. |
| 127 | func (p *openaiCompat) Open(ctx context.Context, b Backend) (*genkit.Genkit, string, error) { |
| 128 | client := &compat_oai.OpenAICompatible{ |
| 129 | Provider: p.name, // model prefix: "<provider>/<id>" |
| 130 | BaseURL: b.BaseURL, |
| 131 | APIKey: b.APIKey, |
| 132 | } |
| 133 | g := genkit.Init(ctx, genkit.WithPlugins(client)) |
| 134 | client.DefineModel(p.name, b.Model, ai.ModelOptions{ |
| 135 | Label: b.Model, |
| 136 | Supports: &compat_oai.Multimodal, // enables Tools, Multiturn, SystemRole… |
| 137 | }) |
| 138 | return g, p.name + "/" + b.Model, nil |
| 139 | } |
| 140 | |
| 141 | // Probe: reachability on /models — the one endpoint every OpenAI-compatible |
| 142 | // server has, and the one the fallback logic already relied on — then whatever |
| 143 | // the server-specific probe adds. The config's context window always wins over |
| 144 | // the server's: the operator knows what they started it with. |
| 145 | func (p *openaiCompat) Probe(ctx context.Context, b Backend) Info { |
| 146 | var info Info |
| 147 | if p.probe != nil { |
| 148 | info = p.probe(ctx, b) |
| 149 | } else { |
| 150 | info.Reachable = reachable(b.BaseURL) |
| 151 | } |
| 152 | if !info.Reachable { |
| 153 | info.Warnings = append(info.Warnings, |
| 154 | fmt.Sprintf("%s: nothing answers at %s — %s", p.name, b.BaseURL, p.startHint)) |
| 155 | } |
| 156 | if b.ContextWindow > 0 { |
| 157 | info.ContextWindow, info.ContextSource = b.ContextWindow, "config" |
| 158 | } |
| 159 | return info |
| 160 | } |
| 161 | |
| 162 | // probeClient: two seconds, like the old reachable(). A local server answers in |
| 163 | // milliseconds; a longer wait only delays the prompt when it is down. |
| 164 | var probeClient = &http.Client{Timeout: 2 * time.Second} |
| 165 | |
| 166 | // reachable reports whether an OpenAI-compatible endpoint answers on /models. |
| 167 | func reachable(baseURL string) bool { |
| 168 | resp, err := probeClient.Get(baseURL + "/models") |
| 169 | if err != nil { |
| 170 | return false |
| 171 | } |
| 172 | defer resp.Body.Close() |
| 173 | return resp.StatusCode == http.StatusOK |
| 174 | } |
| 175 | |
| 176 | // getJSON fetches and decodes one endpoint; false when anything failed. Probes |
| 177 | // do not distinguish "down" from "not that kind of server": both mean "no info". |
| 178 | func getJSON(url string, out any) bool { |
| 179 | resp, err := probeClient.Get(url) |
| 180 | if err != nil { |
| 181 | return false |
| 182 | } |
| 183 | defer resp.Body.Close() |
| 184 | if resp.StatusCode != http.StatusOK { |
| 185 | return false |
| 186 | } |
| 187 | return json.NewDecoder(resp.Body).Decode(out) == nil |
| 188 | } |
| 189 | |
| 190 | // probeModels (DMR): lists /models and warns when the configured model is not |
| 191 | // among them. Measured on DMR: asking for a model that is not pulled returns an |
| 192 | // error only once the request is sent — after the spinner has run a while. |
| 193 | func probeModels(_ context.Context, b Backend) Info { |
| 194 | var list struct { |
| 195 | Data []struct { |
| 196 | ID string `json:"id"` |
| 197 | } `json:"data"` |
| 198 | } |
| 199 | info := Info{Reachable: getJSON(b.BaseURL+"/models", &list)} |
| 200 | if !info.Reachable || len(list.Data) == 0 { |
| 201 | return info |
| 202 | } |
| 203 | for _, m := range list.Data { |
| 204 | if m.ID == b.Model { |
| 205 | return info |
| 206 | } |
| 207 | } |
| 208 | info.Warnings = append(info.Warnings, |
| 209 | fmt.Sprintf("%s: model %q is not in the server's list — docker model pull %s", b.Provider, b.Model, b.Model)) |
| 210 | return info |
| 211 | } |
| 212 | |
| 213 | // probeLlamaCpp: llama-server exposes GET /props, one level above /v1, with the |
| 214 | // context size it actually serves (`default_generation_settings.n_ctx`) — not |
| 215 | // the one the model was trained with, which is what /v1/models reports. The |
| 216 | // served size is the number the context-compression work needs. |
| 217 | func probeLlamaCpp(_ context.Context, b Backend) Info { |
| 218 | var props struct { |
| 219 | Settings struct { |
| 220 | NCtx int `json:"n_ctx"` |
| 221 | } `json:"default_generation_settings"` |
| 222 | ChatTemplate string `json:"chat_template"` |
| 223 | } |
| 224 | root := strings.TrimSuffix(strings.TrimSuffix(b.BaseURL, "/"), "/v1") |
| 225 | info := Info{Reachable: getJSON(root+"/props", &props)} |
| 226 | if !info.Reachable { |
| 227 | // /props missing but /models answering: some other OpenAI-compatible |
| 228 | // server is on that port. Reachable, but nothing llama.cpp to report. |
| 229 | info.Reachable = reachable(b.BaseURL) |
| 230 | return info |
| 231 | } |
| 232 | if props.Settings.NCtx > 0 { |
| 233 | info.ContextWindow, info.ContextSource = props.Settings.NCtx, "/props" |
| 234 | } |
| 235 | if props.ChatTemplate == "" { |
| 236 | info.Warnings = append(info.Warnings, |
| 237 | "llamacpp: the server reports no chat template; tool calls need one (and --jinja)") |
| 238 | } |
| 239 | return info |
| 240 | } |