@sesamespace/hivemind 0.2.0 → 0.3.0

package/packages/memory/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "hivemind-memory"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+axum = "0.7"
+tokio = { version = "1", features = ["full"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+lancedb = "0.15"
+arrow-array = "53"
+arrow-schema = "53"
+reqwest = { version = "0.12", features = ["json"] }
+uuid = { version = "1", features = ["v4"] }
+chrono = { version = "0.4", features = ["serde"] }
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+tower-http = { version = "0.5", features = ["cors", "trace"] }
+anyhow = "1"
+futures = "0.3"

package/packages/memory/src/context.rs

@@ -0,0 +1,179 @@
+use anyhow::Result;
+use arrow_array::{RecordBatch, RecordBatchIterator, StringArray};
+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::sync::Arc;
+use uuid;
+
+const CONTEXTS_TABLE: &str = "contexts";
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct ContextMetadata {
+    pub name: String,
+    pub description: String,
+    pub created_at: String,
+}
+
+pub struct ContextStore {
+    db: Connection,
+}
+
+impl ContextStore {
+    pub async fn new(db: Connection) -> Result<Self> {
+        let store = Self { db };
+        store.ensure_table().await?;
+        Ok(store)
+    }
+
+    fn schema() -> Arc<Schema> {
+        Arc::new(Schema::new(vec![
+            Field::new("name", DataType::Utf8, false),
+            Field::new("description", DataType::Utf8, false),
+            Field::new("created_at", DataType::Utf8, false),
+        ]))
+    }
+
+    async fn ensure_table(&self) -> Result<()> {
+        let names = self.db.table_names().execute().await?;
+        if !names.contains(&CONTEXTS_TABLE.to_string()) {
+            let schema = Self::schema();
+            let batch = RecordBatch::new_empty(schema.clone());
+            let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+            self.db
+                .create_table(CONTEXTS_TABLE, Box::new(batches))
+                .execute()
+                .await?;
+            tracing::info!("Created contexts metadata table");
+        }
+        Ok(())
+    }
+
+    pub async fn create_context(&self, name: &str, description: &str) -> Result<ContextMetadata> {
+        // Check if context already exists
+        if let Some(existing) = self.get_context(name).await? {
+            return Ok(existing);
+        }
+
+        let created_at = Utc::now().to_rfc3339();
+
+        let metadata = ContextMetadata {
+            name: name.to_string(),
+            description: description.to_string(),
+            created_at: created_at.clone(),
+        };
+
+        let schema = Self::schema();
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![
+                Arc::new(StringArray::from(vec![name])),
+                Arc::new(StringArray::from(vec![description])),
+                Arc::new(StringArray::from(vec![created_at.as_str()])),
+            ],
+        )?;
+
+        let table = self.db.open_table(CONTEXTS_TABLE).execute().await?;
+        let batches = RecordBatchIterator::new(vec![Ok(batch)], schema);
+        table.add(Box::new(batches)).execute().await?;
+
+        tracing::info!("Created context: {}", name);
+        Ok(metadata)
+    }
+
+    pub async fn get_context(&self, name: &str) -> Result<Option<ContextMetadata>> {
+        let table = self.db.open_table(CONTEXTS_TABLE).execute().await?;
+        let results = table
+            .query()
+            .only_if(format!("name = '{}'", name))
+            .execute()
+            .await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        for batch in &batches {
+            if batch.num_rows() > 0 {
+                let names = batch
+                    .column_by_name("name")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                let descriptions = batch
+                    .column_by_name("description")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+                let created_ats = batch
+                    .column_by_name("created_at")
+                    .unwrap()
+                    .as_any()
+                    .downcast_ref::<StringArray>()
+                    .unwrap();
+
+                return Ok(Some(ContextMetadata {
+                    name: names.value(0).to_string(),
+                    description: descriptions.value(0).to_string(),
+                    created_at: created_ats.value(0).to_string(),
+                }));
+            }
+        }
+
+        Ok(None)
+    }
+
+    pub async fn list_contexts(&self) -> Result<Vec<ContextMetadata>> {
+        let table = self.db.open_table(CONTEXTS_TABLE).execute().await?;
+        let results = table.query().execute().await?;
+
+        let batches: Vec<RecordBatch> = results.try_collect().await?;
+        let mut contexts = Vec::new();
+
+        for batch in &batches {
+            let names = batch
+                .column_by_name("name")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+            let descriptions = batch
+                .column_by_name("description")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+            let created_ats = batch
+                .column_by_name("created_at")
+                .unwrap()
+                .as_any()
+                .downcast_ref::<StringArray>()
+                .unwrap();
+
+            for i in 0..batch.num_rows() {
+                contexts.push(ContextMetadata {
+                    name: names.value(i).to_string(),
+                    description: descriptions.value(i).to_string(),
+                    created_at: created_ats.value(i).to_string(),
+                });
+            }
+        }
+
+        Ok(contexts)
+    }
+
+    pub async fn delete_context(&self, name: &str) -> Result<bool> {
+        if name == "global" {
+            anyhow::bail!("Cannot delete global context");
+        }
+
+        let table = self.db.open_table(CONTEXTS_TABLE).execute().await?;
+        table
+            .delete(&format!("name = '{}'", name))
+            .await?;
+
+        tracing::info!("Deleted context metadata: {}", name);
+        Ok(true)
+    }
+}

package/packages/memory/src/embeddings.rs

@@ -0,0 +1,51 @@
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
+
+pub struct OllamaClient {
+    base_url: String,
+    model: String,
+    http: reqwest::Client,
+}
+
+#[derive(Serialize)]
+struct EmbedRequest {
+    model: String,
+    input: String,
+}
+
+#[derive(Deserialize)]
+struct EmbedResponse {
+    embeddings: Vec<Vec<f32>>,
+}
+
+impl OllamaClient {
+    pub fn new(base_url: &str, model: &str) -> Self {
+        Self {
+            base_url: base_url.to_string(),
+            model: model.to_string(),
+            http: reqwest::Client::new(),
+        }
+    }
+
+    pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
+        let url = format!("{}/api/embed", self.base_url);
+        let req = EmbedRequest {
+            model: self.model.clone(),
+            input: text.to_string(),
+        };
+
+        let resp = self
+            .http
+            .post(&url)
+            .json(&req)
+            .send()
+            .await?
+            .error_for_status()?;
+
+        let body: EmbedResponse = resp.json().await?;
+        body.embeddings
+            .into_iter()
+            .next()
+            .ok_or_else(|| anyhow::anyhow!("No embedding returned from Ollama"))
+    }
+}