From 3660f1ce0cc74dc63029d269811026b56796bc84 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 12 Oct 2026 00:22:55 -0700 Subject: [PATCH] test(e2e): cover Vertex AI Agent Engine chat and RAG ingest/query through the proxy --- tests/e2e/conftest.py | 7 + .../coverage_registry/llm_conversational.yaml | 2 + .../llm_nonconversational.yaml | 3 + tests/e2e/coverage_registry/schema.py | 3 + tests/e2e/e2e_config.py | 1 + tests/e2e/e2e_metadata.py | 1 + tests/e2e/llm_translation/test_rag_e2e.py | 195 ++++++++++++++++++ .../test_vertex_agent_engine_e2e.py | 78 +++++++ tests/e2e/models.py | 104 ++++++++++ tests/e2e/proxy_client.py | 27 +++ tests/e2e/pytest.ini | 1 + 11 files changed, 422 insertions(+) create mode 100644 tests/e2e/llm_translation/test_rag_e2e.py create mode 100644 tests/e2e/llm_translation/test_vertex_agent_engine_e2e.py diff --git a/tests/e2e/conftest.py b/tests/e2e/conftest.py index 4477b9584aa..0a6947d3c5d 100644 --- a/tests/e2e/conftest.py +++ b/tests/e2e/conftest.py @@ -39,6 +39,7 @@ from e2e_config import ( REDIS_CHAOS_OPT_IN_ENV, 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, @@ -79,6 +80,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, } ) @@ -206,6 +208,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: diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index b11ae907cfb..cad95ec4aea 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/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/ 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"} diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index b05e9c916d2..55118764104 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -68,6 +68,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"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index abef3266aa8..081af0bdd9f 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/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", diff --git a/tests/e2e/e2e_config.py b/tests/e2e/e2e_config.py index fd5c0bae27b..0a57e76cee5 100644 --- a/tests/e2e/e2e_config.py +++ b/tests/e2e/e2e_config.py @@ -215,6 +215,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")) diff --git a/tests/e2e/e2e_metadata.py b/tests/e2e/e2e_metadata.py index a8eadb1637e..ea6c161305c 100644 --- a/tests/e2e/e2e_metadata.py +++ b/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" diff --git a/tests/e2e/llm_translation/test_rag_e2e.py b/tests/e2e/llm_translation/test_rag_e2e.py new file mode 100644 index 00000000000..b823b4f1f7a --- /dev/null +++ b/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}" diff --git a/tests/e2e/llm_translation/test_vertex_agent_engine_e2e.py b/tests/e2e/llm_translation/test_vertex_agent_engine_e2e.py new file mode 100644 index 00000000000..0fddba6412c --- /dev/null +++ b/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}" + ) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 42a59af45d2..2aabe9d5f06 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -861,6 +861,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 ---------- diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 23ab6487889..8ea54f1cc5a 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -83,6 +83,10 @@ from models import ( ModelUpdateBody, OcrBody, OcrResponse, + RagIngestBody, + RagIngestResponse, + RagQueryBody, + RagQueryResponse, RerankBody, RerankResponse, ResponsesStreamBody, @@ -1037,6 +1041,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 diff --git a/tests/e2e/pytest.ini b/tests/e2e/pytest.ini index ff81a593fe8..e1fcfa3bcd6 100644 --- a/tests/e2e/pytest.ini +++ b/tests/e2e/pytest.ini @@ -22,3 +22,4 @@ markers = secret_manager: needs a proxy booted from gateway/secret_manager__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