Conversation
94b104d to
38093a0
Compare
38093a0 to
dd4d8ad
Compare
📝 WalkthroughWalkthroughRefactors inference authentication to support external API keys alongside the existing token. Adds a new Estimated code review effort🎯 3 (Moderate) | ⏱️ ~30 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
src/actix/api/update_api.rs (1)
331-360:update_batchis missing API key extraction.The
update_batchhandler does not extract or propagateApiKeys, unlikeupsert_pointsandupdate_vectorsin this file, and unlike the gRPCupdate_batchhandler insrc/tonic/api/points_api.rs(lines 290-292) which does include API key extraction. Since batch operations can include upsert operations that may require external API keys for inference, this appears to be an oversight.Suggested fix
#[post("/collections/{name}/points/batch")] +#[allow(clippy::too_many_arguments)] async fn update_batch( dispatcher: web::Data<Dispatcher>, collection: Path<CollectionPath>, operations: Json<UpdateOperations>, params: Query<UpdateParams>, service_config: web::Data<ServiceConfig>, ActixAccess(access): ActixAccess, + api_keys: ApiKeys, inference_token: InferenceToken, ) -> impl Responder { let operations = operations.into_inner(); let request_hw_counter = get_request_hardware_counter( &dispatcher, collection.name.clone(), service_config.hardware_reporting(), Some(params.wait), ); - let inference_params = InferenceParams::new(inference_token.clone(), params.timeout); + let inference_params = + InferenceParams::new(inference_token, params.timeout).with_ext_api_keys(api_keys); let timing = Instant::now();src/tonic/api/points_api.rs (1)
668-670: Add API key extraction toquery_groupshandler.The
query_groupshandler at lines 668-670 is missing the external API key extraction that bothqueryandquery_batchhandlers include. Add:let api_keys = extract_api_key(request.metadata());and chain it to InferenceParams:
let inference_params = InferenceParams::new(inference_token, timeout).with_ext_api_keys(api_keys);src/actix/api/query_api.rs (1)
109-132: Inconsistent API keys support across query endpoints.
query_points_batchdoes not extract or forward API keys, whilequery_pointsdoes. If external API keys are needed for inference during query operations, this endpoint should also support them for consistency.The same issue applies to
query_points_groups(lines 196-221).Suggested fix for query_points_batch
#[post("/collections/{name}/points/query/batch")] +#[allow(clippy::too_many_arguments)] async fn query_points_batch( dispatcher: web::Data<Dispatcher>, collection: Path<CollectionPath>, request: Json<QueryRequestBatch>, params: Query<ReadParams>, service_config: web::Data<ServiceConfig>, ActixAccess(access): ActixAccess, + api_keys: ApiKeys, inference_token: InferenceToken, ) -> impl Responder { ... - let inference_params = InferenceParams::new(inference_token, params.timeout()); + let inference_params = + InferenceParams::new(inference_token, params.timeout()).with_ext_api_keys(api_keys);
🤖 Fix all issues with AI agents
In `@src/common/inference/ext_api_keys.rs`:
- Around line 56-67: The current impl From<ApiKeys> for
reqwest::header::HeaderMap panics on invalid header names/values due to expect()
in the from(api_keys: ApiKeys) method; change this to a fallible conversion by
implementing TryFrom<ApiKeys> for reqwest::header::HeaderMap (or add a new
ApiKeys::to_header_map() -> Result<HeaderMap, Error>) so invalid
HeaderName::from_str / HeaderValue::from_str errors are returned instead of
panicking; update callers to handle the Result, or alternatively if you prefer
best-effort behavior, replace the expect() calls inside from() with match/if let
to skip invalid entries and optionally log a warning (reference ApiKeys,
api_keys.keys, and the from method when making the change).
🧹 Nitpick comments (3)
src/common/inference/service.rs (1)
218-223: Prefer explicit type conversion over.into().As per coding guidelines, prefer explicit
SomeType::from(x)over implicitx.into()in Rust.Suggested change
let mut request = request.json(&request_body); if let Some(api_keys) = ext_api_keys { - request = request.headers(api_keys.into()); + request = request.headers(reqwest::header::HeaderMap::from(api_keys)); }src/common/inference/ext_api_keys.rs (2)
35-37: Consider avoiding the clone of MetadataMap.
metadata.clone().into_headers()allocates a full copy of the metadata map. If performance is a concern for high-throughput gRPC calls, consider iterating directly over the metadata entries instead.Alternative implementation without clone
pub fn from_grpc_metadata(metadata: &tonic::metadata::MetadataMap) -> Self { - Self::from_headers(&metadata.clone().into_headers().into()) + let mut api_keys = Self::default(); + for key_and_value in metadata.iter() { + if let tonic::metadata::KeyAndValueRef::Ascii(key, value) = key_and_value { + let k = key.as_str(); + if k.ends_with(EMBEDDING_API_KEY_HEADER_SUFFIX) { + if let Ok(v) = value.to_str() { + api_keys.keys.insert(k.to_string(), v.to_string()); + } + } + } + } + api_keys }
52-54: Consider removing this thin wrapper function.
extract_api_keysimply delegates toApiKeys::from_grpc_metadata. Unless there's a specific reason for this indirection (e.g., future extension), calling the method directly would be more idiomatic.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@src/common/inference/api_keys.rs`:
- Around line 114-126: The From<InferenceApiKeys> impl for
reqwest::header::HeaderMap currently uses expect() on HeaderName::from_str and
HeaderValue::from_str which will panic on invalid user input; modify the
conversion to handle invalid headers gracefully by either changing the API to
return a Result<reqwest::header::HeaderMap, E> from a new TryFrom/TryInto
implementation for InferenceApiKeys (propagate parse errors) or by
filtering/logging invalid entries inside fn from (skip entries where
HeaderName::from_str or HeaderValue::from_str returns Err and record a warning),
updating usage sites accordingly; reference the impl block for
From<InferenceApiKeys>, the from(fn) method, and the api_keys.keys iteration to
locate the code to change.
🧹 Nitpick comments (1)
src/actix/api/update_api.rs (1)
326-327: Minor: Inconsistent attribute ordering.The attribute ordering here (
#[allow]before#[post]) differs fromupsert_pointsandupdate_vectors, which both place the route attribute (#[put]/#[post]) before#[allow(clippy::too_many_arguments)]. Consider reordering for consistency:🔧 Suggested reorder for consistency
-#[allow(clippy::too_many_arguments)] #[post("/collections/{name}/points/batch")] +#[allow(clippy::too_many_arguments)] async fn update_batch(
* feat: forward api keys header * fix: linter * feat: add api_keys to grpc * fix: imports * fix: extract api keys from grpc headers * refactor: make apikeys extracor more compact * fix: linter issue * feat: update header key logic * fix: clippy issues * make external keys a part of InferenceParams::new argument, so we wont forget using it * avoid cloning metadata on header extraction * combine inference token with ext api key for simpler usage * clippy --------- Co-authored-by: generall <andrey@vasnetsov.com>
All Submissions:
devbranch. Did you create your branch fromdev?New Feature Submissions:
cargo +nightly fmt --allcommand prior to submission?cargo clippy --workspace --all-featurescommand?Changes to Core Features:
Before
After