diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 0ab87472748..c3816928564 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -30,6 +30,28 @@ class MessagesRequest(BaseModel): messages: list[ChatMessage] +class CacheControl(BaseModel): + type: str = "ephemeral" + + +class TextBlock(BaseModel): + type: str = "text" + text: str + cache_control: CacheControl | None = None + + +class RichMessage(BaseModel): + role: str + content: list[TextBlock] + + +class RichMessagesRequest(BaseModel): + model: str + max_tokens: int = 64 + system: list[TextBlock] + messages: list[RichMessage] + + class EmbeddingsRequest(BaseModel): model: str input: str @@ -83,11 +105,19 @@ class AnthropicContentBlock(BaseModel): text: str | None = None +class MessagesUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + class MessagesResult(BaseModel): id: str | None = None role: str | None = None model: str | None = None content: list[AnthropicContentBlock] = [] + usage: MessagesUsage = MessagesUsage() @property def text(self) -> str: diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 51ae6ccbc1d..ba701f4869c 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -15,7 +15,7 @@ service_tier lives in test_provider_features_e2e.py. The provider-native cache_control request shape is not expressible with the shared ``ChatBody`` (whose content is a plain string), so the cacheable body is -modelled locally with typed content blocks. +built from the typed content blocks shared in ``endpoints_client.py``. """ from __future__ import annotations @@ -27,6 +27,7 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, unwrap +from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager from models import ChatResponse, LiteLLMParamsBody, Usage from passthrough_client import PassthroughClient @@ -38,21 +39,6 @@ BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" VERTEX_MODEL = "vertex_ai/gemini-2.5-flash" -class CacheControl(BaseModel): - type: str = "ephemeral" - - -class TextBlock(BaseModel): - type: str = "text" - text: str - cache_control: CacheControl | None = None - - -class RichMessage(BaseModel): - role: str - content: list[TextBlock] - - class CacheChatBody(BaseModel): model: str messages: list[RichMessage] diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py index be6cb67eb2d..f5da3744fab 100644 --- a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -28,7 +28,14 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, unwrap -from endpoints_client import AnthropicContentBlock, EndpointsClient +from endpoints_client import ( + CacheControl, + EndpointsClient, + MessagesResult, + RichMessage, + RichMessagesRequest, + TextBlock, +) from lifecycle import ResourceManager from models import LiteLLMParamsBody @@ -41,45 +48,6 @@ CACHE_PRIMING_DEADLINE_SECONDS = 60.0 CACHE_PRIMING_INTERVAL_SECONDS = 3.0 -class CacheControl(BaseModel): - type: str = "ephemeral" - - -class TextBlock(BaseModel): - type: str = "text" - text: str - cache_control: CacheControl | None = None - - -class MessageTurn(BaseModel): - role: str - content: list[TextBlock] - - -class MessagesBody(BaseModel): - model: str - max_tokens: int = 64 - system: list[TextBlock] - messages: list[MessageTurn] - - -class MessagesUsage(BaseModel): - input_tokens: int = 0 - output_tokens: int = 0 - cache_creation_input_tokens: int = 0 - cache_read_input_tokens: int = 0 - - -class MessagesCompletion(BaseModel): - role: str | None = None - content: list[AnthropicContentBlock] = [] - usage: MessagesUsage = MessagesUsage() - - @property - def text(self) -> str: - return "".join(block.text or "" for block in self.content) - - def _cacheable_system_block(marker: str) -> TextBlock: """A system prompt comfortably above Sonnet's 1024-token minimum cacheable size, unique per run so no other run's cache entry can satisfy the read.""" @@ -89,13 +57,13 @@ def _cacheable_system_block(marker: str) -> TextBlock: return TextBlock(text=text, cache_control=CacheControl()) -def _user_turn(text: str, *, cached: bool = False) -> MessageTurn: +def _user_turn(text: str, *, cached: bool = False) -> RichMessage: block = TextBlock(text=text, cache_control=CacheControl() if cached else None) - return MessageTurn(role="user", content=[block]) + return RichMessage(role="user", content=[block]) -def _system_reminder_turn() -> MessageTurn: - return MessageTurn( +def _system_reminder_turn() -> RichMessage: + return RichMessage( role="system", content=[ TextBlock( @@ -106,13 +74,13 @@ def _system_reminder_turn() -> MessageTurn: def _post_messages( - client: EndpointsClient, key: str, body: MessagesBody -) -> Result[MessagesCompletion]: + client: EndpointsClient, key: str, body: RichMessagesRequest +) -> Result[MessagesResult]: return client.gateway.transport.post( "/v1/messages", headers=client.gateway.transport.bearer(key), json=body, - response_type=MessagesCompletion, + response_type=MessagesResult, ) @@ -157,7 +125,7 @@ def _prime_prompt_cache( deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS while True: user_text = _first_turn_user_text(unique_marker()) - body = MessagesBody( + body = RichMessagesRequest( model=model, system=[system_block], messages=[_user_turn(user_text, cached=True)], @@ -193,13 +161,13 @@ class TestBedrockInvokeMidConversationSystem: primed = _prime_prompt_cache(endpoints_client, key, model, system_block) - reminder_turn_body = MessagesBody( + reminder_turn_body = RichMessagesRequest( model=model, system=[system_block], messages=[ _user_turn(primed.first_user_text, cached=True), _system_reminder_turn(), - MessageTurn(role="assistant", content=[TextBlock(text="OK.")]), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), _user_turn("Reply with one word again.", cached=True), ], ) @@ -230,13 +198,13 @@ class TestBedrockInvokeMidConversationSystem: ) key = resources.key() - body = MessagesBody( + body = RichMessagesRequest( model=model, system=[TextBlock(text="You are terse.")], messages=[ _user_turn(f"Say hi. Run {unique_marker()}."), _system_reminder_turn(), - MessageTurn(role="assistant", content=[TextBlock(text="Hi.")]), + RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), _user_turn("Say bye."), ], )