use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use crate::{
ChatCompletionRequest, ChatCompletionResponse, ChatMessage, ChatProvider, EmbeddingInput,
EmbeddingRequest, Result, SdkError, ZhipuClient,
};
static MEMORY_SEQUENCE: AtomicU64 = AtomicU64::new(1);
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MemoryDocument {
pub id: String,
pub text: String,
#[serde(default)]
pub metadata: BTreeMap<String, String>,
}
impl MemoryDocument {
pub fn new(id: impl Into<String>, text: impl Into<String>) -> Self {
Self {
id: id.into(),
text: text.into(),
metadata: BTreeMap::new(),
}
}
pub fn metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.metadata.insert(key.into(), value.into());
self
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct QuantizedVector {
pub scale: f32,
pub values: Vec<i8>,
}
impl QuantizedVector {
pub fn compress(values: &[f32]) -> Result<Self> {
if values.is_empty() || values.iter().any(|value| !value.is_finite()) {
return Err(SdkError::Validation(
"embedding must contain finite values".into(),
));
}
let max = values
.iter()
.map(|value| value.abs())
.fold(0.0_f32, f32::max);
let scale = if max == 0.0 { 1.0 } else { max / 127.0 };
let values = values
.iter()
.map(|value| (value / scale).round().clamp(-127.0, 127.0) as i8)
.collect();
Ok(Self { scale, values })
}
pub fn dimensions(&self) -> usize {
self.values.len()
}
pub fn decompress(&self) -> Vec<f32> {
self.values
.iter()
.map(|value| *value as f32 * self.scale)
.collect()
}
pub fn cosine_similarity(&self, query: &[f32]) -> Result<f32> {
if query.len() != self.values.len()
|| query.is_empty()
|| query.iter().any(|value| !value.is_finite())
{
return Err(SdkError::Validation(
"query embedding dimensions must match stored vector".into(),
));
}
let mut dot = 0.0_f32;
let mut stored_norm = 0.0_f32;
let mut query_norm = 0.0_f32;
for (stored, query) in self.values.iter().zip(query) {
let stored = *stored as f32 * self.scale;
dot += stored * query;
stored_norm += stored * stored;
query_norm += query * query;
}
if stored_norm == 0.0 || query_norm == 0.0 {
return Ok(0.0);
}
Ok(dot / (stored_norm.sqrt() * query_norm.sqrt()))
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct StoredMemory {
pub document: MemoryDocument,
pub embedding: QuantizedVector,
}
#[derive(Debug, Clone, PartialEq)]
pub struct MemoryMatch {
pub document: MemoryDocument,
pub score: f32,
}
#[async_trait]
pub trait EmbeddingProvider: Send + Sync {
async fn embed(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>>;
}
#[async_trait]
pub trait VectorStore: Send + Sync {
async fn upsert(&self, memory: StoredMemory) -> Result<()>;
async fn search(&self, query: &[f32], limit: usize) -> Result<Vec<MemoryMatch>>;
async fn clear(&self) -> Result<()>;
async fn len(&self) -> Result<usize>;
async fn is_empty(&self) -> Result<bool> {
Ok(self.len().await? == 0)
}
}
#[derive(Default)]
pub struct InMemoryVectorStore {
records: RwLock<Vec<StoredMemory>>,
}
impl InMemoryVectorStore {
pub fn new() -> Self {
Self::default()
}
pub fn snapshot(&self) -> Result<Vec<StoredMemory>> {
self.records
.read()
.map(|records| records.clone())
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))
}
pub fn restore(&self, records: Vec<StoredMemory>) -> Result<()> {
let mut current = self
.records
.write()
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))?;
*current = records;
Ok(())
}
pub fn snapshot_json(&self) -> Result<String> {
serde_json::to_string(&self.snapshot()?)
.map_err(|error| SdkError::Configuration(error.to_string().into()))
}
pub fn restore_json(&self, value: &str) -> Result<()> {
let records = serde_json::from_str(value)
.map_err(|error| SdkError::Configuration(error.to_string().into()))?;
self.restore(records)
}
}
#[async_trait]
impl VectorStore for InMemoryVectorStore {
async fn upsert(&self, memory: StoredMemory) -> Result<()> {
let mut records = self
.records
.write()
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))?;
if let Some(existing) = records
.iter_mut()
.find(|existing| existing.document.id == memory.document.id)
{
*existing = memory;
} else {
records.push(memory);
}
Ok(())
}
async fn search(&self, query: &[f32], limit: usize) -> Result<Vec<MemoryMatch>> {
if limit == 0 {
return Ok(Vec::new());
}
let records = self
.records
.read()
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))?;
let mut matches = records
.iter()
.map(|memory| {
memory
.embedding
.cosine_similarity(query)
.map(|score| MemoryMatch {
document: memory.document.clone(),
score,
})
})
.collect::<Result<Vec<_>>>()?;
matches.sort_by(|left, right| right.score.total_cmp(&left.score));
matches.truncate(limit);
Ok(matches)
}
async fn clear(&self) -> Result<()> {
self.records
.write()
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))?
.clear();
Ok(())
}
async fn len(&self) -> Result<usize> {
self.records
.read()
.map(|records| records.len())
.map_err(|_| SdkError::Configuration("vector store lock is poisoned".into()))
}
}
#[async_trait]
pub trait ConversationMemory: Send + Sync {
async fn remember(&self, document: MemoryDocument) -> Result<()>;
async fn recall(&self, query: &str, limit: usize) -> Result<Vec<MemoryMatch>>;
async fn clear(&self) -> Result<()>;
}
pub struct SemanticMemory {
embeddings: Arc<dyn EmbeddingProvider>,
store: Arc<dyn VectorStore>,
}
impl SemanticMemory {
pub fn new(embeddings: Arc<dyn EmbeddingProvider>, store: Arc<dyn VectorStore>) -> Self {
Self { embeddings, store }
}
}
#[async_trait]
impl ConversationMemory for SemanticMemory {
async fn remember(&self, document: MemoryDocument) -> Result<()> {
if document.text.trim().is_empty() {
return Err(SdkError::Validation("memory text cannot be empty".into()));
}
let vectors = self
.embeddings
.embed(std::slice::from_ref(&document.text))
.await?;
let vector = vectors.into_iter().next().ok_or_else(|| SdkError::Decode {
message: "embedding response did not contain a vector".into(),
body: String::new(),
})?;
self.store
.upsert(StoredMemory {
document,
embedding: QuantizedVector::compress(&vector)?,
})
.await
}
async fn recall(&self, query: &str, limit: usize) -> Result<Vec<MemoryMatch>> {
if query.trim().is_empty() || limit == 0 {
return Ok(Vec::new());
}
let vectors = self.embeddings.embed(&[query.to_owned()]).await?;
let vector = vectors.first().ok_or_else(|| SdkError::Decode {
message: "embedding response did not contain a vector".into(),
body: String::new(),
})?;
self.store.search(vector, limit).await
}
async fn clear(&self) -> Result<()> {
self.store.clear().await
}
}
#[derive(Clone)]
pub struct ZhipuEmbeddingProvider {
client: ZhipuClient,
model: String,
dimensions: Option<u32>,
}
impl ZhipuEmbeddingProvider {
pub fn new(client: ZhipuClient, model: impl Into<String>) -> Self {
Self {
client,
model: model.into(),
dimensions: None,
}
}
pub fn dimensions(mut self, value: u32) -> Self {
self.dimensions = Some(value);
self
}
}
#[async_trait]
impl EmbeddingProvider for ZhipuEmbeddingProvider {
async fn embed(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>> {
if inputs.is_empty() {
return Ok(Vec::new());
}
let response = self
.client
.embedding(&EmbeddingRequest {
model: self.model.clone(),
input: EmbeddingInput::Texts(inputs.to_vec()),
dimensions: self.dimensions,
encoding_format: Some("float".into()),
user_id: None,
request_id: None,
extra: Default::default(),
})
.await?;
let mut data = response.data;
data.sort_by_key(|item| item.index);
let vectors = data
.into_iter()
.map(|item| item.embedding)
.collect::<Vec<_>>();
if vectors.len() != inputs.len() || vectors.iter().any(Vec::is_empty) {
return Err(SdkError::Decode {
message: "embedding response count or dimensions are invalid".into(),
body: String::new(),
});
}
Ok(vectors)
}
}
#[derive(Clone)]
pub struct ConversationConfig {
pub model: String,
pub system_prompt: Option<String>,
pub retain_history: bool,
pub max_history_messages: usize,
pub memory: Option<Arc<dyn ConversationMemory>>,
pub recall_limit: usize,
}
impl ConversationConfig {
pub fn new(model: impl Into<String>) -> Self {
Self {
model: model.into(),
system_prompt: None,
retain_history: false,
max_history_messages: 20,
memory: None,
recall_limit: 4,
}
}
pub fn system_prompt(mut self, value: impl Into<String>) -> Self {
self.system_prompt = Some(value.into());
self
}
pub fn retain_history(mut self, value: bool) -> Self {
self.retain_history = value;
self
}
pub fn max_history_messages(mut self, value: usize) -> Self {
self.max_history_messages = value;
self
}
pub fn semantic_memory(
mut self,
memory: Arc<dyn ConversationMemory>,
recall_limit: usize,
) -> Self {
self.memory = Some(memory);
self.recall_limit = recall_limit;
self
}
}
pub struct Conversation {
provider: Arc<dyn ChatProvider>,
config: ConversationConfig,
history: Vec<ChatMessage>,
}
impl Conversation {
pub fn new(provider: Arc<dyn ChatProvider>, config: ConversationConfig) -> Result<Self> {
if config.model.trim().is_empty() {
return Err(SdkError::Configuration(
"conversation model cannot be empty".into(),
));
}
if config.retain_history && config.max_history_messages == 0 {
return Err(SdkError::Configuration(
"history limit must be greater than zero".into(),
));
}
if config.memory.is_some() && config.recall_limit == 0 {
return Err(SdkError::Configuration(
"semantic recall limit must be greater than zero".into(),
));
}
Ok(Self {
provider,
config,
history: Vec::new(),
})
}
pub fn history(&self) -> &[ChatMessage] {
&self.history
}
pub fn clear_history(&mut self) {
self.history.clear();
}
pub async fn clear_memory(&self) -> Result<()> {
if let Some(memory) = &self.config.memory {
memory.clear().await?;
}
Ok(())
}
pub async fn send(&mut self, input: impl Into<String>) -> Result<ChatCompletionResponse> {
let input = input.into();
if input.trim().is_empty() {
return Err(SdkError::Validation(
"conversation input cannot be empty".into(),
));
}
let mut messages = Vec::new();
if let Some(system_prompt) = &self.config.system_prompt {
messages.push(ChatMessage::system(system_prompt));
}
if let Some(memory) = &self.config.memory {
let recalled = memory.recall(&input, self.config.recall_limit).await?;
if !recalled.is_empty() {
let context = recalled
.into_iter()
.map(|item| item.document.text)
.collect::<Vec<_>>();
messages.push(ChatMessage::system(format!(
"Relevant prior conversation context:\n{}",
serde_json::to_string(&context)
.map_err(|error| SdkError::Validation(error.to_string().into()))?
)));
}
}
if self.config.retain_history {
let start = self
.history
.len()
.saturating_sub(self.config.max_history_messages);
messages.extend_from_slice(&self.history[start..]);
}
messages.push(ChatMessage::user(&input));
let response = self
.provider
.complete(ChatCompletionRequest::new(&self.config.model).messages(messages))
.await?;
if let Some(text) = response.text() {
if self.config.retain_history {
self.history.push(ChatMessage::user(&input));
self.history.push(ChatMessage::assistant(text));
let overflow = self
.history
.len()
.saturating_sub(self.config.max_history_messages);
if overflow > 0 {
self.history.drain(..overflow);
}
}
if let Some(memory) = &self.config.memory {
memory
.remember(
MemoryDocument::new(
memory_id()?,
format!("User: {input}\nAssistant: {text}"),
)
.metadata("source", "conversation"),
)
.await?;
}
}
Ok(response)
}
}
fn memory_id() -> Result<String> {
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_err(|_| SdkError::Configuration("system clock is before Unix epoch".into()))?
.as_nanos();
let sequence = MEMORY_SEQUENCE.fetch_add(1, Ordering::Relaxed);
Ok(format!("memory-{timestamp}-{sequence}"))
}
#[cfg(test)]
mod tests {
use std::pin::Pin;
use std::sync::Mutex;
use futures_util::Stream;
use serde_json::json;
use super::*;
use crate::{ChatCompletionChunk, ChatStream, ProviderCapabilities};
struct FakeEmbeddings;
#[async_trait]
impl EmbeddingProvider for FakeEmbeddings {
async fn embed(&self, inputs: &[String]) -> Result<Vec<Vec<f32>>> {
Ok(inputs
.iter()
.map(|input| {
if input.to_ascii_lowercase().contains("rust") {
vec![1.0, 0.0, 0.0]
} else if input.contains("天气") {
vec![0.0, 1.0, 0.0]
} else {
vec![0.0, 0.0, 1.0]
}
})
.collect())
}
}
struct FakeProvider {
requests: Arc<Mutex<Vec<ChatCompletionRequest>>>,
}
#[async_trait]
impl ChatProvider for FakeProvider {
fn name(&self) -> &str {
"fake"
}
fn capabilities(&self) -> ProviderCapabilities {
ProviderCapabilities::openai_compatible()
}
async fn complete(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse> {
let index = {
let mut requests = self.requests.lock().unwrap();
requests.push(request);
requests.len()
};
Ok(serde_json::from_value(json!({
"choices":[{"message":{"content":format!("answer-{index}")}}]
}))
.unwrap())
}
async fn stream(&self, _: ChatCompletionRequest) -> Result<ChatStream> {
let stream: Pin<Box<dyn Stream<Item = Result<ChatCompletionChunk>> + Send>> =
Box::pin(futures_util::stream::empty());
Ok(stream)
}
}
#[test]
fn quantizes_and_compares_vectors() {
let vector = QuantizedVector::compress(&[1.0, -0.5, 0.25]).unwrap();
assert_eq!(vector.dimensions(), 3);
let restored = vector.decompress();
assert!((restored[0] - 1.0).abs() < 0.01);
assert!(vector.cosine_similarity(&[1.0, -0.5, 0.25]).unwrap() > 0.99);
assert_eq!(
QuantizedVector::compress(&[0.0, 0.0])
.unwrap()
.cosine_similarity(&[0.0, 0.0])
.unwrap(),
0.0
);
assert!(QuantizedVector::compress(&[]).is_err());
assert!(QuantizedVector::compress(&[f32::NAN]).is_err());
assert!(vector.cosine_similarity(&[1.0]).is_err());
}
#[tokio::test]
async fn stores_searches_snapshots_and_restores_memories() {
let store = InMemoryVectorStore::new();
assert!(store.is_empty().await.unwrap());
store
.upsert(StoredMemory {
document: MemoryDocument::new("rust", "Rust ownership").metadata("topic", "rust"),
embedding: QuantizedVector::compress(&[1.0, 0.0]).unwrap(),
})
.await
.unwrap();
store
.upsert(StoredMemory {
document: MemoryDocument::new("weather", "北京天气"),
embedding: QuantizedVector::compress(&[0.0, 1.0]).unwrap(),
})
.await
.unwrap();
assert_eq!(store.len().await.unwrap(), 2);
assert_eq!(
store.search(&[0.9, 0.1], 1).await.unwrap()[0].document.id,
"rust"
);
let json = store.snapshot_json().unwrap();
store.clear().await.unwrap();
assert!(store.is_empty().await.unwrap());
store.restore_json(&json).unwrap();
assert_eq!(store.snapshot().unwrap().len(), 2);
store
.upsert(StoredMemory {
document: MemoryDocument::new("rust", "Updated Rust"),
embedding: QuantizedVector::compress(&[1.0, 0.0]).unwrap(),
})
.await
.unwrap();
assert_eq!(store.len().await.unwrap(), 2);
}
#[tokio::test]
async fn semantic_memory_uses_real_embedding_interface() {
let store = Arc::new(InMemoryVectorStore::new());
let memory = SemanticMemory::new(Arc::new(FakeEmbeddings), store.clone());
memory
.remember(MemoryDocument::new("one", "Rust 所有权"))
.await
.unwrap();
memory
.remember(MemoryDocument::new("two", "北京天气"))
.await
.unwrap();
let recalled = memory.recall("Rust 生命周期", 1).await.unwrap();
assert_eq!(recalled[0].document.id, "one");
assert!(memory.recall("", 3).await.unwrap().is_empty());
assert!(memory.recall("Rust", 0).await.unwrap().is_empty());
assert!(
memory
.remember(MemoryDocument::new("empty", " "))
.await
.is_err()
);
memory.clear().await.unwrap();
assert!(store.is_empty().await.unwrap());
}
#[tokio::test]
async fn conversation_supports_stateless_history_and_semantic_modes() {
let stateless_requests = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(FakeProvider {
requests: stateless_requests.clone(),
});
let mut stateless = Conversation::new(
provider,
ConversationConfig::new("model").system_prompt("system"),
)
.unwrap();
assert!(stateless.send(" ").await.is_err());
stateless.clear_memory().await.unwrap();
stateless.send("first").await.unwrap();
stateless.send("second").await.unwrap();
assert_eq!(stateless.history().len(), 0);
assert_eq!(stateless_requests.lock().unwrap()[1].messages.len(), 2);
let history_requests = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(FakeProvider {
requests: history_requests.clone(),
});
let mut history = Conversation::new(
provider,
ConversationConfig::new("model")
.retain_history(true)
.max_history_messages(2),
)
.unwrap();
history.send("first").await.unwrap();
history.send("second").await.unwrap();
assert_eq!(history.history().len(), 2);
assert_eq!(history_requests.lock().unwrap()[1].messages.len(), 3);
history.clear_history();
assert!(history.history().is_empty());
let semantic_requests = Arc::new(Mutex::new(Vec::new()));
let provider = Arc::new(FakeProvider {
requests: semantic_requests.clone(),
});
let memory = Arc::new(SemanticMemory::new(
Arc::new(FakeEmbeddings),
Arc::new(InMemoryVectorStore::new()),
));
let mut semantic = Conversation::new(
provider,
ConversationConfig::new("model").semantic_memory(memory, 2),
)
.unwrap();
semantic.send("Rust first").await.unwrap();
semantic.send("Rust second").await.unwrap();
let requests = semantic_requests.lock().unwrap();
assert_eq!(requests[1].messages.len(), 2);
assert!(
serde_json::to_string(&requests[1].messages[0])
.unwrap()
.contains("Rust first")
);
}
#[test]
fn conversation_rejects_invalid_configuration() {
let provider = Arc::new(FakeProvider {
requests: Arc::new(Mutex::new(Vec::new())),
});
assert!(Conversation::new(provider.clone(), ConversationConfig::new("")).is_err());
assert!(
Conversation::new(
provider.clone(),
ConversationConfig::new("model")
.retain_history(true)
.max_history_messages(0)
)
.is_err()
);
let memory = Arc::new(SemanticMemory::new(
Arc::new(FakeEmbeddings),
Arc::new(InMemoryVectorStore::new()),
));
assert!(
Conversation::new(
provider,
ConversationConfig::new("model").semantic_memory(memory, 0)
)
.is_err()
);
}
#[tokio::test]
async fn zhipu_embedding_adapter_handles_empty_batches_without_network() {
let client = crate::ZhipuConfig::new("key")
.base_url("http://127.0.0.1:1")
.build()
.unwrap();
let embeddings = ZhipuEmbeddingProvider::new(client, "embedding-3").dimensions(128);
assert!(embeddings.embed(&[]).await.unwrap().is_empty());
}
}