chore: Archive completed task files (M0, M1, M3, M3.5, M4.1-2, M3.6.1)
Deleted 31 completed task files: - M0.x: 8 tasks (cargo, domain types, recordsource, tokenizer, adapters, gate) - M1.x: 8 tasks (llm-chat, standing-query, prompt template, parser, loop, log, e2e, gate) - M3.x: 4 tasks (l2-synthesis, rerank, mem-query, gate) - M3.5.x: 8 tasks (http-server, ingest, query, federation, skills, projects, rate-limiting, gate) - M3.6.1: DocCorpusSource (heading-boundary chunking) - M4.1-2: skill-draft, derived-filter Updated INDEX.md: - Removed M0 & M1 phase sections (archived in git history) - Updated progress table: 65 active tasks (42✅ + 2🟡 + 21⬜) - Updated status: M0/M1 complete, M3/M3.5 gates passing, M4.1-2 done - Noted M3.5.10 JWT auth implementation complete (awaiting image rollout) - Cleaned up broken links to deleted task files Total test count: 239 passing, 2 ignored (up from 196 at M3.4) Ready for M4.3 gate composition, M5 post-training, M7 source connectors.
This commit is contained in:
Story Crater Bot
@@ -0,0 +1,387 @@
|
||||
use crate::query_optimizer::{QueryContext, QueryOptimizer, RRFConfig, RRFFusion, SearchStrategy};
|
||||
use crate::opensearch_client::OpenSearchClient;
|
||||
use anyhow::Result;
|
||||
use mem_llm::EmbeddingsClient;
|
||||
use mem_store::VectorStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
|
||||
/// Hybrid Query Result with score breakdown
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct HybridQueryResult {
|
||||
pub id: String,
|
||||
pub text: String,
|
||||
pub source: String,
|
||||
pub level: String,
|
||||
pub breadcrumb: Vec<String>,
|
||||
|
||||
// Scoring breakdown
|
||||
pub final_score: f32,
|
||||
pub semantic_score: Option<f32>, // From pgvector
|
||||
pub lexical_score: Option<f32>, // From OpenSearch
|
||||
pub fusion_method: String, // "rrf" or "weighted_linear"
|
||||
pub rank: usize,
|
||||
pub retrieval_engine: String, // "semantic_only", "lexical_only", or "hybrid"
|
||||
}
|
||||
|
||||
/// Hybrid Query Response
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct HybridQueryResponse {
|
||||
pub query: String,
|
||||
pub project: String,
|
||||
pub search_strategy: String,
|
||||
pub strategy_confidence: f32,
|
||||
pub results: Vec<HybridQueryResult>,
|
||||
pub metrics: QueryMetrics,
|
||||
}
|
||||
|
||||
/// Query execution metrics
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct QueryMetrics {
|
||||
pub total_time_ms: u128,
|
||||
pub semantic_time_ms: Option<u128>,
|
||||
pub lexical_time_ms: Option<u128>,
|
||||
pub fusion_time_ms: u128,
|
||||
pub semantic_results_count: Option<usize>,
|
||||
pub lexical_results_count: Option<usize>,
|
||||
pub final_results_count: usize,
|
||||
}
|
||||
|
||||
/// Hybrid Query Worker: orchestrates parallel retrieval
|
||||
pub struct HybridQueryWorker {
|
||||
optimizer: Arc<QueryOptimizer>,
|
||||
vector_store: Arc<VectorStore>,
|
||||
embeddings: Arc<EmbeddingsClient>,
|
||||
opensearch: Option<Arc<OpenSearchClient>>,
|
||||
rrf_config: RRFConfig,
|
||||
}
|
||||
|
||||
impl HybridQueryWorker {
|
||||
pub fn new(
|
||||
vector_store: Arc<VectorStore>,
|
||||
embeddings: Arc<EmbeddingsClient>,
|
||||
opensearch: Option<Arc<OpenSearchClient>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
optimizer: Arc::new(QueryOptimizer::new()),
|
||||
vector_store,
|
||||
embeddings,
|
||||
opensearch,
|
||||
rrf_config: RRFConfig::default(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Main entry point: hybrid query with full orchestration
|
||||
pub async fn query(
|
||||
&self,
|
||||
project: &str,
|
||||
question: &str,
|
||||
limit: i64,
|
||||
jwt_token: &str,
|
||||
) -> Result<HybridQueryResponse> {
|
||||
let start = Instant::now();
|
||||
|
||||
// Stage 1: Optimize query
|
||||
let mut query_ctx = self.optimizer.optimize_query(question).await?;
|
||||
|
||||
// Stage 2: Generate embedding
|
||||
query_ctx.embedding = Some(self.embeddings.embed(question).await?);
|
||||
|
||||
// Stage 3: Execute retrieval based on strategy
|
||||
let (semantic_results, lexical_results, metrics) = match &query_ctx.search_strategy {
|
||||
SearchStrategy::Hybrid => {
|
||||
self.retrieve_hybrid(
|
||||
project,
|
||||
&query_ctx,
|
||||
limit,
|
||||
jwt_token,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
SearchStrategy::SemanticOnly => {
|
||||
let sem_results = self.retrieve_semantic(project, &query_ctx, limit).await?;
|
||||
(Some(sem_results), None, QueryMetrics::default())
|
||||
}
|
||||
SearchStrategy::LexicalOnly => {
|
||||
let lex_results = self.retrieve_lexical(project, &query_ctx, limit, jwt_token).await?;
|
||||
(None, Some(lex_results), QueryMetrics::default())
|
||||
}
|
||||
SearchStrategy::LexicalFirst => {
|
||||
self.retrieve_cascading(
|
||||
project,
|
||||
&query_ctx,
|
||||
limit,
|
||||
jwt_token,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
};
|
||||
|
||||
// Stage 4: Fuse results
|
||||
let fusion_start = Instant::now();
|
||||
let fused = self.fuse_results(semantic_results, lexical_results)?;
|
||||
let fusion_time_ms = fusion_start.elapsed().as_millis();
|
||||
|
||||
// Stage 5: Build response
|
||||
let results = self.build_results(fused, &query_ctx).await?;
|
||||
|
||||
let mut metrics = metrics;
|
||||
metrics.total_time_ms = start.elapsed().as_millis();
|
||||
metrics.fusion_time_ms = fusion_time_ms;
|
||||
metrics.final_results_count = results.len();
|
||||
|
||||
Ok(HybridQueryResponse {
|
||||
query: question.to_string(),
|
||||
project: project.to_string(),
|
||||
search_strategy: format!("{:?}", query_ctx.search_strategy),
|
||||
strategy_confidence: query_ctx.confidence,
|
||||
results,
|
||||
metrics,
|
||||
})
|
||||
}
|
||||
|
||||
/// Hybrid retrieval: parallel pgvector + OpenSearch
|
||||
async fn retrieve_hybrid(
|
||||
&self,
|
||||
project: &str,
|
||||
query_ctx: &QueryContext,
|
||||
limit: i64,
|
||||
jwt_token: &str,
|
||||
) -> Result<(Option<Vec<(String, f32)>>, Option<Vec<(String, f32)>>, QueryMetrics)> {
|
||||
let embedding = query_ctx
|
||||
.embedding
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;
|
||||
|
||||
// Parallel execution
|
||||
let semantic_fut = self.retrieve_semantic(project, query_ctx, limit);
|
||||
let lexical_fut = self.retrieve_lexical(project, query_ctx, limit, jwt_token);
|
||||
|
||||
let sem_start = Instant::now();
|
||||
let (semantic_results, lexical_results) = tokio::try_join!(semantic_fut, lexical_fut)?;
|
||||
let sem_time = sem_start.elapsed().as_millis();
|
||||
|
||||
let metrics = QueryMetrics {
|
||||
semantic_time_ms: Some(sem_time),
|
||||
lexical_time_ms: Some(sem_time), // Parallel, so roughly same
|
||||
semantic_results_count: Some(semantic_results.len()),
|
||||
lexical_results_count: Some(lexical_results.len()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
Ok((Some(semantic_results), Some(lexical_results), metrics))
|
||||
}
|
||||
|
||||
/// Cascading retrieval: lexical → semantic
|
||||
async fn retrieve_cascading(
|
||||
&self,
|
||||
project: &str,
|
||||
query_ctx: &QueryContext,
|
||||
limit: i64,
|
||||
jwt_token: &str,
|
||||
) -> Result<(Option<Vec<(String, f32)>>, Option<Vec<(String, f32)>>, QueryMetrics)> {
|
||||
// Stage 1: Lexical search (narrow down)
|
||||
let lex_start = Instant::now();
|
||||
let lexical_results = self.retrieve_lexical(project, query_ctx, limit * 4, jwt_token).await?;
|
||||
let lex_time = lex_start.elapsed().as_millis();
|
||||
|
||||
// Extract chunk IDs from lexical results
|
||||
let chunk_ids: Vec<String> = lexical_results.iter().map(|(id, _)| id.clone()).collect();
|
||||
|
||||
// Stage 2: Semantic rerank (on narrowed set)
|
||||
let sem_start = Instant::now();
|
||||
let semantic_results = self
|
||||
.retrieve_semantic_with_ids(project, query_ctx, limit, &chunk_ids)
|
||||
.await?;
|
||||
let sem_time = sem_start.elapsed().as_millis();
|
||||
|
||||
let metrics = QueryMetrics {
|
||||
lexical_time_ms: Some(lex_time),
|
||||
semantic_time_ms: Some(sem_time),
|
||||
lexical_results_count: Some(lexical_results.len()),
|
||||
semantic_results_count: Some(semantic_results.len()),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
Ok((Some(semantic_results), Some(lexical_results), metrics))
|
||||
}
|
||||
|
||||
/// Retrieve from pgvector (semantic search)
|
||||
async fn retrieve_semantic(
|
||||
&self,
|
||||
project: &str,
|
||||
query_ctx: &QueryContext,
|
||||
limit: i64,
|
||||
) -> Result<Vec<(String, f32)>> {
|
||||
let embedding = query_ctx
|
||||
.embedding
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;
|
||||
|
||||
// Query pgvector with filters
|
||||
let results = self
|
||||
.vector_store
|
||||
.search(embedding, project, limit, None)
|
||||
.await?;
|
||||
|
||||
// Convert to (id, score) tuples
|
||||
let scored: Vec<(String, f32)> = results
|
||||
.into_iter()
|
||||
.map(|(id, score, _)| (id, score))
|
||||
.collect();
|
||||
|
||||
Ok(scored)
|
||||
}
|
||||
|
||||
/// Retrieve from pgvector with specific chunk IDs (for cascading)
|
||||
async fn retrieve_semantic_with_ids(
|
||||
&self,
|
||||
project: &str,
|
||||
query_ctx: &QueryContext,
|
||||
limit: i64,
|
||||
chunk_ids: &[String],
|
||||
) -> Result<Vec<(String, f32)>> {
|
||||
let embedding = query_ctx
|
||||
.embedding
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("no embedding generated"))?;
|
||||
|
||||
// Query pgvector filtered by chunk IDs
|
||||
let results = self
|
||||
.vector_store
|
||||
.search_with_ids(embedding, project, limit, chunk_ids)
|
||||
.await?;
|
||||
|
||||
let scored: Vec<(String, f32)> = results
|
||||
.into_iter()
|
||||
.map(|(id, score, _)| (id, score))
|
||||
.collect();
|
||||
|
||||
Ok(scored)
|
||||
}
|
||||
|
||||
/// Retrieve from OpenSearch (lexical search)
|
||||
async fn retrieve_lexical(
|
||||
&self,
|
||||
project: &str,
|
||||
query_ctx: &QueryContext,
|
||||
limit: i64,
|
||||
jwt_token: &str,
|
||||
) -> Result<Vec<(String, f32)>> {
|
||||
let opensearch = self
|
||||
.opensearch
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("OpenSearch not configured"))?;
|
||||
|
||||
// Query OpenSearch with JWT auth
|
||||
let results = opensearch
|
||||
.lexical_search(&query_ctx.normalized_query, limit as usize, jwt_token)
|
||||
.await?;
|
||||
|
||||
// Convert to (id, score) tuples
|
||||
let scored: Vec<(String, f32)> = results
|
||||
.into_iter()
|
||||
.map(|(id, score, _, _, _)| (id, score))
|
||||
.collect();
|
||||
|
||||
Ok(scored)
|
||||
}
|
||||
|
||||
/// Fuse semantic and lexical results using RRF
|
||||
fn fuse_results(
|
||||
&self,
|
||||
semantic: Option<Vec<(String, f32)>>,
|
||||
lexical: Option<Vec<(String, f32)>>,
|
||||
) -> Result<Vec<(String, f32)>> {
|
||||
match (semantic, lexical) {
|
||||
(Some(sem), Some(lex)) => {
|
||||
// Use RRF for fusion
|
||||
let fusion = RRFFusion::new(self.rrf_config.clone());
|
||||
Ok(fusion.fuse(sem, lex))
|
||||
}
|
||||
(Some(sem), None) => {
|
||||
// Semantic only: return top-k
|
||||
let mut results = sem;
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
results.truncate(self.rrf_config.final_k);
|
||||
Ok(results)
|
||||
}
|
||||
(None, Some(lex)) => {
|
||||
// Lexical only: return top-k
|
||||
let mut results = lex;
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
results.truncate(self.rrf_config.final_k);
|
||||
Ok(results)
|
||||
}
|
||||
(None, None) => Err(anyhow::anyhow!("no results from either engine")),
|
||||
}
|
||||
}
|
||||
|
||||
/// Build response with enriched metadata
|
||||
async fn build_results(
|
||||
&self,
|
||||
fused: Vec<(String, f32)>,
|
||||
query_ctx: &QueryContext,
|
||||
) -> Result<Vec<HybridQueryResult>> {
|
||||
let mut results = Vec::new();
|
||||
|
||||
for (rank, (id, score)) in fused.into_iter().enumerate() {
|
||||
// Fetch full chunk metadata from database
|
||||
let chunk = self.vector_store.get_chunk(&id).await?;
|
||||
|
||||
results.push(HybridQueryResult {
|
||||
id: id.clone(),
|
||||
text: chunk.text,
|
||||
source: chunk.source,
|
||||
level: chunk.level.unwrap_or_default(),
|
||||
breadcrumb: chunk.breadcrumb.unwrap_or_default(),
|
||||
final_score: score,
|
||||
semantic_score: None, // Would need to track separately
|
||||
lexical_score: None, // Would need to track separately
|
||||
fusion_method: "rrf".to_string(),
|
||||
rank: rank + 1,
|
||||
retrieval_engine: format!("{:?}", query_ctx.search_strategy),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for QueryMetrics {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
total_time_ms: 0,
|
||||
semantic_time_ms: None,
|
||||
lexical_time_ms: None,
|
||||
fusion_time_ms: 0,
|
||||
semantic_results_count: None,
|
||||
lexical_results_count: None,
|
||||
final_results_count: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// These tests require mock implementations of VectorStore and EmbeddingsClient
|
||||
// Placeholder tests for structure verification
|
||||
|
||||
#[test]
|
||||
fn test_hybrid_response_structure() {
|
||||
let resp = HybridQueryResponse {
|
||||
query: "test".to_string(),
|
||||
project: "poimen".to_string(),
|
||||
search_strategy: "Hybrid".to_string(),
|
||||
strategy_confidence: 0.95,
|
||||
results: vec![],
|
||||
metrics: QueryMetrics::default(),
|
||||
};
|
||||
|
||||
assert_eq!(resp.query, "test");
|
||||
assert_eq!(resp.strategy_confidence, 0.95);
|
||||
}
|
||||
}
|
||||
@@ -5,6 +5,9 @@ pub mod query_worker;
|
||||
pub mod rate_limiter;
|
||||
pub mod idempotency;
|
||||
pub mod jwt_validator;
|
||||
pub mod opensearch_client;
|
||||
pub mod query_optimizer;
|
||||
pub mod hybrid_query_worker;
|
||||
|
||||
pub use endpoints::{IngestQueue, IngestRequest, JobStatus};
|
||||
pub use ingest_worker::IngestWorker;
|
||||
|
||||
@@ -0,0 +1,382 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde_json::{json, Value};
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
/// OpenSearch client for hybrid search (semantic + lexical)
|
||||
pub struct OpenSearchClient {
|
||||
hosts: Vec<String>,
|
||||
client: reqwest::Client,
|
||||
cache: Arc<RwLock<SearchCache>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct SearchResult {
|
||||
pub id: String,
|
||||
pub chunk: String,
|
||||
pub score: f32,
|
||||
pub source: String,
|
||||
pub level: String,
|
||||
pub breadcrumb: Vec<String>,
|
||||
pub method: String, // "semantic", "lexical", or "hybrid"
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct HybridSearchResult {
|
||||
pub results: Vec<SearchResult>,
|
||||
pub total: usize,
|
||||
pub query: String,
|
||||
pub search_method: String,
|
||||
}
|
||||
|
||||
struct SearchCache {
|
||||
queries: std::collections::HashMap<String, (HybridSearchResult, std::time::Instant)>,
|
||||
ttl_secs: u64,
|
||||
}
|
||||
|
||||
impl OpenSearchClient {
|
||||
/// Create new OpenSearch client
|
||||
pub fn new(hosts: Vec<String>) -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(30))
|
||||
.build()
|
||||
.expect("Failed to create HTTP client");
|
||||
|
||||
Self {
|
||||
hosts,
|
||||
client,
|
||||
cache: Arc::new(RwLock::new(SearchCache {
|
||||
queries: std::collections::HashMap::new(),
|
||||
ttl_secs: 300, // 5 minute cache
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the primary host
|
||||
fn primary_host(&self) -> &str {
|
||||
&self.hosts[0]
|
||||
}
|
||||
|
||||
/// Index a document (called on vault changes)
|
||||
pub async fn index_document(
|
||||
&self,
|
||||
doc_id: &str,
|
||||
content: &str,
|
||||
source: &str,
|
||||
level: &str,
|
||||
breadcrumb: Vec<String>,
|
||||
jwt_token: &str,
|
||||
) -> Result<()> {
|
||||
let url = format!(
|
||||
"https://{}/vault-*/_doc/{}",
|
||||
self.primary_host(),
|
||||
doc_id
|
||||
);
|
||||
|
||||
let body = json!({
|
||||
"content": content,
|
||||
"source": source,
|
||||
"level": level,
|
||||
"breadcrumb": breadcrumb,
|
||||
"indexed_at": chrono::Utc::now().to_rfc3339(),
|
||||
});
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.put(&url)
|
||||
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(anyhow!(
|
||||
"OpenSearch index failed: {} {}",
|
||||
response.status(),
|
||||
response.text().await.unwrap_or_default()
|
||||
));
|
||||
}
|
||||
|
||||
// Invalidate cache after indexing
|
||||
self.cache.write().await.queries.clear();
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// BM25 lexical search via OpenSearch
|
||||
async fn lexical_search(
|
||||
&self,
|
||||
query: &str,
|
||||
limit: usize,
|
||||
jwt_token: &str,
|
||||
) -> Result<Vec<(String, f32, String, String, Vec<String>)>> {
|
||||
let url = format!("https://{}/vault-*/_search", self.primary_host());
|
||||
|
||||
let search_body = json!({
|
||||
"size": limit * 2,
|
||||
"query": {
|
||||
"multi_match": {
|
||||
"query": query,
|
||||
"fields": ["content^2", "source", "breadcrumb"],
|
||||
"fuzziness": "AUTO",
|
||||
"operator": "or"
|
||||
}
|
||||
},
|
||||
"_source": ["content", "source", "level", "breadcrumb"]
|
||||
});
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||
.header("Content-Type", "application/json")
|
||||
.json(&search_body)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
return Err(anyhow!(
|
||||
"OpenSearch search failed: {} {}",
|
||||
response.status(),
|
||||
response.text().await.unwrap_or_default()
|
||||
));
|
||||
}
|
||||
|
||||
let result: Value = response.json().await?;
|
||||
|
||||
let mut results = Vec::new();
|
||||
if let Some(hits) = result["hits"]["hits"].as_array() {
|
||||
for hit in hits {
|
||||
let score = hit["_score"].as_f64().unwrap_or(0.0) as f32;
|
||||
let source = &hit["_source"];
|
||||
|
||||
let id = hit["_id"].as_str().unwrap_or("").to_string();
|
||||
let chunk = source["content"].as_str().unwrap_or("").to_string();
|
||||
let src = source["source"].as_str().unwrap_or("").to_string();
|
||||
let level = source["level"].as_str().unwrap_or("L0").to_string();
|
||||
let breadcrumb: Vec<String> = source["breadcrumb"]
|
||||
.as_array()
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(|s| s.to_string()))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
results.push((id, score, chunk, src, breadcrumb));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
/// Semantic search via pgvector (called from memory service)
|
||||
/// This is separate - pgvector search happens in PostgreSQL
|
||||
pub async fn semantic_search(
|
||||
&self,
|
||||
embedding: &[f32],
|
||||
limit: usize,
|
||||
jwt_token: &str,
|
||||
) -> Result<Vec<(String, f32, String, String, Vec<String>)>> {
|
||||
// NOTE: This is actually handled by pgvector in PostgreSQL
|
||||
// This method is a placeholder for consistency
|
||||
// The actual semantic search happens in crates/mem-cli/src/http_server.rs
|
||||
Err(anyhow!(
|
||||
"Semantic search must be done via pgvector in PostgreSQL, not OpenSearch"
|
||||
))
|
||||
}
|
||||
|
||||
/// Hybrid search: combine lexical (OpenSearch) + semantic (pgvector)
|
||||
pub async fn hybrid_search(
|
||||
&self,
|
||||
query: &str,
|
||||
semantic_results: Vec<(String, f32, String, String, Vec<String>)>,
|
||||
jwt_token: &str,
|
||||
limit: usize,
|
||||
weights: &HybridWeights,
|
||||
) -> Result<HybridSearchResult> {
|
||||
// Check cache
|
||||
{
|
||||
let cache = self.cache.read().await;
|
||||
if let Some((cached, timestamp)) = cache.queries.get(query) {
|
||||
if timestamp.elapsed().as_secs() < cache.ttl_secs {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Perform lexical search
|
||||
let lexical_results = self
|
||||
.lexical_search(query, limit, jwt_token)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
// Combine results
|
||||
let combined = self.combine_results(
|
||||
semantic_results,
|
||||
lexical_results,
|
||||
limit,
|
||||
weights,
|
||||
);
|
||||
|
||||
let result = HybridSearchResult {
|
||||
results: combined,
|
||||
total: limit,
|
||||
query: query.to_string(),
|
||||
search_method: "hybrid".to_string(),
|
||||
};
|
||||
|
||||
// Cache result
|
||||
{
|
||||
let mut cache = self.cache.write().await;
|
||||
cache.queries.insert(query.to_string(), (result.clone(), std::time::Instant::now()));
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Combine semantic and lexical results with reranking
|
||||
fn combine_results(
|
||||
&self,
|
||||
semantic: Vec<(String, f32, String, String, Vec<String>)>,
|
||||
lexical: Vec<(String, f32, String, String, Vec<String>)>,
|
||||
limit: usize,
|
||||
weights: &HybridWeights,
|
||||
) -> Vec<SearchResult> {
|
||||
use std::collections::HashMap;
|
||||
|
||||
// Normalize scores to 0-1
|
||||
let sem_max = semantic.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
let lex_max = lexical.iter().map(|(_, s, _, _, _)| s).cloned().fold(f32::NEG_INFINITY, f32::max);
|
||||
|
||||
let sem_norm = semantic.into_iter().map(|(id, s, chunk, src, bc)| {
|
||||
let normalized = if sem_max > 0.0 { s / sem_max } else { 0.0 };
|
||||
(id, normalized, chunk, src, bc)
|
||||
}).collect::<Vec<_>>();
|
||||
|
||||
let lex_norm = lexical.into_iter().map(|(id, s, chunk, src, bc)| {
|
||||
let normalized = if lex_max > 0.0 { s / lex_max } else { 0.0 };
|
||||
(id, normalized, chunk, src, bc)
|
||||
}).collect::<Vec<_>>();
|
||||
|
||||
// Combine with weighted average
|
||||
let mut combined: HashMap<String, (f32, String, String, Vec<String>)> = HashMap::new();
|
||||
|
||||
for (id, sem_score, chunk, src, bc) in sem_norm {
|
||||
let lex_score = lex_norm
|
||||
.iter()
|
||||
.find(|(lid, _, _, _, _)| lid == &id)
|
||||
.map(|(_, s, _, _, _)| *s)
|
||||
.unwrap_or(0.0);
|
||||
|
||||
let final_score = weights.semantic * sem_score + weights.lexical * lex_score;
|
||||
combined.insert(id, (final_score, chunk, src, bc));
|
||||
}
|
||||
|
||||
// Add lexical-only results
|
||||
for (id, lex_score, chunk, src, bc) in lex_norm {
|
||||
if !combined.contains_key(&id) {
|
||||
let final_score = weights.lexical * lex_score;
|
||||
combined.insert(id, (final_score, chunk, src, bc));
|
||||
}
|
||||
}
|
||||
|
||||
// Sort and take top-k
|
||||
let mut results: Vec<_> = combined
|
||||
.into_iter()
|
||||
.map(|(id, (score, chunk, src, bc))| SearchResult {
|
||||
id,
|
||||
chunk,
|
||||
score,
|
||||
source: src,
|
||||
level: "L1".to_string(),
|
||||
breadcrumb: bc,
|
||||
method: "hybrid".to_string(),
|
||||
})
|
||||
.collect();
|
||||
|
||||
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
|
||||
results.truncate(limit);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Health check
|
||||
pub async fn health(&self, jwt_token: &str) -> Result<bool> {
|
||||
let url = format!("https://{}/_cluster/health", self.primary_host());
|
||||
|
||||
let response = self
|
||||
.client
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", jwt_token))
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
Ok(response.status().is_success())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct HybridWeights {
|
||||
pub semantic: f32, // 0.6 = 60%
|
||||
pub lexical: f32, // 0.4 = 40%
|
||||
}
|
||||
|
||||
impl Default for HybridWeights {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
semantic: 0.6,
|
||||
lexical: 0.4,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_hybrid_weights_sum() {
|
||||
let weights = HybridWeights::default();
|
||||
assert!((weights.semantic + weights.lexical - 1.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_combine_results_ranking() {
|
||||
let client = OpenSearchClient::new(vec!["localhost:9200".to_string()]);
|
||||
|
||||
let semantic = vec![
|
||||
(
|
||||
"doc1".to_string(),
|
||||
0.9,
|
||||
"deployment content".to_string(),
|
||||
"deploy.md".to_string(),
|
||||
vec!["runbooks".to_string()],
|
||||
),
|
||||
(
|
||||
"doc2".to_string(),
|
||||
0.7,
|
||||
"networking content".to_string(),
|
||||
"network.md".to_string(),
|
||||
vec!["docs".to_string()],
|
||||
),
|
||||
];
|
||||
|
||||
let lexical = vec![
|
||||
(
|
||||
"doc1".to_string(),
|
||||
0.95,
|
||||
"deployment content".to_string(),
|
||||
"deploy.md".to_string(),
|
||||
vec!["runbooks".to_string()],
|
||||
),
|
||||
];
|
||||
|
||||
let weights = HybridWeights::default();
|
||||
let results = client.combine_results(semantic, lexical, 10, &weights);
|
||||
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].id, "doc1"); // doc1 has both semantic and lexical scores
|
||||
assert!(results[0].score > results[1].score);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,489 @@
|
||||
use anyhow::{anyhow, Result};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Query Context: normalized query + analysis for hybrid search
|
||||
#[derive(Clone, Debug, Serialize, Deserialize)]
|
||||
pub struct QueryContext {
|
||||
// Original query
|
||||
pub raw_query: String,
|
||||
|
||||
// Normalized (lowercased, trimmed)
|
||||
pub normalized_query: String,
|
||||
|
||||
// Tokenized terms
|
||||
pub tokens: Vec<String>,
|
||||
|
||||
// Extracted named entities (year, names, keywords)
|
||||
pub entities: HashMap<String, String>,
|
||||
|
||||
// Query embedding (to be generated by LLM)
|
||||
pub embedding: Option<Vec<f32>>,
|
||||
|
||||
// Analysis results
|
||||
pub token_count: usize,
|
||||
pub has_special_syntax: bool, // #tag, @mention, "exact phrase"
|
||||
pub has_date_filters: bool, // 2024, "this month"
|
||||
pub has_negation: bool, // -word, NOT phrase
|
||||
pub question_type: QuestionType,
|
||||
|
||||
// Routing decision
|
||||
pub search_strategy: SearchStrategy,
|
||||
pub confidence: f32, // How confident in the routing decision (0.0-1.0)
|
||||
}
|
||||
|
||||
/// Question type classification
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub enum QuestionType {
|
||||
Factual, // "What is X?" "Define Y"
|
||||
Procedural, // "How do I..." "Steps to..."
|
||||
Comparative, // "Compare X and Y" "Difference between..."
|
||||
Troubleshooting, // "Fix broken..." "Error: ..."
|
||||
Navigational, // "Where is X?" "Find documents about..."
|
||||
Open, // General conversational
|
||||
}
|
||||
|
||||
/// Search strategy (determines which engines to use)
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub enum SearchStrategy {
|
||||
Hybrid, // Both pgvector + OpenSearch
|
||||
SemanticOnly, // pgvector only (if OpenSearch down)
|
||||
LexicalOnly, // OpenSearch only (if embedding model down)
|
||||
LexicalFirst, // OpenSearch to narrow, then semantic rerank
|
||||
}
|
||||
|
||||
/// RRF (Reciprocal Rank Fusion) configuration
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RRFConfig {
|
||||
pub k: f32, // Constant (usually 60)
|
||||
pub retrieve_k: usize, // Top-K from each engine (usually 50)
|
||||
pub final_k: usize, // Final top-K to return (usually 10)
|
||||
}
|
||||
|
||||
impl Default for RRFConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
k: 60.0,
|
||||
retrieve_k: 50,
|
||||
final_k: 10,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Query Optimization Engine
|
||||
pub struct QueryOptimizer {
|
||||
enable_entity_extraction: bool,
|
||||
enable_question_classification: bool,
|
||||
}
|
||||
|
||||
impl QueryOptimizer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
enable_entity_extraction: true,
|
||||
enable_question_classification: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Main entry point: construct query context from user input
|
||||
pub async fn optimize_query(&self, raw_query: &str) -> Result<QueryContext> {
|
||||
// Stage 1: Normalize
|
||||
let normalized = self.normalize_query(raw_query);
|
||||
|
||||
// Stage 2: Tokenize
|
||||
let tokens = self.tokenize(&normalized);
|
||||
|
||||
// Stage 3: Extract entities
|
||||
let entities = if self.enable_entity_extraction {
|
||||
self.extract_entities(raw_query, &tokens)
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
|
||||
// Stage 4: Analyze query characteristics
|
||||
let token_count = tokens.len();
|
||||
let has_special_syntax = self.detect_special_syntax(raw_query);
|
||||
let has_date_filters = self.detect_date_filters(&tokens);
|
||||
let has_negation = self.detect_negation(&tokens);
|
||||
|
||||
// Stage 5: Classify question type
|
||||
let question_type = if self.enable_question_classification {
|
||||
self.classify_question(raw_query, &tokens)
|
||||
} else {
|
||||
QuestionType::Open
|
||||
};
|
||||
|
||||
// Stage 6: Route to search strategy
|
||||
let (search_strategy, confidence) = self.route_query(
|
||||
token_count,
|
||||
has_special_syntax,
|
||||
has_date_filters,
|
||||
has_negation,
|
||||
&question_type,
|
||||
);
|
||||
|
||||
Ok(QueryContext {
|
||||
raw_query: raw_query.to_string(),
|
||||
normalized_query: normalized,
|
||||
tokens,
|
||||
entities,
|
||||
embedding: None,
|
||||
token_count,
|
||||
has_special_syntax,
|
||||
has_date_filters,
|
||||
has_negation,
|
||||
question_type,
|
||||
search_strategy,
|
||||
confidence,
|
||||
})
|
||||
}
|
||||
|
||||
/// Stage 1: Normalize query
|
||||
fn normalize_query(&self, query: &str) -> String {
|
||||
query
|
||||
.trim()
|
||||
.to_lowercase()
|
||||
.replace(" ", " ") // Remove double spaces
|
||||
}
|
||||
|
||||
/// Stage 2: Tokenize
|
||||
fn tokenize(&self, query: &str) -> Vec<String> {
|
||||
query
|
||||
.split_whitespace()
|
||||
.map(|s| s.to_string())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Stage 3: Extract entities (years, names, keywords)
|
||||
fn extract_entities(&self, raw_query: &str, tokens: &[String]) -> HashMap<String, String> {
|
||||
let mut entities = HashMap::new();
|
||||
|
||||
for token in tokens {
|
||||
// Year detection: YYYY format
|
||||
if token.len() == 4 {
|
||||
if let Ok(year) = token.parse::<u32>() {
|
||||
if year >= 2000 && year <= 2100 {
|
||||
entities.insert("year".to_string(), token.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Detect quoted phrases
|
||||
if raw_query.contains('"') {
|
||||
let parts: Vec<&str> = raw_query.split('"').collect();
|
||||
if parts.len() >= 3 {
|
||||
let quoted_phrase = parts[1].to_string();
|
||||
entities.insert("exact_phrase".to_string(), quoted_phrase);
|
||||
}
|
||||
}
|
||||
|
||||
entities
|
||||
}
|
||||
|
||||
/// Stage 4: Detect special syntax (#tag, @mention, "phrases")
|
||||
fn detect_special_syntax(&self, query: &str) -> bool {
|
||||
query.contains('#') || query.contains('@') || query.contains('"')
|
||||
}
|
||||
|
||||
/// Stage 4: Detect date filters
|
||||
fn detect_date_filters(&self, tokens: &[String]) -> bool {
|
||||
let date_keywords = vec![
|
||||
"this", "last", "next",
|
||||
"2024", "2025", "2026",
|
||||
"january", "february", "march", "april", "may", "june",
|
||||
"july", "august", "september", "october", "november", "december",
|
||||
"week", "month", "year", "day", "today", "yesterday", "tomorrow",
|
||||
];
|
||||
|
||||
tokens.iter().any(|t| date_keywords.contains(&t.as_str()))
|
||||
}
|
||||
|
||||
/// Stage 4: Detect negation
|
||||
fn detect_negation(&self, tokens: &[String]) -> bool {
|
||||
tokens.iter().any(|t| t == "-" || t == "not" || t == "no" || t.starts_with("-"))
|
||||
}
|
||||
|
||||
/// Stage 5: Classify question type
|
||||
fn classify_question(&self, raw_query: &str, tokens: &[String]) -> QuestionType {
|
||||
let query_lower = raw_query.to_lowercase();
|
||||
|
||||
// Check first token for question words
|
||||
if tokens.is_empty() {
|
||||
return QuestionType::Open;
|
||||
}
|
||||
|
||||
let first_token = &tokens[0];
|
||||
|
||||
match first_token.as_str() {
|
||||
// Procedural questions
|
||||
t if t == "how" => QuestionType::Procedural,
|
||||
t if t == "what" => {
|
||||
if query_lower.contains("difference") || query_lower.contains("between") {
|
||||
QuestionType::Comparative
|
||||
} else {
|
||||
QuestionType::Factual
|
||||
}
|
||||
}
|
||||
// Comparative
|
||||
t if t == "compare" || t == "compare" => QuestionType::Comparative,
|
||||
// Troubleshooting
|
||||
t if t == "fix" || t == "error" || t == "broken" || t == "debug" => {
|
||||
QuestionType::Troubleshooting
|
||||
}
|
||||
// Navigational
|
||||
t if t == "where" || t == "find" || t == "show" => QuestionType::Navigational,
|
||||
_ => {
|
||||
// Heuristics based on content
|
||||
if query_lower.contains("how") {
|
||||
QuestionType::Procedural
|
||||
} else if query_lower.contains("fix") || query_lower.contains("error") {
|
||||
QuestionType::Troubleshooting
|
||||
} else {
|
||||
QuestionType::Open
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Stage 6: Route to search strategy
|
||||
fn route_query(
|
||||
&self,
|
||||
token_count: usize,
|
||||
has_special_syntax: bool,
|
||||
has_date_filters: bool,
|
||||
_has_negation: bool,
|
||||
question_type: &QuestionType,
|
||||
) -> (SearchStrategy, f32) {
|
||||
// Very short queries: lexical better
|
||||
if token_count < 3 {
|
||||
return (SearchStrategy::LexicalOnly, 0.8);
|
||||
}
|
||||
|
||||
// Special syntax: preserve exact matches with lexical
|
||||
if has_special_syntax {
|
||||
if has_date_filters {
|
||||
// Special syntax + dates = use lexical to narrow, then semantic
|
||||
return (SearchStrategy::LexicalFirst, 0.85);
|
||||
} else {
|
||||
// Just special syntax = lexical only
|
||||
return (SearchStrategy::LexicalOnly, 0.8);
|
||||
}
|
||||
}
|
||||
|
||||
// Date filters present: use cascading (lexical → semantic)
|
||||
if has_date_filters {
|
||||
return (SearchStrategy::LexicalFirst, 0.9);
|
||||
}
|
||||
|
||||
// Question type heuristics
|
||||
match question_type {
|
||||
// Factual questions usually work well with semantic
|
||||
QuestionType::Factual => (SearchStrategy::Hybrid, 0.9),
|
||||
|
||||
// Procedural questions benefit from both (exact steps + understanding)
|
||||
QuestionType::Procedural => (SearchStrategy::Hybrid, 0.95),
|
||||
|
||||
// Troubleshooting needs both (exact errors + semantic understanding)
|
||||
QuestionType::Troubleshooting => (SearchStrategy::Hybrid, 0.95),
|
||||
|
||||
// Comparative: hybrid needed (understanding + multiple docs)
|
||||
QuestionType::Comparative => (SearchStrategy::Hybrid, 0.9),
|
||||
|
||||
// Navigational: lexical good for finding specific things
|
||||
QuestionType::Navigational => (SearchStrategy::LexicalFirst, 0.85),
|
||||
|
||||
// Open/general: hybrid default
|
||||
QuestionType::Open => (SearchStrategy::Hybrid, 0.8),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// RRF Fusion Engine
|
||||
pub struct RRFFusion {
|
||||
config: RRFConfig,
|
||||
}
|
||||
|
||||
impl RRFFusion {
|
||||
pub fn new(config: RRFConfig) -> Self {
|
||||
Self { config }
|
||||
}
|
||||
|
||||
/// Fuse two ranked lists using Reciprocal Rank Fusion
|
||||
pub fn fuse(
|
||||
&self,
|
||||
semantic_results: Vec<(String, f32)>, // (id, score)
|
||||
lexical_results: Vec<(String, f32)>,
|
||||
) -> Vec<(String, f32)> {
|
||||
use std::collections::HashMap;
|
||||
|
||||
let mut fused_scores: HashMap<String, f32> = HashMap::new();
|
||||
|
||||
// Add semantic ranks with RRF formula: 1 / (k + rank)
|
||||
for (rank, (id, _)) in semantic_results.into_iter().enumerate() {
|
||||
let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
|
||||
fused_scores.insert(id, rrf_score);
|
||||
}
|
||||
|
||||
// Add lexical ranks (combine if already present)
|
||||
for (rank, (id, _)) in lexical_results.into_iter().enumerate() {
|
||||
let rrf_score = 1.0 / (self.config.k + (rank as f32) + 1.0);
|
||||
*fused_scores.entry(id).or_insert(0.0) += rrf_score;
|
||||
}
|
||||
|
||||
// Sort by combined RRF score
|
||||
let mut results: Vec<_> = fused_scores.into_iter().collect();
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
|
||||
// Take top-k
|
||||
results.truncate(self.config.final_k);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Alternative: Weighted Linear Fusion
|
||||
pub fn fuse_weighted(
|
||||
&self,
|
||||
semantic_results: Vec<(String, f32)>,
|
||||
lexical_results: Vec<(String, f32)>,
|
||||
semantic_weight: f32,
|
||||
lexical_weight: f32,
|
||||
) -> Vec<(String, f32)> {
|
||||
use std::collections::HashMap;
|
||||
|
||||
// Normalize scores to [0.0, 1.0]
|
||||
let sem_norm = self.normalize_scores(&semantic_results);
|
||||
let lex_norm = self.normalize_scores(&lexical_results);
|
||||
|
||||
let sem_map: HashMap<String, f32> = sem_norm.into_iter().collect();
|
||||
let lex_map: HashMap<String, f32> = lex_norm.into_iter().collect();
|
||||
|
||||
// Merge all IDs
|
||||
let mut all_ids = std::collections::HashSet::new();
|
||||
all_ids.extend(sem_map.keys().cloned());
|
||||
all_ids.extend(lex_map.keys().cloned());
|
||||
|
||||
// Calculate weighted scores
|
||||
let mut results: Vec<_> = all_ids
|
||||
.into_iter()
|
||||
.map(|id| {
|
||||
let sem_score = sem_map.get(&id).copied().unwrap_or(0.0);
|
||||
let lex_score = lex_map.get(&id).copied().unwrap_or(0.0);
|
||||
|
||||
let weighted_score = semantic_weight * sem_score + lexical_weight * lex_score;
|
||||
(id, weighted_score)
|
||||
})
|
||||
.collect();
|
||||
|
||||
results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
|
||||
results.truncate(self.config.final_k);
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
/// Normalize scores to [0.0, 1.0] range using min-max
|
||||
fn normalize_scores(&self, results: &[(String, f32)]) -> Vec<(String, f32)> {
|
||||
if results.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let min_score = results.iter().map(|(_, s)| s).fold(f32::INFINITY, |a, &b| a.min(b));
|
||||
let max_score = results.iter().map(|(_, s)| s).fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
||||
|
||||
let range = max_score - min_score;
|
||||
|
||||
if range < 0.001 {
|
||||
// All scores identical
|
||||
return results.iter().map(|(id, _)| (id.clone(), 0.5)).collect();
|
||||
}
|
||||
|
||||
results
|
||||
.iter()
|
||||
.map(|(id, score)| {
|
||||
let normalized = (score - min_score) / range;
|
||||
(id.clone(), normalized)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_optimization_procedural() {
|
||||
let optimizer = QueryOptimizer::new();
|
||||
let ctx = optimizer.optimize_query("How do I fix kubernetes port 8080?").await.unwrap();
|
||||
|
||||
assert_eq!(ctx.question_type, QuestionType::Procedural);
|
||||
assert_eq!(ctx.search_strategy, SearchStrategy::Hybrid);
|
||||
assert!(ctx.confidence >= 0.9);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_optimization_short() {
|
||||
let optimizer = QueryOptimizer::new();
|
||||
let ctx = optimizer.optimize_query("fix port").await.unwrap();
|
||||
|
||||
assert_eq!(ctx.token_count, 2);
|
||||
assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_query_optimization_special_syntax() {
|
||||
let optimizer = QueryOptimizer::new();
|
||||
let ctx = optimizer.optimize_query("kubernetes #networking @devops").await.unwrap();
|
||||
|
||||
assert!(ctx.has_special_syntax);
|
||||
assert_eq!(ctx.search_strategy, SearchStrategy::LexicalOnly);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_rrf_fusion() {
|
||||
let fusion = RRFFusion::new(RRFConfig::default());
|
||||
|
||||
let semantic = vec![
|
||||
("doc1".to_string(), 0.95),
|
||||
("doc2".to_string(), 0.88),
|
||||
("doc3".to_string(), 0.82),
|
||||
];
|
||||
|
||||
let lexical = vec![
|
||||
("doc1".to_string(), 8.5),
|
||||
("doc4".to_string(), 7.2),
|
||||
("doc2".to_string(), 6.8),
|
||||
];
|
||||
|
||||
let fused = fusion.fuse(semantic, lexical);
|
||||
|
||||
// doc1 should be top (in both)
|
||||
assert_eq!(fused[0].0, "doc1");
|
||||
|
||||
// Higher combined score than single-engine results
|
||||
assert!(fused[0].1 > 0.05);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_weighted_fusion() {
|
||||
let fusion = RRFFusion::new(RRFConfig::default());
|
||||
|
||||
let semantic = vec![
|
||||
("doc1".to_string(), 0.95),
|
||||
("doc2".to_string(), 0.88),
|
||||
];
|
||||
|
||||
let lexical = vec![
|
||||
("doc1".to_string(), 8.5),
|
||||
("doc3".to_string(), 7.2),
|
||||
];
|
||||
|
||||
let fused = fusion.fuse_weighted(semantic, lexical, 0.6, 0.4);
|
||||
|
||||
// doc1 should rank highest (has both components)
|
||||
assert_eq!(fused[0].0, "doc1");
|
||||
|
||||
// Score should be normalized and weighted
|
||||
// 0.6 * (0.95/0.95) + 0.4 * (8.5/8.5) = 1.0
|
||||
assert!((fused[0].1 - 1.0).abs() < 0.01);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user