use parking_lot::RwLock;
use crate::agent_memory::{
AgentMemory, CodeSnippet, Language, MemoryConfig, MemoryRecall, TaskEpisode, TaskOutcome,
};
use crate::error::Result;
use crate::memory_traits::TransferLevel;
use crate::types::{MetadataValue, VectorId};
pub use crate::memory_traits::{KnowledgeDomain, LanguageCompatibility, ProjectContext};
#[derive(Debug)]
pub struct TransferableRecall {
pub recall: MemoryRecall,
pub transfer_score: f32,
pub combined_score: f32,
pub applicable: bool,
pub adaptation_needed: Option<String>,
pub transfer_level: TransferLevel,
pub concepts: Vec<String>,
}
pub struct ConceptExtractor {
design_patterns: Vec<(&'static str, &'static str)>,
domain_concepts: Vec<(&'static str, &'static str)>,
principles: Vec<(&'static str, &'static str)>,
}
impl Default for ConceptExtractor {
fn default() -> Self {
Self::new()
}
}
impl ConceptExtractor {
pub fn new() -> Self {
Self {
design_patterns: vec![
("factory", "Factory Pattern"),
("singleton", "Singleton Pattern"),
("observer", "Observer Pattern"),
("strategy", "Strategy Pattern"),
("decorator", "Decorator Pattern"),
("adapter", "Adapter Pattern"),
("facade", "Facade Pattern"),
("proxy", "Proxy Pattern"),
("builder", "Builder Pattern"),
("prototype", "Prototype Pattern"),
("middleware", "Middleware Pattern"),
("repository", "Repository Pattern"),
("unit of work", "Unit of Work Pattern"),
("dependency injection", "Dependency Injection"),
("event sourcing", "Event Sourcing"),
("cqrs", "CQRS Pattern"),
("saga", "Saga Pattern"),
("circuit breaker", "Circuit Breaker"),
],
domain_concepts: vec![
("authentication", "Authentication"),
("authorization", "Authorization"),
("jwt", "JWT Tokens"),
("oauth", "OAuth"),
("session", "Session Management"),
("rate limit", "Rate Limiting"),
("throttl", "Throttling"),
("cache", "Caching"),
("pagination", "Pagination"),
("validation", "Input Validation"),
("sanitiz", "Input Sanitization"),
("error handling", "Error Handling"),
("logging", "Logging"),
("monitoring", "Monitoring"),
("testing", "Testing"),
("mocking", "Mocking"),
("serializ", "Serialization"),
("deserializ", "Deserialization"),
("encryption", "Encryption"),
("hashing", "Hashing"),
("compression", "Compression"),
("streaming", "Data Streaming"),
("websocket", "WebSockets"),
("graphql", "GraphQL"),
("rest", "REST API"),
("grpc", "gRPC"),
("queue", "Message Queue"),
("pub sub", "Pub/Sub"),
("batch", "Batch Processing"),
("concurrency", "Concurrency"),
("parallelism", "Parallelism"),
("async", "Async Programming"),
],
principles: vec![
("solid", "SOLID Principles"),
("dry", "DRY Principle"),
("kiss", "KISS Principle"),
("yagni", "YAGNI Principle"),
("separation of concern", "Separation of Concerns"),
("single responsibility", "Single Responsibility"),
("open closed", "Open/Closed Principle"),
("liskov", "Liskov Substitution"),
("interface segregation", "Interface Segregation"),
("dependency inversion", "Dependency Inversion"),
(
"composition over inheritance",
"Composition over Inheritance",
),
("fail fast", "Fail Fast"),
("defensive programming", "Defensive Programming"),
],
}
}
pub fn extract(&self, task: &str, code: &str) -> Vec<String> {
let text = format!("{} {}", task, code).to_lowercase();
let mut concepts = Vec::new();
for (keyword, concept) in &self.design_patterns {
if text.contains(keyword) {
concepts.push(concept.to_string());
}
}
for (keyword, concept) in &self.domain_concepts {
if text.contains(keyword) {
concepts.push(concept.to_string());
}
}
for (keyword, concept) in &self.principles {
if text.contains(keyword) {
concepts.push(concept.to_string());
}
}
concepts.sort();
concepts.dedup();
concepts
}
pub fn infer_transfer_level(&self, concepts: &[String]) -> TransferLevel {
let has_principle = concepts
.iter()
.any(|c| self.principles.iter().any(|(_, name)| c == *name));
if has_principle {
return TransferLevel::Universal;
}
let has_pattern = concepts
.iter()
.any(|c| self.design_patterns.iter().any(|(_, name)| c == *name));
if has_pattern {
return TransferLevel::Domain;
}
if !concepts.is_empty() {
return TransferLevel::Domain;
}
TransferLevel::Context
}
}
impl crate::memory_traits::ConceptExtractor for ConceptExtractor {
fn extract(&self, description: &str, content: &str) -> Vec<String> {
ConceptExtractor::extract(self, description, content)
}
fn is_universal(&self, concept: &str) -> bool {
self.principles.iter().any(|(_, name)| *name == concept)
}
fn universal_concepts(&self) -> Vec<&'static str> {
self.principles.iter().map(|(_, name)| *name).collect()
}
}
pub struct TransferableMemory {
memory: AgentMemory,
current_context: RwLock<Option<ProjectContext>>,
extractor: ConceptExtractor,
relevance_weight: f32,
transfer_threshold: f32,
}
impl TransferableMemory {
pub fn new(config: MemoryConfig) -> Result<Self> {
Ok(Self {
memory: AgentMemory::new(config)?,
current_context: RwLock::new(None),
extractor: ConceptExtractor::new(),
relevance_weight: 0.6,
transfer_threshold: 0.3,
})
}
pub fn from_memory(memory: AgentMemory) -> Self {
Self {
memory,
current_context: RwLock::new(None),
extractor: ConceptExtractor::new(),
relevance_weight: 0.6,
transfer_threshold: 0.3,
}
}
pub fn with_relevance_weight(mut self, weight: f32) -> Self {
self.relevance_weight = weight.clamp(0.0, 1.0);
self
}
pub fn with_transfer_threshold(mut self, threshold: f32) -> Self {
self.transfer_threshold = threshold.clamp(0.0, 1.0);
self
}
pub fn set_embed_fn<F>(&mut self, f: F)
where
F: Fn(&str) -> Vec<f32> + Send + Sync + 'static,
{
self.memory.set_embed_fn(f);
}
pub fn set_project_context(&self, context: ProjectContext) {
self.memory.with_working_context(|ctx| {
ctx.set_project(&context.name);
});
*self.current_context.write() = Some(context);
}
pub fn project_context(&self) -> Option<ProjectContext> {
self.current_context.read().clone()
}
pub fn clear_project_context(&self) {
*self.current_context.write() = None;
}
pub fn memory(&self) -> &AgentMemory {
&self.memory
}
pub fn learn_task_transferable(
&self,
task: &str,
code: &str,
outcome: TaskOutcome,
learnings: Vec<&str>,
transfer_level: Option<TransferLevel>,
domain: Option<KnowledgeDomain>,
) -> Result<VectorId> {
let concepts = self.extractor.extract(task, code);
let level =
transfer_level.unwrap_or_else(|| self.extractor.infer_transfer_level(&concepts));
let domain = domain.unwrap_or_else(|| {
self.current_context
.read()
.as_ref()
.map(|c| c.domain.clone())
.unwrap_or(KnowledgeDomain::General)
});
let mut episode = TaskEpisode {
task: task.to_string(),
code: code.to_string(),
outcome,
steps: Vec::new(),
learnings: learnings.iter().map(|s| s.to_string()).collect(),
errors: Vec::new(),
language: self
.current_context
.read()
.as_ref()
.map(|c| Language::from_str(&c.language))
.unwrap_or(Language::Other("unknown".into())),
project: self.current_context.read().as_ref().map(|c| c.name.clone()),
duration_secs: None,
tags: Vec::new(),
};
episode.tags.push(format!("transfer:{}", level.as_str()));
episode.tags.push(format!("domain:{}", domain.as_str()));
for concept in &concepts {
episode.tags.push(format!(
"concept:{}",
concept.to_lowercase().replace(' ', "_")
));
}
self.memory.learn_episode(episode)
}
pub fn learn_code_transferable(
&self,
snippet: CodeSnippet,
transfer_level: Option<TransferLevel>,
domain: Option<KnowledgeDomain>,
) -> Result<VectorId> {
let concepts = self.extractor.extract(&snippet.description, &snippet.code);
let level =
transfer_level.unwrap_or_else(|| self.extractor.infer_transfer_level(&concepts));
let domain = domain.unwrap_or(KnowledgeDomain::General);
let mut enriched = snippet;
enriched.tags.push(format!("transfer:{}", level.as_str()));
enriched.tags.push(format!("domain:{}", domain.as_str()));
for concept in &concepts {
enriched.tags.push(format!(
"concept:{}",
concept.to_lowercase().replace(' ', "_")
));
}
self.memory.learn_code(enriched)
}
pub fn recall_transferable(&self, query: &str, k: usize) -> Result<Vec<TransferableRecall>> {
let all_results = self.memory.recall_similar(query, k * 3)?;
let mut ranked: Vec<TransferableRecall> = all_results
.into_iter()
.map(|recall| {
let concepts = self.extract_concepts_from_recall(&recall);
let transfer_level = self.infer_level_from_recall(&recall, &concepts);
let transfer_score = self.calculate_transfer_score(&recall, &concepts);
let combined_score = recall.relevance_score * self.relevance_weight
+ transfer_score * (1.0 - self.relevance_weight);
TransferableRecall {
adaptation_needed: self.get_adaptation_needed(&recall),
applicable: transfer_score >= self.transfer_threshold,
recall,
transfer_score,
combined_score,
transfer_level,
concepts,
}
})
.collect();
ranked.sort_by(|a, b| {
b.combined_score
.partial_cmp(&a.combined_score)
.unwrap_or(std::cmp::Ordering::Equal)
});
ranked.truncate(k);
Ok(ranked)
}
pub fn recall_universal(&self, query: &str, k: usize) -> Result<Vec<TransferableRecall>> {
let results = self.recall_transferable(query, k * 2)?;
Ok(results
.into_iter()
.filter(|r| r.transfer_level == TransferLevel::Universal)
.take(k)
.collect())
}
pub fn recall_same_domain(&self, query: &str, k: usize) -> Result<Vec<TransferableRecall>> {
let current_domain = self
.current_context
.read()
.as_ref()
.map(|c| c.domain.clone());
let results = self.recall_transferable(query, k * 2)?;
Ok(results
.into_iter()
.filter(|r| {
if let Some(ref domain) = current_domain {
self.is_domain_compatible(&r.recall, domain)
} else {
true
}
})
.take(k)
.collect())
}
pub fn recall_same_stack(&self, query: &str, k: usize) -> Result<Vec<TransferableRecall>> {
let current_lang = self
.current_context
.read()
.as_ref()
.map(|c| c.language.clone());
let results = self.recall_transferable(query, k * 2)?;
Ok(results
.into_iter()
.filter(|r| {
if let Some(ref lang) = current_lang {
self.is_language_compatible(&r.recall, lang)
} else {
true
}
})
.take(k)
.collect())
}
fn extract_concepts_from_recall(&self, recall: &MemoryRecall) -> Vec<String> {
let mut concepts = Vec::new();
if let Some(ref meta) = recall.metadata {
if let Some(MetadataValue::String(tags)) = meta.get("tags") {
for tag in tags.split(',') {
if tag.starts_with("concept:") {
concepts.push(tag.trim_start_matches("concept:").replace('_', " "));
}
}
}
}
if concepts.is_empty() {
concepts = self.extractor.extract(&recall.content, "");
}
concepts
}
fn infer_level_from_recall(&self, recall: &MemoryRecall, concepts: &[String]) -> TransferLevel {
if let Some(ref meta) = recall.metadata {
if let Some(MetadataValue::String(tags)) = meta.get("tags") {
for tag in tags.split(',') {
if tag.starts_with("transfer:") {
return TransferLevel::from_str(tag.trim_start_matches("transfer:"))
.unwrap_or(TransferLevel::Instance);
}
}
}
}
self.extractor.infer_transfer_level(concepts)
}
fn calculate_transfer_score(&self, recall: &MemoryRecall, concepts: &[String]) -> f32 {
let current = self.current_context.read();
let Some(ref ctx) = *current else {
return 0.5; };
let mut score = 0.0;
let level = self.infer_level_from_recall(recall, concepts);
score += level.transfer_score() * 0.3;
if self.is_domain_compatible(recall, &ctx.domain) {
score += 0.25;
} else if self.is_related_domain(recall, &ctx.domain) {
score += 0.12;
}
let lang_compat = self.get_language_compatibility(recall, &ctx.language);
score += lang_compat * 0.3;
let concept_overlap = self.calculate_concept_overlap(concepts, ctx);
score += concept_overlap * 0.15;
score.min(1.0)
}
fn is_domain_compatible(&self, recall: &MemoryRecall, domain: &KnowledgeDomain) -> bool {
if let Some(ref meta) = recall.metadata {
if let Some(MetadataValue::String(tags)) = meta.get("tags") {
let domain_tag = format!("domain:{}", domain.as_str());
return tags.contains(&domain_tag);
}
}
false
}
fn is_related_domain(&self, recall: &MemoryRecall, domain: &KnowledgeDomain) -> bool {
let related = domain.related_domains();
for rel_domain in related {
if self.is_domain_compatible(recall, &rel_domain) {
return true;
}
}
false
}
fn get_language_compatibility(&self, recall: &MemoryRecall, current_lang: &str) -> f32 {
if let Some(ref meta) = recall.metadata {
if let Some(MetadataValue::String(lang)) = meta.get("language") {
return LanguageCompatibility::compatibility(lang, current_lang);
}
}
0.5 }
fn is_language_compatible(&self, recall: &MemoryRecall, current_lang: &str) -> bool {
self.get_language_compatibility(recall, current_lang) >= 0.6
}
fn calculate_concept_overlap(&self, concepts: &[String], ctx: &ProjectContext) -> f32 {
if concepts.is_empty() {
return 0.0;
}
let matches = concepts
.iter()
.filter(|c| {
ctx.patterns.iter().any(|p| {
p.to_lowercase().contains(&c.to_lowercase())
|| c.to_lowercase().contains(&p.to_lowercase())
})
})
.count();
matches as f32 / concepts.len().max(ctx.patterns.len()).max(1) as f32
}
fn get_adaptation_needed(&self, recall: &MemoryRecall) -> Option<String> {
let current = self.current_context.read();
let Some(ref ctx) = *current else {
return None;
};
if let Some(ref meta) = recall.metadata {
if let Some(MetadataValue::String(source_lang)) = meta.get("language") {
let target_lang = ctx.language.as_str();
return LanguageCompatibility::adaptation_description(source_lang, target_lang);
}
}
None
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_transfer_level_ordering() {
assert!(TransferLevel::Universal > TransferLevel::Domain);
assert!(TransferLevel::Domain > TransferLevel::Context);
assert!(TransferLevel::Context > TransferLevel::Instance);
}
#[test]
fn test_transfer_level_score() {
assert!((TransferLevel::Universal.transfer_score() - 1.0).abs() < 0.01);
assert!((TransferLevel::Instance.transfer_score() - 0.25).abs() < 0.01);
}
#[test]
fn test_language_compatibility() {
assert!((LanguageCompatibility::compatibility("rust", "rust") - 1.0).abs() < 0.01);
assert!(LanguageCompatibility::compatibility("typescript", "javascript") > 0.5);
assert!(LanguageCompatibility::compatibility("python", "ruby") > 0.5);
assert!(LanguageCompatibility::compatibility("rust", "python") < 0.5);
}
#[test]
fn test_concept_extraction() {
let extractor = ConceptExtractor::new();
let concepts = extractor.extract(
"Implement JWT authentication with rate limiting",
"middleware auth jwt token verify",
);
assert!(concepts.contains(&"JWT Tokens".to_string()));
assert!(concepts.contains(&"Authentication".to_string()));
assert!(concepts.contains(&"Rate Limiting".to_string()));
assert!(concepts.contains(&"Middleware Pattern".to_string()));
}
#[test]
fn test_infer_transfer_level() {
let extractor = ConceptExtractor::new();
let concepts = vec!["SOLID Principles".to_string()];
assert_eq!(
extractor.infer_transfer_level(&concepts),
TransferLevel::Universal
);
let concepts = vec!["Factory Pattern".to_string()];
assert_eq!(
extractor.infer_transfer_level(&concepts),
TransferLevel::Domain
);
let concepts = vec!["Authentication".to_string()];
assert_eq!(
extractor.infer_transfer_level(&concepts),
TransferLevel::Domain
);
let concepts: Vec<String> = vec![];
assert_eq!(
extractor.infer_transfer_level(&concepts),
TransferLevel::Context
);
}
#[test]
fn test_domain_related() {
let web = KnowledgeDomain::WebBackend;
let related = web.related_domains();
assert!(related.contains(&KnowledgeDomain::Database));
assert!(related.contains(&KnowledgeDomain::Security));
assert!(!related.contains(&KnowledgeDomain::GameDev));
}
#[test]
fn test_project_context_builder() {
let ctx = ProjectContext::new("my-api", "rust", KnowledgeDomain::WebBackend)
.with_frameworks(vec!["Axum".into(), "SQLx".into()])
.with_patterns(vec!["REST".into(), "Clean Architecture".into()]);
assert_eq!(ctx.name, "my-api");
assert_eq!(ctx.frameworks.len(), 2);
assert_eq!(ctx.patterns.len(), 2);
}
#[test]
fn test_transferable_memory_creation() {
let config = MemoryConfig::small();
let memory = TransferableMemory::new(config);
assert!(memory.is_ok());
}
#[test]
fn test_set_project_context() {
let config = MemoryConfig::small();
let memory = TransferableMemory::new(config).unwrap();
assert!(memory.project_context().is_none());
memory.set_project_context(ProjectContext::new(
"test-project",
"rust",
KnowledgeDomain::CLI,
));
let ctx = memory.project_context();
assert!(ctx.is_some());
assert_eq!(ctx.unwrap().name, "test-project");
}
}