[Feature Request]: Vector query on specific set of documents ids

[ismail-mebsout]

Describe the problem

When querying a collection using an embedding vector, I am unable to set a list of Ids to query on

Describe the proposed solution

Include ids arguments in collection.query() function

Alternatives considered

No response

Importance

nice to have

Additional Information

No response

[tazarov]

@ismail-mebsout, thanks for this request. This is seems to be technically already supported in v1.0 Rust backend. However it is not yet enabled in the Python interfaces as evidenced by the None in L608 (second snippet).

chroma/rust/types/src/api_types.rs

Lines 1239 to 1250 in 5832679

	#[non_exhaustive]
	#[derive(Clone, Validate)]
	pub struct QueryRequest {
	pub tenant_id: String,
	pub database_name: String,
	pub collection_id: CollectionUuid,
	pub ids: Option<Vec<String>>,
	pub r#where: Option<Where>,
	pub embeddings: Vec<Vec<f32>>,
	pub n_results: u32,
	pub include: IncludeList,
	}

chroma/rust/python_bindings/src/bindings.rs

Lines 575 to 613 in 7569eff

	#[pyo3(
	signature = (collection_id, query_embeddings, n_results, r#where = None, where_document = None, include = ["metadatas".to_string(), "documents".to_string()].to_vec(), tenant = DEFAULT_TENANT.to_string(), database = DEFAULT_DATABASE.to_string())
	)]
	#[allow(clippy::too_many_arguments)]
	fn query(
	&self,
	collection_id: String,
	query_embeddings: Vec<Vec<f32>>,
	n_results: u32,
	r#where: Option<String>,
	where_document: Option<String>,
	include: Vec<String>,
	tenant: String,
	database: String,
	py: Python<'_>,
	) -> ChromaPyResult<QueryResponse> {
	let r#where = chroma_types::RawWhereFields::from_json_str(
	r#where.as_deref(),
	where_document.as_deref(),
	)?
	.parse()?;
	let collection_id = chroma_types::CollectionUuid(
	uuid::Uuid::parse_str(&collection_id).map_err(WrappedUuidError)?,
	);
	let include = IncludeList::try_from(include)?;
	let request = chroma_types::QueryRequest::try_new(
	tenant,
	database,
	collection_id,
	None,
	r#where,
	query_embeddings,
	n_results,
	include,
	)?;

[jairad26]

this has been completed in #4250