mirror of
https://gitea.ingwaz.work/Ingwaz/openbrain-mcp.git
synced 2026-03-31 14:49:06 +00:00
Add hybrid text plus vector memory search
This commit is contained in:
@@ -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": {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user