Repository: TabbyML/tabby
File: ee/tabby-schema/src/schema/mod.rs (Rust)

Code before:
            },
        )
        .await
    }

    /// List user groups.
    async fn user_groups(ctx: &Context) -> Result<Vec<UserGroup>> {
        check_user(ctx).await?;
        ctx.locator.user_group().list().await
    }

    async fn source_id_access_policies(
        ctx: &Context,
        source_id: String,
    ) -> Result<SourceIdAccessPolicy> {
        check_admin(ctx).await?;
        let read = ctx
            .locator
            .access_policy()
            .list_source_id_read_access(&source_id)
            .await?;

        Ok(SourceIdAccessPolicy { source_id, read })
    }

    async fn test_model_connection(ctx: &Context, backend: ModelHealthBackend) -> Result<i32> {
        check_user_allow_auth_token(ctx).await?;
        // count request time in milliseconds
        let start = Instant::now();
        match backend {
            ModelHealthBackend::Completion => {
                if let Some(completion) = ctx.locator.completion() {
                    let options = CodeGenerationOptionsBuilder::default()
                        .build()
                        .expect("Failed to build completion options");

                    if !completion.generate("hello Tabby", options).await.is_empty() {
                        return Ok(start.elapsed().as_millis() as i32);
                    }

                    // ast-grep-ignore: use-schema-result
                    Err(CoreError::Other(anyhow::anyhow!(
                        "Failed to connect to the completion model"
                    )))
                } else {
                    Err(CoreError::NotFound("Completion model is not enabled"))
                }
            }
            ModelHealthBackend::Chat => {
                if let Some(chat) = ctx.locator.chat() {
                    let request = CreateChatCompletionRequest::default();

Review comment (bug):
would it be better to create a new error type for Result (Result<i32, TestModelConnectionError>)? e.g we might wanna report different error type to user in future (e.g credential error, timeout, connection failed, ...)

Code after:
            after,
            before,
            first,
            last,
            |after, before, first, last| async move {
                ctx.locator
                    .web_documents()
                    .list_custom_web_documents(ids, after, before, first, last)
                    .await
            },
        )
        .await
    }
    async fn preset_web_documents(
        ctx: &Context,
        ids: Option<Vec<ID>>,
        after: Option<String>,
        before: Option<String>,
        first: Option<i32>,
        last: Option<i32>,
        is_active: Option<bool>,
    ) -> Result<Connection<PresetWebDocument>> {
        check_admin(ctx).await?;
        query_async(
            after,
            before,
            first,
            last,
            |after, before, first, last| async move {
                ctx.locator
                    .web_documents()
                    .list_preset_web_documents(ids, after, before, first, last, is_active)
                    .await
            },
        )
        .await
    }

    /// List user groups.
    async fn user_groups(ctx: &Context) -> Result<Vec<UserGroup>> {
        check_user(ctx).await?;
        ctx.locator.user_group().list().await
    }

    async fn source_id_access_policies(
        ctx: &Context,
        source_id: String,
    ) -> Result<SourceIdAccessPolicy> {
        check_admin(ctx).await?;
        let read = ctx
            .locator