From 143ac3a256a220d519e583dfeb3e7e31763efd08 Mon Sep 17 00:00:00 2001
From: pelikhan <jhalleux@microsoft.com>
Date: Fri, 13 Dec 2024 08:29:20 -0800
Subject: [PATCH] =?UTF-8?q?feat:=20enhance=20model=20pull=20logic=20with?=
 =?UTF-8?q?=20tag=20check=20=F0=9F=9B=A0=EF=B8=8F?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/cli/src/nodehost.ts | 43 +++++++++++++++++++++++++++---------
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/packages/cli/src/nodehost.ts b/packages/cli/src/nodehost.ts
index 2d3531dfcc..ca7fd552e2 100644
--- a/packages/cli/src/nodehost.ts
+++ b/packages/cli/src/nodehost.ts
@@ -71,6 +71,7 @@ import {
 } from "../../core/src/azurecontentsafety"
 import { resolveGlobalConfiguration } from "../../core/src/config"
 import { HostConfiguration } from "../../core/src/hostconfiguration"
+import { YAMLStringify } from "../../core/src/yaml"
 
 class NodeServerManager implements ServerManager {
     async start(): Promise<void> {
@@ -102,13 +103,31 @@ class ModelManager implements ModelService {
         if (this.pulled.includes(modelid)) return { ok: true }
 
         if (provider === MODEL_PROVIDER_OLLAMA) {
-            logVerbose(`${provider}: pull ${model}`)
             try {
+                logVerbose(`${provider}: show ${model}`)
                 const conn = await this.getModelToken(modelid)
 
-                let res: Response
-                // OLLAMA
-                res = await fetch(`${conn.base}/api/pull`, {
+                // test if model is present
+                const resTags = await fetch(`${conn.base}/api/tags`, {
+                    method: "GET",
+                    headers: {
+                        "User-Agent": TOOL_ID,
+                        "Content-Type": "application/json",
+                    },
+                })
+                if (resTags.ok) {
+                    const { models }: { models: { model: string }[] } =
+                        await resTags.json()
+                    if (models.find((m) => m.model === model))
+                        return { ok: true }
+                    logVerbose(
+                        `${provider}: ${model} not found in\n${YAMLStringify(models.map((m) => m.model))}`
+                    )
+                }
+
+                // pull
+                logVerbose(`${provider}: pull ${model}`)
+                const resPull = await fetch(`${conn.base}/api/pull`, {
                     method: "POST",
                     headers: {
                         "User-Agent": TOOL_ID,
@@ -120,15 +139,19 @@ class ModelManager implements ModelService {
                         2
                     ),
                 })
-                if (res.ok) {
-                    const resj = await res.json()
-                    //logVerbose(JSON.stringify(resj, null, 2))
+                if (resPull.ok) {
+                    const resj = await resPull.json()
+                    logVerbose(JSON.stringify(resj, null, 2))
+                }
+                if (resPull.ok) this.pulled.push(modelid)
+                else {
+                    logError(`${provider}: failed to pull model ${model}`)
+                    trace?.error(`${provider}: pull model ${model} failed`)
                 }
-                if (res.ok) this.pulled.push(modelid)
-                return { ok: res.ok, status: res.status }
+                return { ok: resPull.ok, status: resPull.status }
             } catch (e) {
                 logError(`${provider}: failed to pull model ${model}`)
-                trace?.error(`${provider}: pull model failed`, e)
+                trace?.error(`${provider}: pull model ${model} failed`, e)
                 return { ok: false, status: 500, error: serializeError(e) }
             }
         }