use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
use super::EncoderConfig;
const MAX_FILE_SIZE_BYTES: usize = 100 * 1024;
const APPROX_CHARS_PER_TOKEN: f32 = 4.0;
pub struct SemanticEncoder {
config: EncoderConfig,
tokenizer: Arc<RwLock<Option<CodeTokenizer>>>,
model: Arc<RwLock<Option<CodeModel>>>,
cache: Arc<RwLock<EmbeddingCache>>,
}
pub struct CodeTokenizer {
model_type: String,
max_tokens: usize,
vocab_size: usize,
special_tokens: HashMap<String, u32>,
}
pub struct CodeModel {
model_type: String,
embedding_dim: usize,
device: String,
initialized: bool,
}
#[derive(Default)]
pub struct EmbeddingCache {
cache: HashMap<String, CachedEmbedding>,
access_order: Vec<String>,
max_size: usize,
}
#[derive(Clone)]
pub struct CachedEmbedding {
embedding: Vec<f32>,
timestamp: std::time::Instant,
access_count: u32,
}
#[derive(Debug, Clone)]
pub struct TokenizedCode {
pub tokens: Vec<u32>,
pub attention_mask: Vec<u32>,
pub token_type_ids: Vec<u32>,
pub original_length: usize,
pub truncated: bool,
pub language: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeEmbedding {
pub embedding: Vec<f32>,
pub dimension: usize,
pub model_version: String,
pub content_hash: String,
pub language: Option<String>,
pub token_count: usize,
pub inference_time_ms: u64,
}
impl SemanticEncoder {
pub async fn new(config: EncoderConfig) -> Result<Self> {
info!("Creating semantic encoder: {}", config.model_type);
info!("Max tokens: {}, embedding dim: {}", config.max_tokens, config.embedding_dim);
Ok(Self {
config,
tokenizer: Arc::new(RwLock::new(None)),
model: Arc::new(RwLock::new(None)),
cache: Arc::new(RwLock::new(EmbeddingCache::new(1000))), })
}
pub async fn initialize(&self) -> Result<()> {
info!("Initializing semantic encoder: {}", self.config.model_type);
let tokenizer = self.create_tokenizer().await
.context("Failed to create tokenizer")?;
*self.tokenizer.write().await = Some(tokenizer);
let model = self.create_model().await
.context("Failed to create model")?;
*self.model.write().await = Some(model);
info!("Semantic encoder initialized successfully");
Ok(())
}
pub async fn encode(&self, content: &str, language: Option<&str>) -> Result<CodeEmbedding> {
let start_time = std::time::Instant::now();
if content.len() > MAX_FILE_SIZE_BYTES {
warn!("Content size {}KB exceeds limit {}KB, truncating",
content.len() / 1024, MAX_FILE_SIZE_BYTES / 1024);
}
let content_hash = self.hash_content(content);
if let Some(cached) = self.get_cached_embedding(&content_hash).await {
debug!("Cache hit for content hash: {}", content_hash);
return Ok(cached);
}
let tokenized = self.tokenize_content(content, language).await
.context("Failed to tokenize content")?;
let embedding_vec = self.generate_embedding(&tokenized).await
.context("Failed to generate embedding")?;
let inference_time = start_time.elapsed().as_millis() as u64;
let embedding = CodeEmbedding {
embedding: embedding_vec,
dimension: self.config.embedding_dim,
model_version: self.config.model_type.clone(),
content_hash: content_hash.clone(),
language: language.map(String::from),
token_count: tokenized.tokens.len(),
inference_time_ms: inference_time,
};
self.cache_embedding(&content_hash, &embedding).await;
if inference_time > 50 {
warn!("Slow inference: {}ms > 50ms target", inference_time);
}
debug!("Encoded {} tokens in {}ms", tokenized.tokens.len(), inference_time);
Ok(embedding)
}
pub async fn encode_batch(&self, contents: &[(&str, Option<&str>)]) -> Result<Vec<CodeEmbedding>> {
let start_time = std::time::Instant::now();
if contents.len() > self.config.batch_size {
warn!("Batch size {} exceeds limit {}, processing in chunks",
contents.len(), self.config.batch_size);
}
let mut results = Vec::with_capacity(contents.len());
for chunk in contents.chunks(self.config.batch_size) {
let chunk_results = self.encode_chunk(chunk).await?;
results.extend(chunk_results);
}
let total_time = start_time.elapsed().as_millis() as u64;
let avg_time_per_item = total_time / contents.len() as u64;
info!("Batch encoded {} items in {}ms (avg {}ms/item)",
contents.len(), total_time, avg_time_per_item);
Ok(results)
}
pub fn can_handle_content(&self, content: &str) -> bool {
let estimated_tokens = (content.len() as f32 / APPROX_CHARS_PER_TOKEN) as usize;
estimated_tokens <= self.config.max_tokens && content.len() <= MAX_FILE_SIZE_BYTES
}
pub async fn get_metrics(&self) -> EncoderMetrics {
let cache_guard = self.cache.read().await;
EncoderMetrics {
cache_size: cache_guard.cache.len(),
cache_hit_rate: cache_guard.calculate_hit_rate(),
max_tokens: self.config.max_tokens,
embedding_dim: self.config.embedding_dim,
model_type: self.config.model_type.clone(),
}
}
async fn create_tokenizer(&self) -> Result<CodeTokenizer> {
info!("Loading tokenizer for {}", self.config.model_type);
let mut special_tokens = HashMap::new();
special_tokens.insert("<pad>".to_string(), 0);
special_tokens.insert("<unk>".to_string(), 1);
special_tokens.insert("<s>".to_string(), 2);
special_tokens.insert("</s>".to_string(), 3);
Ok(CodeTokenizer {
model_type: self.config.model_type.clone(),
max_tokens: self.config.max_tokens,
vocab_size: 50000, special_tokens,
})
}
async fn create_model(&self) -> Result<CodeModel> {
info!("Loading model {} on device {}", self.config.model_type, self.config.device);
Ok(CodeModel {
model_type: self.config.model_type.clone(),
embedding_dim: self.config.embedding_dim,
device: self.config.device.clone(),
initialized: true,
})
}
async fn tokenize_content(&self, content: &str, language: Option<&str>) -> Result<TokenizedCode> {
let tokenizer_guard = self.tokenizer.read().await;
let tokenizer = tokenizer_guard.as_ref()
.ok_or_else(|| anyhow::anyhow!("Tokenizer not initialized"))?;
let max_chars = (self.config.max_tokens as f32 * APPROX_CHARS_PER_TOKEN) as usize;
let content = if content.len() > max_chars {
&content[..max_chars]
} else {
content
};
let tokens = self.mock_tokenize(content, tokenizer).await?;
let token_count = tokens.len().min(self.config.max_tokens);
let truncated = tokens.len() > self.config.max_tokens;
Ok(TokenizedCode {
tokens: tokens[..token_count].to_vec(),
attention_mask: vec![1; token_count],
token_type_ids: vec![0; token_count],
original_length: content.len(),
truncated,
language: language.map(String::from),
})
}
async fn generate_embedding(&self, tokenized: &TokenizedCode) -> Result<Vec<f32>> {
let model_guard = self.model.read().await;
let model = model_guard.as_ref()
.ok_or_else(|| anyhow::anyhow!("Model not initialized"))?;
let embedding = self.mock_inference(tokenized, model).await?;
Ok(embedding)
}
async fn encode_chunk(&self, chunk: &[(&str, Option<&str>)]) -> Result<Vec<CodeEmbedding>> {
let mut results = Vec::with_capacity(chunk.len());
for (content, language) in chunk {
let embedding = self.encode(content, *language).await?;
results.push(embedding);
}
Ok(results)
}
fn hash_content(&self, content: &str) -> String {
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
let mut hasher = DefaultHasher::new();
content.hash(&mut hasher);
self.config.model_type.hash(&mut hasher);
format!("{:x}", hasher.finish())
}
async fn get_cached_embedding(&self, hash: &str) -> Option<CodeEmbedding> {
let mut cache_guard = self.cache.write().await;
cache_guard.get(hash).map(|cached| CodeEmbedding {
embedding: cached.embedding.clone(),
dimension: self.config.embedding_dim,
model_version: self.config.model_type.clone(),
content_hash: hash.to_string(),
language: None, token_count: 0, inference_time_ms: 0, })
}
async fn cache_embedding(&self, hash: &str, embedding: &CodeEmbedding) {
let mut cache_guard = self.cache.write().await;
cache_guard.insert(hash.to_string(), CachedEmbedding {
embedding: embedding.embedding.clone(),
timestamp: std::time::Instant::now(),
access_count: 1,
});
}
async fn mock_tokenize(&self, content: &str, _tokenizer: &CodeTokenizer) -> Result<Vec<u32>> {
let words: Vec<&str> = content.split_whitespace().collect();
let mut tokens = Vec::with_capacity(words.len() + 2);
tokens.push(2);
for word in words {
let hash = word.bytes().fold(0u32, |acc, b| acc.wrapping_mul(31).wrapping_add(b as u32));
tokens.push(hash % 50000); }
tokens.push(3);
Ok(tokens)
}
async fn mock_inference(&self, tokenized: &TokenizedCode, _model: &CodeModel) -> Result<Vec<f32>> {
let mut embedding = vec![0.0f32; self.config.embedding_dim];
for (i, &token) in tokenized.tokens.iter().enumerate() {
let idx = (token as usize + i) % self.config.embedding_dim;
embedding[idx] += (token as f32) / 50000.0 - 0.5;
}
let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
if norm > 0.0 {
for x in &mut embedding {
*x /= norm;
}
}
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
Ok(embedding)
}
}
impl EmbeddingCache {
fn new(max_size: usize) -> Self {
Self {
cache: HashMap::new(),
access_order: Vec::new(),
max_size,
}
}
fn get(&mut self, key: &str) -> Option<&CachedEmbedding> {
if let Some(embedding) = self.cache.get_mut(key) {
embedding.access_count += 1;
if let Some(pos) = self.access_order.iter().position(|x| x == key) {
self.access_order.remove(pos);
}
self.access_order.push(key.to_string());
Some(embedding)
} else {
None
}
}
fn insert(&mut self, key: String, value: CachedEmbedding) {
if self.cache.len() >= self.max_size && !self.cache.contains_key(&key) {
if let Some(lru_key) = self.access_order.first().cloned() {
self.cache.remove(&lru_key);
self.access_order.retain(|k| k != &lru_key);
}
}
self.cache.insert(key.clone(), value);
self.access_order.push(key);
}
fn calculate_hit_rate(&self) -> f32 {
if self.cache.is_empty() {
return 0.0;
}
let total_accesses: u32 = self.cache.values().map(|e| e.access_count).sum();
let hits = total_accesses.saturating_sub(self.cache.len() as u32);
if total_accesses == 0 {
0.0
} else {
hits as f32 / total_accesses as f32
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EncoderMetrics {
pub cache_size: usize,
pub cache_hit_rate: f32,
pub max_tokens: usize,
pub embedding_dim: usize,
pub model_type: String,
}
pub async fn initialize_encoder(config: &EncoderConfig) -> Result<()> {
info!("Initializing semantic encoder with config: {:?}", config);
if config.max_tokens > 4096 {
warn!("Max tokens {} exceeds recommended 4096", config.max_tokens);
}
if config.max_tokens < 512 {
anyhow::bail!("Max tokens {} too small, minimum 512", config.max_tokens);
}
info!("Encoder initialization complete");
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_encoder_creation() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 16,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config).await.unwrap();
assert!(encoder.can_handle_content("test")); }
#[tokio::test]
async fn test_content_limits() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 16,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config).await.unwrap();
let small_content = "def hello(): return 'world'";
assert!(encoder.can_handle_content(small_content));
let large_content = "x".repeat(200_000);
assert!(!encoder.can_handle_content(&large_content));
}
#[tokio::test]
async fn test_cache_operations() {
let mut cache = EmbeddingCache::new(2);
let embedding1 = CachedEmbedding {
embedding: vec![1.0, 2.0],
timestamp: std::time::Instant::now(),
access_count: 1,
};
let embedding2 = CachedEmbedding {
embedding: vec![3.0, 4.0],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert("key1".to_string(), embedding1.clone());
assert_eq!(cache.cache.len(), 1);
cache.insert("key2".to_string(), embedding2.clone());
assert_eq!(cache.cache.len(), 2);
let embedding3 = CachedEmbedding {
embedding: vec![5.0, 6.0],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert("key3".to_string(), embedding3);
assert_eq!(cache.cache.len(), 2);
assert!(!cache.cache.contains_key("key1")); }
#[tokio::test]
async fn test_encoder_error_handling() {
let invalid_config = EncoderConfig {
model_type: "invalid-model".to_string(),
max_tokens: 0, embedding_dim: 0, batch_size: 0, device: "invalid-device".to_string(),
};
let result = SemanticEncoder::new(invalid_config).await;
assert!(result.is_ok() || result.is_err());
}
#[tokio::test]
async fn test_batch_encoding() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 4,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config).await.unwrap();
encoder.initialize().await.unwrap();
let batch_content = vec![
"function add(a, b) { return a + b; }",
"class User { constructor(name) { this.name = name; } }",
"def calculate_sum(numbers): return sum(numbers)",
"public static void main(String[] args) { System.out.println(\"Hello\"); }",
];
for content in batch_content {
let result = encoder.encode(content, None).await;
assert!(result.is_ok(), "Batch encoding should work for: {}", content);
}
}
#[tokio::test]
async fn test_cache_hit_rate_calculation() {
let mut cache = EmbeddingCache::new(10);
assert_eq!(cache.calculate_hit_rate(), 0.0);
let embedding = CachedEmbedding {
embedding: vec![1.0, 2.0, 3.0],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert("key1".to_string(), embedding.clone());
cache.insert("key2".to_string(), CachedEmbedding {
access_count: 3,
..embedding.clone()
});
cache.insert("key3".to_string(), CachedEmbedding {
access_count: 5,
..embedding
});
let hit_rate = cache.calculate_hit_rate();
assert!(hit_rate >= 0.0 && hit_rate <= 1.0);
}
#[tokio::test]
async fn test_content_tokenization_estimation() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 1024,
embedding_dim: 512,
batch_size: 8,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config).await.unwrap();
encoder.initialize().await.unwrap();
let large_content = "x".repeat(4000); let comment_content = "// ".repeat(100);
let test_cases = vec![
("x", true), ("def test(): pass", true), (comment_content.as_str(), true), (large_content.as_str(), true), ];
for (content, should_handle) in test_cases {
let can_handle = encoder.can_handle_content(content);
assert_eq!(can_handle, should_handle,
"Content handling mismatch for length: {}", content.len());
}
}
#[tokio::test]
async fn test_encoder_with_different_languages() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 16,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config).await.unwrap();
encoder.initialize().await.unwrap();
let language_samples = vec![
("rust", "fn main() { println!(\"Hello, world!\"); }"),
("python", "def hello_world():\n print('Hello, world!')"),
("javascript", "function helloWorld() { console.log('Hello, world!'); }"),
("go", "package main\n\nfunc main() { fmt.Println(\"Hello, world!\") }"),
("java", "public class HelloWorld {\n public static void main(String[] args) {\n System.out.println(\"Hello, world!\");\n }\n}"),
];
for (language, code) in language_samples {
let result = encoder.encode(code, Some(language)).await;
assert!(result.is_ok(), "Should handle {} code: {}", language, code);
if let Ok(embedding) = result {
assert!(!embedding.embedding.is_empty());
assert!(embedding.dimension > 0);
}
}
}
#[tokio::test]
async fn test_concurrent_encoding() {
let config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 4,
device: "cpu".to_string(),
};
let encoder = Arc::new(SemanticEncoder::new(config).await.unwrap());
encoder.initialize().await.unwrap();
let mut handles = vec![];
for i in 0..8 {
let encoder_clone = encoder.clone();
let handle = tokio::spawn(async move {
let content = format!("function test_{}() {{ return {}; }}", i, i);
encoder_clone.encode(&content, Some("javascript")).await
});
handles.push(handle);
}
let mut successful = 0;
for handle in handles {
if let Ok(result) = handle.await {
if result.is_ok() {
successful += 1;
}
}
}
assert!(successful >= 6, "Most concurrent encodings should succeed");
}
#[tokio::test]
async fn test_encoder_metrics() {
let config = EncoderConfig {
model_type: "test-model".to_string(),
max_tokens: 1024,
embedding_dim: 256,
batch_size: 8,
device: "cpu".to_string(),
};
let encoder = SemanticEncoder::new(config.clone()).await.unwrap();
let metrics = encoder.get_metrics().await;
assert_eq!(metrics.max_tokens, config.max_tokens);
assert_eq!(metrics.embedding_dim, config.embedding_dim);
assert_eq!(metrics.model_type, config.model_type);
assert!(metrics.cache_hit_rate >= 0.0);
}
#[tokio::test]
async fn test_cache_eviction_policy() {
let mut cache = EmbeddingCache::new(3);
for i in 0..3 {
let embedding = CachedEmbedding {
embedding: vec![i as f32; 5],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert(format!("key{}", i), embedding);
}
assert_eq!(cache.cache.len(), 3);
if let Some(entry) = cache.cache.get_mut("key1") {
entry.access_count += 1;
}
let new_embedding = CachedEmbedding {
embedding: vec![99.0; 5],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert("new_key".to_string(), new_embedding);
assert_eq!(cache.cache.len(), 3);
assert!(cache.cache.contains_key("new_key"));
}
#[test]
fn test_tokenizer_creation() {
let tokenizer = CodeTokenizer {
model_type: "test-model".to_string(),
max_tokens: 1024,
vocab_size: 50000,
special_tokens: HashMap::new(),
};
assert_eq!(tokenizer.model_type, "test-model");
assert_eq!(tokenizer.max_tokens, 1024);
assert_eq!(tokenizer.vocab_size, 50000);
}
#[test]
fn test_code_model_initialization() {
let model = CodeModel {
model_type: "codet5-small".to_string(),
embedding_dim: 512,
device: "cpu".to_string(),
initialized: false,
};
assert_eq!(model.embedding_dim, 512);
assert_eq!(model.device, "cpu");
assert!(!model.initialized);
}
#[tokio::test]
async fn test_initialization_validation() {
let valid_config = EncoderConfig {
model_type: "codet5-base".to_string(),
max_tokens: 2048,
embedding_dim: 768,
batch_size: 16,
device: "cpu".to_string(),
};
let result = initialize_encoder(&valid_config).await;
assert!(result.is_ok(), "Valid configuration should initialize successfully");
let invalid_config = EncoderConfig {
max_tokens: 256, ..valid_config.clone()
};
let result = initialize_encoder(&invalid_config).await;
assert!(result.is_err(), "Configuration with too few tokens should fail");
let warning_config = EncoderConfig {
max_tokens: 8192, ..valid_config
};
let result = initialize_encoder(&warning_config).await;
assert!(result.is_ok(), "Large token configuration should succeed with warning");
}
#[test]
fn test_cache_edge_cases() {
let mut cache = EmbeddingCache::new(1);
assert_eq!(cache.cache.len(), 0);
assert_eq!(cache.calculate_hit_rate(), 0.0);
let embedding = CachedEmbedding {
embedding: vec![1.0, 2.0],
timestamp: std::time::Instant::now(),
access_count: 1,
};
cache.insert("single".to_string(), embedding);
assert_eq!(cache.cache.len(), 1);
let zero_cache = EmbeddingCache::new(0);
assert_eq!(zero_cache.max_size, 0);
}
}