qdrant-edge 0.8.0

use crate::common::types::ScoreType;
#[cfg(feature = "api")]
use itertools::Itertools as _;
use crate::segment::data_types::load_profile::LoadProfile;
#[cfg(feature = "api")]
use crate::segment::data_types::vectors::NamedQuery;
use crate::segment::types::{Filter, SearchParams, WithPayloadInterface, WithVector};
#[cfg(feature = "api")]
use crate::segment::{data_types::vectors::VectorInternal, vector_storage::query::ContextPair};

use crate::shard::query::query_enum::QueryEnum;

/// DEPRECATED: Search method should be removed and replaced with `ShardQueryRequest`
#[derive(Clone, Debug, PartialEq)]
pub struct CoreSearchRequest {
    /// Every kind of query that can be performed on segment level
    pub query: QueryEnum,
    /// Look only for points which satisfies this conditions
    pub filter: Option<Filter>,
    /// Additional search params
    pub params: Option<SearchParams>,
    /// Max number of result to return
    pub limit: usize,
    /// Offset of the first result to return.
    /// May be used to paginate results.
    /// Note: large offset values may cause performance issues.
    pub offset: usize,
    /// Select which payload to return with the response. Default is false.
    pub with_payload: Option<WithPayloadInterface>,
    /// Options for specifying which vectors to include into response. Default is false.
    pub with_vector: Option<WithVector>,
    pub score_threshold: Option<ScoreType>,
}

impl CoreSearchRequest {
    /// Request-specific [`LoadProfile`] for opening a read-only shard to serve exactly
    /// this search: only the queried vector's components and the filter's field indexes
    /// keep their configured placement.
    pub fn load_profile(&self) -> LoadProfile {
        let Self {
            query,
            filter,
            params: _,
            limit: _,
            offset: _,
            with_payload,
            with_vector: _,
            score_threshold: _,
        } = self;

        // The `with_payload` default of a search is `false`.
        let with_payload = with_payload.as_ref().is_some_and(|wp| wp.is_required());

        LoadProfile::for_search(query.get_vector_name(), filter.as_ref(), with_payload)
    }

    pub fn search_rate_cost(&self) -> usize {
        let mut cost = self.query.search_cost();

        if let Some(filter) = &self.filter {
            cost += filter.total_conditions_count();
        }

        cost
    }
}

#[cfg(feature = "api")]
impl From<api::rest::SearchRequestInternal> for CoreSearchRequest {
    fn from(request: api::rest::SearchRequestInternal) -> Self {
        #[cfg(feature = "api")]
        use crate::segment::data_types::vectors::NamedVectorStruct;

        let api::rest::SearchRequestInternal {
            vector,
            filter,
            score_threshold,
            limit,
            offset,
            params,
            with_vector,
            with_payload,
        } = request;
        Self {
            query: QueryEnum::Nearest(NamedQuery::from(NamedVectorStruct::from(vector))),
            filter,
            params,
            limit,
            offset: offset.unwrap_or_default(),
            with_payload,
            with_vector,
            score_threshold,
        }
    }
}

#[cfg(feature = "api")]
impl TryFrom<api::grpc::qdrant::CoreSearchPoints> for CoreSearchRequest {
    type Error = tonic::Status;

    fn try_from(value: api::grpc::qdrant::CoreSearchPoints) -> Result<Self, Self::Error> {
        use crate::segment::data_types::vectors::VectorInternal;
        use crate::segment::vector_storage::query::{ContextQuery, DiscoverQuery, RecoQuery};

        let query = value
            .query
            .and_then(|query| query.query)
            .map(|query| {
                Ok(match query {
                    api::grpc::qdrant::query_enum::Query::NearestNeighbors(vector) => {
                        let vector_internal = VectorInternal::try_from(vector)?;
                        QueryEnum::Nearest(NamedQuery::from(
                            api::grpc::conversions::into_named_vector_struct(
                                value.vector_name,
                                vector_internal,
                            )?,
                        ))
                    }
                    api::grpc::qdrant::query_enum::Query::RecommendBestScore(query) => {
                        QueryEnum::RecommendBestScore(NamedQuery {
                            query: RecoQuery::try_from(query)?,
                            using: value.vector_name,
                        })
                    }
                    api::grpc::qdrant::query_enum::Query::RecommendSumScores(query) => {
                        QueryEnum::RecommendSumScores(NamedQuery {
                            query: RecoQuery::try_from(query)?,
                            using: value.vector_name,
                        })
                    }
                    api::grpc::qdrant::query_enum::Query::Discover(query) => {
                        let Some(target) = query.target else {
                            return Err(tonic::Status::invalid_argument("Target is not specified"));
                        };

                        let pairs = query
                            .context
                            .into_iter()
                            .map(try_context_pair_from_grpc)
                            .try_collect()?;

                        QueryEnum::Discover(NamedQuery {
                            query: DiscoverQuery::new(target.try_into()?, pairs),
                            using: value.vector_name,
                        })
                    }
                    api::grpc::qdrant::query_enum::Query::Context(query) => {
                        let pairs = query
                            .context
                            .into_iter()
                            .map(try_context_pair_from_grpc)
                            .try_collect()?;

                        QueryEnum::Context(NamedQuery {
                            query: ContextQuery::new(pairs),
                            using: value.vector_name,
                        })
                    }
                })
            })
            .transpose()?
            .ok_or_else(|| tonic::Status::invalid_argument("Query is not specified"))?;

        Ok(Self {
            query,
            filter: value.filter.map(|f| f.try_into()).transpose()?,
            params: value.params.map(TryInto::try_into).transpose()?,
            limit: value.limit as usize,
            offset: value.offset.unwrap_or_default() as usize,
            with_payload: value.with_payload.map(|wp| wp.try_into()).transpose()?,
            with_vector: Some(value.with_vectors.map(Into::into).unwrap_or_default()),
            score_threshold: value.score_threshold,
        })
    }
}

#[cfg(feature = "api")]
fn try_context_pair_from_grpc(
    pair: api::grpc::qdrant::ContextPair,
) -> Result<ContextPair<VectorInternal>, tonic::Status> {
    let api::grpc::qdrant::ContextPair { positive, negative } = pair;
    match (positive, negative) {
        (Some(positive), Some(negative)) => Ok(ContextPair {
            positive: positive.try_into()?,
            negative: negative.try_into()?,
        }),
        _ => Err(tonic::Status::invalid_argument(
            "All context pairs must have both positive and negative parts",
        )),
    }
}

#[cfg(feature = "api")]
impl TryFrom<api::grpc::qdrant::SearchPoints> for CoreSearchRequest {
    type Error = tonic::Status;

    fn try_from(value: api::grpc::qdrant::SearchPoints) -> Result<Self, Self::Error> {
        use crate::sparse::common::sparse_vector::validate_sparse_vector_impl;

        let api::grpc::qdrant::SearchPoints {
            collection_name: _,
            vector,
            filter,
            limit,
            with_payload,
            params,
            score_threshold,
            offset,
            vector_name,
            with_vectors,
            read_consistency: _,
            timeout: _,
            shard_key_selector: _,
            sparse_indices,
        } = value;

        if let Some(sparse_indices) = &sparse_indices {
            let api::grpc::qdrant::SparseIndices { data } = sparse_indices;
            validate_sparse_vector_impl(data, &vector).map_err(|e| {
                tonic::Status::invalid_argument(format!(
                    "Sparse indices does not match sparse vector conditions: {e}"
                ))
            })?;
        }

        let vector_internal =
            VectorInternal::from_vector_and_indices(vector, sparse_indices.map(|v| v.data));

        let vector_struct =
            api::grpc::conversions::into_named_vector_struct(vector_name, vector_internal)?;

        Ok(Self {
            query: QueryEnum::Nearest(NamedQuery::from(vector_struct)),
            filter: filter.map(Filter::try_from).transpose()?,
            params: params.map(SearchParams::try_from).transpose()?,
            limit: limit as usize,
            offset: offset.map(|v| v as usize).unwrap_or_default(),
            with_payload: with_payload
                .map(WithPayloadInterface::try_from)
                .transpose()?,
            with_vector: with_vectors.map(WithVector::from),
            score_threshold: score_threshold.map(|s| s as ScoreType),
        })
    }
}

#[derive(Debug, Clone)]
pub struct CoreSearchRequestBatch {
    pub searches: Vec<CoreSearchRequest>,
}