lib.rs

use std::sync::Arc;

use chrono::{SecondsFormat, Utc};
use pyo3::exceptions::PyRuntimeError;
use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyType};
use pyo3::IntoPyObject;
use tokio::runtime::Runtime;

use lance_context::serde::CONTENT_TYPE_TEXT;
use lance_context::{Context as RustContext, ContextRecord, ContextStore, SearchResult};

const DEFAULT_BINARY_CONTENT_TYPE: &str = "application/octet-stream";
const BINARY_PLACEHOLDER: &str = "[binary]";

#[pyfunction]
fn version() -> &'static str {
    env!("CARGO_PKG_VERSION")
}

#[pyclass]
struct Context {
    inner: RustContext,
    store: ContextStore,
    runtime: Arc<Runtime>,
    run_id: String,
}

#[pymethods]
impl Context {
    #[classmethod]
    fn create(_cls: &Bound<'_, PyType>, uri: &str) -> PyResult<Self> {
        let runtime = Arc::new(Runtime::new().map_err(to_py_err)?);
        let store = runtime
            .block_on(ContextStore::open(uri))
            .map_err(to_py_err)?;
        let run_id = new_run_id();
        Ok(Self {
            inner: RustContext::new(uri),
            store,
            runtime,
            run_id,
        })
    }

    fn uri(&self) -> &str {
        self.inner.uri()
    }

    fn branch(&self) -> &str {
        self.inner.branch()
    }

    fn entries(&self) -> u64 {
        self.inner.entries()
    }

    fn version(&self) -> u64 {
        self.store.version()
    }

    #[pyo3(signature = (role, content, data_type = None))]
    fn add(
        &mut self,
        role: &str,
        content: &Bound<'_, PyAny>,
        data_type: Option<&str>,
    ) -> PyResult<()> {
        let (content_type, text_payload, binary_payload, inner_content) =
            match content.extract::<&[u8]>() {
                Ok(bytes) => (
                    data_type.unwrap_or(DEFAULT_BINARY_CONTENT_TYPE).to_string(),
                    None,
                    Some(bytes.to_vec()),
                    BINARY_PLACEHOLDER.to_string(),
                ),
                Err(_) => {
                    let content_str = content.str()?.to_string();
                    (
                        data_type.unwrap_or(CONTENT_TYPE_TEXT).to_string(),
                        Some(content_str.clone()),
                        None,
                        content_str,
                    )
                }
            };

        let record_id = format!("{}-{}", self.run_id, self.inner.entries() + 1);
        let record = ContextRecord {
            id: record_id,
            run_id: self.run_id.clone(),
            created_at: Utc::now(),
            role: role.to_string(),
            state_metadata: None,
            content_type,
            text_payload,
            binary_payload,
            embedding: None,
        };

        self.runtime
            .block_on(self.store.add(std::slice::from_ref(&record)))
            .map_err(to_py_err)?;
        self.inner.add(role, &inner_content, data_type);
        Ok(())
    }

    #[pyo3(signature = (label = None))]
    fn snapshot(&mut self, label: Option<&str>) -> String {
        self.inner.snapshot(label)
    }

    fn fork(&self, branch_name: &str) -> Self {
        Self {
            inner: self.inner.fork(branch_name),
            store: self.store.clone(),
            runtime: Arc::clone(&self.runtime),
            run_id: new_run_id(),
        }
    }

    fn checkout(&mut self, version_id: u64) -> PyResult<()> {
        self.runtime
            .block_on(self.store.checkout(version_id))
            .map_err(to_py_err)?;
        self.run_id = new_run_id();
        Ok(())
    }

    #[pyo3(signature = (query, limit = None))]
    fn search(
        &self,
        py: Python<'_>,
        query: Vec<f32>,
        limit: Option<usize>,
    ) -> PyResult<Vec<PyObject>> {
        let hits = self
            .runtime
            .block_on(self.store.search(&query, limit))
            .map_err(to_py_err)?;
        hits.into_iter()
            .map(|hit| search_hit_to_py(py, hit))
            .collect()
    }
}

fn new_run_id() -> String {
    format!(
        "run-{}-{}",
        Utc::now().timestamp_micros(),
        std::process::id()
    )
}

fn search_hit_to_py(py: Python<'_>, hit: SearchResult) -> PyResult<PyObject> {
    let SearchResult { record, distance } = hit;
    let ContextRecord {
        id,
        run_id,
        created_at,
        role,
        state_metadata,
        content_type,
        text_payload,
        binary_payload,
        embedding,
    } = record;

    let dict = PyDict::new(py);
    dict.set_item("id", id)?;
    dict.set_item("run_id", run_id)?;
    dict.set_item(
        "created_at",
        created_at.to_rfc3339_opts(SecondsFormat::Micros, true),
    )?;
    dict.set_item("role", role)?;

    let state_obj: PyObject = match state_metadata {
        Some(metadata) => {
            let state_dict = PyDict::new(py);
            state_dict.set_item("step", metadata.step)?;
            state_dict.set_item("active_plan_id", metadata.active_plan_id)?;
            state_dict.set_item("tokens_used", metadata.tokens_used)?;
            state_dict.set_item("custom", metadata.custom)?;
            state_dict.into_pyobject(py)?.unbind().into()
        }
        None => py.None().into_pyobject(py)?.unbind().into(),
    };
    dict.set_item("state_metadata", state_obj)?;
    dict.set_item("content_type", content_type)?;
    dict.set_item("text_payload", text_payload)?;
    match binary_payload {
        Some(payload) => dict.set_item("binary_payload", PyBytes::new(py, &payload))?,
        None => dict.set_item("binary_payload", py.None())?,
    }
    dict.set_item("embedding", embedding)?;
    dict.set_item("distance", distance)?;
    Ok(dict.into_pyobject(py)?.unbind().into())
}

fn to_py_err<E: std::fmt::Display>(err: E) -> PyErr {
    PyRuntimeError::new_err(err.to_string())
}

#[pymodule]
fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(version, m)?)?;
    m.add_class::<Context>()?;
    Ok(())
}