Files
bantam-cpp/llm.cpp
T

484 lines
16 KiB
C++

/* llm.cpp — libcurl HTTP client, Fibonacci backoff,
* streaming SSE parser and the llm() dispatch function. */
#include "bantam.hpp"
#include <array>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <curl/curl.h>
#include <memory>
#include <unistd.h>
/* ----------------------------- memory callback ------------------------ */
static size_t mem_cb(char *ptr, size_t sz, size_t nm, void *ud) {
std::string *m = (std::string *)ud;
size_t total = sz * nm;
m->append((const char *)ptr, total);
return total;
}
/* ----------------------------- curl RAII ------------------------------ */
struct CurlHandle {
CURL *c;
CurlHandle() : c(curl_easy_init()) {}
~CurlHandle() { if (c) curl_easy_cleanup(c); }
CurlHandle(const CurlHandle &) = delete;
CurlHandle &operator=(const CurlHandle &) = delete;
explicit operator bool() const { return c != nullptr; }
};
struct CurlHeaders {
curl_slist *list = nullptr;
CurlHeaders() = default;
~CurlHeaders() { if (list) curl_slist_free_all(list); }
CurlHeaders(const CurlHeaders &) = delete;
CurlHeaders &operator=(const CurlHeaders &) = delete;
void append(std::string_view s) {
std::string tmp(s);
curl_slist *next = curl_slist_append(list, tmp.c_str());
if (next) list = next;
}
};
/* ----------------------------- proxy + headers ------------------------ */
static void set_proxy(CURL *c) {
const char *socks = std::getenv("SOCKS_PROXY");
if (!socks || !*socks) socks = std::getenv("socks_proxy");
if (socks && *socks) {
std::string buf;
if (std::strstr(socks, "://")) buf = socks;
else buf = "socks5://" + std::string(socks);
curl_easy_setopt(c, CURLOPT_PROXY, buf.c_str());
}
}
static void apply_llm_headers(CurlHeaders &h, const Cfg &cfg) {
h.append("Content-Type: application/json");
if (!cfg.api_key.empty() && cfg.api_key != "-") {
h.append(sfmt("Authorization: Bearer %s", cfg.api_key.c_str()));
}
}
/* ----------------------------- http_get ------------------------------- */
int http_get(Cfg &cfg, std::string_view url, std::string &body, long &code,
std::string &errmsg) {
CurlHandle curl;
if (!curl) { errmsg = "curl init failed"; return -1; }
std::string url_s(url);
curl_easy_setopt(curl.c, CURLOPT_URL, url_s.c_str());
curl_easy_setopt(curl.c, CURLOPT_WRITEFUNCTION, mem_cb);
curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &body);
set_proxy(curl.c);
long t = cfg.timeout > 0 ? cfg.timeout : 30;
curl_easy_setopt(curl.c, CURLOPT_CONNECTTIMEOUT, t);
curl_easy_setopt(curl.c, CURLOPT_TIMEOUT, t);
curl_easy_setopt(curl.c, CURLOPT_USERAGENT, "Mozilla/5.0 Bantam/nightly-cpp");
CURLcode res = curl_easy_perform(curl.c);
curl_easy_getinfo(curl.c, CURLINFO_RESPONSE_CODE, &code);
if (res != CURLE_OK) {
errmsg = curl_easy_strerror(res);
return -1;
}
return 0;
}
/* ----------------------------- token estimate ------------------------- */
static int est_tokens_msgs(const Messages &msgs) {
int chars = 0;
for (auto &m : msgs.items) {
if (!m.content.empty()) chars += (int)m.content.size();
if (!m.reasoning.empty()) chars += (int)m.reasoning.size();
for (auto &tc : m.calls) {
if (!tc.name.empty()) chars += (int)tc.name.size();
if (!tc.arguments.empty()) chars += (int)tc.arguments.size();
}
}
if (chars == 0) return 0;
int t = chars / 4;
return t > 0 ? t : 1;
}
static int est_tokens_msg(const Message &m) {
int chars = 0;
if (!m.content.empty()) chars += (int)m.content.size();
if (!m.reasoning.empty()) chars += (int)m.reasoning.size();
for (auto &tc : m.calls) {
if (!tc.name.empty()) chars += (int)tc.name.size();
if (!tc.arguments.empty()) chars += (int)tc.arguments.size();
}
if (chars == 0) return 0;
int t = chars / 4;
return t > 0 ? t : 1;
}
/* ----------------------------- streaming parser ----------------------- */
struct StreamCtx {
std::string content, reas, line_buf, live_printed, buf, raw;
std::vector<std::string> tbl;
Usage last_usage;
std::vector<ToolCall> tcs;
std::vector<int> tcs_idx;
int tool_calls_shown = 0;
int tty = 0;
int live_idx = -1;
MdState md;
int in_reasoning = 0;
int done = 0;
int erase_pending = 0;
volatile int *interrupted = nullptr;
};
static void stream_upsert_tc(StreamCtx *sc, int idx, std::string_view id,
std::string_view name, std::string_view args) {
int slot = -1;
for (size_t i = 0; i < sc->tcs_idx.size(); i++)
if (sc->tcs_idx[i] == idx) { slot = (int)i; break; }
if (slot < 0) {
slot = (int)sc->tcs.size();
sc->tcs.emplace_back();
sc->tcs_idx.push_back(idx);
}
if (!id.empty()) sc->tcs[slot].id = std::string(id);
if (!name.empty()) sc->tcs[slot].name += std::string(name);
if (!args.empty()) sc->tcs[slot].arguments += std::string(args);
}
static void live_target(StreamCtx *sc, int slot, bool closed, std::string &out) {
ToolCall &t = sc->tcs[slot];
str_appendf(out, "[tool call: %s", t.name.c_str());
if (closed || !t.arguments.empty()) {
std::string fa = filter_text(t.arguments);
str_appendf(out, "(%s", fa.c_str());
if (closed) out += ")]";
}
}
static void finish_live(StreamCtx *sc, bool closed) {
if (sc->live_idx < 0) return;
std::string target;
live_target(sc, sc->live_idx, closed, target);
if (target.size() > sc->live_printed.size()) {
cprint(target.substr(sc->live_printed.size()), std::array{33});
}
sc->live_printed.clear();
sc->live_idx = -1;
/* The tool call marker must always terminate its own line so the
* following [tool result: ...] header is separated from it by a
* newline (in streaming mode the marker is shown here and not
* reprinted by the caller). */
printf("\n");
fflush(stdout);
}
static void step_live(StreamCtx *sc) {
if (!sc->tty || sc->tcs.empty()) return;
int cur = (int)sc->tcs.size() - 1;
if (sc->live_idx != -1 && sc->live_idx != cur) finish_live(sc, true);
if (sc->live_idx == -1) { sc->live_idx = cur; sc->live_printed.clear(); }
std::string target;
live_target(sc, cur, false, target);
if (target.size() > sc->live_printed.size()) {
cprint(target.substr(sc->live_printed.size()), std::array{33});
sc->live_printed += target.substr(sc->live_printed.size());
}
sc->tool_calls_shown = 1;
}
static void flush_tbl(StreamCtx *sc) {
if (!sc->tbl.empty()) {
render_table_print(sc->tbl);
sc->tbl.clear();
}
}
static void stream_maybe_erase(StreamCtx *sc) {
if (sc->erase_pending) {
if (g_col) printf("\r\033[K");
sc->erase_pending = 0;
}
}
static void stream_handle_data(StreamCtx *sc, const std::string &data) {
if (data == "[DONE]") { sc->done = 1; return; }
Json d_mut;
if (!Json::parse(data, d_mut) || !d_mut.is_object()) return;
const Json &d = d_mut;
const Json &u = d["usage"];
if (u.is_object()) {
int pt = u["prompt_tokens"].as_int();
int tt = u["total_tokens"].as_int();
if (pt > 0 || tt > 0) {
sc->last_usage.prompt_tokens = pt;
sc->last_usage.completion_tokens = u["completion_tokens"].as_int();
sc->last_usage.total_tokens = tt;
const Json &ptd = u["prompt_tokens_details"];
if (ptd.is_object()) {
int ct_val = ptd["cached_tokens"].as_int();
if (ct_val > 0) sc->last_usage.cached_tokens = ct_val;
}
if (sc->last_usage.cached_tokens == 0) {
int ct_val = u["cached_tokens"].as_int();
if (ct_val > 0) sc->last_usage.cached_tokens = ct_val;
}
}
}
if (d["model"].is_string()) {
sc->last_usage.model = d["model"].str();
}
const Json &choices = d["choices"];
if (!choices.is_array() || choices.empty()) return;
const Json &ch = choices.at(0);
if (!ch.is_object()) return;
const Json &delta = ch["delta"];
if (!delta.is_object()) return;
std::string rc = delta["reasoning_content"].str();
if (rc.empty()) rc = delta["reasoning"].str();
if (rc.empty()) rc = delta["thought"].str();
if (!rc.empty()) {
stream_maybe_erase(sc);
if (!sc->in_reasoning) {
if (!sc->content.empty()) {
flush_tbl(sc);
if (!sc->line_buf.empty()) { render_md_line_print(sc->line_buf, sc->md); sc->line_buf.clear(); }
printf("\n");
}
cprint("--- reasoning start ---\n", std::array{36});
sc->in_reasoning = true;
}
cprint(rc, std::array{2});
sc->reas += rc;
}
std::string ct = delta["content"].str();
if (!ct.empty()) {
stream_maybe_erase(sc);
if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
sc->content += ct;
sc->line_buf += ct;
while (true) {
auto nl = sc->line_buf.find('\n');
if (nl == std::string::npos) break;
std::string cur = sc->line_buf.substr(0, nl);
sc->line_buf.erase(0, nl + 1);
if (!sc->md.in_code && !is_table_line(cur)) {
flush_tbl(sc);
render_md_line_print(cur, sc->md);
} else {
sc->tbl.push_back(cur);
}
}
}
const Json &tcs = delta["tool_calls"];
if (tcs.is_array()) {
stream_maybe_erase(sc);
if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
if (sc->live_idx == -1) {
flush_tbl(sc);
if (!sc->line_buf.empty()) { render_md_line_print(sc->line_buf, sc->md); sc->line_buf.clear(); }
}
for (auto &tc : tcs.aval) {
if (!tc.is_object()) continue;
int idx = tc["index"].as_int();
std::string id = tc["id"].str();
std::string name, args;
const Json &fn = tc["function"];
if (fn.is_object()) {
name = fn["name"].str();
args = fn["arguments"].str();
}
stream_upsert_tc(sc, idx, id, name, args);
}
step_live(sc);
}
}
static int stream_feed(StreamCtx *sc, const char *data, size_t len,
volatile int *interrupted) {
sc->interrupted = interrupted;
sc->buf.append(data, len);
while (true) {
auto nl = sc->buf.find('\n');
if (nl == std::string::npos) break;
std::string line = sc->buf.substr(0, nl);
while (!line.empty() && (line.back() == '\r' || line.back() == ' ' || line.back() == '\t'))
line.pop_back();
if (line.rfind("data:", 0) == 0) {
std::string_view dd = line;
dd.remove_prefix(5);
while (!dd.empty() && dd.front() == ' ') dd.remove_prefix(1);
stream_handle_data(sc, std::string(dd));
}
sc->buf.erase(0, nl + 1);
if (sc->done) break;
if (*interrupted) return -1;
}
return 0;
}
static void stream_finish(StreamCtx *sc) {
flush_tbl(sc);
if (!sc->line_buf.empty()) {
std::string cur = sc->line_buf;
if (!sc->md.in_code && is_table_line(cur)) {
sc->tbl.push_back(cur);
flush_tbl(sc);
} else {
render_md_line_print(cur, sc->md);
}
sc->line_buf.clear();
}
if (sc->live_idx >= 0) finish_live(sc, true);
if (sc->in_reasoning) { cprint("\n--- reasoning end ---\n", std::array{36}); sc->in_reasoning = false; }
}
static void stream_result(StreamCtx *sc, Message &out, Usage &usage) {
out = Message{};
out.role = "assistant";
out.calls_shown = sc->tool_calls_shown;
if (!sc->content.empty()) out.content = sc->content;
if (!sc->reas.empty()) out.reasoning = sc->reas;
for (auto &t : sc->tcs) out.calls.push_back(t);
usage = sc->last_usage;
if (usage.total_tokens == 0 && usage.prompt_tokens > 0)
usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
}
/* ----------------------------- streaming write cb --------------------- */
static size_t stream_write_cb(char *ptr, size_t sz, size_t nm, void *ud) {
StreamCtx *sc = (StreamCtx *)ud;
size_t total = sz * nm;
if (*sc->interrupted) return 0;
sc->raw.append((const char *)ptr, total);
int r = stream_feed(sc, ptr, total, sc->interrupted);
if (r != 0 || *sc->interrupted) return 0;
return total;
}
/* ----------------------------- llm() ---------------------------------- */
static void sleep_interruptible(int secs, volatile int *interrupted) {
struct timespec ts{0, 100000000L};
for (int i = 0; i < secs * 10; i++) {
if (*interrupted) return;
nanosleep(&ts, nullptr);
}
}
static std::string http_err_body(const StreamCtx *sc, const std::string *m, long code) {
std::string_view b = sc ? std::string_view(sc->raw) : (m ? std::string_view(*m) : "");
b = utf8_truncate(b, 4000);
return sfmt("HTTP %ld: %.*s", code, (int)b.size(), b.data());
}
int llm(Cfg &cfg, Messages &msgs, bool use_tools, Message &out, Usage &usage,
std::string &errmsg, volatile int *interrupted) {
out = Message{};
usage = Usage{};
if (*interrupted) { errmsg = "interrupted"; return -1; }
bool stream = (cfg.force_stream >= 0)
? (cfg.force_stream != 0)
: (cfg.stream != 0);
std::string payload = build_request_json(cfg, msgs, use_tools, stream);
std::string url = cfg.endpoint;
while (!url.empty() && url.back() == '/') url.pop_back();
url += "/chat/completions";
int fib[] = {1, 1, 2, 3, 5, 8, 13, 21, 34};
int nf = 9;
std::string pend = cstr("...requesting...", std::array{1, 2});
CurlHandle curl;
if (!curl) { errmsg = "curl init failed"; return -1; }
curl_easy_setopt(curl.c, CURLOPT_USERAGENT, "Mozilla/5.0 Bantam/nightly-cpp");
CurlHeaders headers;
apply_llm_headers(headers, cfg);
curl_easy_setopt(curl.c, CURLOPT_URL, url.c_str());
curl_easy_setopt(curl.c, CURLOPT_POST, 1L);
curl_easy_setopt(curl.c, CURLOPT_POSTFIELDS, payload.c_str());
curl_easy_setopt(curl.c, CURLOPT_POSTFIELDSIZE, (long)payload.size());
curl_easy_setopt(curl.c, CURLOPT_HTTPHEADER, headers.list);
set_proxy(curl.c);
curl_easy_setopt(curl.c, CURLOPT_CONNECTTIMEOUT, (long)(cfg.timeout > 0 ? cfg.timeout : 30));
if (!stream) curl_easy_setopt(curl.c, CURLOPT_TIMEOUT, (long)(cfg.timeout > 0 ? cfg.timeout : 300));
int attempt, rc = -1;
for (attempt = 0; attempt <= nf; attempt++) {
if (*interrupted) { errmsg = "interrupted"; rc = -1; break; }
if (g_col) { printf("\r%s", pend.c_str()); fflush(stdout); }
else printf("%s\n", pend.c_str());
StreamCtx sc;
std::string mbody;
CURLcode res; long code = 0;
if (stream) {
sc.tty = is_terminal(STDOUT_FILENO);
sc.interrupted = interrupted;
sc.erase_pending = 1;
curl_easy_setopt(curl.c, CURLOPT_WRITEFUNCTION, stream_write_cb);
curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &sc);
} else {
curl_easy_setopt(curl.c, CURLOPT_WRITEFUNCTION, mem_cb);
curl_easy_setopt(curl.c, CURLOPT_WRITEDATA, &mbody);
}
res = curl_easy_perform(curl.c);
if (*interrupted) {
if (g_col && (!stream || sc.erase_pending)) printf("\r\033[K");
errmsg = "interrupted";
rc = -1;
break;
}
curl_easy_getinfo(curl.c, CURLINFO_RESPONSE_CODE, &code);
int transport_err = (res != CURLE_OK);
int is4xx = (!transport_err && code >= 400 && code < 500 && code != 408 && code != 429);
if (g_col && !stream) printf("\r\033[K");
if (!transport_err && code < 400) {
if (stream) {
stream_finish(&sc);
if (sc.erase_pending) { if (g_col) printf("\r\033[K"); sc.erase_pending = 0; }
stream_result(&sc, out, usage);
} else {
parse_nonstream_response(mbody, out, usage);
}
rc = 0;
break;
}
/* error path */
if (stream) {
if (transport_err) errmsg = curl_easy_strerror(res);
else errmsg = http_err_body(&sc, nullptr, code);
} else {
if (transport_err) errmsg = curl_easy_strerror(res);
else errmsg = http_err_body(nullptr, &mbody, code);
}
if (is4xx) { rc = -1; break; }
if (attempt < nf) {
if (g_col) printf("\r\033[K");
std::string msg = sfmt("[network error: %s, retrying in %ds...]", errmsg.c_str(), fib[attempt]);
cprint(msg, std::array{31});
printf("\n");
sleep_interruptible(fib[attempt], interrupted);
}
}
if (rc == 0) {
/* The gateway may report zero usage (it does not count tokens). In that
* case fall back to a rough bytes/4 estimate so the user still sees a
* number, but never let a bogus estimate override real gateway numbers. */
if (usage.prompt_tokens == 0) {
usage.prompt_tokens = est_tokens_msgs(msgs);
usage.completion_tokens = est_tokens_msg(out);
usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
}
/* Context size = what the model actually processed this turn
* (prompt + completion). Prefer real gateway numbers; only fall back to
* the estimate when none were provided. The old code set this to the
* prompt-only estimate and ignored both the completion and any real
* prompt_tokens the gateway returned. */
int ctx = usage.prompt_tokens + usage.completion_tokens;
if (ctx > usage.context_tokens) usage.context_tokens = ctx;
}
return rc;
}