types: preserve cached prompt token counts

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.'
 

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',

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