From c70f74f33b404393947bcfc5e55f3588afb76e15 Mon Sep 17 00:00:00 2001
From: shuaiyafei <16618929+shuaiyafei@users.noreply.github.com>
Date: Thu, 22 Oct 2026 11:42:33 +0800
Subject: [PATCH] types: preserve cached prompt token counts

---
 ollama/_types.py                 |  3 +++
 tests/test_client.py             | 44 ++++++++++++++++++++++++++++++++
 tests/test_type_serialization.py | 24 ++++++++++++++++-
 3 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/ollama/_types.py b/ollama/_types.py
index ad041013..4bed1df1 100644
--- a/ollama/_types.py
+++ b/ollama/_types.py
@@ -249,6 +249,9 @@ class BaseGenerateResponse(SubscriptableBaseModel):
   prompt_eval_count: Optional[int] = None
   'Number of tokens evaluated in the prompt.'
 
+  prompt_eval_cached_count: Optional[int] = None
+  'Number of prompt tokens read from the cache.'
+
   prompt_eval_duration: Optional[int] = None
   'Duration of evaluating the prompt in nanoseconds.'
 
diff --git a/tests/test_client.py b/tests/test_client.py
index c212a11d..671dc2ec 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -62,6 +62,50 @@ def test_client_chat(httpserver: HTTPServer):
   assert response['message']['content'] == "I don't know."
 
 
+def _serve_cached_prompt_response(httpserver: HTTPServer, method: str, stream: bool):
+  content = {'message': {'role': 'assistant', 'content': 'Hello'}} if method == 'chat' else {'response': 'Hello'}
+  payload = {'model': 'dummy', **content, 'done': True, 'prompt_eval_count': 50, 'prompt_eval_cached_count': 42}
+  request = httpserver.expect_request(f'/api/{method}', method='POST')
+  if stream:
+    request.respond_with_data(json.dumps({'model': 'dummy', **content, 'done': False}) + '\n' + json.dumps(payload) + '\n', content_type='application/x-ndjson')
+  else:
+    request.respond_with_json(payload)
+
+
+@pytest.mark.parametrize('method', ['generate', 'chat'])
+@pytest.mark.parametrize('stream', [False, True])
+def test_client_prompt_eval_cached_count(httpserver: HTTPServer, method: str, stream: bool):
+  _serve_cached_prompt_response(httpserver, method, stream)
+
+  with Client(httpserver.url_for('/')) as client:
+    response = getattr(client, method)('dummy', stream=stream)
+    if stream:
+      chunks = list(response)
+      assert chunks[0].prompt_eval_cached_count is None
+      response = chunks[-1]
+
+    assert response.prompt_eval_count == 50
+    assert response.prompt_eval_cached_count == 42
+    assert response.model_dump()['prompt_eval_cached_count'] == 42
+
+
+@pytest.mark.parametrize('method', ['generate', 'chat'])
+@pytest.mark.parametrize('stream', [False, True])
+async def test_async_client_prompt_eval_cached_count(httpserver: HTTPServer, method: str, stream: bool):
+  _serve_cached_prompt_response(httpserver, method, stream)
+
+  async with AsyncClient(httpserver.url_for('/')) as client:
+    response = await getattr(client, method)('dummy', stream=stream)
+    if stream:
+      chunks = [chunk async for chunk in response]
+      assert chunks[0].prompt_eval_cached_count is None
+      response = chunks[-1]
+
+    assert response.prompt_eval_count == 50
+    assert response.prompt_eval_cached_count == 42
+    assert response.model_dump()['prompt_eval_cached_count'] == 42
+
+
 def test_client_chat_with_logprobs(httpserver: HTTPServer):
   httpserver.expect_ordered_request(
     '/api/chat',
diff --git a/tests/test_type_serialization.py b/tests/test_type_serialization.py
index e1af3512..5589e66d 100644
--- a/tests/test_type_serialization.py
+++ b/tests/test_type_serialization.py
@@ -4,7 +4,7 @@
 
 import pytest
 
-from ollama._types import ChatRequest, CreateRequest, GenerateRequest, Image
+from ollama._types import ChatRequest, ChatResponse, CreateRequest, GenerateRequest, GenerateResponse, Image
 
 
 def test_image_serialization_bytes():
@@ -122,3 +122,25 @@ def test_think_boolean_serialization():
   assert ChatRequest(model='test-model', think=True).model_dump(exclude_none=True)['think'] is True
   assert ChatRequest(model='test-model', think=False).model_dump(exclude_none=True)['think'] is False
   assert 'think' not in ChatRequest(model='test-model', think=None).model_dump(exclude_none=True)
+
+
+@pytest.mark.parametrize('response_type', [GenerateResponse, ChatResponse])
+@pytest.mark.parametrize('cached_count', [None, 0, 42])
+def test_prompt_eval_cached_count_serialization(response_type, cached_count):
+  payload = {'model': 'dummy', 'prompt_eval_count': 50}
+  if response_type is ChatResponse:
+    payload['message'] = {'role': 'assistant', 'content': 'Hello'}
+  if cached_count is not None:
+    payload['prompt_eval_cached_count'] = cached_count
+
+  response = response_type.model_validate(payload)
+
+  assert response.prompt_eval_count == 50
+  assert response.prompt_eval_cached_count == cached_count
+  if cached_count is None:
+    assert 'prompt_eval_cached_count' not in response.model_dump(exclude_none=True)
+    with pytest.raises(KeyError):
+      response['prompt_eval_cached_count']
+  else:
+    assert response['prompt_eval_cached_count'] == cached_count
+    assert response.model_dump(exclude_none=True)['prompt_eval_cached_count'] == cached_count