use std::sync::Arc;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use crate::auth::{ApiKeyAuth, Auth, AuthConfig, Target};
use crate::client::{ApiKeyPosition, ClientRequest, OramaClient};
use crate::error::Result;
use crate::types::*;
use crate::utils::create_random_string;
#[derive(Debug, Clone)]
pub struct OramaCoreManagerConfig {
pub url: String,
pub master_api_key: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct CreateCollectionParams {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub write_api_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub read_api_key: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub language: Option<Language>,
#[serde(skip_serializing_if = "Option::is_none")]
pub embeddings_model: Option<EmbeddingsModel>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct NewCollectionResponse {
pub id: String,
pub description: Option<String>,
#[serde(rename = "writeAPIKey")]
pub write_api_key: String,
#[serde(rename = "readonlyAPIKey")]
pub readonly_api_key: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CollectionIndexField {
pub field_id: String,
pub field_path: String,
pub is_array: bool,
pub field_type: serde_json::Value,
}
#[derive(Debug, Clone, Deserialize)]
pub struct CollectionIndex {
pub id: String,
pub document_count: u32,
pub fields: Vec<CollectionIndexField>,
pub automatically_chosen_properties: serde_json::Value,
}
#[derive(Debug, Clone, Deserialize)]
pub struct GetCollectionsResponse {
pub id: String,
pub description: Option<String>,
pub document_count: u32,
pub indexes: Vec<CollectionIndex>,
}
#[derive(Debug, Clone)]
pub struct CollectionNamespace {
client: OramaClient,
}
impl CollectionNamespace {
pub(crate) fn new(client: OramaClient) -> Self {
Self { client }
}
pub async fn create(&self, config: CreateCollectionParams) -> Result<NewCollectionResponse> {
let mut body = serde_json::json!({
"id": config.id,
"description": config.description,
"write_api_key": config.write_api_key.unwrap_or_else(|| create_random_string(32)),
"read_api_key": config.read_api_key.unwrap_or_else(|| create_random_string(32)),
});
if let Some(embeddings_model) = config.embeddings_model {
body["embeddings_model"] = serde_json::to_value(embeddings_model)?;
}
let request = ClientRequest::post(
"/v1/collections/create".to_string(),
Target::Writer,
ApiKeyPosition::Header,
body,
);
let response: serde_json::Value = self.client.request(request).await?;
Ok(NewCollectionResponse {
id: response["id"].as_str().unwrap_or_default().to_string(),
description: response["description"].as_str().map(|s| s.to_string()),
write_api_key: response["write_api_key"]
.as_str()
.unwrap_or_default()
.to_string(),
readonly_api_key: response["read_api_key"]
.as_str()
.unwrap_or_default()
.to_string(),
})
}
pub async fn list(&self) -> Result<Vec<GetCollectionsResponse>> {
let request = ClientRequest::<()>::get(
"/v1/collections".to_string(),
Target::Writer,
ApiKeyPosition::Header,
);
self.client.request(request).await
}
pub async fn get(&self, collection_id: &str) -> Result<GetCollectionsResponse> {
let request = ClientRequest::<()>::get(
format!("/v1/collections/{collection_id}"),
Target::Writer,
ApiKeyPosition::Header,
);
self.client.request(request).await
}
pub async fn delete(&self, collection_id: &str) -> Result<()> {
let body = serde_json::json!({
"collection_id_to_delete": collection_id
});
let request = ClientRequest::post(
"/v1/collections/delete".to_string(),
Target::Writer,
ApiKeyPosition::Header,
body,
);
let _: serde_json::Value = self.client.request(request).await?;
Ok(())
}
}
#[derive(Debug, Clone)]
pub struct OramaCoreManager {
pub collection: CollectionNamespace,
}
impl OramaCoreManager {
pub async fn new(config: OramaCoreManagerConfig) -> Result<Self> {
let auth_config =
AuthConfig::ApiKey(ApiKeyAuth::new(config.master_api_key).with_writer_url(config.url));
let client = Client::new();
let auth = Auth::new(auth_config, Arc::new(client));
let orama_client = OramaClient::new(auth)?;
Ok(Self {
collection: CollectionNamespace::new(orama_client),
})
}
}
impl CreateCollectionParams {
pub fn new<S: Into<String>>(id: S) -> Self {
Self {
id: id.into(),
description: None,
write_api_key: None,
read_api_key: None,
language: None,
embeddings_model: None,
}
}
pub fn with_description<S: Into<String>>(mut self, description: S) -> Self {
self.description = Some(description.into());
self
}
pub fn with_write_api_key<S: Into<String>>(mut self, key: S) -> Self {
self.write_api_key = Some(key.into());
self
}
pub fn with_read_api_key<S: Into<String>>(mut self, key: S) -> Self {
self.read_api_key = Some(key.into());
self
}
pub fn with_language(mut self, language: Language) -> Self {
self.language = Some(language);
self
}
pub fn with_embeddings_model(mut self, model: EmbeddingsModel) -> Self {
self.embeddings_model = Some(model);
self
}
}