Merge pull request #39920 from BerriAI/litellm_e2e_prompt_cache_cohere_passthrough_coverage

test(e2e): cover Anthropic and OpenAI prompt caching, Cohere embeddings, and costed /openai chat passthrough

tests/e2e/llm_translation/test_cache_control.py

@@ -10,6 +10,11 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix
   intentionally not covered here.
 - Vertex (gemini-2.5-flash): prompt caching via ``cache_control`` context
   caching; the second identical call must report cached prompt tokens > 0.
+- Anthropic (claude-haiku-4-5, direct): the same ``cache_control`` prefix over
+  the OpenAI-compatible route; the second call must report cache-read tokens > 0.
+- OpenAI (gpt-5.6): automatic prompt caching needs no request marker, so the
+  cacheable prefix goes out as a plain system string with a ``prompt_cache_key``
+  and the second call must report ``prompt_tokens_details.cached_tokens`` > 0.
 
 service_tier lives in test_provider_features_e2e.py.
 
@@ -21,6 +26,7 @@ built from the typed content blocks shared in ``endpoints_client.py``.
 from __future__ import annotations
 
 import time
+from collections.abc import Callable
 
 import pytest
 from pydantic import BaseModel
@@ -29,7 +35,7 @@ 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 models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody, Usage
 from passthrough_client import PassthroughClient
 import os
 
@@ -37,6 +43,8 @@ pytestmark = pytest.mark.e2e
 
 BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"
 VERTEX_MODEL = "vertex_ai/gemini-2.5-flash"
+ANTHROPIC_MODEL = "anthropic/claude-haiku-4-5-20251001"
+OPENAI_MODEL = "openai/gpt-5.6"
 
 
 class CacheChatBody(BaseModel):
@@ -89,17 +97,36 @@ def _cache_chat(
     )
 
 
+def _plain_cache_chat(
+    client: PassthroughClient, key: str, model: str, prefix: str, cache_key: str
+) -> Result[ChatResponse]:
+    """The same cacheable prefix as a plain system string, for providers that cache
+    automatically and take no per-block marker (OpenAI)."""
+    return client.proxy.chat(
+        key,
+        ChatBody(
+            model=model,
+            messages=[
+                ChatMessage(role="system", content=prefix),
+                ChatMessage(role="user", content="Reply with one word."),
+            ],
+            max_tokens=64,
+            prompt_cache_key=cache_key,
+        ),
+    )
+
+
 def _assert_cache_read_on_second_call(
-    client: PassthroughClient, key: str, model: str
+    model: str, send: Callable[[str], Result[ChatResponse]]
 ) -> None:
     prefix = _cacheable_prefix()
 
-    first = unwrap(_cache_chat(client, key, model, prefix))
+    first = unwrap(send(prefix))
     assert first.choices, f"{model}: first cache-priming call returned no choices: {first}"
 
     deadline = time.monotonic() + 30.0
     while True:
-        second = unwrap(_cache_chat(client, key, model, prefix))
+        second = unwrap(send(prefix))
         read_tokens = _cached_read_tokens(second.usage)
         if read_tokens > 0 or time.monotonic() >= deadline:
             break
@@ -125,7 +152,8 @@ class TestCacheControl:
             LiteLLMParamsBody(model=BEDROCK_MODEL, aws_region_name="us-east-1"),
         )
         resources.defer(lambda: client.proxy.delete_model(model_id))
-        _assert_cache_read_on_second_call(client, resources.key(), model)
+        key = resources.key()
+        _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix))
 
     @pytest.mark.covers(
         "llm.chat_completions.vertex.prompt_cache_5m.nonstream.works",
@@ -145,4 +173,40 @@ class TestCacheControl:
             ),
         )
         resources.defer(lambda: client.proxy.delete_model(model_id))
-        _assert_cache_read_on_second_call(client, resources.key(), model)
+        key = resources.key()
+        _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix))
+
+    @pytest.mark.covers(
+        "llm.chat_completions.anthropic.prompt_cache_5m.nonstream.works",
+        exercised_on=[],
+    )
+    def test_anthropic_prompt_caching_reads_cache(
+        self, client: PassthroughClient, resources: ResourceManager
+    ) -> None:
+        model = f"e2e-anthropic-cache-{unique_marker()}"
+        model_id = client.proxy.create_model(
+            model,
+            LiteLLMParamsBody(model=ANTHROPIC_MODEL, api_key="os.environ/ANTHROPIC_API_KEY"),
+        )
+        resources.defer(lambda: client.proxy.delete_model(model_id))
+        key = resources.key()
+        _assert_cache_read_on_second_call(model, lambda prefix: _cache_chat(client, key, model, prefix))
+
+    @pytest.mark.covers(
+        "llm.chat_completions.openai.prompt_cache_5m.nonstream.works",
+        exercised_on=[],
+    )
+    def test_openai_prompt_caching_reads_cache(
+        self, client: PassthroughClient, resources: ResourceManager
+    ) -> None:
+        model = f"e2e-openai-cache-{unique_marker()}"
+        model_id = client.proxy.create_model(
+            model,
+            LiteLLMParamsBody(model=OPENAI_MODEL, api_key="os.environ/OPENAI_API_KEY"),
+        )
+        resources.defer(lambda: client.proxy.delete_model(model_id))
+        key = resources.key()
+        cache_key = f"e2e-openai-cache-{unique_marker()}"
+        _assert_cache_read_on_second_call(
+            model, lambda prefix: _plain_cache_chat(client, key, model, prefix, cache_key)
+        )

tests/e2e/llm_translation/test_embeddings_endpoint_e2e.py

@@ -1,4 +1,4 @@
-"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex.
+"""Live e2e: POST /embeddings returns a real vector across OpenAI, Bedrock, Vertex, Cohere.
 
 Each test registers the deployment it needs at runtime (deleted on teardown) and
 asserts a non-empty, non-zero vector came back. The LIT-3167 guard in
@@ -86,6 +86,26 @@ class TestEmbeddingsEndpoint:
             f"embedding vector is all zeros: {result.body[:300]}"
         )
 
+    @pytest.mark.covers("llm.embeddings.cohere.basic.nonstream.works")
+    def test_cohere_embeddings_returns_vector(
+        self, endpoints_client: EndpointsClient, resources: ResourceManager
+    ) -> None:
+        model = f"e2e-embeddings-cohere-{unique_marker()}"
+        model_id = endpoints_client.create_model(
+            model,
+            LiteLLMParamsBody(model="cohere/embed-v4.0", api_key="os.environ/COHERE_API_KEY"),
+        )
+        resources.defer(lambda: endpoints_client.delete_model(model_id))
+        key = resources.key()
+
+        result = endpoints_client.embeddings(key, model, "Say this is a test!")
+        require_successful_call(result)
+        parsed = EmbeddingsResult.model_validate_json(result.body)
+        assert parsed.first_vector, f"/embeddings returned no vector: {result.body[:300]}"
+        assert any(component != 0.0 for component in parsed.first_vector), (
+            f"embedding vector is all zeros: {result.body[:300]}"
+        )
+
     @pytest.mark.covers("llm.embeddings.vertex.basic.nonstream.works")
     def test_vertex_embeddings_returns_vector(
         self, endpoints_client: EndpointsClient, resources: ResourceManager

tests/e2e/llm_translation/test_passthrough_e2e.py

@@ -17,7 +17,7 @@ import pytest
 from e2e_config import CHEAP_OPENAI_MODEL, unique_marker
 from e2e_http import require_successful_call, unwrap
 from lifecycle import ResourceManager
-from models import KeyGenerateBody, SpendLogRow
+from models import ChatResponse, KeyGenerateBody, SpendLogRow
 from passthrough_client import (
     AnthropicTool,
     GeminiFunctionDeclaration,
@@ -344,6 +344,44 @@ class TestOpenAIPassthroughSpend:
         )
 
 
+class TestOpenAIProviderPrefixChat:
+    """OpenAI-format chat through the raw `/openai/{endpoint}` passthrough (LIT-4752).
+
+    The body goes to OpenAI untranslated with the proxy's own OPENAI_API_KEY swapped
+    in, so the customer gets OpenAI's real completion back, and the gateway must
+    still write a costed pass_through_endpoint row for it.
+    """
+
+    @pytest.mark.covers("llm.chat_completions.openai.passthrough.nonstream.cost_logged")
+    def test_openai_prefix_chat_returns_completion_and_logs_its_cost(
+        self, client: PassthroughClient, scoped_key: str
+    ) -> None:
+        result = client.openai_chat(scoped_key, CHEAP_OPENAI_MODEL, f"Say hi in one word. {unique_marker()}")
+        require_successful_call(result)
+
+        completion = ChatResponse.model_validate_json(result.body)
+        assert completion.id, f"/openai/v1/chat/completions relayed no completion id: {result.body[:300]}"
+        content = (
+            completion.choices[0].message.content
+            if completion.choices and completion.choices[0].message
+            else None
+        )
+        assert content and content.strip(), (
+            f"/openai/v1/chat/completions relayed an empty completion: {result.body[:300]}"
+        )
+        assert completion.usage is not None, f"the completion carried no usage to price from: {completion}"
+
+        row = _fetch_cost_breakdown(client, completion.id)
+        assert row.prompt_tokens == completion.usage.prompt_tokens, (
+            f"logged {row.prompt_tokens} prompt tokens, the completion the customer read "
+            f"reported {completion.usage.prompt_tokens}"
+        )
+        assert row.completion_tokens == completion.usage.completion_tokens, (
+            f"logged {row.completion_tokens} completion tokens, the completion the customer read "
+            f"reported {completion.usage.completion_tokens}"
+        )
+
+
 class TestOpenAIPassthroughWebsocket:
     """The OpenAI passthrough prefixes must answer a websocket upgrade, not only a POST.