From 008a6a6b00d6df0536d1595b85e7b6dae1028273 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Wed, 10 Jan 2024 17:20:23 -0800
Subject: [PATCH] add embeddings

---
 ollama/__init__.py |  2 ++
 ollama/_client.py  | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/ollama/__init__.py b/ollama/__init__.py
index c4ebe2b..048bb14 100644
--- a/ollama/__init__.py
+++ b/ollama/__init__.py
@@ -21,6 +21,7 @@ __all__ = [
   'ResponseError',
   'generate',
   'chat',
+  'embeddings',
   'pull',
   'push',
   'create',
@@ -34,6 +35,7 @@ _client = Client()
 
 generate = _client.generate
 chat = _client.chat
+embeddings = _client.embeddings
 pull = _client.pull
 push = _client.push
 create = _client.create
diff --git a/ollama/_client.py b/ollama/_client.py
index 6c2c9b3..e4aceac 100644
--- a/ollama/_client.py
+++ b/ollama/_client.py
@@ -170,6 +170,17 @@ class Client(BaseClient):
       stream=stream,
     )
 
+  def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]:
+    return self._request(
+      'POST',
+      '/api/embeddings',
+      json={
+        'model': model,
+        'prompt': prompt,
+        'options': options or {},
+      },
+    ).json()
+
   def pull(
     self,
     model: str,
@@ -425,6 +436,19 @@ class AsyncClient(BaseClient):
       stream=stream,
     )
 
+  async def embeddings(self, model: str = '', prompt: str = '', options: Optional[Options] = None) -> Sequence[float]:
+    response = await self._request(
+      'POST',
+      '/api/embeddings',
+      json={
+        'model': model,
+        'prompt': prompt,
+        'options': options or {},
+      },
+    )
+
+    return response.json()
+
   async def pull(
     self,
     model: str,