Repository navigation
feat(llm): add rerank support - #2135
Conversation
2d0b09b to
85ca9fd
Compare
| } | ||
|
|
||
| #[test] | ||
| fn test_bedrock_rerank_forwards_max_tokens_and_passthrough_via_additional_fields() { |
There was a problem hiding this comment.
can we use the standardized golden test file approach
There was a problem hiding this comment.
Done — moved the request/response translation tests to the golden harness in llm/tests.rs (requests/rerank/* and response/{bedrock,vertex,cohere}/rerank*.json fixtures with snapshots). rerank_tests.rs now only keeps the error-translation, empty-documents, and host-resolution tests that don't fit the fixture shape.
| ) | ||
| .map_err(ProxyError::Processing)?; | ||
|
|
||
| // `setup_request` rewrote only the request authority/Host header via |
There was a problem hiding this comment.
😕 this is not great. we already run apply_backend_policies at this point based on this info...
AWS hacks around this by nature of AWS signing happening later anyways but its not ideal.
Can we just have this sequence:
default_connector returns the policies (auth + tls)
let effective_policies = ... merge..
let route_type = resolve_route
let tgt = connector_with_router(route_type)
basically split the policiy from the target generation and move up the route type matching?
There was a problem hiding this comment.
Done — restructured per your sketch: default_connector is split into default_connector_policies() (auth + TLS) and default_connector_target(route_type). The AI backend arm now computes effective policies, resolves the route type from them, and only then picks the connection target. The post-hoc target override is gone.
howardjohn
left a comment
howardjohn
left a comment
There was a problem hiding this comment.
overall looks good, main thing is some minor testing comments and the target selection
| (strng::new("/v1/rerank"), crate::llm::RouteType::Rerank), | ||
| (strng::new("/v2/rerank"), crate::llm::RouteType::Rerank), |
There was a problem hiding this comment.
are v1 and v2 the same format?
There was a problem hiding this comment.
Practically yes for everything we model: the request/response shape (query/documents/top_n, results with index + relevance_score) is the same. v2 renamed v1's max_chunks_per_doc to max_tokens_per_doc and dropped object-form documents (we accept both forms, and unknown fields are preserved for passthrough). Cohere-compatible servers (Jina, vLLM, Voyage-style) still expose the shape at /v1/rerank, which is why both paths map to the same route type.
| // A /v1/rerank request flows through the proxy to a plain-HTTP wiremock upstream (a `custom` rerank | ||
| // provider has no default TLS connector, so no credentials are needed) and back in canonical shape. | ||
| #[tokio::test] | ||
| async fn test_rerank_end_to_end() -> anyhow::Result<()> { |
There was a problem hiding this comment.
I think this type of test should be in gateway_tests.rs or use a similar style. There is no need to exercise the flow this completely. The smoke test is just to have 1 basic test doing the entire end to end flow
There was a problem hiding this comment.
Done — dropped the smoke tests (smoke.rs is back to the single basic e2e) and replaced them with one llm_custom_rerank test in gateway_test.rs following the existing custom-provider style. Bedrock/Vertex translation fidelity is covered by the golden tests, and the real-provider path by the AGENTGATEWAY_E2E suite in tests/llm.rs.
Cohere-shaped /v1+/v2/rerank: passthrough for compatible backends, native translation for Bedrock and Vertex. Wired through the xDS proto and controller too, so it works via both local config and Kubernetes. Signed-off-by: apexlnc <[email protected]>
85ca9fd to
3750b3b
Compare
Substrate's egress policy carries a second credential source beside `credential_uri`: `CredentialHeader.actor_jwt` (field 4, `ActorJWTSource` with `audiences` and `expiration_seconds`), agent-substrate/substrate#1960. Substrate's own Envoy egress mints that JWT for the actor behind the tunnel (agentgateway#2135, agentgateway#2141). The agentgateway egress read every entry as an `ate-secret://` URI, so an `actor_jwt` entry was denied with `403 invalid credential URI: `. `apply_effects` now injects, for an `actor_jwt` entry whose header the request carries, `<prefix><token>`, the token minted through `Control.MintActorJwt` on the policy's own control API backend with the entry's audiences and lifetime. It needs no credential provider. Tokens are cached per actor, sorted audiences and lifetime, and reused until a third of their lifetime remains; concurrent misses share one mint. A request without the header is forwarded unchanged and mints nothing. A failed mint denies the request: 403 when the actor is gone (`NotFound`), 503 when the control API is unavailable or times out, 500 otherwise (new `ProxyError::SubstrateEgressInternal`), the classes of Substrate's `mapActorJWTError`. A minted token that is empty or not a valid header value is a 500. The vendored `ateapi.proto` gains `ActorJWTSource`, `MintActorJwt` and its messages, wire-compatible with Substrate's (audiences = 1, expiration_seconds = 8, actor = 5, as after agentgateway#2421). Signed-off-by: Timo Derstappen <[email protected]> (cherry picked from commit 3902417)
Summary
Adds a Cohere-shaped
/v1/rerank(and/v2/rerank) route that reranks documents against a query.bedrock-agent-runtime/discoveryengine, not the chat/embeddings host), SigV4 signing for Bedrock, and synthetic id↔index inversion for Vertex.Wired through every layer so it is configurable via both local files and the Kubernetes controller:
RouteType.RERANK+ProviderFormat.RERANKTesting
AGENTGATEWAY_E2E, plus credential-free wiremock translation tests for the passthrough/Bedrock/Vertex paths.amazon.rerank-v1:0).