diff --git a/src/llama_leap.nim b/src/llama_leap.nim index db71b7b..aaa2631 100644 --- 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 index 3429b26..01b3781 100644 --- 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()