Merge 3660f1ce0c into f8f6a01cae

tests/e2e/conftest.py

@@ -37,6 +37,7 @@ from e2e_config import (
     PROXY_BASE_URL,
     SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
     SAGEMAKER_NOVA_OPT_IN_ENV,
+    VERTEX_AGENT_ENGINE_OPT_IN_ENV,
     SECRET_MANAGER_OPT_IN_ENV,
     WEEKLY_ANOMALY_OPT_IN_ENV,
     unique_marker,
@@ -77,6 +78,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
         "secret_manager": SECRET_MANAGER_OPT_IN_ENV,
         "sagemaker_nova": SAGEMAKER_NOVA_OPT_IN_ENV,
         "sagemaker_nova2_lite": SAGEMAKER_NOVA2_LITE_OPT_IN_ENV,
+        "vertex_agent_engine": VERTEX_AGENT_ENGINE_OPT_IN_ENV,
     }
 )
 
@@ -200,6 +202,11 @@ def pytest_configure(config: pytest.Config) -> None:
         "sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; "
         "deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it",
     )
+    config.addinivalue_line(
+        "markers",
+        "vertex_agent_engine: needs a deployed Vertex AI Agent Engine the proxy's Vertex credentials can query; "
+        "deselected unless VERTEX_AGENT_ENGINE_RESOURCE names its projects/.../reasoningEngines/... resource",
+    )
 
 
 def pytest_sessionstart(session: pytest.Session) -> None:

tests/e2e/coverage_registry/llm_conversational.yaml

@@ -47,6 +47,8 @@
 - {id: llm.chat_completions.cohere.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: cohere, capability: basic, streaming: nonstream, assertions: [works], source: "test_chat_completions_regression_e2e.py", rationale: "Cohere chat via OpenAI-compatible /chat/completions"}
 
 - {id: llm.chat_completions.vertex.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming over Vertex"}
+- {id: llm.chat_completions.vertex_agent_engine.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: vertex_agent_engine, capability: basic, streaming: nonstream, assertions: [works], source: "llms/vertex_ai/agent_engine/transformation.py", rationale: "A vertex_ai/agent_engine/<reasoning engine> deployment answers /chat/completions with text and finish_reason stop"}
+- {id: llm.chat_completions.vertex_agent_engine.basic.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: vertex_agent_engine, capability: basic, streaming: stream, assertions: [works], source: "llms/vertex_ai/agent_engine/sse_iterator.py", rationale: "A streamed Agent Engine reply yields deltas that concatenate to text and end with finish_reason stop"}
 - {id: llm.chat_completions.vertex.tool_use.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: tool_use, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vertex Gemini function_calling"}
 - {id: llm.chat_completions.vertex.web_search.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: vertex, capability: web_search, streaming: stream, assertions: [works], source: "llms/vertex_ai", rationale: "Vertex Gemini streams Google Search grounding metadata"}
 - {id: llm.chat_completions.vertex.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Gemini vision"}

tests/e2e/coverage_registry/llm_nonconversational.yaml

@@ -69,6 +69,9 @@
 - {id: llm.realtime.openai.passthrough.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: openai, capability: basic, streaming: stream, assertions: [works], fail_before_fix: proven, source: "test_passthrough_e2e.py", rationale: "A websocket upgrade on /openai_passthrough/v1/realtime is accepted and relayed to OpenAI; only HTTP routes were registered under the prefix, so realtime clients were refused with a 403 before a socket existed (GitHub issue #36088)"}
 - {id: llm.vector_stores.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: vector_stores, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.17 / LIT-4778", rationale: "Vector store create/list/retrieve/delete lifecycle"}
 - {id: llm.vector_stores.openai.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: vector_stores, route: openai, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.17 / LIT-4778", rationale: "Vector store search and invalid id errors"}
+- {id: llm.rag.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rag, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "proxy/rag_endpoints/endpoints.py", rationale: "A document sent to /v1/rag/ingest lands in a new OpenAI vector store and /v1/rag/query returns it in search_results and answers from it"}
+- {id: llm.rag.openai.rerank.nonstream.works, module: llm, tier: P1, subject_endpoint: rag, route: openai, capability: rerank, streaming: nonstream, assertions: [works], source: "rag/main.py", rationale: "/v1/rag/query with a Cohere rerank config returns rerank_results alongside the ingested document in search_results"}
+- {id: llm.rag.vertex.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: rag, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "rag/ingestion/vertex_ai_ingestion.py", rationale: "/v1/rag/ingest with custom_llm_provider vertex_ai creates a RAG corpus from server-side credentials and uploads the file to it"}
 - {id: llm.bedrock_native.bedrock_converse.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock native converse happy path"}
 - {id: llm.bedrock_native.bedrock_converse.basic.stream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock native converse-stream"}
 - {id: llm.bedrock_native.bedrock_converse.input_validation.nonstream.works, module: llm, tier: P1, subject_endpoint: bedrock_native, route: bedrock_converse, capability: input_validation, streaming: nonstream, assertions: [works], source: "vendor strategy §9.12 / LIT-4778", rationale: "Bedrock converse missing/empty messages and invalid model"}

tests/e2e/coverage_registry/schema.py

@@ -42,6 +42,7 @@ LlmEndpoint = Literal[
     "realtime",
     "google_native",
     "vector_stores",
+    "rag",
     "ocr",
     "bedrock_native",
 ]
@@ -62,6 +63,7 @@ LlmRoute = Literal[
     "sagemaker",
     "together_ai",
     "vertex",
+    "vertex_agent_engine",
     "xiaomi_mimo",
 ]
 
@@ -82,6 +84,7 @@ LlmCapability = Literal[
     "logprobs",
     "prompt_cache_1h",
     "prompt_cache_5m",
+    "rerank",
     "response_headers",
     "service_tier",
     "structured_output",

tests/e2e/e2e_config.py

@@ -210,6 +210,7 @@ OTEL_TLS_OPT_IN_ENV: Final = "E2E_OTEL_EXPORTER_ENDPOINT"
 SECRET_MANAGER_OPT_IN_ENV: Final = "E2E_SECRET_MANAGER"
 SAGEMAKER_NOVA_OPT_IN_ENV: Final = "SAGEMAKER_NOVA_ENDPOINT"
 SAGEMAKER_NOVA2_LITE_OPT_IN_ENV: Final = "SAGEMAKER_NOVA2_LITE_ENDPOINT"
+VERTEX_AGENT_ENGINE_OPT_IN_ENV: Final = "VERTEX_AGENT_ENGINE_RESOURCE"
 ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
 ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
 ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

tests/e2e/e2e_metadata.py

@@ -67,6 +67,7 @@ class Route(str, Enum):
     RERANK = "rerank"
     OCR = "ocr"
     VECTOR_STORES = "vector_stores"
+    RAG = "rag"
     REALTIME = "realtime"
     A2A = "a2a"
     USER_MANAGEMENT = "user_management"

tests/e2e/llm_translation/test_rag_e2e.py

@@ -0,0 +1,195 @@
+"""/v1/rag/ingest and /v1/rag/query through the gateway against real vector stores."""
+
+from __future__ import annotations
+
+import base64
+import re
+import time
+from typing import Final
+
+import pytest
+from e2e_config import CHEAP_OPENAI_MODEL, POLL_INTERVAL, POLL_TIMEOUT, unique_marker
+from e2e_http import NoBody, unwrap
+from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
+from lifecycle import ResourceManager
+from models import (
+    ChatMessage,
+    LiteLLMParamsBody,
+    OpenAIRagVectorStore,
+    RagIngestBody,
+    RagIngestOptions,
+    RagInlineFile,
+    RagQueryBody,
+    RagQueryResponse,
+    RagRerankConfig,
+    RagRetrievalConfig,
+    RagSearchHit,
+    RagVectorStore,
+    VectorStoreDeleteBody,
+    VectorStoreFileStatus,
+    VertexRagVectorStore,
+)
+from proxy_client import ProxyClient
+
+pytestmark = [pytest.mark.e2e, pytest.mark.provider_live]
+
+COHERE_RERANK_MODEL: Final = "cohere/rerank-v4.0-fast"
+VERTEX_CORPUS_ID: Final = re.compile(r"^projects/[^/]+/locations/[^/]+/ragCorpora/[^/]+$")
+
+
+def _document(codename: str) -> RagInlineFile:
+    text: Final = f"LiteLLM e2e RAG document.\nThe secret project codename is {codename}.\n"
+    return RagInlineFile(
+        filename=f"rag-{codename}.txt",
+        content=base64.b64encode(text.encode()).decode(),
+        content_type="text/plain",
+    )
+
+
+def _ingest(
+    proxy: ProxyClient, resources: ResourceManager, key: str, codename: str, vector_store: RagVectorStore
+) -> tuple[str, str]:
+    ingested: Final = unwrap(
+        proxy.rag_ingest(
+            key,
+            RagIngestBody(
+                file=_document(codename),
+                ingest_options=RagIngestOptions(name=f"e2e-rag-{codename}", vector_store=vector_store),
+            ),
+        )
+    )
+    if ingested.vector_store_id:
+        resources.defer(
+            lambda: proxy.transport.post(
+                "/vector_store/delete",
+                headers=proxy.transport.master,
+                json=VectorStoreDeleteBody(vector_store_id=ingested.vector_store_id),
+                response_type=NoBody,
+            )
+        )
+    assert ingested.status == "completed", f"ingest did not complete: {ingested}"
+    assert ingested.file_id, f"ingest returned no file id: {ingested}"
+    return ingested.vector_store_id, ingested.file_id
+
+
+def _ingest_into_openai_store(proxy: ProxyClient, resources: ResourceManager, key: str, codename: str) -> str:
+    store_id, file_id = _ingest(proxy, resources, key, codename, OpenAIRagVectorStore(ttl_days=1))
+    resources.defer(
+        lambda: proxy.transport.delete(
+            f"/v1/vector_stores/{store_id}", headers=proxy.transport.bearer(key), json=NoBody(), response_type=NoBody
+        )
+    )
+    resources.defer(
+        lambda: proxy.transport.delete(
+            f"/v1/files/{file_id}", headers=proxy.transport.bearer(key), json=NoBody(), response_type=NoBody
+        )
+    )
+    _await_indexed(proxy, key, store_id, file_id)
+    return store_id
+
+
+def _await_indexed(proxy: ProxyClient, key: str, store_id: str, file_id: str) -> None:
+    deadline: Final = time.monotonic() + POLL_TIMEOUT
+    while time.monotonic() < deadline:
+        indexed = unwrap(
+            proxy.transport.get(
+                f"/v1/vector_stores/{store_id}/files/{file_id}",
+                headers=proxy.transport.bearer(key),
+                params=NoBody(),
+                response_type=VectorStoreFileStatus,
+            )
+        )
+        if indexed.status == "completed":
+            return
+        assert indexed.status == "in_progress", f"file {file_id} stopped indexing in {store_id}: {indexed}"
+        time.sleep(POLL_INTERVAL)
+    raise AssertionError(f"file {file_id} was not indexed in {store_id} within {POLL_TIMEOUT}s")
+
+
+def _chat_model(proxy: ProxyClient, resources: ResourceManager) -> str:
+    model: Final = f"e2e-rag-{unique_marker()}"
+    model_id: Final = proxy.create_model(
+        model, LiteLLMParamsBody(model=f"openai/{CHEAP_OPENAI_MODEL}", api_key="os.environ/OPENAI_API_KEY")
+    )
+    resources.defer(lambda: proxy.delete_model(model_id))
+    return model
+
+
+def _question(store_id: str, model: str, rerank: RagRerankConfig | None = None) -> RagQueryBody:
+    return RagQueryBody(
+        model=model,
+        messages=(
+            ChatMessage(role="user", content="What is the secret project codename? Reply with only the codename."),
+        ),
+        retrieval_config=RagRetrievalConfig(vector_store_id=store_id, top_k=5),
+        rerank=rerank,
+    )
+
+
+def _retrieved_text(response: RagQueryResponse) -> str:
+    hits: Final = response.choices[0].message.provider_specific_fields.search_results.data
+    return " ".join(_hit_text(hit) for hit in hits)
+
+
+def _hit_text(hit: RagSearchHit) -> str:
+    return " ".join(part.text for part in hit.content)
+
+
+class TestRagOpenAI:
+    @pytest.mark.covers("llm.rag.openai.basic.nonstream.works")
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.RAG,
+            providers=(Provider.OPENAI,),
+            models=(CHEAP_OPENAI_MODEL,),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_query_answers_from_the_ingested_document(self, proxy: ProxyClient, resources: ResourceManager) -> None:
+        key: Final = resources.key()
+        codename: Final = f"heron-{unique_marker()}"
+        store_id: Final = _ingest_into_openai_store(proxy, resources, key, codename)
+
+        answer: Final = unwrap(proxy.rag_query(key, _question(store_id, _chat_model(proxy, resources))))
+
+        assert codename in _retrieved_text(answer), f"search_results do not carry the ingested document: {answer}"
+        assert codename in (answer.choices[0].message.content or ""), f"the answer ignored the context: {answer}"
+
+    @pytest.mark.covers("llm.rag.openai.rerank.nonstream.works")
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.RAG,
+            providers=(Provider.OPENAI, Provider.COHERE),
+            models=(CHEAP_OPENAI_MODEL, COHERE_RERANK_MODEL),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_query_reranks_the_ingested_document(self, proxy: ProxyClient, resources: ResourceManager) -> None:
+        key: Final = resources.key()
+        codename: Final = f"osprey-{unique_marker()}"
+        store_id: Final = _ingest_into_openai_store(proxy, resources, key, codename)
+
+        answer: Final = unwrap(
+            proxy.rag_query(
+                key,
+                _question(store_id, _chat_model(proxy, resources), RagRerankConfig(model=COHERE_RERANK_MODEL, top_n=3)),
+            )
+        )
+
+        fields: Final = answer.choices[0].message.provider_specific_fields
+        assert codename in _retrieved_text(answer), f"search_results do not carry the ingested document: {answer}"
+        assert fields.rerank_results and fields.rerank_results.results, f"no rerank_results returned: {answer}"
+        assert codename in (answer.choices[0].message.content or ""), f"the answer ignored the context: {answer}"
+
+
+class TestRagVertexAI:
+    @pytest.mark.covers("llm.rag.vertex.basic.nonstream.works")
+    @meta(Subject(domain=Domain.LLM_TRANSLATION, route=Route.RAG, providers=(Provider.VERTEX_AI,), mode=Mode.NONSTREAM))
+    def test_ingest_creates_a_corpus_holding_the_file(self, proxy: ProxyClient, resources: ResourceManager) -> None:
+        key: Final = resources.key()
+
+        corpus_id, _ = _ingest(proxy, resources, key, f"kestrel-{unique_marker()}", VertexRagVectorStore())
+
+        assert VERTEX_CORPUS_ID.match(corpus_id), f"ingest returned no Vertex corpus resource name: {corpus_id!r}"

tests/e2e/llm_translation/test_vertex_agent_engine_e2e.py

@@ -0,0 +1,78 @@
+"""/chat/completions through the gateway to a hosted Vertex AI Agent Engine (Reasoning Engine)."""
+
+from __future__ import annotations
+
+import os
+from typing import Final
+
+import pytest
+from e2e_config import VERTEX_AGENT_ENGINE_OPT_IN_ENV, unique_marker
+from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
+from lifecycle import ResourceManager
+from models import LiteLLMParamsBody
+from openai.types.chat import ChatCompletionMessageParam
+from proxy_client import ProxyClient
+from sdk_clients import SdkClients
+
+pytestmark = [pytest.mark.e2e, pytest.mark.vertex_agent_engine]
+
+
+def _agent_engine_deployment(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]:
+    resource: Final = os.environ[VERTEX_AGENT_ENGINE_OPT_IN_ENV]
+    location: Final = resource.split("/")[3]
+    model: Final = f"e2e-agent-engine-{unique_marker()}"
+    model_id: Final = proxy.create_model(
+        model,
+        LiteLLMParamsBody(
+            model=f"vertex_ai/agent_engine/{resource}",
+            vertex_location=location,
+            vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
+        ),
+    )
+    resources.defer(lambda: proxy.delete_model(model_id))
+    return model, resources.key()
+
+
+def _greeting() -> list[ChatCompletionMessageParam]:
+    return [{"role": "user", "content": f"Hello! What can you do? ({unique_marker()})"}]
+
+
+class TestVertexAgentEngine:
+    @pytest.mark.covers("llm.chat_completions.vertex_agent_engine.basic.nonstream.works")
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.CHAT_COMPLETIONS,
+            providers=(Provider.VERTEX_AI,),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_agent_replies(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None:
+        model, key = _agent_engine_deployment(proxy, resources)
+
+        response: Final = sdk.openai(key).chat.completions.create(model=model, messages=_greeting())
+
+        choice: Final = response.choices[0]
+        assert choice.message.content, f"{model}: the agent returned no text: {response}"
+        assert choice.finish_reason == "stop", f"{model}: unexpected finish_reason: {response}"
+
+    @pytest.mark.covers("llm.chat_completions.vertex_agent_engine.basic.stream.works")
+    @meta(
+        Subject(
+            domain=Domain.LLM_TRANSLATION,
+            route=Route.CHAT_COMPLETIONS,
+            providers=(Provider.VERTEX_AI,),
+            mode=Mode.STREAM,
+        )
+    )
+    def test_agent_streams_a_reply(self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients) -> None:
+        model, key = _agent_engine_deployment(proxy, resources)
+
+        chunks: Final = tuple(sdk.openai(key).chat.completions.create(model=model, messages=_greeting(), stream=True))
+
+        text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
+        finish_reasons: Final = tuple(chunk.choices[0].finish_reason for chunk in chunks if chunk.choices)
+        assert text.strip(), f"{model}: the streamed deltas carried no text across {len(chunks)} chunks"
+        assert finish_reasons and finish_reasons[-1] == "stop", (
+            f"{model}: stream did not end with stop: {finish_reasons}"
+        )

tests/e2e/models.py

@@ -857,6 +857,110 @@ class RerankResponse(BaseModel):
     results: list[RerankItem] = []
 
 
+# ---------- rag ----------
+
+
+class RagVectorStoreBase(BaseModel, frozen=True):
+    vector_store_id: str | None = None
+
+
+class OpenAIRagVectorStore(RagVectorStoreBase, frozen=True):
+    custom_llm_provider: Literal["openai"] = "openai"
+    ttl_days: int | None = None
+
+
+class VertexRagVectorStore(RagVectorStoreBase, frozen=True):
+    custom_llm_provider: Literal["vertex_ai"] = "vertex_ai"
+
+
+RagVectorStore = OpenAIRagVectorStore | VertexRagVectorStore
+
+
+class RagIngestOptions(BaseModel, frozen=True):
+    name: str
+    vector_store: RagVectorStore
+
+
+class RagInlineFile(BaseModel, frozen=True):
+    filename: str
+    content: str = Field(repr=False)
+    content_type: str
+
+
+class RagIngestBody(BaseModel, frozen=True):
+    file: RagInlineFile
+    ingest_options: RagIngestOptions
+
+
+class RagIngestResponse(BaseModel, frozen=True):
+    status: Literal["completed", "in_progress", "failed"]
+    vector_store_id: str
+    file_id: str | None = None
+    error: str | None = None
+
+
+class RagRetrievalConfig(BaseModel, frozen=True):
+    vector_store_id: str
+    custom_llm_provider: Literal["openai"] = "openai"
+    top_k: int
+
+
+class RagRerankConfig(BaseModel, frozen=True):
+    enabled: Literal[True] = True
+    model: str
+    top_n: int
+
+
+class RagQueryBody(BaseModel, frozen=True):
+    model: str
+    messages: tuple[ChatMessage, ...]
+    retrieval_config: RagRetrievalConfig
+    rerank: RagRerankConfig | None = None
+
+
+class RagSearchContent(BaseModel, frozen=True):
+    text: str = ""
+
+
+class RagSearchHit(BaseModel, frozen=True):
+    content: tuple[RagSearchContent, ...] = ()
+
+
+class RagSearchResults(BaseModel, frozen=True):
+    data: tuple[RagSearchHit, ...] = ()
+
+
+class RagRerankResults(BaseModel, frozen=True):
+    results: tuple[RerankItem, ...] = ()
+
+
+class RagProviderFields(BaseModel, frozen=True):
+    search_results: RagSearchResults
+    rerank_results: RagRerankResults | None = None
+
+
+class RagMessage(BaseModel, frozen=True):
+    content: str | None = None
+    provider_specific_fields: RagProviderFields
+
+
+class RagChoice(BaseModel, frozen=True):
+    message: RagMessage
+    finish_reason: str | None = None
+
+
+class RagQueryResponse(BaseModel, frozen=True):
+    choices: tuple[RagChoice, ...]
+
+
+class VectorStoreFileStatus(BaseModel, frozen=True):
+    status: Literal["in_progress", "completed", "failed", "cancelled"]
+
+
+class VectorStoreDeleteBody(BaseModel, frozen=True):
+    vector_store_id: str
+
+
 # ---------- ocr ----------
 
 

tests/e2e/proxy_client.py

@@ -83,6 +83,10 @@ from models import (
     ModelUpdateBody,
     OcrBody,
     OcrResponse,
+    RagIngestBody,
+    RagIngestResponse,
+    RagQueryBody,
+    RagQueryResponse,
     RerankBody,
     RerankResponse,
     ResponsesStreamBody,
@@ -1039,6 +1043,29 @@ class ProxyClient:
             response_type=RerankResponse,
         )
 
+    @step(
+        "Ingest {body.file.filename} with /v1/rag/ingest into a "
+        "{body.ingest_options.vector_store.custom_llm_provider} vector store"
+    )
+    def rag_ingest(self, key: str, body: RagIngestBody) -> Result[RagIngestResponse]:
+        return self.transport.post(
+            "/v1/rag/ingest",
+            headers=self.transport.bearer(key),
+            json=body,
+            response_type=RagIngestResponse,
+            timeout=SLOW_PROVIDER_TIMEOUT_SECONDS,
+        )
+
+    @step("Send a /v1/rag/query request to {body.model} over vector store {body.retrieval_config.vector_store_id}")
+    def rag_query(self, key: str, body: RagQueryBody) -> Result[RagQueryResponse]:
+        return self.transport.post(
+            "/v1/rag/query",
+            headers=self.transport.bearer(key),
+            json=body,
+            response_type=RagQueryResponse,
+            timeout=SLOW_PROVIDER_TIMEOUT_SECONDS,
+        )
+
     @step("Count tokens with /v1/messages/count_tokens for {body.model}")
     def count_tokens(self, key: str, body: CountTokensBody) -> Result[CountTokensResponse]:
         """POST /v1/messages/count_tokens (Anthropic-native). Sends the

tests/e2e/pytest.ini

@@ -21,3 +21,4 @@ markers =
     secret_manager: needs a proxy booted from gateway/secret_manager_<system>_ci_config.yml against that live secret manager; deselected unless E2E_SECRET_MANAGER names the backend (see secret_manager/secret_backends.py)
     sagemaker_nova: needs a deployed SageMaker Nova endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA_ENDPOINT names it
     sagemaker_nova2_lite: needs a deployed SageMaker Nova 2 Lite endpoint the proxy's AWS credentials can invoke; deselected unless SAGEMAKER_NOVA2_LITE_ENDPOINT names it
+    vertex_agent_engine: needs a deployed Vertex AI Agent Engine the proxy's Vertex credentials can query; deselected unless VERTEX_AGENT_ENGINE_RESOURCE names its projects/.../reasoningEngines/... resource