Add hybrid text plus vector memory search

This commit is contained in:
Agent Zero
2026-03-22 22:38:14 +00:00
parent 347805cc29
commit 0ba37f8573
6 changed files with 179 additions and 14 deletions

View File

@@ -69,7 +69,7 @@ pub fn get_tool_definitions() -> Vec<Value> {
}),
json!({
"name": "query",
"description": "Query stored memories using semantic similarity search. Returns the most relevant memories based on the query text.",
"description": "Query stored memories using hybrid semantic plus keyword search. Returns the most relevant memories based on vector similarity and exact-text ranking.",
"inputSchema": {
"type": "object",
"properties": {

View File

@@ -36,9 +36,14 @@ pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String>
.and_then(|v| v.as_f64())
.unwrap_or(0.5) as f32;
let (vector_weight, text_weight) = normalized_weights(
state.config.query.vector_weight,
state.config.query.text_weight,
);
info!(
"Querying memories for agent '{}': '{}' (limit={}, threshold={})",
agent_id, query_text, limit, threshold
"Querying memories for agent '{}': '{}' (limit={}, threshold={}, vector_weight={}, text_weight={})",
agent_id, query_text, limit, threshold, vector_weight, text_weight
);
// Generate embedding for query using Arc<EmbeddingEngine>
@@ -49,7 +54,15 @@ pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String>
// Search database
let matches = state
.db
.query_memories(agent_id, &query_embedding, limit, threshold)
.query_memories(
agent_id,
query_text,
&query_embedding,
limit,
threshold,
vector_weight,
text_weight,
)
.await
.context("Failed to query memories")?;
@@ -63,6 +76,9 @@ pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String>
"id": m.record.id.to_string(),
"content": m.record.content,
"similarity": m.similarity,
"vector_score": m.vector_score,
"text_score": m.text_score,
"hybrid_score": m.hybrid_score,
"keywords": m.record.keywords,
"metadata": m.record.metadata,
"created_at": m.record.created_at.to_rfc3339()
@@ -74,8 +90,22 @@ pub async fn execute(state: &Arc<AppState>, arguments: Value) -> Result<String>
"success": true,
"agent_id": agent_id,
"query": query_text,
"vector_weight": vector_weight,
"text_weight": text_weight,
"count": results.len(),
"results": results
})
.to_string())
}
fn normalized_weights(vector_weight: f32, text_weight: f32) -> (f32, f32) {
let vector_weight = vector_weight.max(0.0);
let text_weight = text_weight.max(0.0);
let total = vector_weight + text_weight;
if total <= f32::EPSILON {
(0.6, 0.4)
} else {
(vector_weight / total, text_weight / total)
}
}