RFC: Common Storage Access Layer for Plast Mem

[YangKeao]

Summary

This RFC proposes a phased refactor which introduces a common storage access layer. The objective is to make storage integration predictable and incremental: adding a new storage provider should mainly require implementing traits, while keeping domain logic unchanged.

Motivation

I (personally) have three motivations to do this refactor:

1. The ecosystem of vector database is still a chaos. I personally installed many different pgsql / mysql / ... distributions to meet the requirement for vector database from different services (pgvecto.rs, paradedb, vectorchord, tidb, milvus...) 🤦 . It'd be really helpful for the users to have richer choices, so users can maintain less databases. Maintaining database is really a boring / risky job.
2. PgSQL cannot always cover all cases. Maybe one day we'll need to have a memory based / cloud based / local file based solutions, or have other more powerful storage solutions (e.g. message queue, graph db...). It'll be more helpful to have a clear storage abstraction.
3. Clean code. Having an abstraction between the domain logic and storage is always good. It'll help us to make it easier to add metrics, more complex migration logic and tests in the future.

Goals

1. Define storage capabilities as traits.
2. Keep behavior and all logic unchanged.
3. Allow backend-specific SQL and optimization to stay in the corresponding crates.
4. This will not be a short journey. We should implement this RFC incrementally.

Non-Goals

1. Adding more storage providers besides paradedb.
2. Standardizing every possible backend detail into one universal SQL layer.

Design Principles

1. Core owns interfaces; backend crates own implementation details.
2. Trait boundaries use portable data formats when possible.
3. Backend-specific optimizations remain allowed behind trait implementations.
4. Transactions are explicit in capability traits when atomicity matters.
5. Refactor is done in horizontal slices so each step is independently valid.

Phased Delivery Plan

Phase 1: Message Queue

1. Introduce message queue store trait in core.

Trait design:

#[async_trait]
pub trait MessageQueueStore: Send + Sync {
  async fn get_state(&self, id: Uuid) -> Result<MessageQueueState, AppError>;
  async fn append_message(&self, id: Uuid, message: Message) -> Result<i32, AppError>;
  async fn append_messages(&self, id: Uuid, messages: Vec<Message>) -> Result<i32, AppError>;
  async fn drain(&self, id: Uuid, count: usize) -> Result<(), AppError>;

  async fn try_set_fence(&self, id: Uuid, fence_count: i32) -> Result<bool, AppError>;
  async fn clear_stale_fence(&self, id: Uuid, ttl_minutes: i64) -> Result<bool, AppError>;
  async fn finalize_job(
    &self,
    id: Uuid,
    prev_episode_summary: Option<String>,
  ) -> Result<(), AppError>;
  async fn clear_fence(&self, id: Uuid) -> Result<(), AppError>;
  async fn get_prev_episode_summary(&self, id: Uuid) -> Result<Option<String>, AppError>;

  async fn add_pending_review(
    &self,
    id: Uuid,
    memory_ids: Vec<Uuid>,
    query: String,
  ) -> Result<(), AppError>;
  async fn take_pending_reviews(&self, id: Uuid)
    -> Result<Option<Vec<PendingReview>>, AppError>;

  async fn message_count(&self, id: Uuid) -> Result<i32, AppError>;
  async fn queue_status(&self, id: Uuid) -> Result<QueueStatus, AppError>;
}

2. Implement it in storage_pg.
3. Wire server and worker to consume the trait.

Phase 2: Retrieval

1. Introduce retrieval traits for episodic and semantic retrieval.

Trait design:

#[async_trait]
pub trait EpisodicMemoryRetrieveStore: Send + Sync {
  async fn search_hybrid(
    &self,
    query: &str,
    query_embedding: Vec<f32>,
    conversation_id: Uuid,
    candidate_limit: i64,
  ) -> Result<Vec<(EpisodicMemory, f64)>, AppError>;
}

#[async_trait]
pub trait SemanticMemoryRetrieveStore: Send + Sync {
  async fn retrieve_by_embedding(
    &self,
    query: &str,
    query_embedding: Vec<f32>,
    limit: i64,
    conversation_id: Uuid,
    category: Option<&str>,
    candidate_limit: i64,
  ) -> Result<Vec<(SemanticMemory, f64)>, AppError>;
}

2. Move backend-specific retrieval SQL into storage_pg.

Phase 3: Episodic Write/Read Store

1. Introduce EpisodicMemoryStore trait for inserts, queries, consolidation flags, and review state updates.

Trait design:

#[async_trait]
pub trait EpisodicMemoryStore: Send + Sync {
  async fn insert(&self, memory: &EpisodicMemory) -> Result<(), AppError>;
  async fn get_by_id(&self, id: Uuid) -> Result<Option<EpisodicMemory>, AppError>;
  async fn fetch_recent(
    &self,
    conversation_id: Uuid,
    since: Option<DateTime<Utc>>,
    limit: u64,
  ) -> Result<Vec<EpisodicMemory>, AppError>;
  async fn fetch_unconsolidated(
    &self,
    conversation_id: Uuid,
  ) -> Result<Vec<EpisodicMemory>, AppError>;
  async fn fetch_unconsolidated_by_ids(
    &self,
    conversation_id: Uuid,
    ids: &[Uuid],
  ) -> Result<Vec<EpisodicMemory>, AppError>;
  async fn count_unconsolidated(&self, conversation_id: Uuid) -> Result<u64, AppError>;
  async fn mark_consolidated(&self, ids: &[Uuid]) -> Result<(), AppError>;
  async fn update_review_state(
    &self,
    id: Uuid,
    stability: f32,
    difficulty: f32,
    last_reviewed_at: DateTime<Utc>,
  ) -> Result<(), AppError>;
}

2. Move episodic persistence operations into backend implementation.
3. Update server/worker call sites to depend on this capability.

Phase 4: Semantic Store

1. Introduce SemanticMemoryStore and transactional SemanticMemoryStoreTx.
   Trait design:

#[async_trait]
pub trait SemanticMemoryStore: Send + Sync {
  async fn begin_tx(&self) -> Result<Box<dyn SemanticMemoryStoreTx>, AppError>;
}

#[async_trait]
pub trait SemanticMemoryStoreTx: Send + Sync {
  async fn find_similar(
    &self,
    embedding: Vec<f32>,
    threshold: f64,
    conversation_id: Uuid,
    limit: i64,
  ) -> Result<Vec<SemanticMemory>, AppError>;
  async fn get_by_id(&self, id: Uuid) -> Result<Option<SemanticMemory>, AppError>;
  async fn insert(&self, memory: &SemanticMemory) -> Result<(), AppError>;
  async fn invalidate(&self, id: Uuid) -> Result<(), AppError>;
  async fn append_source_episodic_ids(
    &self,
    fact_id: Uuid,
    existing_source_episodic_ids: &[Uuid],
    new_source_episodic_ids: &[Uuid],
  ) -> Result<(), AppError>;
  async fn mark_episodic_consolidated(&self, ids: &[Uuid]) -> Result<(), AppError>;
  async fn commit(self: Box<Self>) -> Result<(), AppError>;
  async fn rollback(self: Box<Self>) -> Result<(), AppError>;
}

2. Move semantic consolidation write sequence to transactional trait operations.
3. Keep explicit begin_tx, commit, and rollback contract.
4. Keep similarity input at trait boundary portable (Vec<f32>), with backend-local conversion.

Phase 5: Job Queue

1. Introduce worker job queue abstraction.
   Trait design:

#[async_trait]
pub trait JobQueue: Send + Sync {
  async fn enqueue_segmentation(&self, job: EventSegmentationJob) -> Result<(), AppError>;
  async fn enqueue_review(&self, job: MemoryReviewJob) -> Result<(), AppError>;
  async fn enqueue_semantic(&self, job: SemanticConsolidationJob) -> Result<(), AppError>;
  async fn count_active_jobs(&self, conversation_id: Uuid) -> Result<i64, AppError>;
}

3. Provide PostgreSQL-backed implementation in storage_pg.
4. Remove direct queue backend coupling from server/worker composition.

Validation

TODO: I didn;t find any existing tests in this repo. Need discussion / help to validate the modification.

Most of the modification are just Copy-Paste, so the risk is relatively low IMO.

Future Work

1. Add backend conformance test suite shared across providers.
2. Introduce optional feature flags for backend selection in build/runtime composition.

Disclaimer

I'm a developer of TiDB and also a self-hosting maximalist 🧙‍♂️ . This proposal is submitted as a prerequisite for #41 . The final target is to add the support of TiDB (a database with MySQL compatible grammar and vector search capability) as the storage layer.

[kwaa]

1. The ecosystem of vector database is still a chaos. I personally installed many different pgsql / mysql / ... distributions to meet the requirement for vector database from different services (pgvecto.rs, paradedb, vectorchord, tidb, milvus...) 🤦 . It'd be really helpful for the users to have richer choices, so users can maintain less databases. Maintaining database is really a boring / risky job.

It is worth noting that:

• We are currently using SeaORM and intend to use it in all databases that may be supported in the future.
• On top of that, we use Apalis as our background task framework (perhaps a SeaORM adapter could be written for it later).
• Currently, Plast Mem requires a database that supports Vector and BM25 search.

Ultimately, we may support these databases:

• ParadeDB (Of course)
• TiDB (Of course)
• TursoDB (Possibly)

[siddontang]

Thank you for mentioning TiDB here, @kwaa , if you need help, feel free to reach out to us.

One thing that might interest you: TiDB Cloud Zero

Zero signup. Zero provisioning. Your AI companion gets a database instantly. But here's the twist — a Zero instance expires in 30 days unless explicitly claimed.

Think 一週間フレンズ, but make it 一ヶ月フレンズ (One Month Friends) 😄

Your AI companion remembers everything for 30 days. Day 30 arrives — claim her database or she forgets you forever. It's not a paywall. It's a story.

Free → instant memory → 30-day countdown → "Save her memories" → persistent forever.

[Lilia-Chen]

This abstraction plan is fantastic! Combined with the capability vs strategy decoupling (#44 (comment)), I believe we can start considering implementation in the near term.