Genarrative/server-rs/crates/spacetime-client/src/ai.rs

use super::*;
use crate::mapper::*;

impl SpacetimeClient {
    pub async fn create_ai_task(
        &self,
        input: DomainAiTaskCreateInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection.procedures().create_ai_task_and_return_then(
                procedure_input,
                move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                },
            );
        })
        .await
    }

    pub async fn start_ai_task(
        &self,
        input: DomainAiTaskStartInput,
    ) -> Result<(), SpacetimeClientError> {
        let reducer_input = input.into();

        self.call_reducer_after_connect(move |connection, sender| {
            let callback_sender = sender.clone();
            if let Err(error) =
                connection
                    .reducers
                    .start_ai_task_then(reducer_input, move |_, result| {
                        let mapped = result
                            .map_err(SpacetimeClientError::from_sdk_error)
                            .and_then(|inner| inner.map_err(SpacetimeClientError::reducer_failed));
                        send_reducer_once(&callback_sender, mapped);
                    })
            {
                send_reducer_once(&sender, Err(SpacetimeClientError::from_sdk_error(error)));
            }
        })
        .await
    }

    pub async fn start_ai_task_stage(
        &self,
        input: DomainAiTaskStageStartInput,
    ) -> Result<(), SpacetimeClientError> {
        let reducer_input = input.into();

        self.call_reducer_after_connect(move |connection, sender| {
            let callback_sender = sender.clone();
            if let Err(error) =
                connection
                    .reducers
                    .start_ai_task_stage_then(reducer_input, move |_, result| {
                        let mapped = result
                            .map_err(SpacetimeClientError::from_sdk_error)
                            .and_then(|inner| inner.map_err(SpacetimeClientError::reducer_failed));
                        send_reducer_once(&callback_sender, mapped);
                    })
            {
                send_reducer_once(&sender, Err(SpacetimeClientError::from_sdk_error(error)));
            }
        })
        .await
    }

    pub async fn append_ai_text_chunk(
        &self,
        input: DomainAiTextChunkAppendInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection
                .procedures()
                .append_ai_text_chunk_and_return_then(procedure_input, move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                });
        })
        .await
    }

    pub async fn complete_ai_stage(
        &self,
        input: DomainAiStageCompletionInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection.procedures().complete_ai_stage_and_return_then(
                procedure_input,
                move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                },
            );
        })
        .await
    }

    pub async fn attach_ai_result_reference(
        &self,
        input: DomainAiResultReferenceInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection
                .procedures()
                .attach_ai_result_reference_and_return_then(procedure_input, move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                });
        })
        .await
    }

    pub async fn complete_ai_task(
        &self,
        input: DomainAiTaskFinishInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection.procedures().complete_ai_task_and_return_then(
                procedure_input,
                move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                },
            );
        })
        .await
    }

    pub async fn fail_ai_task(
        &self,
        input: DomainAiTaskFailureInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection.procedures().fail_ai_task_and_return_then(
                procedure_input,
                move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                },
            );
        })
        .await
    }

    pub async fn cancel_ai_task(
        &self,
        input: DomainAiTaskCancelInput,
    ) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
        let procedure_input = input.into();

        self.call_after_connect(move |connection, sender| {
            connection.procedures().cancel_ai_task_and_return_then(
                procedure_input,
                move |_, result| {
                    let mapped = result
                        .map_err(SpacetimeClientError::from_sdk_error)
                        .and_then(map_ai_task_procedure_result);
                    send_once(&sender, mapped);
                },
            );
        })
        .await
    }
}