use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;
use crate::embeddings::{NerEntity, NeuralNer};
use crate::graph_memory::GraphMemory;
use crate::memory::{Experience, ExperienceType, MemorySystem, Query as MemoryQuery};
use crate::similarity::cosine_similarity;
#[inline]
pub fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool {
let needle_bytes = needle.as_bytes();
let needle_len = needle_bytes.len();
if needle_len == 0 {
return true;
}
let haystack_bytes = haystack.as_bytes();
if haystack_bytes.len() < needle_len {
return false;
}
'outer: for start in 0..=(haystack_bytes.len() - needle_len) {
for (i, &needle_byte) in needle_bytes.iter().enumerate() {
if haystack_bytes[start + i].to_ascii_lowercase() != needle_byte {
continue 'outer;
}
}
return true;
}
false
}
#[inline]
pub fn content_hash(content: &str) -> u64 {
use std::hash::Hasher;
let mut hasher = std::collections::hash_map::DefaultHasher::new();
let trimmed = content.trim();
let bytes = trimmed.as_bytes();
const CHUNK_SIZE: usize = 64;
let mut buffer: [u8; CHUNK_SIZE] = [0u8; CHUNK_SIZE];
let mut i = 0;
while i + CHUNK_SIZE <= bytes.len() {
for j in 0..CHUNK_SIZE {
buffer[j] = bytes[i + j].to_ascii_lowercase();
}
hasher.write(&buffer);
i += CHUNK_SIZE;
}
let remaining = bytes.len() - i;
if remaining > 0 {
for j in 0..remaining {
buffer[j] = bytes[i + j].to_ascii_lowercase();
}
hasher.write(&buffer[..remaining]);
}
hasher.finish()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StreamMode {
Conversation,
Sensor,
Event,
}
impl Default for StreamMode {
fn default() -> Self {
StreamMode::Conversation
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractionConfig {
#[serde(default = "default_min_importance")]
pub min_importance: f32,
#[serde(default = "default_true")]
pub auto_dedupe: bool,
#[serde(default = "default_dedupe_threshold")]
pub dedupe_threshold: f32,
#[serde(default = "default_checkpoint_interval")]
pub checkpoint_interval_ms: u64,
#[serde(default = "default_max_buffer_size")]
pub max_buffer_size: usize,
#[serde(default = "default_true")]
pub extract_entities: bool,
#[serde(default = "default_true")]
pub create_relationships: bool,
#[serde(default = "default_true")]
pub merge_consecutive: bool,
#[serde(default = "default_trigger_events")]
pub trigger_events: Vec<String>,
#[serde(default = "default_true")]
pub enable_context_injection: bool,
#[serde(default = "default_injection_min_relevance")]
pub injection_min_relevance: f32,
#[serde(default = "default_injection_max_memories")]
pub injection_max_memories: usize,
#[serde(default = "default_injection_cooldown")]
pub injection_cooldown_secs: u64,
}
fn default_min_importance() -> f32 {
0.3
}
fn default_true() -> bool {
true
}
fn default_dedupe_threshold() -> f32 {
0.85
}
fn default_checkpoint_interval() -> u64 {
5000 }
fn default_max_buffer_size() -> usize {
50
}
fn default_trigger_events() -> Vec<String> {
vec![
"error".to_string(),
"decision".to_string(),
"discovery".to_string(),
"learning".to_string(),
]
}
fn default_injection_min_relevance() -> f32 {
0.70 }
fn default_injection_max_memories() -> usize {
3 }
fn default_injection_cooldown() -> u64 {
180 }
const MIN_CHECKPOINT_INTERVAL_MS: u64 = 100; const MAX_CHECKPOINT_INTERVAL_MS: u64 = 3_600_000; const MAX_BUFFER_SIZE: usize = 10_000; const MAX_TRIGGER_EVENTS: usize = 100; const MAX_SEEN_HASHES: usize = 10_000;
impl ExtractionConfig {
pub fn validate_and_clamp(&mut self) {
self.min_importance = self.min_importance.clamp(0.0, 1.0);
self.dedupe_threshold = self.dedupe_threshold.clamp(0.0, 1.0);
if self.checkpoint_interval_ms > 0 {
self.checkpoint_interval_ms = self
.checkpoint_interval_ms
.clamp(MIN_CHECKPOINT_INTERVAL_MS, MAX_CHECKPOINT_INTERVAL_MS);
}
if self.max_buffer_size == 0 {
self.max_buffer_size = default_max_buffer_size();
}
self.max_buffer_size = self.max_buffer_size.min(MAX_BUFFER_SIZE);
if self.trigger_events.len() > MAX_TRIGGER_EVENTS {
self.trigger_events.truncate(MAX_TRIGGER_EVENTS);
}
self.injection_min_relevance = self.injection_min_relevance.clamp(0.0, 1.0);
self.injection_max_memories = self.injection_max_memories.clamp(1, 10);
self.injection_cooldown_secs = self.injection_cooldown_secs.clamp(0, 3600);
}
}
impl Default for ExtractionConfig {
fn default() -> Self {
Self {
min_importance: default_min_importance(),
auto_dedupe: true,
dedupe_threshold: default_dedupe_threshold(),
checkpoint_interval_ms: default_checkpoint_interval(),
max_buffer_size: default_max_buffer_size(),
extract_entities: true,
create_relationships: true,
merge_consecutive: true,
trigger_events: default_trigger_events(),
enable_context_injection: true,
injection_min_relevance: default_injection_min_relevance(),
injection_max_memories: default_injection_max_memories(),
injection_cooldown_secs: default_injection_cooldown(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StreamHandshake {
pub user_id: String,
#[serde(default)]
pub mode: StreamMode,
#[serde(default)]
pub extraction_config: ExtractionConfig,
pub session_id: Option<String>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum StreamMessage {
Content {
content: String,
#[serde(default)]
source: Option<String>,
#[serde(default)]
timestamp: Option<DateTime<Utc>>,
#[serde(default)]
importance: Option<f32>,
#[serde(default)]
tags: Vec<String>,
#[serde(default)]
metadata: HashMap<String, serde_json::Value>,
},
Sensor {
sensor_id: String,
values: HashMap<String, f64>,
#[serde(default)]
timestamp: Option<DateTime<Utc>>,
#[serde(default)]
units: HashMap<String, String>,
},
Event {
event: String,
description: String,
#[serde(default)]
timestamp: Option<DateTime<Utc>>,
#[serde(default)]
severity: Option<String>,
#[serde(default)]
data: HashMap<String, serde_json::Value>,
},
Flush,
Ping,
Close,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ExtractionResult {
Extraction {
memories_created: usize,
memory_ids: Vec<String>,
entities_detected: Vec<DetectedEntity>,
dedupe_skipped: usize,
processing_time_ms: u64,
timestamp: DateTime<Utc>,
},
ContextInjection {
memories: Vec<SurfacedStreamMemory>,
context_hash: u64,
processing_time_ms: u64,
timestamp: DateTime<Utc>,
},
Ack {
message_type: String,
timestamp: DateTime<Utc>,
},
Error {
code: String,
message: String,
fatal: bool,
timestamp: DateTime<Utc>,
},
Closed {
reason: String,
total_memories_created: usize,
timestamp: DateTime<Utc>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SurfacedStreamMemory {
pub id: String,
pub content: String,
pub memory_type: String,
pub relevance: f32,
pub relevance_breakdown: RelevanceBreakdown,
pub created_at: DateTime<Utc>,
pub tags: Vec<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RelevanceBreakdown {
pub semantic: f32,
pub recency: f32,
pub strength: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectedEntity {
pub text: String,
pub entity_type: String,
pub confidence: f32,
pub existing: bool,
}
impl From<&NerEntity> for DetectedEntity {
fn from(ner: &NerEntity) -> Self {
Self {
text: ner.text.clone(),
entity_type: ner.entity_type.as_str().to_string(),
confidence: ner.confidence,
existing: false,
}
}
}
#[derive(Debug, Clone)]
pub struct BufferedMessage {
pub content: String,
pub source: Option<String>,
#[allow(dead_code)]
pub timestamp: DateTime<Utc>,
pub importance: Option<f32>,
pub tags: Vec<String>,
pub metadata: HashMap<String, serde_json::Value>,
}
const MAX_CONCURRENT_SESSIONS: usize = 1000;
const SESSION_TIMEOUT_SECS: i64 = 3600;
pub struct StreamSession {
pub session_id: String,
pub user_id: String,
pub mode: StreamMode,
pub config: ExtractionConfig,
pub metadata: HashMap<String, serde_json::Value>,
buffer: VecDeque<BufferedMessage>,
last_extraction: DateTime<Utc>,
last_activity: DateTime<Utc>,
total_memories_created: usize,
seen_hashes: HashSet<u64>,
#[allow(dead_code)]
recent_embeddings: VecDeque<(String, Vec<f32>)>,
injection_cooldowns: HashMap<String, DateTime<Utc>>,
recent_context_hashes: VecDeque<u64>,
}
impl StreamSession {
pub fn new(handshake: StreamHandshake) -> Self {
let session_id = handshake
.session_id
.unwrap_or_else(|| Uuid::new_v4().to_string());
let mut config = handshake.extraction_config;
config.validate_and_clamp();
let now = Utc::now();
Self {
session_id,
user_id: handshake.user_id,
mode: handshake.mode,
config,
metadata: handshake.metadata,
buffer: VecDeque::with_capacity(64),
last_extraction: now,
last_activity: now,
total_memories_created: 0,
seen_hashes: HashSet::with_capacity(1024),
recent_embeddings: VecDeque::with_capacity(100),
injection_cooldowns: HashMap::new(),
recent_context_hashes: VecDeque::with_capacity(20),
}
}
fn mark_injected(&mut self, memory_id: &str) {
self.injection_cooldowns
.insert(memory_id.to_string(), Utc::now());
}
fn cleanup_injection_cooldowns(&mut self) {
let threshold = self.config.injection_cooldown_secs as i64 * 2;
let cutoff = Utc::now() - chrono::Duration::seconds(threshold);
self.injection_cooldowns.retain(|_, ts| *ts > cutoff);
}
fn should_extract_by_time(&self) -> bool {
if self.config.checkpoint_interval_ms == 0 {
return false;
}
let elapsed = Utc::now()
.signed_duration_since(self.last_extraction)
.num_milliseconds() as u64;
elapsed >= self.config.checkpoint_interval_ms
}
fn should_extract_by_size(&self) -> bool {
self.buffer.len() >= self.config.max_buffer_size
}
#[inline]
fn hash_content(content: &str) -> u64 {
content_hash(content)
}
fn is_exact_duplicate(&self, content: &str) -> bool {
let hash = Self::hash_content(content);
self.seen_hashes.contains(&hash)
}
fn mark_seen(&mut self, content: &str) {
if self.seen_hashes.len() >= MAX_SEEN_HASHES {
let target = MAX_SEEN_HASHES / 2;
let mut kept = 0usize;
self.seen_hashes.retain(|_| {
kept += 1;
kept <= target
});
}
let hash = Self::hash_content(content);
self.seen_hashes.insert(hash);
}
pub fn buffer_message(&mut self, msg: BufferedMessage) -> bool {
if self.config.auto_dedupe && self.is_exact_duplicate(&msg.content) {
return false;
}
if self.config.merge_consecutive && !self.buffer.is_empty() {
if let Some(last) = self.buffer.back_mut() {
if last.source == msg.source {
last.content.push('\n');
last.content.push_str(&msg.content);
last.tags.extend(msg.tags);
for (k, v) in msg.metadata {
last.metadata.insert(k, v);
}
return true;
}
}
}
self.mark_seen(&msg.content);
self.buffer.push_back(msg);
true
}
fn drain_buffer(&mut self) -> Vec<BufferedMessage> {
self.last_extraction = Utc::now();
self.buffer.drain(..).collect()
}
fn touch(&mut self) {
self.last_activity = Utc::now();
}
fn is_stale(&self) -> bool {
let elapsed = Utc::now()
.signed_duration_since(self.last_activity)
.num_seconds();
elapsed > SESSION_TIMEOUT_SECS
}
}
pub struct StreamingMemoryExtractor {
neural_ner: Arc<NeuralNer>,
sessions: Arc<RwLock<HashMap<String, StreamSession>>>,
}
impl StreamingMemoryExtractor {
pub fn new(neural_ner: Arc<NeuralNer>) -> Self {
Self {
neural_ner,
sessions: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn create_session(&self, handshake: StreamHandshake) -> Result<String, String> {
self.cleanup_stale_sessions().await;
let mut sessions = self.sessions.write().await;
if sessions.len() >= MAX_CONCURRENT_SESSIONS {
return Err(format!(
"Maximum concurrent sessions ({}) reached. Try again later.",
MAX_CONCURRENT_SESSIONS
));
}
let session = StreamSession::new(handshake);
let session_id = session.session_id.clone();
sessions.insert(session_id.clone(), session);
Ok(session_id)
}
pub async fn cleanup_stale_sessions(&self) -> usize {
let mut sessions = self.sessions.write().await;
let before_count = sessions.len();
sessions.retain(|_id, session| !session.is_stale());
let removed = before_count - sessions.len();
if removed > 0 {
tracing::info!("Cleaned up {} stale streaming sessions", removed);
}
removed
}
pub async fn session_count(&self) -> usize {
self.sessions.read().await.len()
}
pub async fn process_message(
&self,
session_id: &str,
message: StreamMessage,
memory_system: Arc<parking_lot::RwLock<MemorySystem>>,
) -> ExtractionResult {
let mut sessions = self.sessions.write().await;
let session = match sessions.get_mut(session_id) {
Some(s) => s,
None => {
return ExtractionResult::Error {
code: "SESSION_NOT_FOUND".to_string(),
message: format!("Session {} not found", session_id),
fatal: true,
timestamp: Utc::now(),
}
}
};
session.touch();
match message {
StreamMessage::Content {
content,
source,
timestamp,
importance,
tags,
metadata,
} => {
let msg = BufferedMessage {
content,
source,
timestamp: timestamp.unwrap_or_else(Utc::now),
importance,
tags,
metadata,
};
let buffered = session.buffer_message(msg);
let should_extract = session.should_extract_by_time()
|| session.should_extract_by_size()
|| !buffered;
if should_extract {
drop(sessions);
return self.extract_memories(session_id, memory_system).await;
}
ExtractionResult::Ack {
message_type: "content".to_string(),
timestamp: Utc::now(),
}
}
StreamMessage::Event {
event,
description,
timestamp,
severity,
data,
} => {
let is_trigger = {
let sessions = self.sessions.read().await;
sessions
.get(session_id)
.map(|s| {
s.config
.trigger_events
.iter()
.any(|t| t.eq_ignore_ascii_case(&event))
})
.unwrap_or(false)
};
let content = format!(
"[{}] {}: {}",
severity.unwrap_or_default(),
event,
description
);
let mut metadata: HashMap<String, serde_json::Value> = data;
metadata.insert("event_type".to_string(), serde_json::json!(event));
let msg = BufferedMessage {
content,
source: Some("event".to_string()),
timestamp: timestamp.unwrap_or_else(Utc::now),
importance: if is_trigger { Some(0.8) } else { None },
tags: vec![event.clone()],
metadata,
};
{
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.buffer_message(msg);
}
}
if is_trigger {
return self.extract_memories(session_id, memory_system).await;
}
ExtractionResult::Ack {
message_type: "event".to_string(),
timestamp: Utc::now(),
}
}
StreamMessage::Sensor {
sensor_id,
values,
timestamp,
units,
} => {
let mut parts: Vec<String> = Vec::new();
for (key, value) in &values {
let unit = units.get(key).map(|u| u.as_str()).unwrap_or("");
parts.push(format!("{}={}{}", key, value, unit));
}
let content = format!("[{}] {}", sensor_id, parts.join(", "));
let msg = BufferedMessage {
content,
source: Some(format!("sensor:{}", sensor_id)),
timestamp: timestamp.unwrap_or_else(Utc::now),
importance: None,
tags: vec!["sensor".to_string(), sensor_id],
metadata: HashMap::new(),
};
let should_extract = {
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.buffer_message(msg);
session.should_extract_by_time() || session.should_extract_by_size()
} else {
return ExtractionResult::Error {
code: "SESSION_NOT_FOUND".to_string(),
message: format!("Session '{}' not found", session_id),
fatal: true,
timestamp: Utc::now(),
};
}
};
if should_extract {
return self.extract_memories(session_id, memory_system).await;
}
ExtractionResult::Ack {
message_type: "sensor".to_string(),
timestamp: Utc::now(),
}
}
StreamMessage::Flush => {
drop(sessions);
self.extract_memories(session_id, memory_system).await
}
StreamMessage::Ping => ExtractionResult::Ack {
message_type: "ping".to_string(),
timestamp: Utc::now(),
},
StreamMessage::Close => {
drop(sessions);
let final_result = self.extract_memories(session_id, memory_system).await;
let mut sessions = self.sessions.write().await;
let total = sessions
.get(session_id)
.map(|s| s.total_memories_created)
.unwrap_or(0);
sessions.remove(session_id);
ExtractionResult::Closed {
reason: "client_requested".to_string(),
total_memories_created: total
+ match &final_result {
ExtractionResult::Extraction {
memories_created, ..
} => *memories_created,
_ => 0,
},
timestamp: Utc::now(),
}
}
}
}
async fn extract_memories(
&self,
session_id: &str,
memory_system: Arc<parking_lot::RwLock<MemorySystem>>,
) -> ExtractionResult {
let start = std::time::Instant::now();
let (messages, config, user_metadata, mode) = {
let mut sessions = self.sessions.write().await;
let session = match sessions.get_mut(session_id) {
Some(s) => s,
None => {
return ExtractionResult::Error {
code: "SESSION_NOT_FOUND".to_string(),
message: format!("Session {} not found", session_id),
fatal: true,
timestamp: Utc::now(),
}
}
};
let messages = session.drain_buffer();
let config = session.config.clone();
let metadata = session.metadata.clone();
let mode = session.mode;
(messages, config, metadata, mode)
};
if messages.is_empty() {
return ExtractionResult::Extraction {
memories_created: 0,
memory_ids: vec![],
entities_detected: vec![],
dedupe_skipped: 0,
processing_time_ms: start.elapsed().as_millis() as u64,
timestamp: Utc::now(),
};
}
let mut memory_ids = Vec::new();
let mut all_entities = Vec::new();
let mut dedupe_skipped = 0;
for msg in messages {
let importance = msg
.importance
.unwrap_or_else(|| Self::calculate_importance(&msg.content, mode, &config));
if importance < config.min_importance {
dedupe_skipped += 1;
continue;
}
let entities: Vec<NerEntity> = if config.extract_entities {
match self.neural_ner.extract(&msg.content) {
Ok(ents) => ents,
Err(e) => {
tracing::debug!("NER extraction failed: {}", e);
Vec::new()
}
}
} else {
Vec::new()
};
for ent in &entities {
all_entities.push(DetectedEntity::from(ent));
}
let experience_type = Self::determine_experience_type(mode, &msg);
let mut string_metadata: HashMap<String, String> = HashMap::new();
for (k, v) in user_metadata.iter() {
string_metadata.insert(k.clone(), v.to_string());
}
for (k, v) in msg.metadata {
string_metadata.insert(k, v.to_string());
}
let tags: Vec<String> = msg.tags.clone();
for tag in &tags {
string_metadata.insert(format!("tag:{}", tag), "true".to_string());
}
let mut all_entity_names: Vec<String> =
entities.iter().map(|e| e.text.clone()).collect();
for tag in &tags {
if !all_entity_names.iter().any(|e| e.eq_ignore_ascii_case(tag)) {
all_entity_names.push(tag.clone());
}
}
let experience = Experience {
content: msg.content,
experience_type,
entities: all_entity_names,
metadata: string_metadata,
embeddings: None, tags,
..Default::default()
};
let memory_sys = memory_system.read();
match memory_sys.remember(experience, Some(msg.timestamp)) {
Ok(memory_id) => {
memory_ids.push(memory_id.0.to_string());
}
Err(e) => {
tracing::warn!("Failed to store streaming memory: {}", e);
}
}
}
{
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
session.total_memories_created += memory_ids.len();
}
}
ExtractionResult::Extraction {
memories_created: memory_ids.len(),
memory_ids,
entities_detected: all_entities,
dedupe_skipped,
processing_time_ms: start.elapsed().as_millis() as u64,
timestamp: Utc::now(),
}
}
fn calculate_importance(content: &str, mode: StreamMode, _config: &ExtractionConfig) -> f32 {
let mut importance: f32 = 0.5;
let word_count = content.split_whitespace().count();
if word_count > 50 {
importance += 0.1;
} else if word_count < 10 {
importance -= 0.1;
}
match mode {
StreamMode::Conversation => {
if content.contains('?') {
importance += 0.15;
}
if content.contains("```") || content.contains("fn ") || content.contains("def ") {
importance += 0.2;
}
if contains_ignore_ascii_case(content, "error")
|| contains_ignore_ascii_case(content, "failed")
{
importance += 0.2;
}
}
StreamMode::Sensor => {
importance = 0.4;
}
StreamMode::Event => {
if contains_ignore_ascii_case(content, "error") {
importance += 0.3;
} else if contains_ignore_ascii_case(content, "warning") {
importance += 0.15;
}
}
}
importance.clamp(0.0, 1.0)
}
fn determine_experience_type(mode: StreamMode, msg: &BufferedMessage) -> ExperienceType {
for tag in &msg.tags {
if contains_ignore_ascii_case(tag, "error") {
return ExperienceType::Error;
}
if contains_ignore_ascii_case(tag, "decision") {
return ExperienceType::Decision;
}
if contains_ignore_ascii_case(tag, "learning") {
return ExperienceType::Learning;
}
if contains_ignore_ascii_case(tag, "discovery") {
return ExperienceType::Discovery;
}
}
match mode {
StreamMode::Conversation => ExperienceType::Conversation,
StreamMode::Sensor => ExperienceType::Observation,
StreamMode::Event => ExperienceType::Observation,
}
}
pub async fn close_session(&self, session_id: &str) -> Option<usize> {
let mut sessions = self.sessions.write().await;
sessions
.remove(session_id)
.map(|s| s.total_memories_created)
}
pub async fn inject_context(
&self,
session_id: &str,
content: &str,
memory_system: Arc<parking_lot::RwLock<MemorySystem>>,
graph_memory: Arc<parking_lot::RwLock<GraphMemory>>,
) -> Option<ExtractionResult> {
let start = std::time::Instant::now();
let (config, _user_id) = {
let sessions = self.sessions.read().await;
let session = sessions.get(session_id)?;
if !session.config.enable_context_injection {
return None;
}
(session.config.clone(), session.user_id.clone())
};
let context_hash = content_hash(content);
{
let sessions = self.sessions.read().await;
if let Some(session) = sessions.get(session_id) {
if session.recent_context_hashes.contains(&context_hash) {
return None; }
}
}
let content_for_embed = content.to_string();
let memory_for_embed = memory_system.clone();
let context_embedding: Vec<f32> = tokio::task::spawn_blocking(move || {
let guard = memory_for_embed.read();
guard
.compute_embedding(&content_for_embed)
.unwrap_or_else(|_| vec![0.0; 384])
})
.await
.ok()?;
let min_relevance = config.injection_min_relevance;
let max_per_message = config.injection_max_memories;
let cooldown_seconds = config.injection_cooldown_secs;
let content_for_query = content.to_string();
let max_results = max_per_message * 2; let context_emb = context_embedding.clone();
let cooldown_snapshot: HashSet<String> = {
let sessions_guard = self.sessions.read().await;
if let Some(session) = sessions_guard.get(session_id) {
session
.injection_cooldowns
.iter()
.filter(|(_, ts)| {
let elapsed = Utc::now().signed_duration_since(**ts).num_seconds() as u64;
elapsed < cooldown_seconds
})
.map(|(id, _)| id.clone())
.collect()
} else {
HashSet::new()
}
};
let surfaced: Vec<SurfacedStreamMemory> = {
let memory = memory_system.clone();
let graph = graph_memory.clone();
tokio::task::spawn_blocking(move || {
let memory_guard = memory.read();
let graph_guard = graph.read();
let now = Utc::now();
let query = MemoryQuery {
query_text: Some(content_for_query),
max_results,
..Default::default()
};
let results = memory_guard.recall(&query).unwrap_or_default();
const RECENCY_DECAY_RATE: f32 = 0.01;
let mut candidates: Vec<(_, f32, f32, f32, f32)> = results
.into_iter()
.filter_map(|m| {
let memory_embedding = m.experience.embeddings.as_ref()?.clone();
let score = m.get_score().unwrap_or(0.0);
let semantic = cosine_similarity(&memory_embedding, &context_emb);
let hours_old = (now - m.created_at).num_hours().max(0) as f32;
let recency = (-RECENCY_DECAY_RATE * hours_old).exp();
let hebbian_strength = graph_guard
.get_memory_hebbian_strength(&m.id)
.unwrap_or(0.0);
Some((m, score, semantic, recency, hebbian_strength))
})
.collect();
candidates.sort_by(|a, b| b.1.total_cmp(&a.1));
candidates
.into_iter()
.filter(|(m, score, _, _, _)| {
if *score < min_relevance {
return false;
}
if cooldown_snapshot.contains(&m.id.0.to_string()) {
return false;
}
true
})
.take(max_per_message)
.map(
|(m, score, semantic, recency, strength)| SurfacedStreamMemory {
id: m.id.0.to_string(),
content: m.experience.content.clone(),
memory_type: format!("{:?}", m.experience.experience_type),
relevance: score,
relevance_breakdown: RelevanceBreakdown {
semantic,
recency,
strength,
},
created_at: m.created_at,
tags: m.experience.entities.clone(),
},
)
.collect()
})
.await
.ok()?
};
if surfaced.is_empty() {
return None;
}
{
let mut sessions = self.sessions.write().await;
if let Some(session) = sessions.get_mut(session_id) {
for mem in &surfaced {
session.mark_injected(&mem.id);
}
session.recent_context_hashes.push_back(context_hash);
if session.recent_context_hashes.len() > 20 {
session.recent_context_hashes.pop_front();
}
session.cleanup_injection_cooldowns();
}
}
Some(ExtractionResult::ContextInjection {
memories: surfaced,
context_hash,
processing_time_ms: start.elapsed().as_millis() as u64,
timestamp: Utc::now(),
})
}
pub async fn get_session_stats(&self, session_id: &str) -> Option<SessionStats> {
let sessions = self.sessions.read().await;
sessions.get(session_id).map(|s| SessionStats {
session_id: s.session_id.clone(),
user_id: s.user_id.clone(),
mode: s.mode,
buffer_size: s.buffer.len(),
total_memories_created: s.total_memories_created,
last_extraction: s.last_extraction,
})
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionStats {
pub session_id: String,
pub user_id: String,
pub mode: StreamMode,
pub buffer_size: usize,
pub total_memories_created: usize,
pub last_extraction: DateTime<Utc>,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::embeddings::NerEntityType;
#[test]
fn test_extraction_config_defaults() {
let config = ExtractionConfig::default();
assert_eq!(config.min_importance, 0.3);
assert!(config.auto_dedupe);
assert_eq!(config.checkpoint_interval_ms, 5000);
assert_eq!(config.max_buffer_size, 50);
assert!(config.enable_context_injection);
assert_eq!(config.injection_min_relevance, 0.70);
assert_eq!(config.injection_max_memories, 3);
assert_eq!(config.injection_cooldown_secs, 180);
}
#[test]
fn test_stream_mode_default() {
let mode = StreamMode::default();
assert_eq!(mode, StreamMode::Conversation);
}
#[test]
fn test_content_hash_consistency() {
let h1 = content_hash("Hello World");
let h2 = content_hash("hello world");
let h3 = content_hash(" hello world ");
assert_eq!(h1, h2);
assert_eq!(h2, h3);
}
#[test]
fn test_calculate_importance_conversation() {
let config = ExtractionConfig::default();
let short =
StreamingMemoryExtractor::calculate_importance("ok", StreamMode::Conversation, &config);
assert!(short < 0.5);
let question = StreamingMemoryExtractor::calculate_importance(
"How do I implement streaming in Rust?",
StreamMode::Conversation,
&config,
);
assert!(question > 0.5);
let error = StreamingMemoryExtractor::calculate_importance(
"Error: connection failed to database server unexpectedly while processing request",
StreamMode::Conversation,
&config,
);
assert!(error > 0.6);
}
#[test]
fn test_determine_experience_type() {
let msg_error = BufferedMessage {
content: "test".to_string(),
source: None,
timestamp: Utc::now(),
importance: None,
tags: vec!["error".to_string()],
metadata: HashMap::new(),
};
assert_eq!(
StreamingMemoryExtractor::determine_experience_type(
StreamMode::Conversation,
&msg_error
),
ExperienceType::Error
);
let msg_default = BufferedMessage {
content: "test".to_string(),
source: None,
timestamp: Utc::now(),
importance: None,
tags: vec![],
metadata: HashMap::new(),
};
assert_eq!(
StreamingMemoryExtractor::determine_experience_type(
StreamMode::Conversation,
&msg_default
),
ExperienceType::Conversation
);
assert_eq!(
StreamingMemoryExtractor::determine_experience_type(StreamMode::Sensor, &msg_default),
ExperienceType::Observation
);
}
#[test]
fn test_stream_handshake_deserialization() {
let json = r#"{
"user_id": "test-user",
"mode": "conversation",
"extraction_config": {
"min_importance": 0.5,
"checkpoint_interval_ms": 10000
}
}"#;
let handshake: StreamHandshake = serde_json::from_str(json).unwrap();
assert_eq!(handshake.user_id, "test-user");
assert_eq!(handshake.mode, StreamMode::Conversation);
assert_eq!(handshake.extraction_config.min_importance, 0.5);
assert_eq!(handshake.extraction_config.checkpoint_interval_ms, 10000);
assert!(handshake.extraction_config.auto_dedupe);
}
#[test]
fn test_stream_message_variants() {
let content_json = r#"{
"type": "content",
"content": "Hello world",
"source": "user",
"tags": ["greeting"]
}"#;
let msg: StreamMessage = serde_json::from_str(content_json).unwrap();
matches!(msg, StreamMessage::Content { .. });
let event_json = r#"{
"type": "event",
"event": "error",
"description": "Database connection failed",
"severity": "error"
}"#;
let msg: StreamMessage = serde_json::from_str(event_json).unwrap();
matches!(msg, StreamMessage::Event { .. });
let flush_json = r#"{"type": "flush"}"#;
let msg: StreamMessage = serde_json::from_str(flush_json).unwrap();
matches!(msg, StreamMessage::Flush);
}
#[test]
fn test_detected_entity_from_ner() {
let ner_entity = NerEntity {
text: "Microsoft".to_string(),
entity_type: NerEntityType::Organization,
confidence: 0.95,
start: 0,
end: 9,
};
let detected = DetectedEntity::from(&ner_entity);
assert_eq!(detected.text, "Microsoft");
assert_eq!(detected.entity_type, "ORG");
assert_eq!(detected.confidence, 0.95);
assert!(!detected.existing);
}
}