model pulling

monofuel · Jan 13, 2024 · 90d6e48 · 90d6e48
1 parent 4184dee
commit 90d6e48
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/src/llama_leap.nim b/src/llama_leap.nim
@@ -155,8 +155,25 @@ proc generate*(api: OllamaAPI, req: JsonNode): JsonNode =
   result = fromJson(resp.body)
 
 proc listModels*(api: OllamaAPI): ListResp =
+  ## List all the models available
   let url = api.baseUrl / "tags"
   let resp = api.curlPool.get(url, timeout = api.curlTimeout)
   if resp.code != 200:
     raise newException(CatchableError, &"ollama list tags failed: {resp.code} {resp.body}")
   result = fromJson(resp.body, ListResp)
+
+proc pullModel*(api: OllamaAPI, name: string) =
+  ## Ask the ollama server to pull a model
+  let url = api.baseUrl / "pull"
+  let req = %*{"name": name, "stream": false}
+
+  var headers: curly.HttpHeaders
+  headers["Content-Type"] = "application/json"
+
+  let resp = api.curlPool.post(url, headers, toJson(req), api.curlTimeout)
+  if resp.code != 200:
+    raise newException(CatchableError, &"ollama pull failed: {resp.code} {resp.body}")
+  let respJson = fromJson(resp.body)
+  let status = respJson["status"].getStr
+  if status != "success":
+    raise newException(CatchableError, &"ollama pull bad status: {resp.body}")
diff --git a/tests/test_llama_leap.nim b/tests/test_llama_leap.nim
@@ -13,6 +13,10 @@ suite "llama_leap":
   teardown:
     ollama.close()
 
+  suite "pull":
+    test "pull model":
+      ollama.pullModel(TestModel)
+
   suite "list":
     test "list model tags":
       let resp = ollama.listModels()