Hermit/src/provider_switch.c

#include "cst_provider_switch.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

int cst_ps_fetch(const char *provider, const char *api_base, const char *key,
                 char ids[][128], int max, char *err, size_t err_sz) {
  cst_model_meta_t *m = calloc((size_t)(max > 0 ? max : 1), sizeof(*m));
  if (!m) {
    snprintf(err, err_sz, "out of memory");
    return -1;
  }
  int r = cst_llm_list_models_meta(api_base, key, m, max, err, err_sz);
  if (r >= 0)
    cst_models_cache_save(provider, api_base, m, r);
  for (int i = 0; i < r; i++)
    snprintf(ids[i], 128, "%s", m[i].id);
  free(m);
  return r;
}

int cst_ps_cached(const char *provider, char ids[][128], int max, int *fresh) {
  *fresh = 0;
  cst_model_meta_t *m = calloc((size_t)(max > 0 ? max : 1), sizeof(*m));
  if (!m)
    return -1;
  long age = 0;
  int n = cst_models_cache_load(provider, m, max, &age);
  for (int i = 0; i < n; i++)
    snprintf(ids[i], 128, "%s", m[i].id);
  free(m);
  *fresh = n > 0 && age >= 0 && age < CST_MODELS_CACHE_TTL;
  return n > 0 ? n : -1;
}

const char *cst_ps_base(const cst_agent_config_t *cfg, const char *provider) {
  if (cfg && strcmp(cfg->provider, provider) == 0)
    return cfg->api_base;
  const cst_provider_t *p = cst_provider_find(provider);
  return p ? p->api_base : "";
}

int cst_ps_models(const cst_agent_config_t *cfg, int refresh,
                  cst_ps_model_t *out, int max, int *api_rc, char *err,
                  size_t err_sz) {
  err[0] = 0;
  *api_rc = -1;
  char(*ids)[128] = calloc(CST_PS_MAX_MODELS, sizeof(*ids));
  if (!ids) {
    snprintf(err, err_sz, "out of memory");
    return 0;
  }
  int fresh = 0;
  int na = cst_ps_cached(cfg->provider, ids, CST_PS_MAX_MODELS, &fresh);
  *api_rc = na;
  if ((refresh || !fresh) && cfg->api_key[0]) {
    char(*got)[128] = calloc(CST_PS_MAX_MODELS, sizeof(*got));
    if (got) {
      int r = cst_ps_fetch(cfg->provider, cfg->api_base, cfg->api_key, got,
                           CST_PS_MAX_MODELS, err, err_sz);
      *api_rc = r;
      if (r >= 0) {
        memcpy(ids, got, (size_t)r * sizeof(*ids));
        na = r;
      }
      free(got);
    }
  }
  int n = 0;
  for (int i = 0; i < na && n < max; i++, n++) {
    snprintf(out[n].id, sizeof(out[n].id), "%s", ids[i]);
    out[n].api = 1;
    cst_model_info(cfg->provider, ids[i], &out[n].info);
  }
  char builtin[32][128];
  int nb = cst_models_list(cfg->model, builtin, 32);
  for (int i = 0; i < nb && n < max; i++) {
    int dup = 0;
    for (int k = 0; k < na && !dup; k++)
      dup = strcmp(ids[k], builtin[i]) == 0;
    if (dup)
      continue;
    memcpy(out[n].id, builtin[i], sizeof(out[n].id));
    out[n].id[sizeof(out[n].id) - 1] = 0;
    out[n].api = 0;
    cst_model_info(cfg->provider, builtin[i], &out[n].info);
    n++;
  }
  free(ids);
  return n;
}

int cst_ps_verify_save(const char *provider, const char *api_base,
                       const char *key, char ids[][128], int max, int *nids,
                       char *err, size_t err_sz) {
  err[0] = 0;
  char(*own)[128] = NULL;
  if (!ids) {
    own = calloc(CST_PS_MAX_MODELS, sizeof(*own));
    if (!own) {
      snprintf(err, err_sz, "out of memory");
      return CST_PS_UNVERIFIED | CST_PS_UNSAVED;
    }
    ids = own;
    max = CST_PS_MAX_MODELS;
  }
  int r = cst_ps_fetch(provider, api_base, key, ids, max, err, err_sz);
  free(own);
  if (nids)
    *nids = r;
  if (r == CST_MODELS_AUTH)
    return CST_MODELS_AUTH;
  int flags = r < 0 ? CST_PS_UNVERIFIED : 0;
  if (r >= 0)
    err[0] = 0;
  if (cst_creds_save(provider, key) != 0)
    flags |= CST_PS_UNSAVED;
  return flags;
}

void cst_ps_apply_provider(cst_agent_config_t *cfg, const cst_provider_t *p,
                           const char *key) {
  cst_settings_t s;
  cst_settings_load(&s);
  const char *model = cst_settings_model(&s, p->name);
  snprintf(cfg->provider, sizeof(cfg->provider), "%s", p->name);
  snprintf(cfg->api_base, sizeof(cfg->api_base), "%s", p->api_base);
  snprintf(cfg->model, sizeof(cfg->model), "%s", model ? model : p->model);
  snprintf(cfg->api_key, sizeof(cfg->api_key), "%s", key);
  cst_settings_remember(cfg);
}

int cst_ps_switch_provider(cst_agent_config_t *cfg, const char *name,
                           const char *key, char *err, size_t err_sz) {
  err[0] = 0;
  const cst_provider_t *p = name ? cst_provider_find(name) : NULL;
  if (!p) {
    snprintf(err, err_sz, "unknown provider '%.64s'", name ? name : "");
    return CST_PS_UNKNOWN;
  }
  char k[512];
  int flags = 0;
  char(*ids)[128] = calloc(CST_PS_MAX_MODELS, sizeof(*ids));
  int nids = -1;
  if (key && key[0]) {
    snprintf(k, sizeof(k), "%s", key);
    flags = cst_ps_verify_save(p->name, p->api_base, k, ids,
                               ids ? CST_PS_MAX_MODELS : 0, &nids, err,
                               err_sz);
    if (flags < 0) {
      memset(k, 0, sizeof(k));
      free(ids);
      return flags;
    }
  } else if (cst_creds_resolve(p->name, k, sizeof(k), NULL) != 0) {
    free(ids);
    snprintf(err, err_sz, "%s needs an API key", p->name);
    return CST_PS_NEED_KEY;
  }
  cst_ps_apply_provider(cfg, p, k);
  memset(k, 0, sizeof(k));
  if (!cfg->model[0] && ids) {
    /* a router preset: the first model the API lists */
    int fresh = 0;
    if (nids <= 0)
      nids = cst_ps_cached(p->name, ids, CST_PS_MAX_MODELS, &fresh);
    if (nids <= 0) {
      char e2[256];
      nids = cst_ps_fetch(p->name, cfg->api_base, cfg->api_key, ids,
                          CST_PS_MAX_MODELS, e2, sizeof(e2));
    }
    if (nids > 0) {
      snprintf(cfg->model, sizeof(cfg->model), "%s", ids[0]);
      cst_settings_remember(cfg);
    }
  }
  free(ids);
  return flags;
}

int cst_ps_set_key(cst_agent_config_t *cfg, const char *key, char *err,
                   size_t err_sz) {
  if (!key || !key[0]) {
    snprintf(err, err_sz, "empty key");
    return CST_PS_NEED_KEY;
  }
  int r = cst_ps_verify_save(cfg->provider, cfg->api_base, key, NULL, 0, NULL,
                             err, err_sz);
  if (r >= 0)
    snprintf(cfg->api_key, sizeof(cfg->api_key), "%s", key);
  return r;
}