tokenfactory-pi 0.2.4 → 0.2.6

package/dist/index.d.ts

@@ -1,11 +1,11 @@
 /**
- * Nebius Token Factory — pi extension
+ * Nebius Token Factory - pi extension
  *
  * Fetches the current model catalog from the Token Factory API on startup
  * and registers all tool-capable text-generation models as a "nebius" provider.
  *
  * Environment:
- *   NEBIUS_API_KEY — required, Token Factory API key
+ *   NEBIUS_API_KEY - required, Token Factory API key
  *
  * Usage:
  *   pi -e /path/to/tokenfactory-pi

package/dist/index.js

@@ -1,27 +1,30 @@
 /**
- * Nebius Token Factory — pi extension
+ * Nebius Token Factory - pi extension
  *
  * Fetches the current model catalog from the Token Factory API on startup
  * and registers all tool-capable text-generation models as a "nebius" provider.
  *
  * Environment:
- *   NEBIUS_API_KEY — required, Token Factory API key
+ *   NEBIUS_API_KEY - required, Token Factory API key
  *
  * Usage:
  *   pi -e /path/to/tokenfactory-pi
  *   pi -e /path/to/tokenfactory-pi --provider nebius
  *   pi -e /path/to/tokenfactory-pi --provider nebius --model Qwen/Qwen3-32B
  */
+import { gunzipSync } from "node:zlib";
 const PROVIDER_NAME = "nebius";
 const BASE_URL = "https://api.tokenfactory.nebius.com/v1";
 const ENV_VAR = "NEBIUS_API_KEY";
+function isGzip(bytes) {
+    return bytes.length >= 2 && bytes[0] === 0x1f && bytes[1] === 0x8b;
+}
 // ============================================================================
 // Helpers
 // ============================================================================
-function isToolCapableTextModel(m) {
-    const features = m.supported_features || [];
+function isTextModel(m) {
     const modality = m.architecture?.modality || "";
-    return features.includes("tools") && modality.includes("->text");
+    return modality.includes("->text");
 }
 function parseInputModalities(modality) {
     const input = ["text"];
@@ -30,11 +33,25 @@ function parseInputModalities(modality) {
     return input;
 }
 function parseCostPerMillion(raw) {
-    return parseFloat(raw || "0") * 1_000_000;
+    const parsed = parseFloat(raw || "0");
+    return isNaN(parsed) ? 0 : parsed * 1_000_000;
 }
 function isReasoningModel(id) {
     return /(-R1|-Thinking|QwQ)/.test(id);
 }
+async function readTokenFactoryResponse(res) {
+    const bytes = Buffer.from(await res.arrayBuffer());
+    const body = (isGzip(bytes) ? gunzipSync(bytes) : bytes).toString("utf8");
+    try {
+        return JSON.parse(body);
+    }
+    catch {
+        const preview = body.slice(0, 200).replace(/\s+/g, " ");
+        throw new Error(`Invalid Token Factory JSON response (${res.status} ${res.statusText}, ` +
+            `content-type=${res.headers.get("content-type") || "unknown"}, ` +
+            `content-encoding=${res.headers.get("content-encoding") || "none"}): ${preview}`);
+    }
+}
 // ============================================================================
 // Extension entry point
 // ============================================================================
@@ -46,31 +63,51 @@ export default async function (pi) {
     let response;
     try {
         const res = await fetch(`${BASE_URL}/models?verbose=true`, {
-            headers: { Authorization: `Bearer ${apiKey}` },
+            headers: {
+                Authorization: `Bearer ${apiKey}`,
+                "Accept-Encoding": "identity",
+            },
         });
         if (!res.ok) {
             console.warn(`[${PROVIDER_NAME}] API returned ${res.status}: ${res.statusText}`);
+            console.warn(`[${PROVIDER_NAME}] Response headers:`, [...res.headers.entries()]);
             return;
         }
-        response = (await res.json());
+        response = await readTokenFactoryResponse(res);
     }
     catch (error) {
         console.warn(`[${PROVIDER_NAME}] Failed to fetch models:`, error);
         return;
     }
     if (!Array.isArray(response.data)) {
-        console.warn(`[${PROVIDER_NAME}] Unexpected API response shape`);
+        console.warn(`[${PROVIDER_NAME}] Unexpected API response shape. Expected array, got:`, typeof response.data);
+        console.warn(`[${PROVIDER_NAME}] Response data:`, response);
         return;
     }
     const models = [];
     for (const m of response.data) {
-        if (!isToolCapableTextModel(m))
+        // Skip models without valid IDs
+        if (!m.id || m.id.trim() === "") {
+            console.log(`[${PROVIDER_NAME}] Skipping model with empty ID:`, m);
+            continue;
+        }
+        if (!isTextModel(m)) {
             continue;
+        }
         const modality = m.architecture?.modality || "";
+        // Validate and set defaults for critical fields
+        const contextLength = m.context_length && m.context_length > 0 ? m.context_length : 131072;
+        const modelName = m.name || m.id || "unknown-model";
+        const modelId = m.id || modelName;
+        // Skip models with zero context length even after defaults
+        if (contextLength <= 0) {
+            console.log(`[${PROVIDER_NAME}] Skipping model with invalid context length:`, m.id, m.context_length);
+            continue;
+        }
         models.push({
-            id: m.id,
-            name: m.name || m.id,
-            reasoning: isReasoningModel(m.id),
+            id: modelId,
+            name: modelName,
+            reasoning: isReasoningModel(modelId),
             input: parseInputModalities(modality),
             cost: {
                 input: parseCostPerMillion(m.pricing?.prompt),
@@ -78,8 +115,8 @@ export default async function (pi) {
                 cacheRead: 0,
                 cacheWrite: 0,
             },
-            contextWindow: m.context_length || 131072,
-            maxTokens: Math.min(m.context_length || 32768, 32768),
+            contextWindow: contextLength,
+            maxTokens: Math.min(contextLength, 32768),
             compat: {
                 supportsDeveloperRole: false,
                 maxTokensField: "max_tokens",
@@ -90,6 +127,9 @@ export default async function (pi) {
         baseUrl: BASE_URL,
         apiKey: ENV_VAR,
         api: "openai-completions",
+        headers: {
+            "Accept-Encoding": "identity",
+        },
         models,
     });
     // /nebius-models command to list and select a model
@@ -101,7 +141,7 @@ export default async function (pi) {
                 return;
             }
             const items = models
-                .sort((a, b) => a.id.localeCompare(b.id))
+                .sort((a, b) => b.cost.output - a.cost.output)
                 .map((m) => {
                 const tags = [];
                 if (m.reasoning)
@@ -111,7 +151,33 @@ export default async function (pi) {
                 const suffix = tags.length > 0 ? ` (${tags.join(", ")})` : "";
                 return `${m.id}${suffix}`;
             });
-            await ctx.ui.select(`Nebius Token Factory — ${models.length} models`, items);
+            // Show the selection dialog and capture the result
+            const selectedItem = await ctx.ui.select(`Nebius Token Factory - ${models.length} models`, items);
+            // If user cancelled (selectedItem is undefined), do nothing
+            if (!selectedItem) {
+                ctx.ui.notify("Model selection cancelled", "warning");
+                return;
+            }
+            // Extract the model ID from the selection (remove the suffix)
+            const selectedId = selectedItem.split(" (")[0];
+            const selectedModel = models.find(m => m.id === selectedId);
+            if (!selectedModel) {
+                ctx.ui.notify(`Model not found: ${selectedId}`, "error");
+                return;
+            }
+            // Find the model in the registry and switch to it
+            const model = ctx.modelRegistry.find(PROVIDER_NAME, selectedModel.id);
+            if (!model) {
+                ctx.ui.notify(`Model not found in registry: ${PROVIDER_NAME}/${selectedModel.id}`, "error");
+                return;
+            }
+            const success = await pi.setModel(model);
+            if (success) {
+                ctx.ui.notify(`Switched to model: ${model.provider}/${model.id}`, "info");
+            }
+            else {
+                ctx.ui.notify(`Failed to switch to model: ${model.provider}/${model.id}. Check API key?`, "error");
+            }
         },
     });
 }

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "tokenfactory-pi",
-  "version": "0.2.4",
+  "version": "0.2.6",
   "description": "Nebius Token Factory provider extension for pi coding agent. Requires `npm install -g @mariozechner/pi-coding-agent`. Install with `pi install npm:tokenfactory-pi`",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -43,4 +43,4 @@
     "check": "tsc --noEmit",
     "prepare": "npm run build"
   }
-}
+}