use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FrameKind {
DemotedMessage,
CompactionSummary,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FrameMetadata {
pub schema_version: u8,
pub kind: FrameKind,
pub conversation_id: String,
pub chat_role: String,
pub dedup_key: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub scope: Option<String>,
}
impl FrameMetadata {
pub fn try_from_map(
map: &std::collections::BTreeMap<String, String>,
) -> Result<Self, serde_json::Error> {
let mut obj = serde_json::Map::new();
#[allow(clippy::collapsible_if)]
for (k, v) in map {
if let Ok(num) = v.parse::<u8>() {
if k == "schema_version" {
obj.insert(k.clone(), serde_json::Value::Number(num.into()));
continue;
}
}
obj.insert(k.clone(), serde_json::Value::String(v.clone()));
}
serde_json::from_value(serde_json::Value::Object(obj))
}
pub fn into_map(self) -> std::collections::BTreeMap<String, String> {
let mut map = std::collections::BTreeMap::new();
map.insert(
"schema_version".to_string(),
self.schema_version.to_string(),
);
map.insert("kind".to_string(), self.kind.as_str().to_string());
map.insert("conversation_id".to_string(), self.conversation_id);
map.insert("chat_role".to_string(), self.chat_role);
map.insert("dedup_key".to_string(), self.dedup_key);
if let Some(scope) = self.scope {
map.insert("scope".to_string(), scope);
}
map
}
}
impl FrameKind {
pub fn as_str(self) -> &'static str {
match self {
Self::DemotedMessage => "demoted_message",
Self::CompactionSummary => "compaction_summary",
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::panic, clippy::indexing_slicing)]
mod tests {
use super::*;
#[test]
fn round_trips_through_map() {
let original = FrameMetadata {
schema_version: 1,
kind: FrameKind::DemotedMessage,
conversation_id: "conv-1".to_string(),
chat_role: "user".to_string(),
dedup_key: "deadbeef".to_string(),
scope: Some("scope-a".to_string()),
};
let map = original.clone().into_map();
let parsed = FrameMetadata::try_from_map(&map).unwrap();
assert_eq!(parsed, original);
}
#[test]
fn scope_round_trips_when_absent() {
let original = FrameMetadata {
schema_version: 1,
kind: FrameKind::CompactionSummary,
conversation_id: "conv-2".to_string(),
chat_role: "assistant".to_string(),
dedup_key: "cafebabe".to_string(),
scope: None,
};
let map = original.clone().into_map();
assert!(!map.contains_key("scope"));
let parsed = FrameMetadata::try_from_map(&map).unwrap();
assert_eq!(parsed, original);
}
}