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,