diff --git a/README.md b/README.md index 3038d7ce..dd2b6fa6 100644 --- a/README.md +++ b/README.md @@ -249,6 +249,19 @@ ollama.embed(model='gemma4', input='The sky is blue because of rayleigh scatteri ollama.embed(model='gemma4', input=['The sky is blue because of rayleigh scattering', 'Grass is green because of chlorophyll']) ``` +### Embed (structured media) + +With an Ollama server and embedding model that support [structured embedding inputs](https://github.com/ollama/ollama/blob/3295b3e96ba386219f4baa841cdbe402a04b13dc/api/types.go#L647-L674), an input can be an object with `text`, `image`, or `audio` string fields. Media values must be caller-provided base64 strings. Objects and text strings can also be mixed in a batch; each item remains a separate embedding input. + +Given `image_base64` and `audio_base64` containing the encoded media: + +```python +ollama.embed(model='your-multimodal-embedding-model', input={'text': 'Describe this image', 'image': image_base64}) +ollama.embed(model='your-multimodal-embedding-model', input=['plain text', {'audio': audio_base64}]) +``` + +The same inputs work with `Client.embed` and `await AsyncClient.embed`. The SDK does not read media paths, bytes, or PDFs for these fields. Older servers and models without the corresponding embedding capability may reject the request; chat image support alone does not imply multimodal embedding support. + ### Ps ```python diff --git a/ollama/_client.py b/ollama/_client.py index 3654e3c0..4d6921a9 100644 --- a/ollama/_client.py +++ b/ollama/_client.py @@ -50,6 +50,7 @@ DeleteRequest, EmbeddingsRequest, EmbeddingsResponse, + EmbedInput, EmbedRequest, EmbedResponse, GenerateRequest, @@ -437,7 +438,7 @@ def systemone( def embed( self, model: str = '', - input: Union[str, Sequence[str]] = '', + input: Union[EmbedInput, Sequence[EmbedInput]] = '', truncate: Optional[bool] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, @@ -1097,7 +1098,7 @@ async def systemone( async def embed( self, model: str = '', - input: Union[str, Sequence[str]] = '', + input: Union[EmbedInput, Sequence[EmbedInput]] = '', truncate: Optional[bool] = None, options: Optional[Union[Mapping[str, Any], Options]] = None, keep_alive: Optional[Union[float, str]] = None, diff --git a/ollama/_types.py b/ollama/_types.py index ad041013..db2f4739 100644 --- a/ollama/_types.py +++ b/ollama/_types.py @@ -422,9 +422,12 @@ class ChatResponse(BaseGenerateResponse): 'Log probabilities for generated tokens if requested.' +EmbedInput = Union[str, Mapping[str, str]] + + class EmbedRequest(BaseRequest): - input: Union[str, Sequence[str]] - 'Input text to embed.' + input: Union[EmbedInput, Sequence[EmbedInput]] + 'Text or structured text/image/audio input, or a batch of inputs. Media values are base64 strings.' truncate: Optional[bool] = None 'Truncate the input to the maximum token length.' diff --git a/tests/test_embed.py b/tests/test_embed.py new file mode 100644 index 00000000..925519f3 --- /dev/null +++ b/tests/test_embed.py @@ -0,0 +1,110 @@ +import pytest +from pydantic import ValidationError + +from ollama import AsyncClient, Client, EmbedResponse, Options, ResponseError + +pytestmark = pytest.mark.anyio + +# A 1x1 PNG and a one-sample PCM WAV. Responses below are HTTP fixtures, not inference. +PNG_BASE64 = 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzhVAAAAAElFTkSuQmCC' +WAV_BASE64 = 'UklGRiYAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQIAAAAAAA==' + + +@pytest.fixture +def anyio_backend(): + return 'asyncio' + + +async def embed(httpserver, async_mode, **kwargs): + if async_mode: + async with AsyncClient(httpserver.url_for('/')) as client: + return await client.embed(**kwargs) + with Client(httpserver.url_for('/')) as client: + return client.embed(**kwargs) + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +@pytest.mark.parametrize( + 'input', + [ + {'text': 'A tiny image 🖼️', 'image': PNG_BASE64}, + {'image': PNG_BASE64}, + {'text': 'A short sound', 'audio': WAV_BASE64}, + {'audio': WAV_BASE64}, + ['plain text', {'text': '图片', 'image': PNG_BASE64}, {'audio': WAV_BASE64}], + ({'image': PNG_BASE64}, 'another item'), + ], + ids=['text-image', 'image', 'text-audio', 'audio', 'mixed-list', 'mixed-tuple'], +) +async def test_embed_structured_inputs(httpserver, async_mode, input): + wire_input = list(input) if isinstance(input, tuple) else input + count = len(wire_input) if isinstance(wire_input, list) else 1 + embeddings = [[float(i), i + 0.5] for i in range(count)] + httpserver.expect_ordered_request('/api/embed', method='POST', json={'model': 'dummy', 'input': wire_input}).respond_with_json({'model': 'dummy', 'embeddings': embeddings, 'prompt_eval_count': 3}) + + response = await embed(httpserver, async_mode, model='dummy', input=input) + + assert isinstance(response, EmbedResponse) + assert response.embeddings == embeddings + assert response['prompt_eval_count'] == 3 + httpserver.check_assertions() + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +@pytest.mark.parametrize('input', ['hello', ['one', 'two'], ('one', 'two'), '', []], ids=['scalar', 'list', 'tuple', 'empty-string', 'empty-list']) +async def test_embed_text_inputs(httpserver, async_mode, input): + wire_input = list(input) if isinstance(input, tuple) else input + httpserver.expect_ordered_request('/api/embed', method='POST', json={'model': 'dummy', 'input': wire_input}).respond_with_json({'embeddings': [[0.25]]}) + + response = await embed(httpserver, async_mode, model='dummy', input=input) + + assert response.embeddings == [[0.25]] + httpserver.check_assertions() + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +async def test_embed_default_input(httpserver, async_mode): + httpserver.expect_ordered_request('/api/embed', method='POST', json={'model': 'dummy', 'input': ''}).respond_with_json({'embeddings': []}) + + response = await embed(httpserver, async_mode, model='dummy') + + assert response.embeddings == [] + httpserver.check_assertions() + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +async def test_embed_missing_model_still_rejected(httpserver, async_mode): + with pytest.raises(ValidationError, match='model'): + await embed(httpserver, async_mode) + assert not httpserver.log + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +@pytest.mark.parametrize('options_model', [False, True], ids=['mapping', 'Options']) +async def test_embed_structured_options(httpserver, async_mode, options_model): + input = {'image': PNG_BASE64} + options = {'num_ctx': 128} + httpserver.expect_ordered_request( + '/api/embed', + method='POST', + json={'model': 'dummy', 'input': input, 'truncate': False, 'options': options, 'keep_alive': '5m', 'dimensions': 2}, + ).respond_with_json({'embeddings': [[0.25, 0.5]], 'total_duration': 12}) + + response = await embed(httpserver, async_mode, model='dummy', input=input, truncate=False, options=Options(**options) if options_model else options, keep_alive='5m', dimensions=2) + + assert response.embeddings == [[0.25, 0.5]] + assert response.total_duration == 12 + httpserver.check_assertions() + + +@pytest.mark.parametrize('async_mode', [False, True], ids=['sync', 'async']) +async def test_embed_structured_server_error(httpserver, async_mode): + # Media decoding and model support remain server concerns, not SDK file loading. + input = {'image': 'invalid-base64'} + httpserver.expect_ordered_request('/api/embed', method='POST', json={'model': 'dummy', 'input': input}).respond_with_json({'error': 'input.image: invalid base64'}, status=400) + + with pytest.raises(ResponseError, match=r'input\.image: invalid base64') as error: + await embed(httpserver, async_mode, model='dummy', input=input) + + assert error.value.status_code == 400 + httpserver.check_assertions()