letta-code extension: register oMLX as a pi provider using /v1/models/status as the source of truth for per-model capabilities (modality, context window, max tokens, thinking default). Requires letta-code with PR #2552 merged.

omlx-provider.ts

// Registers oMLX (https://github.com/jundot/omlx) as a letta-code pi provider,
// using oMLX's /v1/models/status endpoint as the source of truth for per-model
// capabilities (modality, context window, max tokens, thinking default).
//
// Without this, letta-code falls back to substring heuristics on the model id
// (llava / vision / vl for VLMs, gpt-oss / qwen3 / deepseek-r1 for reasoning,
// 128K context default) — see https://github.com/letta-ai/letta-code/issues/2541.
//
// Requires letta-code with PR #2552 merged (the `letta.registerProvider()` API).

const OMLX_BASE_URL = process.env.OMLX_BASE_URL ?? "http://localhost:8000/v1";
const OMLX_API_KEY = process.env.OMLX_API_KEY ?? "1234";

// Provider name must match the entry in ~/.letta/lc-local-backend/providers/auth.json
// so the registration augments the existing connection rather than introducing
// a parallel one.
const OMLX_PROVIDER_NAME = "lmstudio";

const OMLX_DISCOVERY_TIMEOUT_MS = 3_000;

type OmlxModelType = "llm" | "vlm" | "embeddings";

interface OmlxModelStatus {
  id: string;
  model_type: OmlxModelType;
  thinking_default: boolean | null;
  max_context_window: number;
  max_tokens: number;
}

interface OmlxModelsStatusResponse {
  models?: OmlxModelStatus[];
}

async function fetchOmlxModels(
  signal: AbortSignal | undefined,
): Promise<OmlxModelStatus[]> {
  const controller = new AbortController();
  const timeout = setTimeout(
    () => controller.abort(new Error("oMLX discovery timed out")),
    OMLX_DISCOVERY_TIMEOUT_MS,
  );
  signal?.addEventListener("abort", () => controller.abort(signal.reason), {
    once: true,
  });
  try {
    const response = await fetch(`${OMLX_BASE_URL}/models/status`, {
      headers: { Authorization: `Bearer ${OMLX_API_KEY}` },
      signal: controller.signal,
    });
    if (!response.ok) {
      throw new Error(
        `oMLX returned HTTP ${response.status} ${response.statusText}`,
      );
    }
    const body = (await response.json()) as OmlxModelsStatusResponse;
    return body.models ?? [];
  } finally {
    clearTimeout(timeout);
  }
}

interface ProviderModelRegistration {
  id: string;
  name: string;
  reasoning: boolean;
  input: ("text" | "image")[];
  cost: {
    input: number;
    output: number;
    cacheRead: number;
    cacheWrite: number;
  };
  contextWindow: number;
  maxTokens: number;
}

function toProviderModel(m: OmlxModelStatus): ProviderModelRegistration {
  return {
    id: m.id,
    name: m.id,
    reasoning: m.thinking_default === true,
    input: m.model_type === "vlm" ? ["text", "image"] : ["text"],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    contextWindow: m.max_context_window,
    maxTokens: m.max_tokens,
  };
}

export default async function activate(letta: any) {
  if (!letta?.capabilities?.providers) {
    return;
  }

  let models: OmlxModelStatus[];
  try {
    models = await fetchOmlxModels(letta.signal);
  } catch (error) {
    const reason = error instanceof Error ? error.message : String(error);
    console.warn(
      `[omlx-provider] could not reach oMLX at ${OMLX_BASE_URL}: ${reason}`,
    );
    return;
  }

  const chatModels = models
    .filter((m) => m.model_type !== "embeddings")
    .map(toProviderModel);

  if (chatModels.length === 0) {
    console.warn(
      `[omlx-provider] oMLX returned no chat-capable models at ${OMLX_BASE_URL}`,
    );
    return;
  }

  try {
    letta.registerProvider(OMLX_PROVIDER_NAME, {
      baseUrl: OMLX_BASE_URL,
      apiKey: OMLX_API_KEY,
      api: "openai-completions",
      models: chatModels,
    });
  } catch (error) {
    const reason = error instanceof Error ? error.message : String(error);
    console.warn(`[omlx-provider] registerProvider threw: ${reason}`);
    return;
  }

  return () => {};
}