diff --git a/packages/sie_gateway/openapi.json b/packages/sie_gateway/openapi.json index e06f3e02e..ab2617b9e 100644 --- a/packages/sie_gateway/openapi.json +++ b/packages/sie_gateway/openapi.json @@ -727,6 +727,15 @@ "minimum": 0, "type": "integer" }, + "images": { + "description": "Number of input images observed by the worker after successful execution.", + "format": "int32", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, "prompt_tokens": { "format": "int32", "minimum": 0, @@ -1806,6 +1815,15 @@ "minimum": 0, "type": "integer" }, + "images": { + "description": "Number of input images observed by the worker after successful execution.", + "format": "int32", + "minimum": 1, + "type": [ + "integer", + "null" + ] + }, "prompt_tokens": { "format": "int32", "minimum": 0, diff --git a/packages/sie_gateway/src/handlers/proxy.rs b/packages/sie_gateway/src/handlers/proxy.rs index 98783cfdd..1808a013a 100644 --- a/packages/sie_gateway/src/handlers/proxy.rs +++ b/packages/sie_gateway/src/handlers/proxy.rs @@ -4449,7 +4449,7 @@ async fn queue_mode_streaming_generate( HeaderName::from_static("x-sie-server-version"), HeaderValue::from_static(GATEWAY_VERSION), ); - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( response.headers_mut(), model_revision, bundle_config_hash, @@ -6363,11 +6363,7 @@ fn build_chat_completion_body( // backend-config identifier (see `system_fingerprint`). Present and // identical in shape on both the blocking and streaming responses. "system_fingerprint": system_fingerprint(model), - "usage": { - "prompt_tokens": usage.prompt_tokens, - "completion_tokens": usage.completion_tokens, - "total_tokens": usage.total_tokens, - } + "usage": usage }); Ok(serde_json::to_vec(&body).unwrap_or_default()) } @@ -7323,7 +7319,7 @@ async fn proxy_chat_inner( HeaderName::from_static("x-sie-server-version"), HeaderValue::from_static(GATEWAY_VERSION), ); - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( response.headers_mut(), model_revision.as_deref(), &bundle_config_hash, @@ -7760,11 +7756,7 @@ fn build_text_completion_body( "finish_reason": map_chat_finish_reason(&outcome.finish_reason), }], "system_fingerprint": system_fingerprint(model), - "usage": { - "prompt_tokens": usage.prompt_tokens, - "completion_tokens": usage.completion_tokens, - "total_tokens": usage.total_tokens, - } + "usage": usage }); Ok(serde_json::to_vec(&body).unwrap_or_default()) } @@ -7969,7 +7961,7 @@ pub async fn proxy_completions(State(state): State>, req: Request) HeaderName::from_static("x-sie-version"), HeaderValue::from_static(GATEWAY_VERSION), ); - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( h, model_revision.as_deref(), &bundle_config_hash, @@ -8569,7 +8561,7 @@ pub async fn proxy_responses(State(state): State>, req: Request) - HeaderName::from_static("x-sie-version"), HeaderValue::from_static(GATEWAY_VERSION), ); - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( h, model_revision.as_deref(), &bundle_config_hash, @@ -9368,7 +9360,7 @@ fn insert_model_revision_header( } } -fn insert_stream_model_revision_header( +fn insert_buffered_generation_model_revision_header( headers: &mut HeaderMap, model_revision: Option<&str>, expected_bundle_config_hash: &str, @@ -9778,13 +9770,7 @@ pub(crate) fn build_generate_success_body_v2( outcome: &crate::queue::streaming::StreamOutcome, use_msgpack: bool, ) -> Vec { - let usage_value = outcome.usage.as_ref().map(|u| { - json!({ - "prompt_tokens": u.prompt_tokens, - "completion_tokens": u.completion_tokens, - "total_tokens": u.total_tokens, - }) - }); + let usage_value = outcome.usage.as_ref().map(|usage| json!(usage)); let mut body = serde_json::Map::new(); body.insert("model".to_string(), json!(model)); body.insert("text".to_string(), json!(outcome.text)); @@ -13935,7 +13921,7 @@ mod tests { execution_identity_sha256: None, execution_binding_sha256: None, }; - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( &mut headers, Some(revision), execution_hash.as_str(), @@ -13950,7 +13936,7 @@ mod tests { headers.clear(); stream_outcome.executed_bundle_config_hash = Some("b".repeat(64)); - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( &mut headers, Some(revision), execution_hash.as_str(), @@ -13959,7 +13945,7 @@ mod tests { assert!(headers.get("x-sie-model-revision").is_none()); stream_outcome.executed_bundle_config_hash = None; - insert_stream_model_revision_header( + insert_buffered_generation_model_revision_header( &mut headers, Some(revision), execution_hash.as_str(), @@ -14058,6 +14044,58 @@ mod tests { assert!(headers.get("x-sie-execution-identity-sha256").is_none()); } + #[tokio::test] + async fn buffered_generation_preserves_terminal_images_and_validated_binding() { + use crate::queue::streaming::{ChunkEnvelope, StreamCollector}; + for (identity, binding) in [ + (Some("a".repeat(64)), Some("b".repeat(64))), + (Some("a".repeat(64)), None), + (None, Some("b".repeat(64))), + (Some("bad".to_string()), Some("b".repeat(64))), + (Some("a".repeat(64)), Some("invalid".to_string())), + (Some("a".repeat(64)), Some("c".repeat(64))), + ] { + let (sender, _receiver) = tokio::sync::oneshot::channel(); + let mut collector = StreamCollector::new(sender, "test/model".into(), "default".into()); + for seq in 0..=1 { + let chunk: ChunkEnvelope = serde_json::from_value(json!({ + "kind": "chunk", "request_id": "request-1", "attempt_id": "attempt-1", + "seq": seq, "text_delta": "", "done": seq == 1, + "execution_identity_sha256": if seq == 0 { identity.clone() } else { Some("a".repeat(64)) }, + "execution_binding_sha256": if seq == 0 { binding.clone() } else { Some("b".repeat(64)) }, + "finish_reason": if seq == 1 { Some("stop") } else { None }, + "usage": if seq == 1 { Some(json!({ + "prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7, "images": 1 + })) } else { None }, + })).unwrap(); + collector.apply(chunk); + } + let outcome = collector.build_outcome().unwrap(); + let mut headers = HeaderMap::new(); + insert_stream_execution_binding_header(&mut headers, &outcome); + insert_stream_execution_identity_header(&mut headers, &outcome); + assert_eq!( + headers.contains_key("x-sie-execution-binding-sha256"), + identity.as_deref() == Some(&"a".repeat(64)) + && binding.as_deref() == Some(&"b".repeat(64)) + ); + assert_eq!( + headers.contains_key("x-sie-execution-identity-sha256"), + headers.contains_key("x-sie-execution-binding-sha256") + ); + for msgpack in [false, true] { + let bytes = build_generate_success_body_v2("test/model", &outcome, msgpack); + let body: serde_json::Value = if msgpack { + rmp_serde::from_slice(&bytes).unwrap() + } else { + serde_json::from_slice(&bytes).unwrap() + }; + assert_eq!(body["usage"]["images"], 1); + assert_eq!(body["usage"]["total_tokens"], 7); + } + } + } + #[test] fn test_execution_binding_header_requires_valid_unanimous_worker_proof() { let binding = "e".repeat(64); @@ -14568,6 +14606,7 @@ mod tests { text: "ok".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 1, total_tokens: 2, @@ -17680,6 +17719,7 @@ mod tests { text: "Hello world!".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 3, total_tokens: 8, @@ -22056,6 +22096,7 @@ mod tests { text: String::new(), finish_reason: "stop".to_string(), usage: Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 9, total_tokens: 14, @@ -22111,6 +22152,7 @@ mod tests { text: String::new(), finish_reason: "tool_calls".to_string(), usage: Some(UsageBlock { + images: None, prompt_tokens: 6, completion_tokens: 12, total_tokens: 18, @@ -22178,6 +22220,7 @@ mod tests { text: String::new(), finish_reason: "stop".to_string(), usage: Some(UsageBlock { + images: None, prompt_tokens: 3, completion_tokens: 4, total_tokens: 7, @@ -22560,6 +22603,7 @@ mod tests { text: "a continuation".to_string(), finish_reason: "length".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 4, completion_tokens: 16, total_tokens: 20, @@ -23183,6 +23227,7 @@ mod tests { text: "a joke".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -23240,6 +23285,7 @@ mod tests { text: "Hi there!".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 3, total_tokens: 8, @@ -23330,6 +23376,7 @@ mod tests { text: "Hi".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 1, total_tokens: 2, @@ -23367,6 +23414,7 @@ mod tests { text: "Hi".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 1, total_tokens: 2, @@ -23399,6 +23447,7 @@ mod tests { text: String::new(), finish_reason: "tool_calls".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 7, completion_tokens: 11, total_tokens: 18, @@ -23972,6 +24021,7 @@ mod tests { text: "Hi".to_string(), finish_reason: "stop".to_string(), usage: Some(crate::queue::streaming::UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 1, total_tokens: 2, diff --git a/packages/sie_gateway/src/handlers/sse.rs b/packages/sie_gateway/src/handlers/sse.rs index a9edc7069..ff82770b7 100644 --- a/packages/sie_gateway/src/handlers/sse.rs +++ b/packages/sie_gateway/src/handlers/sse.rs @@ -1348,11 +1348,7 @@ fn build_usage_only_chunk_event( _ => return None, }; let usage = chunk.usage.as_ref()?; - let mut usage_body = json!({ - "prompt_tokens": usage.prompt_tokens, - "completion_tokens": usage.completion_tokens, - "total_tokens": usage.total_tokens, - }); + let mut usage_body = json!(usage); merge_terminal_usage_extras(&mut usage_body, terminal_extras); Some(json!({ "id": id, @@ -1436,11 +1432,7 @@ fn build_generate_chunk_event(chunk: &ChunkEnvelope, terminal_extras: &[(String, obj.insert("finish_reason".to_string(), json!(fr)); } if let Some(u) = chunk.usage.as_ref() { - let mut usage = json!({ - "prompt_tokens": u.prompt_tokens, - "completion_tokens": u.completion_tokens, - "total_tokens": u.total_tokens, - }); + let mut usage = json!(u); merge_terminal_usage_extras(&mut usage, terminal_extras); obj.insert("usage".to_string(), usage); } @@ -1781,6 +1773,7 @@ mod tests { // whose usage is the count-so-far. finish_reason: "cancelled".to_string(), usage: Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens, total_tokens: 5 + completion_tokens, @@ -2023,6 +2016,7 @@ mod tests { let mut terminal = _terminal_chunk("error", None); terminal.seq = 42; terminal.usage = Some(UsageBlock { + images: None, prompt_tokens: 3, completion_tokens: 2, total_tokens: 5, @@ -2044,6 +2038,7 @@ mod tests { let mut terminal = _terminal_chunk("error", None); terminal.seq = 42; terminal.usage = Some(UsageBlock { + images: None, prompt_tokens: 3, completion_tokens: 2, total_tokens: 5, @@ -2725,6 +2720,7 @@ mod tests { let chunk = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 10, completion_tokens: 7, total_tokens: 17, @@ -3142,6 +3138,7 @@ mod tests { collector.apply(_terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 2, completion_tokens: 2, total_tokens: 4, @@ -3177,6 +3174,7 @@ mod tests { let terminal = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -3217,6 +3215,7 @@ mod tests { let terminal = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -3262,6 +3261,7 @@ mod tests { let terminal = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -3327,6 +3327,7 @@ mod tests { let terminal = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -3361,6 +3362,7 @@ mod tests { let terminal = _terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 7, total_tokens: 12, @@ -3493,6 +3495,7 @@ mod tests { collector.apply(_terminal_chunk( "stop", Some(UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 1, total_tokens: 2, diff --git a/packages/sie_gateway/src/openapi.rs b/packages/sie_gateway/src/openapi.rs index 6c2f62c95..7945e4d62 100644 --- a/packages/sie_gateway/src/openapi.rs +++ b/packages/sie_gateway/src/openapi.rs @@ -1979,6 +1979,10 @@ pub struct GenerateUsage { pub prompt_tokens: u32, pub completion_tokens: u32, pub total_tokens: u32, + /// Number of input images observed by the worker after successful execution. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(minimum = 1)] + pub images: Option, } #[derive(Debug, Serialize, Deserialize, ToSchema)] @@ -2288,6 +2292,10 @@ pub struct ChatCompletionUsage { pub prompt_tokens: u32, pub completion_tokens: u32, pub total_tokens: u32, + /// Number of input images observed by the worker after successful execution. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schema(minimum = 1)] + pub images: Option, } #[derive(Debug, Serialize, Deserialize, ToSchema)] diff --git a/packages/sie_gateway/src/queue/streaming.rs b/packages/sie_gateway/src/queue/streaming.rs index edd72a7ce..bed8d98a1 100644 --- a/packages/sie_gateway/src/queue/streaming.rs +++ b/packages/sie_gateway/src/queue/streaming.rs @@ -200,6 +200,8 @@ pub struct UsageBlock { pub completion_tokens: u32, #[serde(default)] pub total_tokens: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub images: Option, } #[derive(Debug, Clone, Deserialize, Serialize)] @@ -1097,6 +1099,10 @@ impl StreamCollector { Some(self.logprobs.clone()) }; + let complete_execution_evidence = self.execution_identity_consistent + && self.execution_binding_consistent + && self.execution_identity_sha256.is_some() + && self.execution_binding_sha256.is_some(); Some(StreamOutcome { text, finish_reason: meta.finish_reason.clone(), @@ -1110,12 +1116,10 @@ impl StreamCollector { logprobs, candidates: meta.candidates.clone(), executed_bundle_config_hash: meta.executed_bundle_config_hash.clone(), - execution_identity_sha256: self - .execution_identity_consistent + execution_identity_sha256: complete_execution_evidence .then(|| self.execution_identity_sha256.clone()) .flatten(), - execution_binding_sha256: self - .execution_binding_consistent + execution_binding_sha256: complete_execution_evidence .then(|| self.execution_binding_sha256.clone()) .flatten(), }) @@ -1194,6 +1198,7 @@ mod tests { finish_reason: if done { Some("stop".to_string()) } else { None }, usage: if done { Some(UsageBlock { + images: None, prompt_tokens: 5, completion_tokens: 3, total_tokens: 8, @@ -1616,6 +1621,7 @@ mod tests { collector.last_output_at = Some(first + std::time::Duration::from_millis(400)); collector.output_event_count = 2; collector.final_meta.as_mut().expect("terminal").usage = Some(UsageBlock { + images: None, prompt_tokens: 1, completion_tokens: 4, total_tokens: 5, diff --git a/packages/sie_sdk/src/sie_sdk/client/async_.py b/packages/sie_sdk/src/sie_sdk/client/async_.py index 1f6f3bf23..d4f80ff33 100644 --- a/packages/sie_sdk/src/sie_sdk/client/async_.py +++ b/packages/sie_sdk/src/sie_sdk/client/async_.py @@ -255,6 +255,9 @@ def _parse_generate_result_async( "completion_tokens": _coerce_token_count(usage.get("completion_tokens")), "total_tokens": _coerce_token_count(usage.get("total_tokens")), } + images = usage.get("images") + if isinstance(images, int) and not isinstance(images, bool) and 0 < images <= 2**32 - 1: + parsed_usage["images"] = images settled = settled_charge_from_usage(usage) if settled is not None: parsed_usage["credits_charged"], parsed_usage["rate_book_version"] = settled diff --git a/packages/sie_sdk/src/sie_sdk/client/sync.py b/packages/sie_sdk/src/sie_sdk/client/sync.py index a74909c6a..1fd93fd67 100644 --- a/packages/sie_sdk/src/sie_sdk/client/sync.py +++ b/packages/sie_sdk/src/sie_sdk/client/sync.py @@ -239,6 +239,9 @@ def _parse_generate_result( "completion_tokens": _coerce_token_count(usage.get("completion_tokens")), "total_tokens": _coerce_token_count(usage.get("total_tokens")), } + images = usage.get("images") + if isinstance(images, int) and not isinstance(images, bool) and 0 < images <= 2**32 - 1: + parsed_usage["images"] = images settled = settled_charge_from_usage(usage) if settled is not None: parsed_usage["credits_charged"], parsed_usage["rate_book_version"] = settled diff --git a/packages/sie_sdk/src/sie_sdk/types.py b/packages/sie_sdk/src/sie_sdk/types.py index 0e0beff56..b9cb4dc38 100644 --- a/packages/sie_sdk/src/sie_sdk/types.py +++ b/packages/sie_sdk/src/sie_sdk/types.py @@ -599,6 +599,7 @@ class GenerationUsage(TypedDict): prompt_tokens: int completion_tokens: int total_tokens: int + images: NotRequired[int] credits_charged: NotRequired[int] rate_book_version: NotRequired[str] @@ -727,6 +728,7 @@ class ChatUsage(TypedDict): prompt_tokens: int completion_tokens: int total_tokens: int + images: NotRequired[int] class ChatChoice(TypedDict, total=False): diff --git a/packages/sie_sdk/tests/client/test_generate.py b/packages/sie_sdk/tests/client/test_generate.py index a6873012f..2af84cc3f 100644 --- a/packages/sie_sdk/tests/client/test_generate.py +++ b/packages/sie_sdk/tests/client/test_generate.py @@ -19,8 +19,9 @@ import pytest from sie_sdk import SIEAsyncClient, SIEClient, SIEConnectionError from sie_sdk.client._shared import MODAL_CONTINUATION_MAX_HOPS, validate_generate_grammar -from sie_sdk.client.async_ import _AioResponse +from sie_sdk.client.async_ import _AioResponse, _parse_generate_result_async from sie_sdk.client.errors import ModelLoadingError, ProvisioningError, RequestError, ServerError +from sie_sdk.client.sync import _parse_generate_result def _ok_response(payload: dict, headers: dict[str, str] | None = None) -> MagicMock: @@ -1239,3 +1240,42 @@ def test_async_parser_raises_on_missing_or_non_string_required_fields(self, enve with pytest.raises(RequestError): _parse_generate_result_async(envelope) + + +@pytest.mark.parametrize("images", [None, 0, -1, True, "1", 1.5, 2**32]) +def test_generate_parsers_omit_invalid_image_usage(images: object) -> None: + for parse in (_parse_generate_result, _parse_generate_result_async): + result = parse({"model": "test/model", "text": "ok", "usage": {"images": images}}) + assert "images" not in result["usage"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_buffered_generate_preserves_image_usage_and_binding(asynchronous: bool) -> None: + body = _ok_envelope() + body["usage"]["images"] = 1 + headers = { + "content-type": "application/json", + "x-sie-execution-identity-sha256": "a" * 64, + "x-sie-execution-binding-sha256": "b" * 64, + "x-sie-units-images": "1", + } + if asynchronous: + source = AsyncMock(return_value=_aio_raw_resp(200, json.dumps(body).encode(), headers)) + session = MagicMock() + session.post = _make_session_post(source) + session.close = AsyncMock() + with patch("sie_sdk.client.async_.aiohttp.ClientSession", return_value=session): + async with SIEAsyncClient("http://localhost:8080") as client: + result = await client.generate("m", prompt="describe", max_new_tokens=8) + else: + response = _ok_response(body) + response.headers = headers + with patch("sie_sdk.client.sync.httpx.Client") as transport: + transport.return_value.post.return_value = response + with SIEClient("http://localhost:8080") as client: + result = client.generate("m", prompt="describe", max_new_tokens=8) + assert result["usage"]["images"] == 1 + assert result["request"]["usage"]["images"] == 1 + assert result["request"]["execution_identity_sha256"] == "a" * 64 + assert result["request"]["execution_binding_sha256"] == "b" * 64 diff --git a/packages/sie_server/openapi.json b/packages/sie_server/openapi.json index 6df18a1eb..58f96fc44 100644 --- a/packages/sie_server/openapi.json +++ b/packages/sie_server/openapi.json @@ -1411,6 +1411,20 @@ "minimum": 0.0, "title": "Total Tokens", "description": "Total prompt and generated tokens" + }, + "images": { + "anyOf": [ + { + "type": "integer", + "maximum": 4294967295.0, + "minimum": 1.0 + }, + { + "type": "null" + } + ], + "title": "Images", + "description": "Observed input images on successful image-conditioned generation; omitted for text-only requests" } }, "type": "object", diff --git a/packages/sie_server/src/sie_server/api/generate.py b/packages/sie_server/src/sie_server/api/generate.py index ac0ae3fda..d37ee5114 100644 --- a/packages/sie_server/src/sie_server/api/generate.py +++ b/packages/sie_server/src/sie_server/api/generate.py @@ -950,6 +950,8 @@ async def _stream_generate_events( "total_tokens": prompt_tokens + completion_tokens, }, } + if images and terminal_error is None and finish_reason not in {"error", "cancelled"}: + terminal["usage"]["images"] = len(images) if ttft_ms is not None: terminal["ttft_ms"] = ttft_ms if terminal_error is not None: @@ -1369,6 +1371,7 @@ async def generate( "prompt_tokens": result.prompt_tokens, "completion_tokens": result.completion_tokens, "total_tokens": result.prompt_tokens + result.completion_tokens, + **({"images": len(images)} if images else {}), }, } ) diff --git a/packages/sie_server/src/sie_server/processors/streaming.py b/packages/sie_server/src/sie_server/processors/streaming.py index 088684c4f..0691340d6 100644 --- a/packages/sie_server/src/sie_server/processors/streaming.py +++ b/packages/sie_server/src/sie_server/processors/streaming.py @@ -2031,6 +2031,8 @@ async def _stream_generate( # Dispatch a copy of that mapping so admission and execution cannot # drift through a second set of keyword-construction rules. gen_kwargs = dict(generation_parameters) + images = gen_kwargs.get("images") + image_count = len(images) if images else None # Worker-side phase boundary (#3136): everything between work receipt # and this adapter dispatch — deserialization, validation, model load, # chat-template render, context checks, and KV-admission wait — is the @@ -2384,6 +2386,7 @@ async def _next_chunk() -> GenerationChunk: finish_reason=chunk.finish_reason or "stop", prompt_tokens=chunk.prompt_tokens, completion_tokens=chunk.completion_tokens, + images=image_count if not terminal_has_error and chunk.finish_reason != "cancelled" else None, ttft_ms=_compute_ttft_ms(publish_at, first_text_at), error_code=(chunk.error_code or "inference_error") if terminal_has_error else None, error_message=( @@ -4067,6 +4070,7 @@ def _encode_chunk( finish_reason: FinishReason | None = None, prompt_tokens: int | None = None, completion_tokens: int | None = None, + images: int | None = None, ttft_ms: float | None = None, error_code: str | None = None, error_message: str | None = None, @@ -4105,6 +4109,8 @@ def _encode_chunk( "completion_tokens": int(completion_tokens or 0), "total_tokens": int((prompt_tokens or 0) + (completion_tokens or 0)), } + if images is not None: + payload["usage"]["images"] = images if ttft_ms is not None: payload["ttft_ms"] = ttft_ms if error_code is not None or error_message is not None: diff --git a/packages/sie_server/src/sie_server/types/openapi.py b/packages/sie_server/src/sie_server/types/openapi.py index 986bb869f..e648f8ac4 100644 --- a/packages/sie_server/src/sie_server/types/openapi.py +++ b/packages/sie_server/src/sie_server/types/openapi.py @@ -462,6 +462,12 @@ class GenerateUsageModel(BaseModel): prompt_tokens: int = Field(..., ge=0, description="Number of prompt tokens") completion_tokens: int = Field(..., ge=0, description="Number of generated tokens") total_tokens: int = Field(..., ge=0, description="Total prompt and generated tokens") + images: int | None = Field( + default=None, + ge=1, + le=(1 << 32) - 1, + description="Observed input images on successful image-conditioned generation; omitted for text-only requests", + ) class GenerateResponseModel(BaseModel): diff --git a/packages/sie_server/tests/api/test_generate.py b/packages/sie_server/tests/api/test_generate.py index b74be5857..59fdd76b0 100644 --- a/packages/sie_server/tests/api/test_generate.py +++ b/packages/sie_server/tests/api/test_generate.py @@ -467,6 +467,11 @@ async def render(_config: object, prompt: str, image_count: int) -> str: assert fake_adapter.last_call is not None assert fake_adapter.last_call["prompt"] == "Read the image" assert fake_adapter.last_call["images"] == [{"data": b"hello", "format": "png"}] + if stream: + chunks = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: {")] + assert chunks[-1]["usage"]["images"] == 1 + else: + assert response.json()["usage"]["images"] == 1 @pytest.mark.parametrize("enable_thinking", [False, True]) def test_native_image_prompt_uses_pinned_trusted_model_tokenizer( diff --git a/packages/sie_server/tests/api/test_generate_stream.py b/packages/sie_server/tests/api/test_generate_stream.py index de35f5b02..16fe0ff62 100644 --- a/packages/sie_server/tests/api/test_generate_stream.py +++ b/packages/sie_server/tests/api/test_generate_stream.py @@ -151,6 +151,7 @@ def _events( *, suppress_thinking: bool = False, preflight_result: GenerationPreflightResult | None = None, + images: list[dict[str, Any]] | None = None, ) -> AsyncIterator[str]: return _stream_generate_events( adapter, @@ -170,6 +171,7 @@ def _events( top_logprobs=3, suppress_thinking=suppress_thinking, preflight_result=preflight_result, + images=images, ) @@ -648,3 +650,19 @@ async def test_done_is_final_sse_line() -> None: ] raw = await _drain(_FakeAdapter(chunks)) assert raw[-1].strip() == "data: [DONE]" + + +@pytest.mark.parametrize("finish_reason", ["stop", "error", "cancelled"]) +async def test_stream_image_usage_requires_success(finish_reason: str) -> None: + adapter = _FakeAdapter( + [ + GenerationChunk( + text_delta="", done=True, finish_reason=finish_reason, prompt_tokens=5, completion_tokens=2 + ), + ] + ) + events, _ = _parse_sse([event async for event in _events(adapter, images=[{"data": b"image", "format": "png"}])]) + if finish_reason == "stop": + assert events[-1]["usage"]["images"] == 1 + else: + assert "images" not in events[-1]["usage"] diff --git a/packages/sie_server/tests/processors/test_streaming.py b/packages/sie_server/tests/processors/test_streaming.py index cad1b9743..d37663769 100644 --- a/packages/sie_server/tests/processors/test_streaming.py +++ b/packages/sie_server/tests/processors/test_streaming.py @@ -3951,3 +3951,78 @@ def encode(self, text, *, add_special_tokens): assert terminal["error"]["code"] == "context_exceeded" assert "~image_tokens (1280)" in terminal["error"]["message"] msg.ack.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("image_count", [0, 1, 2]) +@pytest.mark.parametrize("finish_reason", ["stop", "error", "cancelled"]) +async def test_terminal_image_usage_counts_executed_images_only( + monkeypatch: pytest.MonkeyPatch, image_count: int, finish_reason: str +) -> None: + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = "rendered prompt" + monkeypatch.setattr(StreamingProcessor, "_get_tokenizer", AsyncMock(return_value=tokenizer)) + adapter = _FakeGenAdapter( + [ + GenerationChunk(text_delta="ok", is_first=True), + GenerationChunk( + text_delta="", + done=True, + finish_reason=finish_reason, + prompt_tokens=5, + completion_tokens=2, + ), + ] + ) + nc = AsyncMock() + proc = StreamingProcessor(nc=nc, registry=_make_registry(adapter), worker_id="w1") + wi = _make_work_item( + generate={ + "messages": [ + { + "role": "user", + "content": "describe", + "images": [{"data": b"image", "format": "png"}] * image_count, + } + ], + "max_new_tokens": 8, + } + ) + await proc.process(_make_msg(wi), "test/model") + chunks = [msgpack.unpackb(call.args[1], raw=False) for call in nc.publish.call_args_list] + terminal = next(chunk for chunk in chunks if chunk.get("done")) + usage = terminal["usage"] + assert usage["prompt_tokens"] == 5 + assert usage["completion_tokens"] == 2 + if image_count and finish_reason == "stop": + assert usage["images"] == image_count + else: + assert "images" not in usage + assert all("images" not in chunk.get("usage", {}) for chunk in chunks if not chunk.get("done")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("emit_terminal", [False, True]) +async def test_image_generation_without_token_counts_does_not_synthesize_usage( + monkeypatch: pytest.MonkeyPatch, emit_terminal: bool +) -> None: + tokenizer = MagicMock() + tokenizer.apply_chat_template.return_value = "rendered prompt" + monkeypatch.setattr(StreamingProcessor, "_get_tokenizer", AsyncMock(return_value=tokenizer)) + script = [GenerationChunk(text_delta="ok", is_first=True)] + if emit_terminal: + script.append(GenerationChunk(text_delta="", done=True, finish_reason="stop")) + nc = AsyncMock() + proc = StreamingProcessor(nc=nc, registry=_make_registry(_FakeGenAdapter(script)), worker_id="w1") + wi = _make_work_item( + generate={ + "messages": [{"role": "user", "content": "describe", "images": [{"data": b"image", "format": "png"}]}], + "max_new_tokens": 8, + } + ) + + await proc.process(_make_msg(wi), "test/model") + + terminal = _terminal_chunk(nc) + assert terminal["done"] is True + assert "usage" not in terminal diff --git a/packages/sie_ts_sdk/src/internal/parsing.ts b/packages/sie_ts_sdk/src/internal/parsing.ts index fea72fe6f..56ace67e4 100644 --- a/packages/sie_ts_sdk/src/internal/parsing.ts +++ b/packages/sie_ts_sdk/src/internal/parsing.ts @@ -723,6 +723,7 @@ export function parseExtractResults(data: unknown[]): ExtractResult[] { } interface WireUsageBlock { + images?: unknown; prompt_tokens?: number; completion_tokens?: number; total_tokens?: number; @@ -793,6 +794,12 @@ function coerceTokenCount(v: unknown): number { return typeof v === "number" && Number.isFinite(v) ? Math.trunc(v) : 0; } +function isPositiveSafeInteger(value: unknown): value is number { + return ( + typeof value === "number" && Number.isSafeInteger(value) && value > 0 && value <= 0xffffffff + ); +} + export function parseGenerateResult(data: Record): GenerateResult { const wire = data as WireGenerateResult; if (typeof wire.model !== "string") { @@ -815,6 +822,7 @@ export function parseGenerateResult(data: Record): GenerateResu promptTokens: coerceTokenCount(usage.prompt_tokens), completionTokens: coerceTokenCount(usage.completion_tokens), totalTokens: coerceTokenCount(usage.total_tokens), + ...(isPositiveSafeInteger(usage.images) ? { images: usage.images } : {}), // #2434: the gateway merges the settled charge into this same block, so // rebuilding it field-by-field must carry the charge across. Absence // stays absence — a request that committed no debit gets neither key. diff --git a/packages/sie_ts_sdk/src/types.ts b/packages/sie_ts_sdk/src/types.ts index 14c21d7ad..eaec69798 100644 --- a/packages/sie_ts_sdk/src/types.ts +++ b/packages/sie_ts_sdk/src/types.ts @@ -962,6 +962,7 @@ export interface GenerationUsage { promptTokens: number; completionTokens: number; totalTokens: number; + images?: number; creditsCharged?: number; rateBookVersion?: string; } @@ -1277,6 +1278,7 @@ export interface ChatUsage { prompt_tokens: number; completion_tokens: number; total_tokens: number; + images?: number; credits_charged?: number; rate_book_version?: string; } diff --git a/packages/sie_ts_sdk/tests/internal.test.ts b/packages/sie_ts_sdk/tests/internal.test.ts index deed7fbc1..9cdb8b0aa 100644 --- a/packages/sie_ts_sdk/tests/internal.test.ts +++ b/packages/sie_ts_sdk/tests/internal.test.ts @@ -677,3 +677,17 @@ describe("parseGenerateResult usage coercion (BUG 13c)", () => { }); }); }); + +describe("authoritative generation image usage", () => { + it.each([1, 2])("preserves %i observed images", (images) => { + const result = parseGenerateResult({ model: "m", text: "ok", usage: { images } }); + expect(result.usage.images).toBe(images); + }); + it.each([undefined, null, 0, -1, 1.5, "1", true, Number.MAX_SAFE_INTEGER + 1])( + "omits missing or malformed image usage %s", + (images) => { + const result = parseGenerateResult({ model: "m", text: "ok", usage: { images } }); + expect(result.usage).not.toHaveProperty("images"); + }, + ); +});