use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::rdfstar::{MetadataBuilder, RdfStarProvenanceStore};
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ProvenanceTracker {
pub entity_to_tensor: HashMap<String, usize>,
pub tensor_to_entity: HashMap<usize, String>,
pub shape_to_rule: HashMap<String, String>,
pub node_to_shape: HashMap<usize, String>,
#[serde(skip)]
pub rdfstar_store: Option<RdfStarProvenanceStore>,
}
impl ProvenanceTracker {
pub fn new() -> Self {
ProvenanceTracker {
entity_to_tensor: HashMap::new(),
tensor_to_entity: HashMap::new(),
shape_to_rule: HashMap::new(),
node_to_shape: HashMap::new(),
rdfstar_store: None,
}
}
pub fn with_rdfstar() -> Self {
ProvenanceTracker {
entity_to_tensor: HashMap::new(),
tensor_to_entity: HashMap::new(),
shape_to_rule: HashMap::new(),
node_to_shape: HashMap::new(),
rdfstar_store: Some(RdfStarProvenanceStore::new()),
}
}
pub fn rdfstar_store_mut(&mut self) -> &mut RdfStarProvenanceStore {
if self.rdfstar_store.is_none() {
self.rdfstar_store = Some(RdfStarProvenanceStore::new());
}
self.rdfstar_store
.as_mut()
.expect("rdfstar_store just initialized above")
}
pub fn rdfstar_store(&self) -> Option<&RdfStarProvenanceStore> {
self.rdfstar_store.as_ref()
}
pub fn track_entity(&mut self, entity_iri: String, tensor_idx: usize) {
self.entity_to_tensor.insert(entity_iri.clone(), tensor_idx);
self.tensor_to_entity.insert(tensor_idx, entity_iri);
}
pub fn track_shape(&mut self, shape_iri: String, rule_expr: String, node_idx: usize) {
self.shape_to_rule.insert(shape_iri.clone(), rule_expr);
self.node_to_shape.insert(node_idx, shape_iri);
}
pub fn get_entity(&self, tensor_idx: usize) -> Option<&str> {
self.tensor_to_entity.get(&tensor_idx).map(|s| s.as_str())
}
pub fn get_tensor(&self, entity_iri: &str) -> Option<usize> {
self.entity_to_tensor.get(entity_iri).copied()
}
pub fn to_rdf_star(&self) -> Vec<String> {
let mut statements = Vec::new();
for (tensor_idx, entity_iri) in &self.tensor_to_entity {
statements.push(format!(
"<< <{}> <http://tensorlogic.org/tensor> \"{}\"^^<http://www.w3.org/2001/XMLSchema#integer> >> <http://tensorlogic.org/computedBy> <http://tensorlogic.org/engine> .",
entity_iri, tensor_idx
));
}
for (node_idx, shape_iri) in &self.node_to_shape {
if let Some(rule) = self.shape_to_rule.get(shape_iri) {
statements.push(format!(
"<< <http://tensorlogic.org/node/{}> <http://tensorlogic.org/derivedFrom> <{}> >> <http://tensorlogic.org/rule> \"{}\" .",
node_idx, shape_iri, rule.replace('"', "\\\"")
));
}
}
statements
}
pub fn to_json(&self) -> Result<String> {
Ok(serde_json::to_string_pretty(self)?)
}
pub fn from_json(json: &str) -> Result<Self> {
Ok(serde_json::from_str(json)?)
}
pub fn track_inferred_triple(
&mut self,
subject: String,
predicate: String,
object: String,
rule_id: Option<String>,
confidence: Option<f64>,
) {
let store = self.rdfstar_store_mut();
let mut builder = MetadataBuilder::for_triple(subject, predicate, object)
.generated_by("http://tensorlogic.org/inference".to_string());
if let Some(rule) = rule_id {
builder = builder.rule_id(rule);
}
if let Some(conf) = confidence {
builder = builder.confidence(conf);
}
let now = chrono::Utc::now().to_rfc3339();
builder = builder.generated_at(now);
store.add_metadata(builder.build());
}
pub fn get_high_confidence_inferences(
&self,
min_confidence: f64,
) -> Vec<&crate::rdfstar::StatementMetadata> {
self.rdfstar_store
.as_ref()
.map(|store| store.get_by_min_confidence(min_confidence))
.unwrap_or_default()
}
pub fn to_rdfstar_turtle(&self) -> String {
let mut turtle = String::new();
turtle.push_str("# Combined Provenance Export\n");
turtle.push_str("@prefix tl: <http://tensorlogic.org/> .\n");
turtle.push_str("@prefix prov: <http://www.w3.org/ns/prov#> .\n");
turtle.push_str("@prefix xsd: <http://www.w3.org/2001/XMLSchema#> .\n\n");
turtle.push_str("# Legacy Provenance\n");
for stmt in self.to_rdf_star() {
turtle.push_str(&stmt);
turtle.push('\n');
}
if let Some(store) = &self.rdfstar_store {
turtle.push_str("\n# RDF* Statement-Level Provenance\n");
turtle.push_str(&store.to_turtle());
}
turtle
}
}
impl Default for ProvenanceTracker {
fn default() -> Self {
Self::new()
}
}