use serde::{Deserialize, Serialize};
use surrealdb::Surreal;
use super::edge_view::{self, EdgeView, NameCache};
use super::embed::Embedder;
use super::error::GraphError;
use super::store::Db;
use super::types::{EntityDetail, GraphStats, MatchSource, QueryOptions};
use crate::config::GraphScoringConfig;
pub const STRONG_CONFIDENCE: f64 = 0.8;
pub const DOUBTFUL_CONFIDENCE: f64 = 0.5;
const HIGHLIGHT_EDGES: usize = 5;
const MAX_EDGES_PER_ENTITY: usize = 8;
const TOPIC_GRAPH_DEPTH: u32 = 1;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct MemoryOverview {
pub stats: GraphStats,
pub groups: Vec<TypeGroup>,
pub confidence: ConfidenceSummary,
pub uncertain: Vec<EdgeView>,
pub self_reinforced: Vec<EdgeView>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TypeGroup {
pub entity_type: String,
pub count: u64,
pub top: Vec<KnownEntity>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct KnownEntity {
pub id: String,
pub name: String,
pub entity_type: String,
#[serde(rename = "abstract")]
pub abstract_text: String,
#[serde(default)]
pub access_count: i64,
#[serde(default)]
pub utility_score: f64,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ConfidenceSummary {
pub strong: u64,
pub uncertain: u64,
pub doubtful: u64,
}
impl ConfidenceSummary {
#[must_use]
pub fn total(&self) -> u64 {
self.strong + self.uncertain + self.doubtful
}
fn record(&mut self, confidence: f64) {
if confidence >= STRONG_CONFIDENCE {
self.strong += 1;
} else if confidence >= DOUBTFUL_CONFIDENCE {
self.uncertain += 1;
} else {
self.doubtful += 1;
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TopicReport {
pub topic: String,
pub entities: Vec<TopicEntity>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TopicEntity {
pub entity: EntityDetail,
pub score: f64,
pub source: MatchSource,
pub edges: Vec<EdgeView>,
#[serde(default)]
pub edges_omitted: usize,
}
pub async fn overview(
db: &Surreal<Db>,
stats: GraphStats,
per_type: usize,
) -> Result<MemoryOverview, GraphError> {
let mut groups = Vec::with_capacity(stats.entity_type_counts.len());
for (entity_type, count) in &stats.entity_type_counts {
groups.push(TypeGroup {
entity_type: entity_type.clone(),
count: *count,
top: strongest_of_type(db, entity_type, per_type).await?,
});
}
groups.sort_by(|left, right| {
right
.count
.cmp(&left.count)
.then_with(|| left.entity_type.cmp(&right.entity_type))
});
let mut cache = NameCache::new();
let uncertain = edge_view::views(db, &mut cache, &least_certain(db).await?).await?;
let self_reinforced =
edge_view::views(db, &mut cache, &most_self_reinforced(db).await?).await?;
Ok(MemoryOverview {
stats,
groups,
confidence: confidence_summary(db).await?,
uncertain,
self_reinforced,
})
}
pub async fn about(
db: &Surreal<Db>,
embedder: &dyn Embedder,
scoring: &GraphScoringConfig,
topic: &str,
limit: usize,
) -> Result<TopicReport, GraphError> {
let options = QueryOptions {
limit,
entity_type: None,
keyword: None,
graph_depth: TOPIC_GRAPH_DEPTH,
include_episodes: false,
};
let result = super::query::query(db, embedder, scoring, topic, &options).await?;
let mut cache = NameCache::new();
let mut entities = Vec::with_capacity(result.entities.len());
for scored in result.entities {
let all = edge_view::live_edges_of(db, &scored.entity.id_string()).await?;
let edges_omitted = all.len().saturating_sub(MAX_EDGES_PER_ENTITY);
let shown = &all[..all.len().min(MAX_EDGES_PER_ENTITY)];
entities.push(TopicEntity {
entity: scored.entity,
score: scored.score,
source: scored.source,
edges: edge_view::views(db, &mut cache, shown).await?,
edges_omitted,
});
}
Ok(TopicReport {
topic: topic.to_string(),
entities,
})
}
async fn strongest_of_type(
db: &Surreal<Db>,
entity_type: &str,
limit: usize,
) -> Result<Vec<KnownEntity>, GraphError> {
#[derive(serde::Deserialize)]
struct Row {
id: serde_json::Value,
name: String,
entity_type: String,
#[serde(rename = "abstract")]
abstract_text: String,
#[serde(default, deserialize_with = "super::util::count_or_zero")]
access_count: i64,
#[serde(default)]
utility_score: Option<f64>,
}
let query = format!(
r#"SELECT id, name, entity_type, abstract, access_count, utility_score, updated_at
FROM entity WHERE entity_type = $entity_type
ORDER BY utility_score DESC, access_count DESC, updated_at DESC
LIMIT {}"#,
limit.clamp(1, 100)
);
let mut response = db
.query(&query)
.bind(("entity_type", entity_type.to_string()))
.await?;
let rows: Vec<Row> = super::deserialize_take(&mut response, 0)?;
Ok(rows
.into_iter()
.map(|row| KnownEntity {
id: edge_view::record_id(&row.id),
name: row.name,
entity_type: row.entity_type,
abstract_text: row.abstract_text,
access_count: row.access_count,
utility_score: row.utility_score.unwrap_or(0.5),
})
.collect())
}
async fn confidence_summary(db: &Surreal<Db>) -> Result<ConfidenceSummary, GraphError> {
let mut response = db
.query("SELECT VALUE confidence FROM relates_to WHERE valid_until IS NONE")
.await?;
let confidences: Vec<f64> = super::deserialize_take(&mut response, 0)?;
let mut summary = ConfidenceSummary::default();
for confidence in confidences {
summary.record(confidence);
}
Ok(summary)
}
async fn least_certain(db: &Surreal<Db>) -> Result<Vec<super::types::Relationship>, GraphError> {
let query = format!(
r#"SELECT * FROM relates_to
WHERE valid_until IS NONE AND confidence < {STRONG_CONFIDENCE}
ORDER BY confidence ASC
LIMIT {HIGHLIGHT_EDGES}"#
);
let mut response = db.query(&query).await?;
super::deserialize_take(&mut response, 0)
}
async fn most_self_reinforced(
db: &Surreal<Db>,
) -> Result<Vec<super::types::Relationship>, GraphError> {
let query = format!(
r#"SELECT * FROM relates_to
WHERE self_reinforcements IS NOT NONE AND self_reinforcements > 0
ORDER BY self_reinforcements DESC
LIMIT {HIGHLIGHT_EDGES}"#
);
let mut response = db.query(&query).await?;
super::deserialize_take(&mut response, 0)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn confidence_bands_split_at_the_documented_thresholds() {
let mut summary = ConfidenceSummary::default();
for confidence in [1.0, STRONG_CONFIDENCE, 0.79, DOUBTFUL_CONFIDENCE, 0.49, 0.0] {
summary.record(confidence);
}
assert_eq!(summary.strong, 2);
assert_eq!(summary.uncertain, 2);
assert_eq!(summary.doubtful, 2);
assert_eq!(summary.total(), 6);
}
#[test]
fn an_empty_graph_has_nothing_to_be_sure_of() {
let summary = ConfidenceSummary::default();
assert_eq!(summary.total(), 0);
}
}