lance-context 0.2.4__tar.gz → 0.3.1__tar.gz

{lance_context-0.2.4 → lance_context-0.3.1}/Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-members = ["crates/lance-context-core", "crates/lance-context", "python"]
+members = ["crates/lance-context-core", "crates/lance-context-api", "python"]
 resolver = "2"
 
 [workspace.package]

{lance_context-0.2.4 → lance_context-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lance-context
-Version: 0.2.4
+Version: 0.3.1
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: Apache Software License
@@ -14,15 +14,16 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Rust
 Classifier: Topic :: Scientific/Engineering
-Requires-Dist: pylance==1.0.2
-Requires-Dist: lancedb==0.26.1
-Requires-Dist: lance-namespace==0.4.5
-Requires-Dist: lance-graph==0.4.0
+Requires-Dist: pylance==7.0.0
+Requires-Dist: lancedb==0.27.1
+Requires-Dist: lance-namespace==0.7.7
+Requires-Dist: lance-graph==0.5.4
 Requires-Dist: ruff ; extra == 'dev'
 Requires-Dist: pyright ; extra == 'dev'
 Requires-Dist: pytest ; extra == 'tests'
+Requires-Dist: pytest-asyncio ; extra == 'tests'
 Requires-Dist: ruff ; extra == 'tests'
-Requires-Dist: moto[s3] ; extra == 'tests'
+Requires-Dist: moto[s3,server] ; extra == 'tests'
 Requires-Dist: boto3 ; extra == 'tests'
 Requires-Dist: botocore ; extra == 'tests'
 Provides-Extra: dev

lance_context-0.3.1/crates/lance-context-api/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "lance-context-api"
+version = "0.2.4"
+edition = "2021"
+license = "Apache-2.0"
+authors = ["Lance Devs <dev@lancedb.com>"]
+repository = "https://github.com/lancedb/lance-context"
+description = "Shared request/response types for the lance-context REST API"
+keywords = ["context", "lance", "api"]
+
+[dependencies]
+base64 = "0.22"
+chrono = { version = "0.4", default-features = false, features = ["clock", "serde"] }
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
+thiserror = "2"

lance_context-0.3.1/crates/lance-context-api/src/lib.rs

@@ -0,0 +1,330 @@
+use base64::{engine::general_purpose::STANDARD as BASE64, Engine};
+use chrono::{DateTime, Utc};
+use serde::{Deserialize, Serialize};
+use serde_json::Value;
+use std::future::Future;
+
+// ---------------------------------------------------------------------------
+// Unified error
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, thiserror::Error)]
+pub enum ContextError {
+    #[error("{0}")]
+    NotFound(String),
+    #[error("{0}")]
+    AlreadyExists(String),
+    #[error("{0}")]
+    InvalidRequest(String),
+    #[error("{0}")]
+    Internal(String),
+    #[error("Compaction already in progress")]
+    CompactionInProgress,
+}
+
+pub type ContextResult<T> = Result<T, ContextError>;
+
+// ---------------------------------------------------------------------------
+// Unified trait
+// ---------------------------------------------------------------------------
+
+pub trait ContextStoreApi {
+    fn add(
+        &mut self,
+        records: &[AddRecordRequest],
+    ) -> impl Future<Output = ContextResult<AddRecordsResponse>> + Send;
+
+    fn get(&self, id: &str) -> impl Future<Output = ContextResult<Option<RecordDto>>> + Send;
+
+    fn list(
+        &self,
+        limit: Option<usize>,
+        offset: Option<usize>,
+    ) -> impl Future<Output = ContextResult<Vec<RecordDto>>> + Send;
+
+    fn search(
+        &self,
+        query: &[f32],
+        limit: Option<usize>,
+    ) -> impl Future<Output = ContextResult<Vec<SearchResultDto>>> + Send;
+
+    fn version(&self) -> u64;
+
+    fn checkout(&mut self, version: u64) -> impl Future<Output = ContextResult<()>> + Send;
+
+    fn compact(
+        &mut self,
+        options: Option<CompactRequest>,
+    ) -> impl Future<Output = ContextResult<CompactResponse>> + Send;
+
+    fn compaction_stats(&self) -> impl Future<Output = ContextResult<CompactStatsResponse>> + Send;
+}
+
+// ---------------------------------------------------------------------------
+// Context lifecycle
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CreateContextRequest {
+    pub name: String,
+    #[serde(default)]
+    pub storage_options: Option<std::collections::HashMap<String, String>>,
+    #[serde(default)]
+    pub id_index_type: Option<String>,
+    #[serde(default)]
+    pub blob_columns: Option<Vec<String>>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ContextInfo {
+    pub name: String,
+    pub uri: String,
+    pub version: u64,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ListContextsResponse {
+    pub contexts: Vec<ContextInfo>,
+}
+
+// ---------------------------------------------------------------------------
+// Records
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct StateMetadataDto {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub step: Option<i32>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub active_plan_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub tokens_used: Option<i32>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub custom: Option<String>,
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct AddRecordRequest {
+    #[serde(default = "default_role")]
+    pub role: String,
+    #[serde(default = "default_content_type")]
+    pub content_type: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub text_payload: Option<String>,
+    #[serde(
+        default,
+        skip_serializing_if = "Option::is_none",
+        serialize_with = "serialize_base64_opt",
+        deserialize_with = "deserialize_base64_opt"
+    )]
+    pub binary_payload: Option<Vec<u8>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub embedding: Option<Vec<f32>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub bot_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub session_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub external_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub state_metadata: Option<StateMetadataDto>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub expires_at: Option<DateTime<Utc>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub retention_policy: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub supersedes_id: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct AddRecordsRequest {
+    pub records: Vec<AddRecordRequest>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct AddRecordsResponse {
+    pub version: u64,
+    pub ids: Vec<String>,
+    pub count: usize,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RecordDto {
+    pub id: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub external_id: Option<String>,
+    pub run_id: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub bot_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub session_id: Option<String>,
+    pub created_at: DateTime<Utc>,
+    pub role: String,
+    pub content_type: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub text_payload: Option<String>,
+    #[serde(
+        default,
+        skip_serializing_if = "Option::is_none",
+        serialize_with = "serialize_base64_opt",
+        deserialize_with = "deserialize_base64_opt"
+    )]
+    pub binary_payload: Option<Vec<u8>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub embedding: Option<Vec<f32>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub state_metadata: Option<StateMetadataDto>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<Value>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub expires_at: Option<DateTime<Utc>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub retention_policy: Option<String>,
+    pub lifecycle_status: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub retired_at: Option<DateTime<Utc>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub retired_reason: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub supersedes_id: Option<String>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub superseded_by_id: Option<String>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ListRecordsResponse {
+    pub records: Vec<RecordDto>,
+}
+
+// ---------------------------------------------------------------------------
+// Single record lookup
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct GetRecordResponse {
+    pub record: Option<RecordDto>,
+}
+
+// ---------------------------------------------------------------------------
+// Search
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SearchRequest {
+    pub query: Vec<f32>,
+    #[serde(default = "default_search_limit")]
+    pub limit: usize,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SearchResultDto {
+    pub record: RecordDto,
+    pub distance: f32,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct SearchResponse {
+    pub results: Vec<SearchResultDto>,
+}
+
+// ---------------------------------------------------------------------------
+// Versioning
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct VersionResponse {
+    pub version: u64,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CheckoutRequest {
+    pub version: u64,
+}
+
+// ---------------------------------------------------------------------------
+// Compaction
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Default, Serialize, Deserialize)]
+pub struct CompactRequest {
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub target_rows_per_fragment: Option<usize>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub materialize_deletions: Option<bool>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CompactResponse {
+    pub fragments_removed: usize,
+    pub fragments_added: usize,
+    pub files_removed: usize,
+    pub files_added: usize,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct CompactStatsResponse {
+    pub total_fragments: usize,
+    pub is_compacting: bool,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub last_compaction: Option<DateTime<Utc>>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub last_error: Option<String>,
+    pub total_compactions: u64,
+}
+
+// ---------------------------------------------------------------------------
+// Error
+// ---------------------------------------------------------------------------
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ErrorBody {
+    pub code: String,
+    pub message: String,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ErrorResponse {
+    pub error: ErrorBody,
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+fn default_content_type() -> String {
+    "text/plain".to_string()
+}
+
+fn default_role() -> String {
+    "user".to_string()
+}
+
+fn default_search_limit() -> usize {
+    10
+}
+
+fn serialize_base64_opt<S>(data: &Option<Vec<u8>>, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: serde::Serializer,
+{
+    match data {
+        Some(bytes) => serializer.serialize_some(&BASE64.encode(bytes)),
+        None => serializer.serialize_none(),
+    }
+}
+
+fn deserialize_base64_opt<'de, D>(deserializer: D) -> Result<Option<Vec<u8>>, D::Error>
+where
+    D: serde::Deserializer<'de>,
+{
+    let opt: Option<String> = Option::deserialize(deserializer)?;
+    match opt {
+        Some(s) => BASE64
+            .decode(&s)
+            .map(Some)
+            .map_err(serde::de::Error::custom),
+        None => Ok(None),
+    }
+}

{lance_context-0.2.4 → lance_context-0.3.1}/crates/lance-context-core/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "lance-context-core"
-version = "0.2.4"
+version = "0.3.1"
 edition = "2021"
 license = "Apache-2.0"
 authors = ["Lance Devs <dev@lancedb.com>"]
@@ -11,18 +11,22 @@ keywords = ["context", "multimodal", "lance", "agents", "storage"]
 categories = ["database", "data-structures", "science"]
 
 [dependencies]
-arrow-array = "56.2.0"
-arrow-ipc = "56.2.0"
-arrow-schema = "56.2.0"
+arrow-array = "58"
+arrow-ipc = "58"
+arrow-schema = "58"
 chrono = { version = "0.4", default-features = false, features = ["clock"] }
-lance = "1.0.0"
-lancedb = "0.23.1"
-lance-namespace = "1.0.1"
-lance-graph = "0.4.0"
+lance = "7.0.0"
+lance-context-api = { path = "../lance-context-api" }
+lance-index = "7.0.0"
+lance-namespace = "7.0.0"
+lancedb = "0.30.0"
+lance-graph = "0.5.4"
 serde = { version = "1", features = ["derive"] }
+serde_json = "1"
 futures = "0.3"
 tokio = { version = "1", features = ["sync", "time"] }
 tracing = "0.1"
+uuid = { version = "1.20.0", features = ["v4", "v5"] }
 
 [dev-dependencies]
 tempfile = "3"

lance_context-0.3.1/crates/lance-context-core/src/api_impl.rs

@@ -0,0 +1,180 @@
+use chrono::Utc;
+use uuid::Uuid;
+
+use lance_context_api::{
+    AddRecordRequest, AddRecordsResponse, CompactRequest, CompactResponse, CompactStatsResponse,
+    ContextError, ContextResult, ContextStoreApi, RecordDto, SearchResultDto, StateMetadataDto,
+};
+
+use crate::record::{ContextRecord, StateMetadata, LIFECYCLE_ACTIVE};
+use crate::store::{CompactionConfig, ContextStore};
+
+impl ContextStoreApi for ContextStore {
+    async fn add(&mut self, records: &[AddRecordRequest]) -> ContextResult<AddRecordsResponse> {
+        let run_id = Uuid::new_v4().to_string();
+        let mut ids = Vec::with_capacity(records.len());
+        let mut core_records = Vec::with_capacity(records.len());
+
+        for r in records {
+            let id = Uuid::new_v4().to_string();
+            ids.push(id.clone());
+            core_records.push(ContextRecord {
+                id,
+                external_id: r.external_id.clone(),
+                run_id: run_id.clone(),
+                bot_id: r.bot_id.clone(),
+                session_id: r.session_id.clone(),
+                created_at: Utc::now(),
+                role: r.role.clone(),
+                state_metadata: r.state_metadata.as_ref().map(|sm| StateMetadata {
+                    step: sm.step,
+                    active_plan_id: sm.active_plan_id.clone(),
+                    tokens_used: sm.tokens_used,
+                    custom: sm.custom.clone(),
+                }),
+                metadata: r.metadata.clone(),
+                expires_at: r.expires_at,
+                retention_policy: r.retention_policy.clone(),
+                lifecycle_status: LIFECYCLE_ACTIVE.to_string(),
+                retired_at: None,
+                retired_reason: None,
+                supersedes_id: r.supersedes_id.clone(),
+                superseded_by_id: None,
+                content_type: r.content_type.clone(),
+                text_payload: r.text_payload.clone(),
+                binary_payload: r.binary_payload.clone(),
+                embedding: r.embedding.clone(),
+            });
+        }
+
+        let count = core_records.len();
+        let version = self.add(&core_records).await.map_err(to_ctx_err)?;
+        Ok(AddRecordsResponse {
+            version,
+            ids,
+            count,
+        })
+    }
+
+    async fn get(&self, id: &str) -> ContextResult<Option<RecordDto>> {
+        let record = ContextStore::get(self, id).await.map_err(to_ctx_err)?;
+        Ok(record.map(record_to_dto))
+    }
+
+    async fn list(
+        &self,
+        limit: Option<usize>,
+        offset: Option<usize>,
+    ) -> ContextResult<Vec<RecordDto>> {
+        let records = ContextStore::list(self, limit, offset)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(records.into_iter().map(record_to_dto).collect())
+    }
+
+    async fn search(
+        &self,
+        query: &[f32],
+        limit: Option<usize>,
+    ) -> ContextResult<Vec<SearchResultDto>> {
+        let results = ContextStore::search(self, query, limit)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(results
+            .into_iter()
+            .map(|sr| SearchResultDto {
+                record: record_to_dto(sr.record),
+                distance: sr.distance,
+            })
+            .collect())
+    }
+
+    fn version(&self) -> u64 {
+        ContextStore::version(self)
+    }
+
+    async fn checkout(&mut self, version: u64) -> ContextResult<()> {
+        ContextStore::checkout(self, version)
+            .await
+            .map_err(to_ctx_err)
+    }
+
+    async fn compact(&mut self, options: Option<CompactRequest>) -> ContextResult<CompactResponse> {
+        let config = options.map(|req| {
+            let mut c = CompactionConfig::default();
+            if let Some(v) = req.target_rows_per_fragment {
+                c.target_rows_per_fragment = v;
+            }
+            if let Some(v) = req.materialize_deletions {
+                c.materialize_deletions = v;
+            }
+            c
+        });
+
+        let metrics = ContextStore::compact(self, config)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(CompactResponse {
+            fragments_removed: metrics.fragments_removed,
+            fragments_added: metrics.fragments_added,
+            files_removed: metrics.files_removed,
+            files_added: metrics.files_added,
+        })
+    }
+
+    async fn compaction_stats(&self) -> ContextResult<CompactStatsResponse> {
+        let stats = ContextStore::compaction_stats(self)
+            .await
+            .map_err(to_ctx_err)?;
+        Ok(CompactStatsResponse {
+            total_fragments: stats.total_fragments,
+            is_compacting: stats.is_compacting,
+            last_compaction: stats.last_compaction,
+            last_error: stats.last_error,
+            total_compactions: stats.total_compactions,
+        })
+    }
+}
+
+fn record_to_dto(r: ContextRecord) -> RecordDto {
+    RecordDto {
+        id: r.id,
+        external_id: r.external_id,
+        run_id: r.run_id,
+        bot_id: r.bot_id,
+        session_id: r.session_id,
+        created_at: r.created_at,
+        role: r.role,
+        content_type: r.content_type,
+        text_payload: r.text_payload,
+        binary_payload: r.binary_payload,
+        embedding: r.embedding,
+        state_metadata: r.state_metadata.map(|sm| StateMetadataDto {
+            step: sm.step,
+            active_plan_id: sm.active_plan_id,
+            tokens_used: sm.tokens_used,
+            custom: sm.custom,
+        }),
+        metadata: r.metadata,
+        expires_at: r.expires_at,
+        retention_policy: r.retention_policy,
+        lifecycle_status: r.lifecycle_status,
+        retired_at: r.retired_at,
+        retired_reason: r.retired_reason,
+        supersedes_id: r.supersedes_id,
+        superseded_by_id: r.superseded_by_id,
+    }
+}
+
+fn to_ctx_err(err: lance::Error) -> ContextError {
+    let msg = err.to_string();
+    if msg.contains("already in progress") {
+        ContextError::CompactionInProgress
+    } else if msg.contains("not found") || msg.contains("DatasetNotFound") {
+        ContextError::NotFound(msg)
+    } else if msg.contains("Invalid") {
+        ContextError::InvalidRequest(msg)
+    } else {
+        ContextError::Internal(msg)
+    }
+}

lance_context-0.3.1/crates/lance-context-core/src/lib.rs

@@ -0,0 +1,20 @@
+//! Core types for the lance-context storage layer.
+#![recursion_limit = "256"]
+
+mod api_impl;
+mod context;
+mod record;
+pub mod serde;
+mod store;
+
+pub use context::{Context, ContextEntry, Snapshot};
+pub use record::{
+    ContextRecord, LifecycleQueryOptions, MetadataFilter, RecordFilters, SearchResult,
+    StateMetadata, LIFECYCLE_ACTIVE, LIFECYCLE_CONTRADICTED,
+};
+pub use store::{
+    CompactionConfig, CompactionStats, ContextStore, ContextStoreOptions, IdIndexType,
+};
+
+// Re-export CompactionMetrics from lance for Python bindings
+pub use lance::dataset::optimize::CompactionMetrics;