genius-core-client 0.4.0

use serde::de::DeserializeOwned;

use crate::client::InternalClient;
use crate::{
    retry_await,
    types::{entity::HSMLEntity, error::HstpError},
};

#[derive(serde::Deserialize, Debug)]
pub struct QueryEntitiesReturn {
    pub entities: Vec<HSMLEntity>,
    pub distances: Option<Vec<f32>>,
}

pub async fn query_t<T: DeserializeOwned>(
    client: &mut InternalClient,
    query: String,
    retries: u32,
) -> Result<T, HstpError> {
    let query = kortex_gen_grpc::hstp::v1::QueryRequest { query };

    let response = retry_await!(retries, client.query(query.clone()))
        .map_err(|e| {
            HstpError::new(
                kortex_gen_grpc::hstp::v1::ErrorCode::UnhandledError,
                format!(
                    "Failed to query entities - gRPC error code {}: {:#?}",
                    e.code() as u8,
                    e
                ),
                "".into(),
            )
        })?
        .into_inner();

    match response.response {
        Some(kortex_gen_grpc::hstp::v1::query_response::Response::Result(entity_array)) => {
            let entities: Result<T, _> = serde_json::from_str(&entity_array);
            match entities {
                Ok(entities) => Ok(entities),
                Err(e) => Err(HstpError::new(
                    kortex_gen_grpc::hstp::v1::ErrorCode::UnhandledError,
                    format!(
                        "Error parsing entities, query may have returned the wrong shape: {}",
                        e
                    ),
                    "".into(),
                )),
            }
        }
        Some(kortex_gen_grpc::hstp::v1::query_response::Response::Error(error)) => {
            Err(error.into())
        }
        _ => Err(HstpError::new(
            kortex_gen_grpc::hstp::v1::ErrorCode::UnhandledError,
            "Unknown response".to_string(),
            "".to_string(),
        )),
    }
}