From 59ca37f2572a18b2e90685e3a613303e12f5a2e0 Mon Sep 17 00:00:00 2001
From: monofuel <monofuel@japura.net>
Date: Fri, 12 Jan 2024 21:34:23 -0500
Subject: [PATCH] embeddings

---
 src/llama_leap.nim        | 28 ++++++++++++++++++++++++++++
 tests/test_llama_leap.nim |  4 ++++
 2 files changed, 32 insertions(+)

diff --git a/src/llama_leap.nim b/src/llama_leap.nim
index aaa2631..ca0bcf9 100644
--- a/src/llama_leap.nim
+++ b/src/llama_leap.nim
@@ -65,6 +65,12 @@ type
     details*: ModelDetails
   ListResp* = ref object
     models*: seq[OllamaModel]
+  EmbeddingReq* = ref object
+    model*: string
+    prompt*: string
+    options*: Option[ModelParameters] # bag of model parameters
+  EmbeddingResp* = ref object
+    embedding*: seq[float64]
 
 proc renameHook*(v: var GenerateReq, fieldName: var string) =
   ## `template` is a special keyword in nim, so we need to rename it during serialization
@@ -177,3 +183,25 @@ proc pullModel*(api: OllamaAPI, name: string) =
   let status = respJson["status"].getStr
   if status != "success":
     raise newException(CatchableError, &"ollama pull bad status: {resp.body}")
+
+proc generateEmbeddings*(
+  api: OllamaAPI,
+  model: string,
+  prompt: string,
+  options: Option[ModelParameters] = none(ModelParameters)
+): EmbeddingResp =
+  ## Get the embeddings for a prompt
+  let url = api.baseUrl / "embeddings"
+  var req = EmbeddingReq(
+    model: model,
+    prompt: prompt,
+    options: options
+  )
+
+  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 embedding failed: {resp.code} {resp.body}")
+  result = fromJson(resp.body, EmbeddingResp)
diff --git a/tests/test_llama_leap.nim b/tests/test_llama_leap.nim
index 01b3781..0a3c958 100644
--- a/tests/test_llama_leap.nim
+++ b/tests/test_llama_leap.nim
@@ -82,3 +82,7 @@ suite "llama_leap":
       )
       let resp2 = ollama.generate(req2)
       echo "2> " & resp2.response.strip()
+  suite "embeddings":
+    test "generate embeddings":
+      let resp = ollama.generateEmbeddings(TestModel, "How are you today?")
+      echo "> " & resp.embedding.join(",")