cognoxium_core/
lib.rs

1//! Cognoxium's native data and tokenization core.
2//!
3//! The crate deliberately contains no network client, model invocation, API
4//! key handling, or telemetry. Python bindings translate structured errors at
5//! the language boundary.
6
7use std::collections::HashMap;
8use std::sync::Arc;
9
10use arrow::datatypes::{DataType, Field, Fields, Schema};
11use bytes::Bytes;
12use chrono::{DateTime, Utc};
13use rayon::prelude::*;
14use serde::{Deserialize, Serialize};
15use sha2::{Digest, Sha256};
16use thiserror::Error;
17
18/// Stable schema version written into manifests and serialized records.
19pub const SCHEMA_VERSION: &str = "1";
20
21/// A context payload. Binary storage is supported even when a renderer cannot
22/// yet send that payload to a model provider.
23#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
24#[serde(tag = "kind", rename_all = "snake_case")]
25pub enum Payload {
26    Text {
27        text: String,
28    },
29    Json {
30        json: String,
31    },
32    Binary {
33        binary: Bytes,
34    },
35    Reference {
36        uri: String,
37        digest: Option<String>,
38        size: Option<u64>,
39    },
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
43#[serde(rename_all = "snake_case")]
44pub enum Trust {
45    Trusted,
46    Untrusted,
47    Quarantined,
48}
49
50#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
51#[serde(rename_all = "snake_case")]
52pub enum Sensitivity {
53    Public,
54    Internal,
55    Confidential,
56    Restricted,
57}
58
59#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
60#[serde(rename_all = "snake_case")]
61pub enum Retention {
62    Optional,
63    Preferred,
64    Required,
65}
66
67/// Provider-neutral conversational role.
68#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
69#[serde(rename_all = "snake_case")]
70pub enum Role {
71    System,
72    Developer,
73    User,
74    Assistant,
75    Tool,
76}
77
78/// Stable provenance reference attached to a context item.
79#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
80pub struct SourceRef {
81    pub uri: String,
82    pub kind: String,
83    pub name: Option<String>,
84    pub digest: Option<String>,
85}
86
87/// Canonical in-memory context record shared by native extensions.
88#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
89pub struct ContextItem {
90    pub id: String,
91    pub payload: Payload,
92    pub mime_type: String,
93    pub role: Role,
94    pub sources: Vec<SourceRef>,
95    pub content_hash: String,
96    pub trust: Trust,
97    pub sensitivity: Sensitivity,
98    pub retention: Retention,
99    pub priority: f64,
100    pub created_at: DateTime<Utc>,
101    pub expires_at: Option<DateTime<Utc>>,
102    pub parent_ids: Vec<String>,
103    pub merged_from_ids: Vec<String>,
104    pub metadata: serde_json::Map<String, serde_json::Value>,
105    pub truncatable: bool,
106    pub min_tokens: u64,
107}
108
109/// Locale-independent diagnostic returned across language boundaries.
110#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
111pub struct Diagnostic {
112    pub code: String,
113    pub parameters: serde_json::Map<String, serde_json::Value>,
114    pub help_code: String,
115}
116
117/// Loads context without coupling the core to a storage or agent framework.
118pub trait Source: Send + Sync {
119    type Error;
120
121    fn id(&self) -> &str;
122    fn load(&self) -> Result<Vec<ContextItem>, Self::Error>;
123}
124
125/// Combines deterministic tokenization with provider envelope accounting.
126pub trait TokenProfile: Send + Sync {
127    fn id(&self) -> &str;
128    fn fingerprint(&self) -> &str;
129    fn count_batch(&self, values: &[String]) -> Result<Vec<usize>, CoreError>;
130}
131
132/// Scores context items in one batch without imposing an output order.
133pub trait Ranker: Send + Sync {
134    fn fingerprint(&self) -> &str;
135    fn score_batch(&self, query: &str, items: &[ContextItem]) -> Result<Vec<f64>, CoreError>;
136}
137
138/// Returns a stable reason code when an item violates a boundary policy.
139pub trait Policy: Send + Sync {
140    fn id(&self) -> &str;
141    fn evaluate(&self, item: &ContextItem, boundary: &str) -> Option<String>;
142}
143
144/// Converts selected context into an explicit provider representation.
145pub trait Renderer: Send + Sync {
146    type Output;
147    type Error;
148
149    fn id(&self) -> &str;
150    fn render(&self, items: &[ContextItem]) -> Result<Self::Output, Self::Error>;
151}
152
153#[derive(Debug, Error)]
154pub enum CoreError {
155    #[error("unsupported tokenizer encoding: {0}")]
156    UnsupportedEncoding(String),
157    #[error("invalid Hugging Face tokenizer: {0}")]
158    InvalidTokenizer(String),
159}
160
161/// Compute a stable lowercase SHA-256 digest.
162pub fn sha256_hex(value: &[u8]) -> String {
163    use std::fmt::Write as _;
164
165    let digest = Sha256::digest(value);
166    let mut encoded = String::with_capacity(digest.len() * 2);
167    for byte in digest {
168        write!(&mut encoded, "{byte:02x}").expect("writing to a String cannot fail");
169    }
170    encoded
171}
172
173/// Return the Arrow schema used at the Rust/Python interoperability boundary.
174pub fn context_schema() -> Schema {
175    let mut json_metadata = HashMap::new();
176    json_metadata.insert("ARROW:extension:name".to_owned(), "arrow.json".to_owned());
177    json_metadata.insert("ARROW:extension:metadata".to_owned(), "{}".to_owned());
178
179    let payload_fields = Fields::from(vec![
180        Arc::new(Field::new("kind", DataType::UInt8, false)),
181        Arc::new(Field::new("text", DataType::LargeUtf8, true)),
182        Arc::new(Field::new("json", DataType::LargeUtf8, true).with_metadata(json_metadata)),
183        Arc::new(Field::new("binary", DataType::LargeBinary, true)),
184        Arc::new(Field::new("uri", DataType::LargeUtf8, true)),
185        Arc::new(Field::new("digest", DataType::LargeUtf8, true)),
186        Arc::new(Field::new("size", DataType::UInt64, true)),
187    ]);
188    let source_fields = Fields::from(vec![
189        Arc::new(Field::new("uri", DataType::LargeUtf8, false)),
190        Arc::new(Field::new("kind", DataType::Utf8, false)),
191        Arc::new(Field::new("name", DataType::LargeUtf8, true)),
192        Arc::new(Field::new("digest", DataType::LargeUtf8, true)),
193    ]);
194
195    Schema::new(vec![
196        Field::new("id", DataType::LargeUtf8, false),
197        Field::new("payload", DataType::Struct(payload_fields), false),
198        Field::new("mime_type", DataType::Utf8, false),
199        Field::new("role", DataType::Utf8, false),
200        Field::new(
201            "sources",
202            DataType::List(Arc::new(Field::new(
203                "item",
204                DataType::Struct(source_fields),
205                false,
206            ))),
207            false,
208        ),
209        Field::new("content_hash", DataType::Utf8, false),
210        Field::new("trust", DataType::Utf8, false),
211        Field::new("sensitivity", DataType::Utf8, false),
212        Field::new("retention", DataType::Utf8, false),
213        Field::new("priority", DataType::Float64, false),
214        Field::new(
215            "created_at",
216            DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())),
217            false,
218        ),
219        Field::new(
220            "expires_at",
221            DataType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, Some("UTC".into())),
222            true,
223        ),
224        Field::new(
225            "parent_ids",
226            DataType::List(Arc::new(Field::new("item", DataType::LargeUtf8, false))),
227            false,
228        ),
229        Field::new(
230            "merged_from_ids",
231            DataType::List(Arc::new(Field::new("item", DataType::LargeUtf8, false))),
232            false,
233        ),
234        Field::new("metadata", DataType::LargeUtf8, false),
235        Field::new("truncatable", DataType::Boolean, false),
236        Field::new("min_tokens", DataType::UInt64, false),
237    ])
238}
239
240/// Create a DataFusion session for lazy, in-process execution.
241#[cfg(feature = "query-engine")]
242pub fn session_context() -> datafusion::prelude::SessionContext {
243    datafusion::prelude::SessionContext::new()
244}
245
246/// Count an OpenAI-compatible encoding in parallel without entering Python.
247#[cfg(feature = "native-tokenizers")]
248pub fn count_tiktoken_batch(encoding: &str, values: &[String]) -> Result<Vec<usize>, CoreError> {
249    match encoding {
250        "o200k_base" => {
251            let tokenizer = tiktoken_rs::o200k_base_singleton();
252            Ok(values
253                .par_iter()
254                .map(|value| tokenizer.encode_with_special_tokens(value).len())
255                .collect())
256        }
257        "cl100k_base" => {
258            let tokenizer = tiktoken_rs::cl100k_base_singleton();
259            Ok(values
260                .par_iter()
261                .map(|value| tokenizer.encode_with_special_tokens(value).len())
262                .collect())
263        }
264        other => Err(CoreError::UnsupportedEncoding(other.to_owned())),
265    }
266}
267
268/// Count values with a Hugging Face tokenizer JSON definition.
269#[cfg(feature = "native-tokenizers")]
270pub fn count_huggingface_batch(
271    tokenizer_json: &str,
272    values: &[String],
273) -> Result<Vec<usize>, CoreError> {
274    let tokenizer = tokenizers::Tokenizer::from_bytes(tokenizer_json.as_bytes())
275        .map_err(|error| CoreError::InvalidTokenizer(error.to_string()))?;
276    values
277        .par_iter()
278        .map(|value| {
279            tokenizer
280                .encode(value.as_str(), false)
281                .map(|encoding| encoding.len())
282                .map_err(|error| CoreError::InvalidTokenizer(error.to_string()))
283        })
284        .collect()
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290
291    #[test]
292    fn sha256_is_stable() {
293        assert_eq!(
294            sha256_hex(b"cognoxium"),
295            "ddf0670f67f3e9919a3761ca57ac898de58572974e4f2985e1b95f503af9e8f3"
296        );
297    }
298
299    #[test]
300    fn payload_schema_is_structured_for_future_modalities() {
301        let schema = context_schema();
302        let payload = schema.field_with_name("payload").expect("payload field");
303        assert!(matches!(payload.data_type(), DataType::Struct(_)));
304    }
305
306    #[cfg(feature = "native-tokenizers")]
307    #[test]
308    fn counts_known_openai_encoding() {
309        let counts = count_tiktoken_batch("o200k_base", &["hello world".to_owned()])
310            .expect("supported encoding");
311        assert_eq!(counts.len(), 1);
312        assert!(counts[0] > 0);
313    }
314}
cognoxium_core/lib.rs

cognoxium_core/
lib.rs
