@sesamespace/hivemind 0.1.0

package/packages/memory/src/src/promotion.rs

@@ -0,0 +1,637 @@
+use anyhow::Result;
+use arrow_array::{RecordBatch, RecordBatchIterator, StringArray, UInt64Array};
+use arrow_schema::{DataType, Field, Schema};
+use chrono::Utc;
+use futures::stream::TryStreamExt;
+use lancedb::{connection::Connection, query::ExecutableQuery, query::QueryBase, Table};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+const ACCESS_TABLE: &str = "episode_access";
+const COOCCURRENCE_TABLE: &str = "episode_cooccurrence";
+const L3_TABLE: &str = "l3_knowledge";
+
+// Promotion thresholds
+const ACCESS_THRESHOLD: u64 = 5; // episode must be accessed at least N times
+const COOCCURRENCE_THRESHOLD: u64 = 3; // episode must co-occur with others at least N times
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct AccessRecord {
+    pub episode_id: String,
+    pub access_count: u64,
+    pub last_accessed: String,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct CooccurrenceRecord {
+    pub episode_a: String,
+    pub episode_b: String,
+    pub count: u64,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct L3Entry {
+    pub id: String,
+    pub source_episode_id: String,
+    pub context_name: String,
+    pub content: String,
+    pub promoted_at: String,
+    pub access_count: u64,
+    pub connection_density: u64,
+}
+
+pub struct PromotionEngine {
+    db: Connection,
+}
+
+impl PromotionEngine {
+    pub async fn new(db: Connection) -> Result<Self> {
+        let engine = Self { db };
+        engine.ensure_tables().await?;
+        Ok(engine)
+    }
+
+    fn access_schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![
+            Field::new("episode_id", DataType::Utf8, false),
+            Field::new("access_count", DataType::UInt64, false),
+            Field::new("last_accessed", DataType::Utf8, false),
+        ]))
+    }
+
+    fn cooccurrence_schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![
+            Field::new("episode_a", DataType::Utf8, false),
+            Field::new("episode_b", DataType::Utf8, false),
+            Field::new("count", DataType::UInt64, false),
+        ]))
+    }
+
+    fn l3_schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Utf8, false),
+            Field::new("source_episode_id", DataType::Utf8, false),
+            Field::new("context_name", DataType::Utf8, false),
+            Field::new("content", DataType::Utf8, false),
+            Field::new("promoted_at", DataType::Utf8, false),
+            Field::new("access_count", DataType::UInt64, false),
+            Field::new("connection_density", DataType::UInt64, false),
+        ]))
+    }
+
+    async fn ensure_tables(&self) -> Result<()> {
+        let names = self.db.table_names().execute().await?;
+
+        if !names.contains(&ACCESS_TABLE.to_string()) {
+            let schema = Self::access_schema();
+            let batch = RecordBatch::new_empty(schema.clone());
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            self.db
+                .create_table(ACCESS_TABLE, Box::new(batches))
+                .execute()
+                .await?;
+            tracing::info!("Created episode_access table");
+        }
+
+        if !names.contains(&COOCCURRENCE_TABLE.to_string()) {
+            let schema = Self::cooccurrence_schema();
+            let batch = RecordBatch::new_empty(schema.clone());
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            self.db
+                .create_table(COOCCURRENCE_TABLE, Box::new(batches))
+                .execute()
+                .await?;
+            tracing::info!("Created episode_cooccurrence table");
+        }
+
+        if !names.contains(&L3_TABLE.to_string()) {
+            let schema = Self::l3_schema();
+            let batch = RecordBatch::new_empty(schema.clone());
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            self.db
+                .create_table(L3_TABLE, Box::new(batches))
+                .execute()
+                .await?;
+            tracing::info!("Created l3_knowledge table");
+        }
+
+        Ok(())
+    }
+
+    /// Record a single access for an episode
+    pub async fn record_access(&self, episode_id: &str) -> Result<()> {
+        let table = self.db.open_table(ACCESS_TABLE).execute().await?;
+
+        // Check if record exists
+        let existing = self.get_access_record(episode_id).await?;
+
+        if let Some(record) = existing {
+            // Delete old record and insert updated one
+            table
+                .delete(&format!("episode_id = '{}'", episode_id))
+                .await?;
+
+            let new_count = record.access_count + 1;
+            let now = Utc::now().to_rfc3339();
+            let schema = Self::access_schema();
+            let batch = RecordBatch::try_new(
+                schema.clone(),
+                vec![
+                    Arc::new(StringArray::from(vec![episode_id])),
+                    Arc::new(UInt64Array::from(vec![new_count])),
+                    Arc::new(StringArray::from(vec![now.as_str()])),
+                ],
+            )?;
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            table.add(Box::new(batches)).execute().await?;
+        } else {
+            // Insert new record
+            let now = Utc::now().to_rfc3339();
+            let schema = Self::access_schema();
+            let batch = RecordBatch::try_new(
+                schema.clone(),
+                vec![
+                    Arc::new(StringArray::from(vec![episode_id])),
+                    Arc::new(UInt64Array::from(vec![1u64])),
+                    Arc::new(StringArray::from(vec![now.as_str()])),
+                ],
+            )?;
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            table.add(Box::new(batches)).execute().await?;
+        }
+
+        Ok(())
+    }
+
+    /// Record co-access for a set of episodes retrieved together
+    pub async fn record_co_access(&self, episode_ids: &[String]) -> Result<()> {
+        if episode_ids.len() < 2 {
+            return Ok(());
+        }
+
+        let table = self.db.open_table(COOCCURRENCE_TABLE).execute().await?;
+
+        // For each pair, update or create cooccurrence record
+        for i in 0..episode_ids.len() {
+            for j in (i + 1)..episode_ids.len() {
+                let (a, b) = if episode_ids[i] < episode_ids[j] {
+                    (&episode_ids[i], &episode_ids[j])
+                } else {
+                    (&episode_ids[j], &episode_ids[i])
+                };
+
+                let existing = self.get_cooccurrence(a, b).await?;
+
+                if let Some(record) = existing {
+                    table
+                        .delete(&format!(
+                            "episode_a = '{}' AND episode_b = '{}'",
+                            a, b
+                        ))
+                        .await?;
+
+                    let new_count = record.count + 1;
+                    let schema = Self::cooccurrence_schema();
+                    let batch = RecordBatch::try_new(
+                        schema.clone(),
+                        vec![
+                            Arc::new(StringArray::from(vec![a.as_str()])),
+                            Arc::new(StringArray::from(vec![b.as_str()])),
+                            Arc::new(UInt64Array::from(vec![new_count])),
+                        ],
+                    )?;
+                    let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+                    table.add(Box::new(batches)).execute().await?;
+                } else {
+                    let schema = Self::cooccurrence_schema();
+                    let batch = RecordBatch::try_new(
+                        schema.clone(),
+                        vec![
+                            Arc::new(StringArray::from(vec![a.as_str()])),
+                            Arc::new(StringArray::from(vec![b.as_str()])),
+                            Arc::new(UInt64Array::from(vec![1u64])),
+                        ],
+                    )?;
+                    let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+                    table.add(Box::new(batches)).execute().await?;
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Run the promotion engine: check all episodes and promote eligible ones to L3
+    pub async fn run_promotion(&self, context_filter: Option<&str>) -> Result<Vec<String>> {
+        let mut promoted = Vec::new();
+
+        // Get all access records that meet the threshold
+        let access_records = self.get_high_access_episodes().await?;
+
+        for record in access_records {
+            // Check if already promoted
+            if self.is_promoted(&record.episode_id).await? {
+                continue;
+            }
+
+            // Get connection density for this episode
+            let density = self.get_connection_density(&record.episode_id).await?;
+
+            // Check if both thresholds are met
+            if record.access_count >= ACCESS_THRESHOLD && density >= COOCCURRENCE_THRESHOLD {
+                // Get the episode to determine its context
+                let episode_context = self.get_episode_context(&record.episode_id).await?;
+                if let Some(ref ctx) = episode_context {
+                    // Apply context filter if specified
+                    if let Some(filter) = context_filter {
+                        if ctx != filter {
+                            continue;
+                        }
+                    }
+                }
+
+                let episode_content = self.get_episode_content(&record.episode_id).await?;
+                if let (Some(ctx), Some(content)) = (episode_context, episode_content) {
+                    self.promote_episode(
+                        &record.episode_id,
+                        &ctx,
+                        &content,
+                        record.access_count,
+                        density,
+                    )
+                    .await?;
+                    promoted.push(record.episode_id);
+                }
+            }
+        }
+
+        if !promoted.is_empty() {
+            tracing::info!("Promoted {} episodes to L3", promoted.len());
+        }
+
+        Ok(promoted)
+    }
+
+    async fn get_access_record(&self, episode_id: &str) -> Result<Option<AccessRecord>> {
+        let table = self.db.open_table(ACCESS_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("episode_id = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                let ids = batch
+                    .column_by_name("episode_id")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                let counts = batch
+                    .column_by_name("access_count")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<UInt64Array>()
+                    .unwrap();
+                let last = batch
+                    .column_by_name("last_accessed")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+
+                return Ok(Some(AccessRecord {
+                    episode_id: ids.value(0).to_string(),
+                    access_count: counts.value(0),
+                    last_accessed: last.value(0).to_string(),
+                }));
+            }
+        }
+        Ok(None)
+    }
+
+    async fn get_cooccurrence(&self, a: &str, b: &str) -> Result<Option<CooccurrenceRecord>> {
+        let table = self.db.open_table(COOCCURRENCE_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("episode_a = '{}' AND episode_b = '{}'", a, b))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                let ep_a = batch
+                    .column_by_name("episode_a")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                let ep_b = batch
+                    .column_by_name("episode_b")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                let counts = batch
+                    .column_by_name("count")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<UInt64Array>()
+                    .unwrap();
+
+                return Ok(Some(CooccurrenceRecord {
+                    episode_a: ep_a.value(0).to_string(),
+                    episode_b: ep_b.value(0).to_string(),
+                    count: counts.value(0),
+                }));
+            }
+        }
+        Ok(None)
+    }
+
+    async fn get_high_access_episodes(&self) -> Result<Vec<AccessRecord>> {
+        let table = self.db.open_table(ACCESS_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("access_count >= {}", ACCESS_THRESHOLD))
+            .execute()
+            .await?;
+
+        let mut records = Vec::new();
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+
+        for batch in &batches {
+            let ids = batch
+                .column_by_name("episode_id")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+            let counts = batch
+                .column_by_name("access_count")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap();
+            let last = batch
+                .column_by_name("last_accessed")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+
+            for i in 0..batch.num_rows() {
+                records.push(AccessRecord {
+                    episode_id: ids.value(i).to_string(),
+                    access_count: counts.value(i),
+                    last_accessed: last.value(i).to_string(),
+                });
+            }
+        }
+
+        Ok(records)
+    }
+
+    async fn get_connection_density(&self, episode_id: &str) -> Result<u64> {
+        let table = self.db.open_table(COOCCURRENCE_TABLE).execute().await?;
+
+        // Sum all cooccurrence counts where this episode is involved
+        let results_a = table
+            .query()
+            .only_if(format!("episode_a = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let mut total: u64 = 0;
+        let batches: Vec<RecordBatch> = results_a.try_collect().await?;
+        for batch in &batches {
+            let counts = batch
+                .column_by_name("count")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap();
+            for i in 0..batch.num_rows() {
+                total += counts.value(i);
+            }
+        }
+
+        let results_b = table
+            .query()
+            .only_if(format!("episode_b = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results_b.try_collect().await?;
+        for batch in &batches {
+            let counts = batch
+                .column_by_name("count")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<UInt64Array>()
+                .unwrap();
+            for i in 0..batch.num_rows() {
+                total += counts.value(i);
+            }
+        }
+
+        Ok(total)
+    }
+
+    async fn is_promoted(&self, episode_id: &str) -> Result<bool> {
+        let table = self.db.open_table(L3_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("source_episode_id = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                return Ok(true);
+            }
+        }
+        Ok(false)
+    }
+
+    async fn get_episode_context(&self, episode_id: &str) -> Result<Option<String>> {
+        let table = self.db.open_table("episodes").execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("id = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                let contexts = batch
+                    .column_by_name("context_name")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                return Ok(Some(contexts.value(0).to_string()));
+            }
+        }
+        Ok(None)
+    }
+
+    async fn get_episode_content(&self, episode_id: &str) -> Result<Option<String>> {
+        let table = self.db.open_table("episodes").execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("id = '{}'", episode_id))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                let contents = batch
+                    .column_by_name("content")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                return Ok(Some(contents.value(0).to_string()));
+            }
+        }
+        Ok(None)
+    }
+
+    async fn promote_episode(
+        &self,
+        episode_id: &str,
+        context: &str,
+        content: &str,
+        access_count: u64,
+        connection_density: u64,
+    ) -> Result<()> {
+        let id = uuid::Uuid::new_v4().to_string();
+        let now = Utc::now().to_rfc3339();
+
+        let table = self.db.open_table(L3_TABLE).execute().await?;
+        let schema = Self::l3_schema();
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(StringArray::from(vec![id.as_str()])),
+                Arc::new(StringArray::from(vec![episode_id])),
+                Arc::new(StringArray::from(vec![context])),
+                Arc::new(StringArray::from(vec![content])),
+                Arc::new(StringArray::from(vec![now.as_str()])),
+                Arc::new(UInt64Array::from(vec![access_count])),
+                Arc::new(UInt64Array::from(vec![connection_density])),
+            ],
+        )?;
+
+        let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+        table.add(Box::new(batches)).execute().await?;
+
+        tracing::info!(
+            "Promoted episode {} to L3 (access: {}, density: {})",
+            episode_id,
+            access_count,
+            connection_density
+        );
+
+        Ok(())
+    }
+
+    /// Get all L3 entries for a context
+    pub async fn get_l3_entries(&self, context: &str) -> Result<Vec<L3Entry>> {
+        let table = self.db.open_table(L3_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("context_name = '{}'", context))
+            .execute()
+            .await?;
+
+        let mut entries = Vec::new();
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+
+        for batch in &batches {
+            let ids = batch.column_by_name("id").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
+            let source_ids = batch.column_by_name("source_episode_id").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
+            let contexts = batch.column_by_name("context_name").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
+            let contents = batch.column_by_name("content").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
+            let promoted_ats = batch.column_by_name("promoted_at").unwrap().as_any().downcast_ref::<StringArray>().unwrap();
+            let access_counts = batch.column_by_name("access_count").unwrap().as_any().downcast_ref::<UInt64Array>().unwrap();
+            let densities = batch.column_by_name("connection_density").unwrap().as_any().downcast_ref::<UInt64Array>().unwrap();
+
+            for i in 0..batch.num_rows() {
+                entries.push(L3Entry {
+                    id: ids.value(i).to_string(),
+                    source_episode_id: source_ids.value(i).to_string(),
+                    context_name: contexts.value(i).to_string(),
+                    content: contents.value(i).to_string(),
+                    promoted_at: promoted_ats.value(i).to_string(),
+                    access_count: access_counts.value(i),
+                    connection_density: densities.value(i),
+                });
+            }
+        }
+
+        Ok(entries)
+    }
+
+    /// Check if an episode is referenced across multiple contexts and should be auto-promoted to global
+    pub async fn check_cross_context_promotion(&self, episode_content: &str) -> Result<Vec<String>> {
+        // This is checked by counting how many contexts have similar content
+        // Called externally when a new episode is stored
+        // Returns list of context names where similar content exists
+        let table = self.db.open_table("episodes").execute().await?;
+        let results = table.query().execute().await?;
+
+        let mut context_set: std::collections::HashSet<String> = std::collections::HashSet::new();
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+
+        let lower_content = episode_content.to_lowercase();
+        // Simple substring matching for cross-context detection
+        // Real implementation would use embedding similarity
+        for batch in &batches {
+            let contexts = batch
+                .column_by_name("context_name")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+            let contents = batch
+                .column_by_name("content")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+
+            for i in 0..batch.num_rows() {
+                let content = contents.value(i).to_lowercase();
+                // Check for significant overlap (shared words)
+                let shared = count_shared_significant_words(&lower_content, &content);
+                if shared >= 3 {
+                    context_set.insert(contexts.value(i).to_string());
+                }
+            }
+        }
+
+        Ok(context_set.into_iter().collect())
+    }
+}
+
+/// Count shared significant words (length > 4) between two texts
+fn count_shared_significant_words(a: &str, b: &str) -> usize {
+    let words_a: std::collections::HashSet<&str> = a
+        .split_whitespace()
+        .filter(|w| w.len() > 4)
+        .collect();
+    let words_b: std::collections::HashSet<&str> = b
+        .split_whitespace()
+        .filter(|w| w.len() > 4)
+        .collect();
+    words_a.intersection(&words_b).count()
+}