stasis-rs 0.9.3

Durable AI orchestration framework with runtime jobs, lineage, and memory integration
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

use std::sync::Arc;

use async_trait::async_trait;
use serde_json::json;

use crate::application::orchestration::runtime_job_payloads::{
    MemoryTransformJobPayload, MemoryTransformOperationPayload,
};
use crate::application::runtime::in_memory_runtime::{JobExecutionOutcome, JobHandler};
use crate::application::runtime::memory_job_request_helpers::memory_scope_from_fields;
use crate::application::runtime::memory_operation_job_outcome_helpers::{
    operation_failure, operation_success, policy_violation_failure,
};
use crate::domain::errors::Result;
use crate::domain::runtime::job::Job;
use crate::ports::outbound::memory::memory_models::{
    MemoryTransformOperation, MemoryTransformRequest,
};
use crate::ports::outbound::memory::memory_operations::MemoryOperations;

pub struct MemoryTransformJobHandler {
    operations: Arc<dyn MemoryOperations>,
}

impl MemoryTransformJobHandler {
    pub fn new(operations: Arc<dyn MemoryOperations>) -> Self {
        Self { operations }
    }

    fn parse_payload(raw: &str) -> std::result::Result<MemoryTransformJobPayload, String> {
        serde_json::from_str(raw).map_err(|err| {
            format!("policy violation: invalid memory-transform payload json: {err}")
        })
    }

    fn map_operation(value: Option<MemoryTransformOperationPayload>) -> MemoryTransformOperation {
        match value {
            Some(MemoryTransformOperationPayload::ReindexEmbeddings) => {
                MemoryTransformOperation::ReindexEmbeddings
            }
            Some(MemoryTransformOperationPayload::EmbedTagBackfill) => {
                MemoryTransformOperation::EmbedTagBackfill
            }
            Some(MemoryTransformOperationPayload::ReindexTagEmbeddings) => {
                MemoryTransformOperation::ReindexTagEmbeddings
            }
            _ => MemoryTransformOperation::EmbedBackfill,
        }
    }
}

#[async_trait]
impl JobHandler for MemoryTransformJobHandler {
    fn job_type(&self) -> &'static str {
        "workflow.stasis.memory.transform"
    }

    async fn execute(&self, job: &Job) -> Result<JobExecutionOutcome> {
        let payload = match Self::parse_payload(&job.payload_ref) {
            Ok(payload) => payload,
            Err(message) => return Ok(policy_violation_failure("stasis-memory-transform", message)),
        };

        let request = MemoryTransformRequest {
            scope: memory_scope_from_fields(
                None,
                payload.session_ids,
                payload.tiers,
                payload.from_utc,
                payload.to_utc,
            ),
            filter: Default::default(),
            operation: Self::map_operation(payload.operation),
            dry_run: payload.dry_run.unwrap_or(true),
            batch_size: payload.batch_size.unwrap_or(100),
            max_nodes: payload.max_nodes.unwrap_or(5000),
            provider_id: payload.provider_id,
            model: payload.model,
        };

        match self.operations.transform(&request).await {
            Ok(result) => Ok(operation_success(
                "stasis-memory-transform",
                "memory-transform",
                &job.id,
                json!({
                    "scanned": result.scanned,
                    "selected": result.selected,
                    "updated": result.updated,
                    "failed": result.failed,
                }),
            )),
            Err(err) => Ok(operation_failure("stasis-memory-transform", err.to_string())),
        }
    }
}