crank/apps/admin-api/src/service/samples.rs

use crank_core::{GeneratedDraft, GeneratedDraftStatus, OperationId, SampleId, WorkspaceId};
use crank_mapping::{JsonPathRoot, infer_mapping_from_samples};
use crank_registry::{OperationSampleMetadata, SampleKind, SaveSampleMetadataRequest};
use crank_schema::Schema;
use serde_json::Value;
use time::OffsetDateTime;
use tracing::{info, instrument};

use crate::{
    error::ApiError,
    service::{
        AdminService, DraftGenerationResult, GenerateDraftPayload, new_prefixed_id, now_string,
    },
};

impl AdminService {
    #[instrument(skip(self, payload), fields(operation_id = %operation_id.as_str(), sample_kind = ?sample_kind))]
    pub async fn save_json_sample(
        &self,
        workspace_id: &WorkspaceId,
        operation_id: &OperationId,
        sample_kind: SampleKind,
        payload: &Value,
    ) -> Result<OperationSampleMetadata, ApiError> {
        let summary = self.get_operation(workspace_id, operation_id).await?;
        let version = summary.current_draft_version;
        let sample_id = SampleId::new(new_prefixed_id("sample"));
        let now = OffsetDateTime::now_utc();
        let file_name = match sample_kind {
            SampleKind::InputJson => "input.json",
            SampleKind::OutputJson => "output.json",
            SampleKind::YamlImportSource => "source.yaml",
        };
        let storage_ref = self
            .storage
            .write_json_sample(operation_id, version, sample_kind, &sample_id, payload)
            .await?;
        let metadata = OperationSampleMetadata {
            id: sample_id,
            operation_id: operation_id.clone(),
            version,
            sample_kind,
            storage_ref,
            content_type: "application/json".to_owned(),
            file_name: Some(file_name.to_owned()),
            created_at: now,
        };

        self.registry
            .save_sample_metadata(SaveSampleMetadataRequest { sample: &metadata })
            .await?;
        info!(
            operation_id = %operation_id.as_str(),
            sample_id = %metadata.id.as_str(),
            version,
            "json sample saved"
        );

        Ok(metadata)
    }

    #[instrument(skip(self, payload), fields(operation_id = %operation_id.as_str()))]
    pub async fn generate_draft(
        &self,
        workspace_id: &WorkspaceId,
        operation_id: &OperationId,
        payload: GenerateDraftPayload,
    ) -> Result<DraftGenerationResult, ApiError> {
        let summary = self.get_operation(workspace_id, operation_id).await?;
        let samples = self
            .registry
            .list_sample_metadata(operation_id, summary.current_draft_version)
            .await?;

        let input_sample = latest_sample_ref(&samples, SampleKind::InputJson)
            .ok_or_else(|| ApiError::validation("input_json sample was not found"))?;
        let output_sample = latest_sample_ref(&samples, SampleKind::OutputJson)
            .ok_or_else(|| ApiError::validation("output_json sample was not found"))?;
        let input_value = self.storage.read_json(&input_sample.storage_ref).await?;
        let output_value = self.storage.read_json(&output_sample.storage_ref).await?;

        let input_schema = Schema::from_json_sample(&input_value);
        let output_schema = Schema::from_json_sample(&output_value);
        let input_mapping = infer_mapping_from_samples(
            &input_value,
            JsonPathRoot::Mcp,
            &input_value,
            JsonPathRoot::RequestBody,
        );
        let output_mapping = infer_mapping_from_samples(
            &output_value,
            JsonPathRoot::ResponseBody,
            &output_value,
            JsonPathRoot::Output,
        );
        let source_types = if payload.sources.is_empty() {
            vec![
                "input_json_sample".to_owned(),
                "output_json_sample".to_owned(),
            ]
        } else {
            payload.sources
        };
        let generated_draft = GeneratedDraft {
            status: GeneratedDraftStatus::Available,
            source_types,
            generated_at: Some(now_string()?),
            input_schema_generated: true,
            output_schema_generated: true,
            input_mapping_generated: true,
            output_mapping_generated: true,
            warnings: Vec::new(),
        };

        let result = DraftGenerationResult {
            generated_draft,
            input_schema,
            output_schema,
            input_mapping,
            output_mapping,
        };
        info!(operation_id = %operation_id.as_str(), "draft generated from samples");

        Ok(result)
    }
}

fn latest_sample_ref(
    samples: &[OperationSampleMetadata],
    sample_kind: SampleKind,
) -> Option<OperationSampleMetadata> {
    samples
        .iter()
        .rev()
        .find(|sample| sample.sample_kind == sample_kind)
        .cloned()
}