llmgw/internal/vscodeext/ext/lib.js

'use strict';

// What the extension does that does not need VS Code: reading the gateway's
// model list, turning VS Code's messages into chat-completion messages, and
// reading the streamed answer back. Kept apart so it can be tested with plain
// Node, against a gateway, without an editor running.

const USER = 1;
const ASSISTANT = 2;
// Not in every VS Code that runs this: a system message, where it exists.
const SYSTEM = 3;

// gatewayBase is the address the person typed, as the API is under it: no
// trailing slash, and no /v1, which people copy from client guides.
function gatewayBase(raw) {
  let s = String(raw || '').trim();
  if (!s) return '';
  // A bare host is https; any other scheme than http(s) is not a gateway.
  if (/^[a-z][a-z0-9+.-]*:\/\//i.test(s)) {
    if (!/^https?:\/\//i.test(s)) return '';
  } else {
    s = 'https://' + s;
  }
  let u;
  try {
    u = new URL(s);
  } catch {
    return '';
  }
  if (u.protocol !== 'https:' && u.protocol !== 'http:') return '';
  let path = u.pathname.replace(/\/+$/, '');
  if (path.endsWith('/v1')) path = path.slice(0, -3);
  return u.origin + path;
}

// modelInfo is what VS Code is told of one entry of /v1/models, or null for
// one that does not chat: an embedding model, or one that makes pictures,
// meshes or anything else through a workflow.
//
// The sizes are the catalog's when it gives them. VS Code trims a
// conversation to maxInputTokens, and a request over what a replica holds is
// refused by the gateway, so guessing high is worse than guessing low.
function modelInfo(entry, showWithoutTools) {
  if (!entry || typeof entry.id !== 'string' || !entry.id) return null;
  const caps = entry.capabilities || {};
  if (caps.embeddings || caps.images || caps.meshes || caps.workflows) return null;
  const tools = !!caps.tools;
  if (!tools && !showWithoutTools) return null;
  const context = Number(entry.context_length) > 0 ? Number(entry.context_length) : 8192;
  let output = Number(entry.max_output_tokens) > 0 ? Number(entry.max_output_tokens) : 0;
  output = Math.min(output || 4096, Math.max(256, Math.floor(context / 4)));
  const detail = entry.alias_of ? `llmgw · ${entry.alias_of}` : 'llmgw';
  return {
    id: entry.id,
    name: entry.id,
    family: entry.alias_of || entry.id,
    version: 'llmgw',
    detail,
    tooltip: entry.description || undefined,
    maxInputTokens: Math.max(256, context - output),
    maxOutputTokens: output,
    capabilities: { toolCalling: tools, imageInput: !!caps.vision },
  };
}

function isToolCall(p) {
  return p && typeof p.callId === 'string' && typeof p.name === 'string' && 'input' in p;
}
function isToolResult(p) {
  return p && typeof p.callId === 'string' && Array.isArray(p.content);
}
function isText(p) {
  return p && typeof p.value === 'string' && !isToolCall(p) && !isToolResult(p);
}
function isImage(p) {
  return p && typeof p.mimeType === 'string' && p.mimeType.startsWith('image/') && p.data instanceof Uint8Array;
}

// resultText flattens a tool result to the text a chat-completion tool
// message carries. Parts that are not text (prompt-tsx trees) go as JSON.
function resultText(content) {
  return content
    .map((c) => (isText(c) ? c.value : typeof c === 'string' ? c : JSON.stringify(c)))
    .join('');
}

// chatMessages turns VS Code's request messages into chat-completion ones.
//
// One VS Code message can hold several kinds of part; the API they are going
// to wants each in its own message: an assistant's tool calls in one, each
// tool result in a tool message straight after it, and only then the text a
// person wrote beside the results.
function chatMessages(messages, vision) {
  const out = [];
  for (const m of messages) {
    const parts = Array.isArray(m.content) ? m.content : [];
    if (m.role === ASSISTANT) {
      const text = parts.filter(isText).map((p) => p.value).join('');
      const calls = parts.filter(isToolCall).map((p) => ({
        id: p.callId,
        type: 'function',
        function: { name: p.name, arguments: JSON.stringify(p.input ?? {}) },
      }));
      const msg = { role: 'assistant', content: text || (calls.length ? null : '') };
      if (calls.length) msg.tool_calls = calls;
      out.push(msg);
      continue;
    }
    for (const p of parts.filter(isToolResult)) {
      out.push({ role: 'tool', tool_call_id: p.callId, content: resultText(p.content) });
    }
    const texts = parts.filter(isText).map((p) => p.value);
    const images = vision ? parts.filter(isImage) : [];
    if (!texts.length && !images.length) continue;
    const role = m.role === SYSTEM ? 'system' : 'user';
    if (!images.length) {
      out.push({ role, content: texts.join('') });
      continue;
    }
    const content = [];
    if (texts.length) content.push({ type: 'text', text: texts.join('') });
    for (const img of images) {
      const b64 = Buffer.from(img.data).toString('base64');
      content.push({ type: 'image_url', image_url: { url: `data:${img.mimeType};base64,${b64}` } });
    }
    out.push({ role, content });
  }
  return out;
}

// chatTools turns VS Code's tool list into chat-completion functions. A tool
// without a schema still needs one: an object that takes nothing.
function chatTools(tools) {
  return (tools || []).map((t) => ({
    type: 'function',
    function: {
      name: t.name,
      description: t.description || '',
      parameters: t.inputSchema && Object.keys(t.inputSchema).length ? t.inputSchema : { type: 'object', properties: {} },
    },
  }));
}

// requestBody is the whole chat-completion request. toolMode 2 is VS Code's
// "Required": the model must call one of the tools.
function requestBody(model, messages, options) {
  const body = {
    model: model.id,
    messages: chatMessages(messages, !!(model.capabilities && model.capabilities.imageInput)),
    stream: true,
    max_tokens: model.maxOutputTokens,
  };
  const tools = chatTools(options && options.tools);
  if (tools.length) {
    body.tools = tools;
    if (options.toolMode === 2) body.tool_choice = 'required';
  }
  const mo = (options && options.modelOptions) || {};
  for (const k of ['temperature', 'top_p', 'stop', 'seed']) {
    if (mo[k] !== undefined) body[k] = mo[k];
  }
  return body;
}

// sseData yields the data of each server-sent event in a byte stream, as the
// chat-completion stream sends them: "data: {...}" lines, ending "[DONE]".
async function* sseData(stream) {
  const decoder = new TextDecoder();
  let buf = '';
  for await (const chunk of stream) {
    buf += decoder.decode(chunk, { stream: true });
    let i;
    while ((i = buf.indexOf('\n')) >= 0) {
      const line = buf.slice(0, i).replace(/\r$/, '');
      buf = buf.slice(i + 1);
      if (line.startsWith('data:')) yield line.slice(5).trimStart();
    }
  }
  buf += decoder.decode();
  if (buf.startsWith('data:')) yield buf.slice(5).trimStart();
}

// Answer collects one streamed answer: its text as it comes, and its tool
// calls, whose arguments arrive in pieces and are only whole at the end.
class Answer {
  constructor() {
    this.calls = new Map();
    this.usage = null;
  }

  // add takes one event's data and returns the text it carries, if any.
  add(data) {
    if (data === '[DONE]') return '';
    let ev;
    try {
      ev = JSON.parse(data);
    } catch {
      return '';
    }
    if (ev.error) {
      const msg = typeof ev.error === 'string' ? ev.error : ev.error.message || JSON.stringify(ev.error);
      throw new Error(msg);
    }
    if (ev.usage) this.usage = ev.usage;
    let text = '';
    for (const choice of ev.choices || []) {
      const d = choice.delta || {};
      if (typeof d.content === 'string') text += d.content;
      for (const tc of d.tool_calls || []) {
        const key = tc.index ?? tc.id ?? this.calls.size;
        const c = this.calls.get(key) || { id: '', name: '', args: '' };
        if (tc.id) c.id = tc.id;
        if (tc.function && tc.function.name) c.name += tc.function.name;
        if (tc.function && typeof tc.function.arguments === 'string') c.args += tc.function.arguments;
        this.calls.set(key, c);
      }
    }
    return text;
  }

  // toolCalls is every call the answer made, its arguments parsed. A call
  // whose arguments are not JSON goes with them as text, under "input", so
  // the tool sees what the model wrote rather than nothing.
  toolCalls() {
    const out = [];
    let n = 0;
    for (const c of this.calls.values()) {
      n++;
      if (!c.name) continue;
      let input = {};
      if (c.args.trim()) {
        try {
          input = JSON.parse(c.args);
        } catch {
          input = { input: c.args };
        }
      }
      out.push({ id: c.id || `call_${n}`, name: c.name, input });
    }
    return out;
  }
}

// errorMessage reads what a gateway said when it refused a request.
function errorMessage(status, bodyText) {
  let msg = '';
  try {
    const j = JSON.parse(bodyText);
    const e = j.error;
    msg = (typeof e === 'string' ? e : e && e.message) || j.message || '';
  } catch {
    msg = bodyText.trim().slice(0, 300);
  }
  if (status === 401) return 'llmgw: the gateway did not accept the key. Run "llmgw: Sign In" to make a new one.';
  if (status === 403) return `llmgw: this key may not use that model${msg ? ` (${msg})` : ''}.`;
  if (status === 429) return `llmgw: a limit was reached${msg ? `: ${msg}` : ''}.`;
  return `llmgw: the gateway answered ${status}${msg ? `: ${msg}` : ''}.`;
}

// tokenGuess is how many tokens a text is, near enough for VS Code to decide
// what fits: about four characters to a token in English and in code.
function tokenGuess(text) {
  return Math.ceil(String(text || '').length / 4);
}

// downgrades reports whether next is the gateway at current by plain http
// where current is https: the same host and port, a weaker scheme.
function downgrades(current, next) {
  try {
    const a = new URL(current);
    const b = new URL(next);
    return a.protocol === 'https:' && b.protocol === 'http:' && a.hostname === b.hostname;
  } catch (e) {
    return false;
  }
}

module.exports = {
  USER, ASSISTANT, SYSTEM,
  gatewayBase, modelInfo, chatMessages, chatTools, requestBody, sseData, Answer, errorMessage, tokenGuess, downgrades,
};