Model list retrieval.

nodes/ollama_generate_string.py

@@ -2,7 +2,12 @@ import os
 import json
 from collections import deque
  
-from server import PromptServer
+from server import PromptServer 
+from aiohttp import web
+import aiohttp
+
+OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "127.0.0.1:11434")
+
 class OllamaGenerateString :
     def __init__(self): 
         pass
@@ -35,6 +40,13 @@ class OllamaGenerateString :
         print("  generating new with ", int_value)
         return { "ui": { "next_int": [int_value], "next_value": [f"super cool: {int_value}"]}, "result": [result] } 
 
-@PromptServer.instance.routes.get("/iw/api/ollama_generate")
-async def load_drawing(request):
-    print("/iw/api/ollama_generate") 
+@PromptServer.instance.routes.get("/iw/api/ollama/models")
+async def retrieve_models(request):
+    print("/iw/api/ollama/models") 
+    async with aiohttp.ClientSession() as session:
+        async with session.get(f'http://{OLLAMA_HOST}/api/tags') as response:
+            body = await response.json()
+            model_simple = [v["model"] for v in body["models"]]
+            return web.json_response({"status": "success", "models": model_simple}, status=200)
+
+    return web.json_response({"status": "fail"}, status=500)

web/ollama_generate_string.js

@@ -5,28 +5,45 @@ import { ComfyWidgets } from "../../scripts/widgets.js";
 app.registerExtension({
   name: "iw.ollama.generate",
 
-  async beforeRegisterNodeDef(nodeType, nodeData, app) { 
+  async beforeRegisterNodeDef(nodeType, nodeData, app) {
     let find_widget = (node, name) => node.widgets.find((w) => w.name === name);
+
+
+
     if (nodeData.name === "IW_OllamaGenerateString") {
+
+
+      const resp = await api.fetchApi('/iw/api/ollama/models', { });
+
+      if (resp.status !== 200) {
+        const err = `Error uploading temp file: ${resp.status} - ${resp.statusText}`
+        useToastStore().addAlert(err)
+        throw new Error(err)
+      }
+      const result_in_json = await resp.json();
+      const model_list = result_in_json.models;
+      console.log("model list:", model_list)
+
       console.log("Something ollama thing");
+
       nodeType.prototype.initState = function () {
         console.log("hello");
       };
       nodeType.prototype.initWidgets = function (app) { 
-        this.setupWidgets(); 
+        this.setupWidgets(model_list);  
       };
 
-      nodeType.prototype.setupWidgets = function (app) {
+      nodeType.prototype.setupWidgets = function (app, models) {
         this.addWidget(
           "combo",
           "model",
-          "unknown",
+          model_list[0],
           (value) => {
             this.tool = value;
           },
           {
-            values: ["a", "b", "c"],
-            default: "unknown",
+            values: model_list,
+            default: model_list[0],
           }
         );