feat: Integrate workspace setting (#642)

* chore: update workspace setting

* chore: update workspace setting with changes

* chore: fix test
This commit is contained in:
Nathan.fooo 2024-06-23 15:36:32 +08:00 committed by GitHub
parent c0937bfbeb
commit e33ef45d05
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 87 additions and 33 deletions

View File

@ -198,16 +198,18 @@ impl Display for EmbeddingsModel {
#[repr(u8)]
pub enum AIModel {
#[default]
GPT35 = 0,
GPT4o = 1,
Claude3Sonnet = 2,
Claude3Opus = 3,
Local = 4,
DefaultModel = 0,
GPT35 = 1,
GPT4o = 2,
Claude3Sonnet = 3,
Claude3Opus = 4,
Local = 5,
}
impl AIModel {
pub fn to_str(&self) -> &str {
match self {
AIModel::DefaultModel => "default-model",
AIModel::GPT35 => "gpt-3.5-turbo",
AIModel::GPT4o => "gpt-4o",
AIModel::Claude3Sonnet => "claude-3-sonnet-20240229",
@ -227,7 +229,7 @@ impl FromStr for AIModel {
"claude-3-sonnet" => Ok(AIModel::Claude3Sonnet),
"claude-3-opus" => Ok(AIModel::Claude3Opus),
"local" => Ok(AIModel::Local),
_ => Ok(AIModel::GPT35),
_ => Ok(AIModel::DefaultModel),
}
}
}

View File

@ -72,4 +72,4 @@ test_util = ["scraper"]
template = ["workspace-template"]
sync_verbose_log = ["collab-rt-protocol/verbose_log"]
test_fast_sync = []
enable_brotli = ["brotli"]
enable_brotli = ["brotli"]

View File

@ -4,6 +4,7 @@ use tracing::instrument;
use client_api_entity::AFWorkspaceSettings;
use shared_entity::response::{AppResponse, AppResponseError};
use crate::entity::AFWorkspaceSettingsChange;
use crate::http::log_request_id;
use crate::Client;
@ -32,8 +33,8 @@ impl Client {
pub async fn update_workspace_settings<T: AsRef<str>>(
&self,
workspace_id: T,
settings: &AFWorkspaceSettings,
) -> Result<(), AppResponseError> {
changes: &AFWorkspaceSettingsChange,
) -> Result<AFWorkspaceSettings, AppResponseError> {
let url = format!(
"{}/api/workspace/{}/settings",
self.base_url,
@ -42,10 +43,11 @@ impl Client {
let resp = self
.http_client_with_auth(Method::POST, &url)
.await?
.json(&settings)
.json(&changes)
.send()
.await?;
log_request_id(&resp);
AppResponse::<()>::from_response(resp).await?.into_error()
let resp = AppResponse::<AFWorkspaceSettings>::from_response(resp).await?;
resp.into_data()
}
}

View File

@ -511,10 +511,47 @@ pub struct AFWorkspace {
#[derive(Serialize, Deserialize)]
pub struct AFWorkspaces(pub Vec<AFWorkspace>);
#[derive(Default, Serialize, Deserialize)]
#[derive(Serialize, Deserialize)]
pub struct AFWorkspaceSettings {
#[serde(default)]
pub disable_indexing: bool,
pub disable_search_indexing: bool,
#[serde(default)]
pub ai_model: String,
}
impl Default for AFWorkspaceSettings {
fn default() -> Self {
Self {
disable_search_indexing: false,
ai_model: "".to_string(),
}
}
}
#[derive(Default, Serialize, Deserialize)]
pub struct AFWorkspaceSettingsChange {
#[serde(skip_serializing_if = "Option::is_none")]
pub disable_search_indexing: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ai_model: Option<String>,
}
impl AFWorkspaceSettingsChange {
pub fn new() -> Self {
Self {
disable_search_indexing: None,
ai_model: None,
}
}
pub fn disable_search_indexing(mut self, disable_search_indexing: bool) -> Self {
self.disable_search_indexing = Some(disable_search_indexing);
self
}
pub fn ai_model(mut self, ai_model: String) -> Self {
self.ai_model = Some(ai_model);
self
}
}
#[derive(Serialize, Deserialize)]

View File

@ -324,18 +324,18 @@ async fn post_workspace_settings_handler(
user_uuid: UserUuid,
state: Data<AppState>,
workspace_id: web::Path<Uuid>,
data: Json<AFWorkspaceSettings>,
) -> Result<JsonAppResponse<()>> {
data: Json<AFWorkspaceSettingsChange>,
) -> Result<JsonAppResponse<AFWorkspaceSettings>> {
let uid = state.user_cache.get_user_uid(&user_uuid).await?;
workspace::ops::update_workspace_settings(
let settings = workspace::ops::update_workspace_settings(
&state.pg_pool,
&state.workspace_access_control,
&workspace_id,
&uid,
&data.into_inner(),
data.into_inner(),
)
.await?;
Ok(AppResponse::Ok().into())
Ok(AppResponse::Ok().with_data(settings).into())
}
#[instrument(skip_all, err)]

View File

@ -1,4 +1,4 @@
use database_entity::dto::PublishCollabItem;
use database_entity::dto::{AFWorkspaceSettingsChange, PublishCollabItem};
use std::collections::HashMap;
use database_entity::dto::PublishInfo;
@ -517,8 +517,8 @@ pub async fn update_workspace_settings(
workspace_access_control: &impl WorkspaceAccessControl,
workspace_id: &Uuid,
owner_uid: &i64,
workspace_settings: &AFWorkspaceSettings,
) -> Result<(), AppResponseError> {
change: AFWorkspaceSettingsChange,
) -> Result<AFWorkspaceSettings, AppResponseError> {
let has_access = workspace_access_control
.enforce_role(owner_uid, &workspace_id.to_string(), AFRole::Owner)
.await?;
@ -531,9 +531,21 @@ pub async fn update_workspace_settings(
}
let mut tx = pg_pool.begin().await?;
upsert_workspace_settings(&mut tx, workspace_id, workspace_settings).await?;
let mut setting = select_workspace_settings(tx.deref_mut(), workspace_id)
.await?
.unwrap_or_default();
if let Some(disable_indexing) = change.disable_search_indexing {
setting.disable_search_indexing = disable_indexing;
}
if let Some(ai_model) = change.ai_model {
setting.ai_model = ai_model;
}
// Update the workspace settings in the database
upsert_workspace_settings(&mut tx, workspace_id, &setting).await?;
tx.commit().await?;
Ok(())
Ok(setting)
}
async fn check_workspace_owner(

View File

@ -1,7 +1,7 @@
use app_error::ErrorCode;
use client_api::Client;
use client_api_test::generate_unique_registered_user_client;
use database_entity::dto::{AFRole, AFWorkspaceInvitationStatus, AFWorkspaceSettings};
use database_entity::dto::{AFRole, AFWorkspaceInvitationStatus, AFWorkspaceSettingsChange};
use shared_entity::dto::workspace_dto::WorkspaceMemberInvitation;
use uuid::Uuid;
@ -13,17 +13,20 @@ async fn get_and_set_workspace_by_owner() {
let mut settings = c.get_workspace_settings(&workspace_id).await.unwrap();
assert!(
!settings.disable_indexing,
!settings.disable_search_indexing,
"indexing should be enabled by default"
);
settings.disable_indexing = true;
c.update_workspace_settings(&workspace_id, &settings)
.await
.unwrap();
settings.disable_search_indexing = true;
c.update_workspace_settings(
&workspace_id,
&AFWorkspaceSettingsChange::new().disable_search_indexing(true),
)
.await
.unwrap();
let settings = c.get_workspace_settings(&workspace_id).await.unwrap();
assert!(settings.disable_indexing);
assert!(settings.disable_search_indexing);
}
#[tokio::test]
@ -48,9 +51,7 @@ async fn get_and_set_workspace_by_non_owner() {
let resp = bob_client
.update_workspace_settings(
&alice_workspace_id.to_string(),
&AFWorkspaceSettings {
disable_indexing: true,
},
&AFWorkspaceSettingsChange::new().disable_search_indexing(true),
)
.await;
assert!(