@sesamespace/hivemind 0.2.0 → 0.3.0

package/packages/memory/src/main.rs

@@ -0,0 +1,626 @@
+mod context;
+mod embeddings;
+mod promotion;
+mod scoring;
+mod store;
+mod tasks;
+
+use axum::{
+    extract::{Path, Query, State},
+    http::StatusCode,
+    routing::{delete, get, patch, post},
+    Json, Router,
+};
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use tower_http::trace::TraceLayer;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
+
+use crate::context::ContextStore;
+use crate::embeddings::OllamaClient;
+use crate::promotion::PromotionEngine;
+use crate::store::{Episode, EpisodeInput, MemoryStore};
+use crate::tasks::{TaskInput, TaskRecord, TaskStore, TaskUpdate};
+
+use crate::scoring::ScoringConfig;
+
+#[derive(Clone)]
+struct AppState {
+    store: Arc<MemoryStore>,
+    ollama: Arc<OllamaClient>,
+    context_store: Arc<ContextStore>,
+    promotion_engine: Arc<PromotionEngine>,
+    task_store: Arc<TaskStore>,
+    scoring_config: Arc<ScoringConfig>,
+}
+
+#[derive(Deserialize)]
+struct SearchQuery {
+    q: String,
+    context: Option<String>,
+    limit: Option<usize>,
+}
+
+#[derive(Serialize)]
+struct SearchResult {
+    episodes: Vec<ScoredEpisode>,
+}
+
+#[derive(Serialize)]
+struct ScoredEpisode {
+    #[serde(flatten)]
+    episode: Episode,
+    score: f64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    source_context: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct CreateContextRequest {
+    name: String,
+    #[serde(default)]
+    description: String,
+}
+
+#[derive(Serialize)]
+struct HealthResponse {
+    status: String,
+    version: String,
+}
+
+#[derive(Serialize)]
+struct ContextInfoResponse {
+    name: String,
+    description: String,
+    created_at: String,
+    episode_count: usize,
+}
+
+#[derive(Serialize)]
+struct ListContextsResponse {
+    contexts: Vec<ContextInfoResponse>,
+}
+
+#[derive(Serialize)]
+struct CrossContextResult {
+    context: String,
+    episodes: Vec<ScoredEpisode>,
+}
+
+#[derive(Serialize)]
+struct CrossContextResponse {
+    results: Vec<CrossContextResult>,
+}
+
+#[derive(Deserialize)]
+struct ShareEpisodeRequest {
+    target_context: String,
+}
+
+#[derive(Deserialize)]
+struct CoAccessRequest {
+    episode_ids: Vec<String>,
+}
+
+#[derive(Serialize)]
+struct PromotionResponse {
+    promoted_count: usize,
+    episode_ids: Vec<String>,
+}
+
+#[derive(Deserialize)]
+struct PromotionQuery {
+    context: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct TaskQuery {
+    context: String,
+    status: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct NextTaskQuery {
+    context: String,
+}
+
+#[derive(Serialize)]
+struct ListTasksResponse {
+    tasks: Vec<TaskRecord>,
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::registry()
+        .with(
+            tracing_subscriber::EnvFilter::try_from_default_env()
+                .unwrap_or_else(|_| "hivemind_memory=debug,tower_http=debug".into()),
+        )
+        .with(tracing_subscriber::fmt::layer())
+        .init();
+
+    let ollama_url =
+        std::env::var("OLLAMA_URL").unwrap_or_else(|_| "http://localhost:11434".to_string());
+    let embedding_model =
+        std::env::var("EMBEDDING_MODEL").unwrap_or_else(|_| "nomic-embed-text".to_string());
+    let db_path = std::env::var("DB_PATH").unwrap_or_else(|_| "./data/lancedb".to_string());
+    let port: u16 = std::env::var("PORT")
+        .unwrap_or_else(|_| "3434".to_string())
+        .parse()?;
+
+    let store = MemoryStore::new(&db_path).await?;
+    let ollama = OllamaClient::new(&ollama_url, &embedding_model);
+    let context_store = ContextStore::new(store.get_connection()).await?;
+    let promotion_engine = PromotionEngine::new(store.get_connection()).await?;
+    let task_store = TaskStore::new(store.get_connection()).await?;
+
+    let scoring_config = ScoringConfig::new();
+
+    let state = AppState {
+        store: Arc::new(store),
+        ollama: Arc::new(ollama),
+        context_store: Arc::new(context_store),
+        promotion_engine: Arc::new(promotion_engine),
+        task_store: Arc::new(task_store),
+        scoring_config: Arc::new(scoring_config),
+    };
+
+    let app = Router::new()
+        // Health
+        .route("/health", get(health))
+        // Episodes
+        .route("/episodes", post(create_episode))
+        .route("/episodes/{id}/access", post(record_access))
+        .route("/episodes/{id}/share", post(share_episode))
+        .route("/episodes/co-access", post(record_co_access))
+        // Search
+        .route("/search", get(search_episodes))
+        .route("/search/cross-context", get(search_cross_context))
+        // Contexts
+        .route("/contexts", get(list_contexts))
+        .route("/contexts", post(create_context))
+        .route("/contexts/{name}", get(get_context_episodes))
+        .route("/contexts/{name}", delete(delete_context))
+        // Scoring config
+        .route("/contexts/{name}/scoring", post(set_context_scoring))
+        // Promotion
+        .route("/promotion/run", post(run_promotion))
+        .route("/promotion/l3", get(get_l3_knowledge))
+        // Tasks
+        .route("/tasks", get(list_tasks))
+        .route("/tasks", post(create_task))
+        .route("/tasks/next", get(get_next_task))
+        .route("/tasks/{id}", patch(update_task))
+        .layer(TraceLayer::new_for_http())
+        .with_state(state);
+
+    let listener = tokio::net::TcpListener::bind(format!("0.0.0.0:{}", port)).await?;
+    tracing::info!("Memory daemon listening on port {}", port);
+    axum::serve(listener, app).await?;
+
+    Ok(())
+}
+
+async fn health() -> Json<HealthResponse> {
+    Json(HealthResponse {
+        status: "ok".to_string(),
+        version: env!("CARGO_PKG_VERSION").to_string(),
+    })
+}
+
+// === Episode endpoints ===
+
+async fn create_episode(
+    State(state): State<AppState>,
+    Json(input): Json<EpisodeInput>,
+) -> Result<Json<Episode>, (StatusCode, String)> {
+    let content_for_check = input.content.clone();
+    let context_for_check = input.context_name.clone().unwrap_or_else(|| "global".to_string());
+
+    let embedding = state
+        .ollama
+        .embed(&input.content)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Embedding failed: {e}")))?;
+
+    let episode = state
+        .store
+        .insert_episode(input, embedding.clone())
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Store failed: {e}")))?;
+
+    // Check for auto-promotion to Global: if similar content exists in 3+ contexts
+    if context_for_check != "global" {
+        let state_clone = state.clone();
+        let content = content_for_check;
+        let embed = embedding;
+        tokio::spawn(async move {
+            if let Ok(contexts) = state_clone
+                .promotion_engine
+                .check_cross_context_promotion(&content)
+                .await
+            {
+                let unique_contexts: std::collections::HashSet<&str> =
+                    contexts.iter().map(|s| s.as_str()).collect();
+                if unique_contexts.len() >= 3 {
+                    // Auto-promote to global context
+                    let global_input = EpisodeInput {
+                        context_name: Some("global".to_string()),
+                        role: "system".to_string(),
+                        content: content.clone(),
+                    };
+                    if let Err(e) = state_clone.store.insert_episode(global_input, embed).await {
+                        tracing::warn!("Auto-promotion to global failed: {}", e);
+                    } else {
+                        tracing::info!(
+                            "Auto-promoted episode to global (referenced in {} contexts)",
+                            unique_contexts.len()
+                        );
+                    }
+                }
+            }
+        });
+    }
+
+    Ok(Json(episode))
+}
+
+async fn record_access(
+    State(state): State<AppState>,
+    Path(id): Path<String>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    state
+        .promotion_engine
+        .record_access(&id)
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Record access failed: {e}"),
+            )
+        })?;
+
+    Ok(StatusCode::NO_CONTENT)
+}
+
+async fn share_episode(
+    State(state): State<AppState>,
+    Path(id): Path<String>,
+    Json(req): Json<ShareEpisodeRequest>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    // Get the original episode
+    let episode = state
+        .store
+        .get_by_id(&id)
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Get episode failed: {e}"),
+            )
+        })?
+        .ok_or_else(|| (StatusCode::NOT_FOUND, "Episode not found".to_string()))?;
+
+    // Re-embed and store in target context
+    let embedding = state
+        .ollama
+        .embed(&episode.content)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Embedding failed: {e}")))?;
+
+    let input = EpisodeInput {
+        context_name: Some(req.target_context),
+        role: episode.role,
+        content: episode.content,
+    };
+
+    state
+        .store
+        .insert_episode(input, embedding)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Store failed: {e}")))?;
+
+    Ok(StatusCode::CREATED)
+}
+
+async fn record_co_access(
+    State(state): State<AppState>,
+    Json(req): Json<CoAccessRequest>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    state
+        .promotion_engine
+        .record_co_access(&req.episode_ids)
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Co-access failed: {e}"),
+            )
+        })?;
+
+    Ok(StatusCode::NO_CONTENT)
+}
+
+// === Search endpoints ===
+
+async fn search_episodes(
+    State(state): State<AppState>,
+    Query(query): Query<SearchQuery>,
+) -> Result<Json<SearchResult>, (StatusCode, String)> {
+    let embedding = state
+        .ollama
+        .embed(&query.q)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Embedding failed: {e}")))?;
+
+    let limit = query.limit.unwrap_or(10);
+    let context = query.context.as_deref().unwrap_or("global");
+
+    let results = state
+        .store
+        .search(embedding, context, limit)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Search failed: {e}")))?;
+
+    let half_life = state.scoring_config.get_half_life(context);
+    let scored: Vec<ScoredEpisode> = results
+        .into_iter()
+        .map(|(episode, distance)| {
+            let score = scoring::combined_score(&episode.timestamp, distance, half_life);
+            ScoredEpisode {
+                episode,
+                score,
+                source_context: None,
+            }
+        })
+        .collect();
+
+    Ok(Json(SearchResult { episodes: scored }))
+}
+
+async fn search_cross_context(
+    State(state): State<AppState>,
+    Query(query): Query<SearchQuery>,
+) -> Result<Json<CrossContextResponse>, (StatusCode, String)> {
+    let embedding = state
+        .ollama
+        .embed(&query.q)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Embedding failed: {e}")))?;
+
+    let limit = query.limit.unwrap_or(10);
+
+    // Search across all contexts
+    let results = state
+        .store
+        .search_all(embedding, limit)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Search failed: {e}")))?;
+
+    // Group by context
+    let mut by_context: std::collections::HashMap<String, Vec<ScoredEpisode>> =
+        std::collections::HashMap::new();
+
+    for (episode, distance) in results {
+        let ctx = episode.context_name.clone();
+        let half_life = state.scoring_config.get_half_life(&ctx);
+        let score = scoring::combined_score(&episode.timestamp, distance, half_life);
+        by_context
+            .entry(ctx.clone())
+            .or_default()
+            .push(ScoredEpisode {
+                episode,
+                score,
+                source_context: Some(ctx),
+            });
+    }
+
+    let results: Vec<CrossContextResult> = by_context
+        .into_iter()
+        .map(|(context, episodes)| CrossContextResult { context, episodes })
+        .collect();
+
+    Ok(Json(CrossContextResponse { results }))
+}
+
+// === Context endpoints ===
+
+async fn list_contexts(
+    State(state): State<AppState>,
+) -> Result<Json<ListContextsResponse>, (StatusCode, String)> {
+    let contexts = state
+        .context_store
+        .list_contexts()
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("List failed: {e}")))?;
+
+    let mut infos = Vec::new();
+    for ctx in contexts {
+        let count = state
+            .store
+            .count_by_context(&ctx.name)
+            .await
+            .unwrap_or(0);
+        infos.push(ContextInfoResponse {
+            name: ctx.name,
+            description: ctx.description,
+            created_at: ctx.created_at,
+            episode_count: count,
+        });
+    }
+
+    Ok(Json(ListContextsResponse { contexts: infos }))
+}
+
+async fn create_context(
+    State(state): State<AppState>,
+    Json(req): Json<CreateContextRequest>,
+) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
+    state
+        .context_store
+        .create_context(&req.name, &req.description)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Failed: {e}")))?;
+
+    Ok((
+        StatusCode::CREATED,
+        Json(serde_json::json!({ "name": req.name, "created": true })),
+    ))
+}
+
+async fn get_context_episodes(
+    State(state): State<AppState>,
+    Path(name): Path<String>,
+) -> Result<Json<Vec<Episode>>, (StatusCode, String)> {
+    let episodes = state
+        .store
+        .get_by_context(&name)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Query failed: {e}")))?;
+
+    Ok(Json(episodes))
+}
+
+async fn delete_context(
+    State(state): State<AppState>,
+    Path(name): Path<String>,
+) -> Result<StatusCode, (StatusCode, String)> {
+    state
+        .context_store
+        .delete_context(&name)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Delete failed: {e}")))?;
+
+    // Also delete episodes in that context
+    state
+        .store
+        .delete_by_context(&name)
+        .await
+        .map_err(|e| {
+            (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Delete episodes failed: {e}"),
+            )
+        })?;
+
+    Ok(StatusCode::NO_CONTENT)
+}
+
+// === Scoring config endpoint ===
+
+#[derive(Deserialize)]
+struct ScoringConfigRequest {
+    half_life_hours: f64,
+}
+
+async fn set_context_scoring(
+    State(state): State<AppState>,
+    Path(name): Path<String>,
+    Json(req): Json<ScoringConfigRequest>,
+) -> StatusCode {
+    state.scoring_config.set_half_life(&name, req.half_life_hours);
+    tracing::info!("Set half-life for context '{}' to {} hours", name, req.half_life_hours);
+    StatusCode::NO_CONTENT
+}
+
+// === Promotion endpoint ===
+
+async fn run_promotion(
+    State(state): State<AppState>,
+    Query(query): Query<PromotionQuery>,
+) -> Result<Json<PromotionResponse>, (StatusCode, String)> {
+    let promoted = state
+        .promotion_engine
+        .run_promotion(query.context.as_deref())
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Promotion failed: {e}")))?;
+
+    Ok(Json(PromotionResponse {
+        promoted_count: promoted.len(),
+        episode_ids: promoted,
+    }))
+}
+
+#[derive(Deserialize)]
+struct L3Query {
+    context: String,
+}
+
+#[derive(Serialize)]
+struct L3Response {
+    entries: Vec<promotion::L3Entry>,
+}
+
+async fn get_l3_knowledge(
+    State(state): State<AppState>,
+    Query(query): Query<L3Query>,
+) -> Result<Json<L3Response>, (StatusCode, String)> {
+    let entries = state
+        .promotion_engine
+        .get_l3_entries(&query.context)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("L3 query failed: {e}")))?;
+
+    Ok(Json(L3Response { entries }))
+}
+
+// === Task endpoints ===
+
+async fn create_task(
+    State(state): State<AppState>,
+    Json(input): Json<TaskInput>,
+) -> Result<(StatusCode, Json<TaskRecord>), (StatusCode, String)> {
+    let task = state
+        .task_store
+        .create_task(input)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Create task failed: {e}")))?;
+
+    Ok((StatusCode::CREATED, Json(task)))
+}
+
+async fn list_tasks(
+    State(state): State<AppState>,
+    Query(query): Query<TaskQuery>,
+) -> Result<Json<ListTasksResponse>, (StatusCode, String)> {
+    let tasks = state
+        .task_store
+        .list_tasks(&query.context, query.status.as_deref())
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("List tasks failed: {e}")))?;
+
+    Ok(Json(ListTasksResponse { tasks }))
+}
+
+async fn get_next_task(
+    State(state): State<AppState>,
+    Query(query): Query<NextTaskQuery>,
+) -> Result<Json<TaskRecord>, (StatusCode, String)> {
+    let task = state
+        .task_store
+        .get_next_task(&query.context)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Get next task failed: {e}")))?;
+
+    match task {
+        Some(t) => Ok(Json(t)),
+        None => Err((StatusCode::NOT_FOUND, "No available tasks".to_string())),
+    }
+}
+
+async fn update_task(
+    State(state): State<AppState>,
+    Path(id): Path<String>,
+    Json(update): Json<TaskUpdate>,
+) -> Result<Json<TaskRecord>, (StatusCode, String)> {
+    let task = state
+        .task_store
+        .update_task(&id, update)
+        .await
+        .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("Update task failed: {e}")))?;
+
+    match task {
+        Some(t) => Ok(Json(t)),
+        None => Err((StatusCode::NOT_FOUND, "Task not found".to_string())),
+    }
+}