vector-store-flows 0.1.2

//! Vector store integration for [Flows.network](https://flows.network)
//!
//! # Quick Start
//!
//! The flow function below shows the procedure from
//! creating a collection to upserting points,
//! and then to searching points.
//!
//! ```
//! use std::collections::HashMap;
//!
//! use flowsnet_platform_sdk::logger;
//! use lambda_flows::{request_received, send_response};
//! use serde_json::{json, Value};
//!
//! use vector_store_flows::*;
//!
//! #[no_mangle]
//! #[tokio::main(flavor = "current_thread")]
//! pub async fn run() {
//!     logger::init();
//!     request_received(handler).await;
//! }
//!
//! async fn handler(_qry: HashMap<String, Value>, _body: Vec<u8>) {
//!     let collection_name = "test";
//!
//!     // Delete collection
//!     _ = delete_collection(collection_name).await;
//!
//!     // Create and get collection
//!     {
//!         let p = CollectionCreateParams { vector_size: 4 };
//!         if let Err(_) = create_collection(collection_name, &p).await {
//!             return;
//!         }
//!
//!         match collection_info(collection_name).await {
//!             Ok(ci) => {
//!                 log::debug!(
//!                     "There are {} vectors in collection `{}` just when created",
//!                     ci.points_count,
//!                     collection_name
//!                 );
//!             }
//!             Err(_) => {
//!                 return;
//!             }
//!         }
//!     }
//!
//!     // Upsert points
//!     {
//!         let p = vec![
//!             Point {
//!                 id: PointId::Num(1),
//!                 vector: vec![0.05, 0.61, 0.76, 0.74],
//!                 payload: Some(json!({
//!                     "city": "Berlin",
//!                     "country": "Germany",
//!                     "count": 1000000,
//!                     "square": 12.5,
//!                     "coords": {"lat": 1.0, "lon": 2.0},
//!                 })),
//!             },
//!             Point {
//!                 id: PointId::Num(2),
//!                 vector: vec![0.19, 0.81, 0.75, 0.11],
//!                 payload: Some(json!({
//!                     "city": ["Berlin", "London"],
//!                 })),
//!             },
//!             Point {
//!                 id: PointId::Num(3),
//!                 vector: vec![0.36, 0.55, 0.47, 0.94],
//!                 payload: Some(json!({
//!                     "city": ["Berlin", "Moscow"],
//!                 })),
//!             },
//!             Point {
//!                 id: PointId::Num(4),
//!                 vector: vec![0.18, 0.01, 0.85, 0.8],
//!                 payload: Some(json!({
//!                     "city": ["London", "Moscow"],
//!                 })),
//!             },
//!             Point {
//!                 id: PointId::Uuid(String::from("98a9a4b1-4ef2-46fb-8315-a97d874fe1d7")),
//!                 vector: vec![0.24, 0.18, 0.22, 0.44],
//!                 payload: Some(json!({
//!                     "count": [0],
//!                 })),
//!             },
//!             Point {
//!                 id: PointId::Uuid(String::from("f0e09527-b096-42a8-94e9-ea94d342b925")),
//!                 vector: vec![0.35, 0.08, 0.11, 0.44],
//!                 payload: None,
//!             },
//!         ];
//!
//!         if let Err(_) = upsert_points(collection_name, p).await {
//!             return;
//!         }
//!
//!         log::debug!("Points has been upserted.");
//!     }
//!
//!     // Search points
//!     {
//!         let p = PointsSearchParams {
//!             vector: vec![0.2, 0.1, 0.9, 0.7],
//!             limit: 3,
//!         };
//!
//!         match search_points(collection_name, &p).await {
//!             Ok(sp) => send_response(
//!                 200,
//!                 vec![(
//!                     String::from("content-type"),
//!                     String::from("text/html; charset=UTF-8"),
//!                 )],
//!                 serde_json::to_vec_pretty(&sp).unwrap(),
//!             ),
//!             Err(e) => send_response(
//!                 400,
//!                 vec![(
//!                     String::from("content-type"),
//!                     String::from("text/html; charset=UTF-8"),
//!                 )],
//!                 e.as_bytes().to_vec(),
//!             ),
//!         }
//!     }
//! }
//! ```
//!

use http_req::{
    request::{Method, Request},
    uri::Uri,
};
use lazy_static::lazy_static;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

lazy_static! {
    static ref VECTOR_STORE_API_PREFIX: String = String::from(
        std::option_env!("VECTOR_STORE_API_PREFIX")
            .unwrap_or("https://vector-store.flows.network/api")
    );
}

extern "C" {
    // Return the user id of the flows platform
    fn get_flows_user(p: *mut u8) -> i32;

    // Return the flow id
    fn get_flow_id(p: *mut u8) -> i32;
}

unsafe fn _get_flows_user() -> String {
    let mut flows_user = Vec::<u8>::with_capacity(100);
    let c = get_flows_user(flows_user.as_mut_ptr());
    flows_user.set_len(c as usize);
    String::from_utf8(flows_user).unwrap()
}

unsafe fn _get_flow_id() -> String {
    let mut flow_id = Vec::<u8>::with_capacity(100);
    let c = get_flow_id(flow_id.as_mut_ptr());
    if c == 0 {
        panic!("Failed to get flow id");
    }
    flow_id.set_len(c as usize);
    String::from_utf8(flow_id).unwrap()
}

/// The information of the collection.
/// A collection is a named set of points (vectors with a payload) among which you can search.
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CollectionInfo {
    /// Count of points in the collection
    pub points_count: u64,
}

/// The parameters for creating the collection
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CollectionCreateParams {
    /// Max size of vectors
    pub vector_size: u64,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(untagged)]
pub enum PointId {
    Uuid(String),
    Num(u64),
}

/// The point struct.
/// A point is a record consisting of a vector and an optional payload.
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Point {
    /// Id of the point
    pub id: PointId,

    /// Vectors
    pub vector: Vec<f32>,

    /// Additional information along with vectors
    pub payload: Option<Map<String, Value>>,
}

/// The parameters for searching for points
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct PointsSearchParams {
    /// Vectors
    pub vector: Vec<f32>,

    /// Max number of result to return
    pub limit: u64,
}

/// The point struct with the score returned by searching
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ScoredPoint {
    /// Id of the point
    pub id: PointId,

    /// Vectors
    pub vector: Option<Vec<f32>>,

    /// Additional information along with vectors
    pub payload: Option<Map<String, Value>>,

    /// Points vector distance to the query vector
    pub score: f32,
}

/// Get detailed information about specified existing collection.
pub async fn collection_info(collection_name: &str) -> Result<CollectionInfo, String> {
    unsafe {
        let flows_user = _get_flows_user();
        let flow_id = _get_flow_id();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/{}/collectionInfo",
            VECTOR_STORE_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            collection_name,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        match Request::new(&uri)
            .method(Method::GET)
            .header("Content-Type", "application/json")
            .send(&mut writer)
        {
            Ok(res) => {
                if res.status_code().is_success() {
                    serde_json::from_slice::<CollectionInfo>(&writer)
                        .or_else(|e| Err(e.to_string()))
                } else {
                    let err = String::from_utf8_lossy(&writer);
                    log::error!("{}", err);
                    Err(err.into_owned())
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// Create new collection with given parameters
pub async fn create_collection(
    collection_name: &str,
    params: &CollectionCreateParams,
) -> Result<(), String> {
    unsafe {
        let flows_user = _get_flows_user();
        let flow_id = _get_flow_id();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/{}/createCollection",
            VECTOR_STORE_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            collection_name,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&params).unwrap_or_default();

        match Request::new(&uri)
            .method(Method::PUT)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                if res.status_code().is_success() {
                    Ok(())
                } else {
                    let err = String::from_utf8_lossy(&writer);
                    log::error!("{}", err);
                    Err(err.into_owned())
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// Drop collection and all associated data
pub async fn delete_collection(collection_name: &str) -> Result<(), String> {
    unsafe {
        let flows_user = _get_flows_user();
        let flow_id = _get_flow_id();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/{}/deleteCollection",
            VECTOR_STORE_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            collection_name,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        match Request::new(&uri)
            .method(Method::DELETE)
            .header("Content-Type", "application/json")
            .send(&mut writer)
        {
            Ok(res) => {
                if res.status_code().is_success() {
                    Ok(())
                } else {
                    let err = String::from_utf8_lossy(&writer);
                    log::error!("{}", err);
                    Err(err.into_owned())
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// Perform insert + updates on points. If point with given ID already exists - it will be overwritten.
pub async fn upsert_points(collection_name: &str, points: Vec<Point>) -> Result<(), String> {
    unsafe {
        let flows_user = _get_flows_user();
        let flow_id = _get_flow_id();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/{}/upsertPoints",
            VECTOR_STORE_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            collection_name,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&points).unwrap_or_default();

        match Request::new(&uri)
            .method(Method::PUT)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                if res.status_code().is_success() {
                    Ok(())
                } else {
                    let err = String::from_utf8_lossy(&writer);
                    log::error!("{}", err);
                    Err(err.into_owned())
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}

/// Retrieve closest points based on vector similarity and given filtering conditions
pub async fn search_points(
    collection_name: &str,
    params: &PointsSearchParams,
) -> Result<Vec<ScoredPoint>, String> {
    unsafe {
        let flows_user = _get_flows_user();
        let flow_id = _get_flow_id();

        let mut writer = Vec::new();
        let uri = format!(
            "{}/{}/{}/{}/searchPoints",
            VECTOR_STORE_API_PREFIX.as_str(),
            flows_user,
            flow_id,
            collection_name,
        );
        let uri = Uri::try_from(uri.as_str()).unwrap();
        let body = serde_json::to_vec(&params).unwrap_or_default();

        match Request::new(&uri)
            .method(Method::POST)
            .header("Content-Type", "application/json")
            .header("Content-Length", &body.len())
            .body(&body)
            .send(&mut writer)
        {
            Ok(res) => {
                if res.status_code().is_success() {
                    serde_json::from_slice::<Vec<ScoredPoint>>(&writer)
                        .or_else(|e| Err(e.to_string()))
                } else {
                    let err = String::from_utf8_lossy(&writer);
                    log::error!("{}", err);
                    Err(err.into_owned())
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }
}