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

use super::*;

use crate::mapper::{
    custom_world::{format_ai_result_reference_kind, format_ai_task_kind},
    inventory::map_ai_result_reference_kind,
    npc::map_ai_task_kind,
};

impl From<DomainAiTaskCreateInput> for AiTaskCreateInput {
    fn from(input: DomainAiTaskCreateInput) -> Self {
        Self {
            task_id: input.task_id,
            task_kind: map_ai_task_kind(input.task_kind),
            owner_user_id: input.owner_user_id,
            request_label: input.request_label,
            source_module: input.source_module,
            source_entity_id: input.source_entity_id,
            request_payload_json: input.request_payload_json,
            stages: input.stages.into_iter().map(Into::into).collect(),
            created_at_micros: input.created_at_micros,
        }
    }
}

impl From<DomainAiTaskStartInput> for AiTaskStartInput {
    fn from(input: DomainAiTaskStartInput) -> Self {
        Self {
            task_id: input.task_id,
            started_at_micros: input.started_at_micros,
        }
    }
}

impl From<DomainAiTaskStageStartInput> for AiTaskStageStartInput {
    fn from(input: DomainAiTaskStageStartInput) -> Self {
        Self {
            task_id: input.task_id,
            stage_kind: map_ai_task_stage_kind(input.stage_kind),
            started_at_micros: input.started_at_micros,
        }
    }
}

impl From<DomainAiTextChunkAppendInput> for AiTextChunkAppendInput {
    fn from(input: DomainAiTextChunkAppendInput) -> Self {
        Self {
            task_id: input.task_id,
            stage_kind: map_ai_task_stage_kind(input.stage_kind),
            sequence: input.sequence,
            delta_text: input.delta_text,
            created_at_micros: input.created_at_micros,
        }
    }
}

impl From<DomainAiStageCompletionInput> for AiStageCompletionInput {
    fn from(input: DomainAiStageCompletionInput) -> Self {
        Self {
            task_id: input.task_id,
            stage_kind: map_ai_task_stage_kind(input.stage_kind),
            text_output: input.text_output,
            structured_payload_json: input.structured_payload_json,
            warning_messages: input.warning_messages,
            completed_at_micros: input.completed_at_micros,
        }
    }
}

impl From<DomainAiResultReferenceInput> for AiResultReferenceInput {
    fn from(input: DomainAiResultReferenceInput) -> Self {
        Self {
            task_id: input.task_id,
            reference_kind: map_ai_result_reference_kind(input.reference_kind),
            reference_id: input.reference_id,
            label: input.label,
            created_at_micros: input.created_at_micros,
        }
    }
}

impl From<DomainAiTaskFinishInput> for AiTaskFinishInput {
    fn from(input: DomainAiTaskFinishInput) -> Self {
        Self {
            task_id: input.task_id,
            completed_at_micros: input.completed_at_micros,
        }
    }
}

impl From<DomainAiTaskFailureInput> for AiTaskFailureInput {
    fn from(input: DomainAiTaskFailureInput) -> Self {
        Self {
            task_id: input.task_id,
            failure_message: input.failure_message,
            completed_at_micros: input.completed_at_micros,
        }
    }
}

impl From<DomainAiTaskCancelInput> for AiTaskCancelInput {
    fn from(input: DomainAiTaskCancelInput) -> Self {
        Self {
            task_id: input.task_id,
            completed_at_micros: input.completed_at_micros,
        }
    }
}

impl From<DomainAiTaskStageBlueprint> for AiTaskStageBlueprint {
    fn from(blueprint: DomainAiTaskStageBlueprint) -> Self {
        Self {
            stage_kind: map_ai_task_stage_kind(blueprint.stage_kind),
            label: blueprint.label,
            detail: blueprint.detail,
            order: blueprint.order,
        }
    }
}

pub(crate) fn map_ai_task_procedure_result(
    result: AiTaskProcedureResult,
) -> Result<AiTaskMutationRecord, SpacetimeClientError> {
    if !result.ok {
        return Err(SpacetimeClientError::procedure_failed(result.error_message));
    }

    let task = result
        .task
        .ok_or_else(|| SpacetimeClientError::missing_snapshot("ai_task 快照"))?;

    Ok(AiTaskMutationRecord {
        task: map_ai_task_snapshot(task),
        text_chunk: result.text_chunk.map(map_ai_text_chunk_snapshot),
    })
}

pub(crate) fn map_ai_task_snapshot(snapshot: AiTaskSnapshot) -> AiTaskRecord {
    AiTaskRecord {
        task_id: snapshot.task_id,
        task_kind: format_ai_task_kind(snapshot.task_kind).to_string(),
        owner_user_id: snapshot.owner_user_id,
        request_label: snapshot.request_label,
        source_module: snapshot.source_module,
        source_entity_id: snapshot.source_entity_id,
        request_payload_json: snapshot.request_payload_json,
        status: format_ai_task_status(snapshot.status).to_string(),
        failure_message: snapshot.failure_message,
        stages: snapshot
            .stages
            .into_iter()
            .map(map_ai_task_stage_snapshot)
            .collect(),
        result_references: snapshot
            .result_references
            .into_iter()
            .map(map_ai_result_reference_snapshot)
            .collect(),
        latest_text_output: snapshot.latest_text_output,
        latest_structured_payload_json: snapshot.latest_structured_payload_json,
        version: snapshot.version,
        created_at: format_timestamp_micros(snapshot.created_at_micros),
        started_at: snapshot.started_at_micros.map(format_timestamp_micros),
        completed_at: snapshot.completed_at_micros.map(format_timestamp_micros),
        updated_at: format_timestamp_micros(snapshot.updated_at_micros),
    }
}

pub(crate) fn map_ai_task_stage_snapshot(snapshot: AiTaskStageSnapshot) -> AiTaskStageRecord {
    AiTaskStageRecord {
        stage_kind: format_ai_task_stage_kind(snapshot.stage_kind).to_string(),
        label: snapshot.label,
        detail: snapshot.detail,
        order: snapshot.order,
        status: format_ai_task_stage_status(snapshot.status).to_string(),
        text_output: snapshot.text_output,
        structured_payload_json: snapshot.structured_payload_json,
        warning_messages: snapshot.warning_messages,
        started_at: snapshot.started_at_micros.map(format_timestamp_micros),
        completed_at: snapshot.completed_at_micros.map(format_timestamp_micros),
    }
}

pub(crate) fn map_ai_text_chunk_snapshot(snapshot: AiTextChunkSnapshot) -> AiTextChunkRecord {
    AiTextChunkRecord {
        chunk_id: snapshot.chunk_id,
        task_id: snapshot.task_id,
        stage_kind: format_ai_task_stage_kind(snapshot.stage_kind).to_string(),
        sequence: snapshot.sequence,
        delta_text: snapshot.delta_text,
        created_at: format_timestamp_micros(snapshot.created_at_micros),
    }
}

pub(crate) fn map_ai_result_reference_snapshot(
    snapshot: AiResultReferenceSnapshot,
) -> AiResultReferenceRecord {
    AiResultReferenceRecord {
        result_ref_id: snapshot.result_ref_id,
        task_id: snapshot.task_id,
        reference_kind: format_ai_result_reference_kind(snapshot.reference_kind).to_string(),
        reference_id: snapshot.reference_id,
        label: snapshot.label,
        created_at: format_timestamp_micros(snapshot.created_at_micros),
    }
}

pub(crate) fn map_ai_task_stage_kind(value: DomainAiTaskStageKind) -> AiTaskStageKind {
    match value {
        DomainAiTaskStageKind::PreparePrompt => AiTaskStageKind::PreparePrompt,
        DomainAiTaskStageKind::RequestModel => AiTaskStageKind::RequestModel,
        DomainAiTaskStageKind::RepairResponse => AiTaskStageKind::RepairResponse,
        DomainAiTaskStageKind::NormalizeResult => AiTaskStageKind::NormalizeResult,
        DomainAiTaskStageKind::PersistResult => AiTaskStageKind::PersistResult,
    }
}

pub(crate) fn format_ai_task_status(value: AiTaskStatus) -> &'static str {
    match value {
        AiTaskStatus::Pending => "pending",
        AiTaskStatus::Running => "running",
        AiTaskStatus::Completed => "completed",
        AiTaskStatus::Failed => "failed",
        AiTaskStatus::Cancelled => "cancelled",
    }
}

pub(crate) fn format_ai_task_stage_kind(value: AiTaskStageKind) -> &'static str {
    match value {
        AiTaskStageKind::PreparePrompt => "prepare_prompt",
        AiTaskStageKind::RequestModel => "request_model",
        AiTaskStageKind::RepairResponse => "repair_response",
        AiTaskStageKind::NormalizeResult => "normalize_result",
        AiTaskStageKind::PersistResult => "persist_result",
    }
}

pub(crate) fn format_ai_task_stage_status(value: AiTaskStageStatus) -> &'static str {
    match value {
        AiTaskStageStatus::Pending => "pending",
        AiTaskStageStatus::Running => "running",
        AiTaskStageStatus::Completed => "completed",
        AiTaskStageStatus::Skipped => "skipped",
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AiTaskStageRecord {
    pub stage_kind: String,
    pub label: String,
    pub detail: String,
    pub order: u32,
    pub status: String,
    pub text_output: Option<String>,
    pub structured_payload_json: Option<String>,
    pub warning_messages: Vec<String>,
    pub started_at: Option<String>,
    pub completed_at: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AiResultReferenceRecord {
    pub result_ref_id: String,
    pub task_id: String,
    pub reference_kind: String,
    pub reference_id: String,
    pub label: Option<String>,
    pub created_at: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AiTextChunkRecord {
    pub chunk_id: String,
    pub task_id: String,
    pub stage_kind: String,
    pub sequence: u32,
    pub delta_text: String,
    pub created_at: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AiTaskRecord {
    pub task_id: String,
    pub task_kind: String,
    pub owner_user_id: String,
    pub request_label: String,
    pub source_module: String,
    pub source_entity_id: Option<String>,
    pub request_payload_json: Option<String>,
    pub status: String,
    pub failure_message: Option<String>,
    pub stages: Vec<AiTaskStageRecord>,
    pub result_references: Vec<AiResultReferenceRecord>,
    pub latest_text_output: Option<String>,
    pub latest_structured_payload_json: Option<String>,
    pub version: u32,
    pub created_at: String,
    pub started_at: Option<String>,
    pub completed_at: Option<String>,
    pub updated_at: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct AiTaskMutationRecord {
    pub task: AiTaskRecord,
    pub text_chunk: Option<AiTextChunkRecord>,
}