From e499fef26bbda4b9c9c05fe6426fb848f03f8404 Mon Sep 17 00:00:00 2001 From: RKS Date: Thu, 1 Oct 2026 00:07:14 -0400 Subject: [PATCH] Python: fix(ollama): preserve zero embedding keep_alive --- .../_embedding_client.py | 2 +- .../ollama/test_ollama_embedding_client.py | 27 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/packages/ollama/agent_framework_ollama/_embedding_client.py b/python/packages/ollama/agent_framework_ollama/_embedding_client.py index 03f7078c5c6..bbefa11ba56 100644 --- a/python/packages/ollama/agent_framework_ollama/_embedding_client.py +++ b/python/packages/ollama/agent_framework_ollama/_embedding_client.py @@ -148,7 +148,7 @@ async def get_embeddings( kwargs: dict[str, Any] = {"model": model, "input": list(values)} if (truncate := opts.get("truncate")) is not None: kwargs["truncate"] = truncate - if keep_alive := opts.get("keep_alive"): + if (keep_alive := opts.get("keep_alive")) is not None: kwargs["keep_alive"] = keep_alive if dimensions := opts.get("dimensions"): kwargs["dimensions"] = dimensions diff --git a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py b/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py index 9cd12bbc3f8..6cf2b67405d 100644 --- a/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py +++ b/python/packages/ollama/tests/ollama/test_ollama_embedding_client.py @@ -1,10 +1,14 @@ # Copyright (c) Microsoft. All rights reserved. +import json import os +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from agent_framework import Embedding, GeneratedEmbeddings +from ollama import AsyncClient from agent_framework_ollama import OllamaEmbeddingClient, OllamaEmbeddingOptions from agent_framework_ollama._feature_usage import FeatureIndex @@ -117,6 +121,29 @@ async def test_ollama_embedding_get_embeddings_with_options() -> None: ) +@pytest.mark.parametrize("keep_alive", [0, 0.0, -1, 300, "0", "5m", None]) +async def test_ollama_embedding_keep_alive_request(keep_alive: float | str | None) -> None: + """Preserve explicit keep-alive durations, including numeric zero, on the wire.""" + requests: list[dict[str, Any]] = [] + + def handle_request(request: httpx.Request) -> httpx.Response: + requests.append(json.loads(request.content)) + return httpx.Response(200, json={"model": "test-model", "embeddings": [[0.1, 0.2]]}) + + sdk_client = AsyncClient(host="http://ollama.test", transport=httpx.MockTransport(handle_request)) + async with sdk_client._client: + client = OllamaEmbeddingClient(model="test-model", client=sdk_client) + options: OllamaEmbeddingOptions = {} if keep_alive is None else {"keep_alive": keep_alive} + result = await client.get_embeddings(["hello"], options=options) + + assert result[0].vector == [0.1, 0.2] + assert len(requests) == 1 + if keep_alive is None: + assert "keep_alive" not in requests[0] + else: + assert requests[0]["keep_alive"] == keep_alive + + async def test_ollama_embedding_get_embeddings_no_model_raises() -> None: """Test that missing model at call time raises ValueError.""" with patch("agent_framework_ollama._embedding_client.AsyncClient") as mock_client_cls: