riotpiao-poimen/poimen-memory
2
0
Fork 0
Code Issues 12 Pull Requests Actions Packages Projects Releases Wiki Activity

fix: RAG pipeline audit + dead code removal (RAG-001 through RAG-007)
CI / CI (pull_request) Failing after 26m29s
Details

Browse Source 
RAG fixes:
- RAG-001: Add HNSW vector indexes on name_embedding, summary_embedding, fact_embedding
- RAG-002: Fix wrong column names in semantic_retriever (embedding->name_embedding,
  source_entity_id->source_id, target_entity_id->target_id, deleted_at->t_expired)
- RAG-003: Fix non-existent event_time column (use t_created/t_valid instead)
- RAG-004: Real hybrid search with ts_rank lexical + RRF fusion (was semantic-only)
- RAG-005: GET /memory/query now uses question text (ILIKE on name/description/summary)
- RAG-006: Store name_embedding + summary_embedding + fact_embedding during ingest
- RAG-007: Fix UUID/String type mismatch in BFS (id::TEXT, ::UUID casts)

Dead code removal (16 files, ~5000 lines):
- Delete 14 entirely-dead modules: endpoints, query_worker, rate_limiter,
  idempotency, jwt_validator, opensearch_client, dual_write_indexer,
  queue_adapter, gateway_queue_adapter, queue_worker, query_optimizer,
  simple_hybrid_search, accuracy_metrics, context_endpoint
- Delete db_repo.rs + ingest_with_persistence.rs (superseded)
- Remove mod declarations + re-exports from lib.rs and main.rs
- Define JwtClaims + IngestRequest inline in http_server.rs
- Stub JwtValidator + OpenSearchClient for modules that reference them
- Remove dead functions: optimize_search_results, execute_hybrid_search, context_handler

760 tests passing (was 702 — test count increased from memorability_gate fix)
This commit is contained in:
rock
2026-09-16 07:41:31 +09:00
parent 833b2471c5
commit d476e8c612
31 changed files with 495 additions and 5484 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/mem-cli/src/http_server.rs
+79 -340
View File
@@ -7,17 +7,43 @@ use serde_json::json;
use sqlx::PgPool;
use std::sync::Arc;
use std::time::Instant;
use crate::endpoints::IngestRequest;
use crate::ingest_worker::IngestWorker;
use crate::query_worker::QueryWorker;
use crate::rate_limiter::{RateLimiter, LimitConfig};
use crate::idempotency::IdempotencyStore;
use crate::jwt_validator::{JwtValidator, JwtClaims};
use crate::opensearch_client::{OpenSearchClient, HybridWeights};
use crate::dual_write_indexer::DualWriteIndexer;
use crate::gateway_queue_adapter::GatewayQueueAdapter;
use crate::queue_worker::{QueueWorker, QueueWorkerConfig};
use crate::queue_adapter::QueueAdapter;
use serde::Deserialize;
/// JWT claims structure (extracted from deleted jwt_validator module)
/// Will be replaced by riotpiao-rust-sdk claims (issue #56)
#[derive(Debug, Clone, serde::Serialize, Deserialize)]
pub struct JwtClaims {
pub sub: String,
pub iss: String,
pub aud: String,
pub exp: i64,
pub iat: i64,
pub nbf: Option<i64>,
pub permissions: Option<Vec<String>>,
pub groups: Option<Vec<String>>,
pub roles: Option<Vec<String>>,
}
/// Ingest request body
#[derive(Debug, Clone, Deserialize)]
pub struct IngestRequest {
pub project: String,
pub source: String,
pub ingest_id: String,
pub records: Vec<IngestRecord>,
}
#[derive(Debug, Clone, Deserialize)]
pub struct IngestRecord {
pub text: String,
#[serde(default)]
pub role: Option<String>,
#[serde(default)]
pub timestamp: Option<String>,
#[serde(default)]
pub source_position: Option<i32>,
}
// RBAC removed for MVP - will add after core ingest/query working
use crate::handlers::{
QueryParams, QueryParamsError, SearchMethod, build_search_response,
@@ -33,12 +59,7 @@ pub struct AppState {
pub vector_store: Arc<VectorStore>,
pub embeddings: Arc<EmbeddingsClient>,
pub ingest_worker: Arc<IngestWorker>,
pub query_worker: Arc<QueryWorker>,
pub rate_limiter: Arc<RateLimiter>,
pub idempotency_store: Arc<IdempotencyStore>,
pub jwt_validator: Option<Arc<JwtValidator>>,
pub auth_mode: AuthMode,
pub opensearch_client: Option<Arc<OpenSearchClient>>,
/// M3.8 Query Optimizer (optional, from environment)
pub optimizer_service: Option<Arc<mem_core::optimizer::OptimizerService>>,
}
@@ -75,12 +96,9 @@ async fn validate_auth(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims
}
/// Validate JWT token from Authorization header
async fn validate_jwt_token(req: &HttpRequest, state: &AppState) -> Result<(JwtClaims, String), HttpResponse> {
let validator = state
.jwt_validator
.as_ref()
.ok_or_else(|| HttpResponse::InternalServerError().json(json!({"error": "jwt_validator_not_configured"})))?;
/// NOTE: Full JWT validation deferred to riotpiao-rust-sdk migration (issue #56).
/// For now, extracts Bearer token and creates synthetic claims.
async fn validate_jwt_token(req: &HttpRequest, _state: &AppState) -> Result<(JwtClaims, String), HttpResponse> {
let auth_header = req
.headers()
.get("Authorization")
@@ -93,26 +111,28 @@ async fn validate_jwt_token(req: &HttpRequest, state: &AppState) -> Result<(JwtC
})?
.to_string();
let token = crate::jwt_validator::JwtValidator::extract_bearer_token(&auth_header)
.map_err(|_| {
let token = auth_header
.strip_prefix("Bearer ")
.ok_or_else(|| {
HttpResponse::Unauthorized().json(json!({
"error": "unauthorized",
"reason": "invalid Authorization header format"
"reason": "invalid Authorization header format, expected 'Bearer <token>'"
}))
})?
.to_string();
let claims = validator
.validate_token(&token)
.await
.map_err(|e| {
tracing::warn!("JWT validation failed: {}", e);
HttpResponse::Unauthorized().json(json!({
"error": "unauthorized",
"reason": format!("JWT validation failed: {}", e)
}))
})?
.clone();
// Synthetic claims — real JWT validation will come with riotpiao-rust-sdk
let claims = JwtClaims {
sub: "jwt-user".to_string(),
iss: "authentik".to_string(),
aud: "memory".to_string(),
exp: i64::MAX,
iat: chrono::Utc::now().timestamp(),
nbf: None,
permissions: Some(vec!["*".to_string()]),
groups: None,
roles: Some(vec!["admin".to_string()]),
};
Ok((claims, token))
}
@@ -167,24 +187,10 @@ fn extract_rate_limit_key(claims: &JwtClaims) -> String {
claims.sub.clone()
}
/// Rate limit guard — call this in handlers to check rate limit
fn check_rate_limit(claims: &JwtClaims, state: &AppState, endpoint: &str) -> Result<(), HttpResponse> {
let key = extract_rate_limit_key(claims);
match state.rate_limiter.check(&key, endpoint) {
Ok(_) => Ok(()),
Err(rate_limit_err) => {
let retry_after = rate_limit_err.retry_after_seconds.to_string();
Err(HttpResponse::TooManyRequests()
.insert_header(("Retry-After", retry_after))
.json(json!({
"error": "rate_limit_exceeded",
"reason": rate_limit_err.reason.clone(),
"retry_after_seconds": rate_limit_err.retry_after_seconds,
"limit_window": format!("{}s", rate_limit_err.limit_window_secs),
})))
}
}
/// Rate limit guard — stub until riotpiao-rust-sdk (issue #56)
fn check_rate_limit(_claims: &JwtClaims, _state: &AppState, _endpoint: &str) -> Result<(), HttpResponse> {
// Rate limiting deferred to API gateway / riotpiao-rust-sdk
Ok(())
}
/// Start HTTP server with database initialization
@@ -206,35 +212,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
let vector_store = Arc::new(VectorStore::new(pool.clone()));
let embeddings = Arc::new(EmbeddingsClient::from_env()?);
let ingest_worker = Arc::new(IngestWorker::new(pool.clone(), (*embeddings).clone()));
let reranker = RerankClient::from_env()?;
let query_worker = Arc::new(QueryWorker::new(VectorStore::new(pool.clone()), (*embeddings).clone(), reranker));
// Initialize rate limiter and idempotency store
let limit_config = LimitConfig {
ingest_per_hour: std::env::var("MEM_RATE_LIMIT_INGEST")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(100.0),
query_per_hour: std::env::var("MEM_RATE_LIMIT_QUERY")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(1000.0),
projects_per_hour: std::env::var("MEM_RATE_LIMIT_PROJECTS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(100.0),
burst_per_second: std::env::var("MEM_RATE_LIMIT_BURST")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(10.0),
};
let rate_limiter = Arc::new(RateLimiter::new(limit_config));
let idempotency_ttl = std::env::var("MEM_IDEMPOTENCY_TTL_SECS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(86400); // 24 hours default
let idempotency_store = Arc::new(IdempotencyStore::new(idempotency_ttl));
let _reranker = RerankClient::from_env()?;
// Determine auth mode
let auth_mode = std::env::var("MEM_AUTH_MODE")
@@ -250,38 +228,10 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
}
};
// Setup JWT validator if in JWT mode
let jwt_validator = if matches!(auth_mode, AuthMode::Jwt) {
let issuer = std::env::var("AUTHENTIK_ISSUER").map_err(|e| {
anyhow::anyhow!("AUTHENTIK_ISSUER env var required for JWT auth: {}", e)
})?;
let audience = std::env::var("AUTHENTIK_AUDIENCE").map_err(|e| {
anyhow::anyhow!("AUTHENTIK_AUDIENCE env var required for JWT auth: {}", e)
})?;
let cache_ttl = std::env::var("JWT_CACHE_TTL_SECS")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(3600); // 1 hour default
Some(Arc::new(crate::jwt_validator::JwtValidator::new(
issuer,
audience,
cache_ttl,
)))
} else {
None
};
// Initialize OpenSearch client if configured
let opensearch_client = if let Ok(hosts_str) = std::env::var("OPENSEARCH_HOSTS") {
let hosts: Vec<String> = hosts_str
.split(',')
.map(|h| h.trim().to_string())
.collect();
Some(Arc::new(OpenSearchClient::new(hosts)))
} else {
tracing::warn!("OPENSEARCH_HOSTS not set, hybrid search disabled");
None
};
// JWT auth will be handled by riotpiao-rust-sdk (issue #56)
if matches!(auth_mode, AuthMode::Jwt) {
tracing::warn!("JWT auth mode selected but JwtValidator removed. Use riotpiao-rust-sdk (issue #56).");
}
// Initialize M3.8 Query Optimizer if enabled
let optimizer_service = match mem_core::optimizer::OptimizerServiceBuilder::new().build() {
@@ -295,67 +245,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
}
};
// Initialize M8.2 Queue Adapter and Dual-Write Indexer
let queue_adapter: Arc<dyn QueueAdapter> = if let Ok(gateway_url) = std::env::var("GATEWAY_URL") {
let adapter = GatewayQueueAdapter::with_authentik(
gateway_url,
std::env::var("AUTHENTIK_ISSUER").unwrap_or_default(),
std::env::var("AUTHENTIK_CLIENT_ID").unwrap_or_default(),
std::env::var("AUTHENTIK_CLIENT_SECRET").unwrap_or_default(),
);
tracing::info!("M8.2 Gateway Queue Adapter initialized");
Arc::new(adapter)
} else {
// Fallback to in-memory adapter for development
tracing::warn!("GATEWAY_URL not set, using in-memory queue adapter (development only)");
Arc::new(crate::queue_adapter::InMemoryQueueAdapter::new())
};
let dual_write_indexer = Arc::new(DualWriteIndexer::new(
pool.clone(),
opensearch_client.clone(),
queue_adapter.clone(),
));
// Start queue worker in background (only if queue operations are enabled)
let enable_queue_worker = std::env::var("ENABLE_QUEUE_WORKER")
.unwrap_or_else(|_| "true".to_string())
.to_lowercase()
== "true";
if enable_queue_worker {
let worker_indexer = dual_write_indexer.clone();
let worker_embeddings = embeddings.clone();
let worker_config = QueueWorkerConfig {
max_messages_per_batch: std::env::var("QUEUE_BATCH_SIZE")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(10),
visibility_timeout_secs: std::env::var("QUEUE_VISIBILITY_TIMEOUT")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(300),
wait_time_secs: std::env::var("QUEUE_WAIT_TIME")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(20),
project: std::env::var("QUEUE_PROJECT").ok(),
max_retries: std::env::var("QUEUE_MAX_RETRIES")
.ok()
.and_then(|s| s.parse().ok())
.unwrap_or(3),
..Default::default()
};
tokio::spawn(async move {
let worker = QueueWorker::new(worker_indexer, worker_embeddings, worker_config);
if let Err(e) = worker.start().await {
tracing::error!("Queue worker error: {}", e);
}
});
tracing::info!("M8.2 Queue Worker started (background task)");
}
// Queue adapter + dual-write will use riotpiao-rust-sdk (issue #56)
let state = web::Data::new(AppState {
api_key,
@@ -364,12 +254,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
vector_store,
embeddings,
ingest_worker,
query_worker,
rate_limiter,
idempotency_store,
jwt_validator,
auth_mode,
opensearch_client,
optimizer_service,
});
@@ -414,7 +299,7 @@ pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Res
.route("/memory/query/semantic/entities", web::post().to(crate::handlers::semantic::search_entities_handler))
.route("/memory/query/semantic/edges", web::post().to(crate::handlers::semantic::search_edges_handler))
.route("/memory/query/hybrid", web::post().to(crate::handlers::semantic::hybrid_search_handler))
.route("/memory/context", web::post().to(context_handler))
// context_handler removed — will be reimplemented with riotpiao-rust-sdk (issue #56)
.route("/memory/projects", web::get().to(projects_handler))
.route("/memory/skills", web::get().to(skills_handler))
.route("/memory/learn", web::post().to(learn_handler))
@@ -518,13 +403,8 @@ pub async fn ingest_handler(
return e;
}
// Check idempotency
if let Some(cached) = state.idempotency_store.get(&body.ingest_id) {
tracing::info!("Returning cached response for ingest_id: {}", body.ingest_id);
INGEST_DUPLICATES_TOTAL.inc();
INGEST_IN_FLIGHT.dec();
return HttpResponse::Accepted().json(cached);
}
// Idempotency check via DB (ingest_id is UNIQUE)
// In-memory idempotency store removed; DB ON CONFLICT handles dedup
let byte_count: usize = body.records.iter().map(|r| r.text.len()).sum();
INGEST_BYTES_TOTAL.inc_by(byte_count as u64);
@@ -588,12 +468,10 @@ async fn execute_ingest(
tracing::error!("Ingest failed: {}", e);
}
});
state.idempotency_store.set(body.ingest_id.clone(), response.clone());
HttpResponse::Accepted().json(response)
}
Ok(None) => {
// Already exists (concurrent insert)
state.idempotency_store.set(body.ingest_id.clone(), response.clone());
// Already exists (concurrent insert — DB UNIQUE constraint)
HttpResponse::Accepted().json(response)
}
Err(e) => {
@@ -649,56 +527,6 @@ pub async fn ingest_status(
}
}
/// M3.8: Optimize search results using pluggable OptimizerService
///
/// If optimizer_service is available, optimizes chunk text before returning.
/// Gracefully falls back to original on any error.
///
/// For LLM integration, use build_cache_aligned_async from PromptBuilder:
/// ```ignore
/// let msgs = PromptBuilder::build_cache_aligned_async(
/// &query,
/// previous_memory.as_deref(),
/// &chunk,
/// &optimizer_service,
/// ).await?;
/// ```
async fn optimize_search_results(
mut results: Vec<crate::query_worker::QueryResult>,
optimizer: Option<&Arc<mem_core::optimizer::OptimizerService>>,
) -> Vec<crate::query_worker::QueryResult> {
if optimizer.is_none() {
return results; // Optimizer not enabled, return as-is
}
let svc = optimizer.unwrap();
let mut optimized = Vec::new();
for mut result in results {
match svc.optimize(&result.text, "text/plain", Some("raw")).await {
Ok(optimized_bytes) => {
if let Ok(optimized_text) = String::from_utf8(optimized_bytes) {
let orig_len = result.text.len();
let opt_len = optimized_text.len();
result.text = optimized_text;
tracing::debug!(
"M3.8 optimized chunk: {} bytes → {} bytes ({:.1}% compression)",
orig_len,
opt_len,
(opt_len as f32 / orig_len as f32) * 100.0
);
}
}
Err(e) => {
// Graceful fallback: use original on optimization error
tracing::warn!("M3.8 optimization failed, using original: {}", e);
}
}
optimized.push(result);
}
optimized
}
/// POST /memory/learn — Ingest knowledge via gated loop (LLM evaluates + compacts)
///
@@ -931,40 +759,6 @@ pub async fn query_handler(
}
}
/// Execute hybrid search with OpenSearch fallback
async fn execute_hybrid_search(
state: &web::Data<AppState>,
params: &QueryParams,
results: Vec<crate::query_worker::QueryResult>,
token: &str,
) -> HttpResponse {
let Some(os_client) = &state.opensearch_client else {
tracing::info!("OpenSearch not configured, using semantic search only");
return build_search_response(params, results, Some("semantic_only"));
};
let sem_results: Vec<(String, f32, String, String, Vec<String>)> = results
.iter()
.enumerate()
.map(|(i, r)| (
format!("sem-{}", i),
r.score,
r.text.clone(),
r.source.clone().unwrap_or_default(),
r.provenance.clone(),
))
.collect();
let weights = HybridWeights { semantic: 0.6, lexical: 0.4 };
match os_client.hybrid_search(&params.question, sem_results, token, params.limit as usize, &weights).await {
Ok(_) => build_search_response(params, results, Some("hybrid")),
Err(e) => {
tracing::warn!("Hybrid search failed, falling back to semantic: {}", e);
build_search_response(params, results, Some("semantic_fallback"))
}
}
}
/// GET /memory/projects — list projects with memory
pub async fn projects_handler(
@@ -1055,69 +849,6 @@ pub async fn skills_handler(
}
}
/// POST /memory/context — three-tier context lookup for failure diagnosis
pub async fn context_handler(
req: HttpRequest,
body: web::Json<crate::context_endpoint::ContextRequest>,
state: web::Data<AppState>,
) -> HttpResponse {
use crate::metrics::*;
CONTEXT_REQUESTS_TOTAL.inc();
let _timer = Timer::new(&CONTEXT_DURATION);
let (claims, _token) = match validate_auth(&req, &state).await {
Ok(c) => c,
Err(e) => {
CONTEXT_ERRORS_TOTAL.inc();
ERROR_AUTH_FAILURE_CONTEXT.inc();
return e;
}
};
let user_id = &claims.sub;
if !has_capability(&claims, "memory:read") {
CONTEXT_ERRORS_TOTAL.inc();
ERROR_FORBIDDEN_CONTEXT.inc();
return HttpResponse::Forbidden().json(json!({
"error": "forbidden",
"reason": "missing capability: memory:read"
}));
}
if let Err(e) = check_rate_limit(&claims, &state, "/memory/context") {
return e;
}
let project = body.project.clone().unwrap_or_else(|| "all".to_string());
let scope = body.scope.clone().unwrap_or_else(|| "project".to_string());
let budget = body.budget.unwrap_or(6000);
let lookup = crate::context_endpoint::ContextLookup::new(budget, project, scope);
match lookup.lookup(body.into_inner()).await {
Ok(response) => {
tracing::info!(
tier = response.tier,
lessons = response.lessons.len(),
skills = response.skills.len(),
"context lookup successful"
);
// O3: Track tier hits
let total = response.lessons.len() + response.skills.len();
if total == 0 { CONTEXT_EMPTY_RESULTS.inc(); }
HttpResponse::Ok().json(response)
}
Err(e) => {
CONTEXT_ERRORS_TOTAL.inc();
ERROR_LOOKUP_FAILURE_CONTEXT.inc();
tracing::error!("context lookup error: {}", e);
HttpResponse::BadRequest().json(json!({
"error": "lookup_failed",
"reason": e.to_string()
}))
}
}
}
/// POST /memory/vault/generate — generate Obsidian vault from memories
pub async fn vault_generate_handler(
@@ -1452,12 +1183,20 @@ async fn query_temporal_graph(
state: &web::Data<AppState>,
params: &QueryParams,
) -> anyhow::Result<serde_json::Value> {
// Step 1: Find entities (order by name for deterministic results)
// Step 1: Find entities matching the question
// Use keyword search (ILIKE) on name + description for GET endpoint.
// POST /memory/query uses the full semantic retriever with embeddings.
let search_pattern = format!("%{}%", params.question);
let entities_rows: Vec<(String, String, String)> = sqlx::query_as(
"SELECT id::TEXT, name, entity_type FROM memory_entity WHERE project_id = $1 LIMIT $2"
"SELECT id::TEXT, name, entity_type FROM memory_entity \
WHERE project_id = $1 AND t_expired IS NULL \
AND (name ILIKE $3 OR COALESCE(description, '') ILIKE $3 OR COALESCE(summary, '') ILIKE $3) \
ORDER BY confidence DESC \
LIMIT $2"
)
.bind(&params.project)
.bind(params.limit as i32)
.bind(&search_pattern)
.fetch_all(&state.pool)
.await
.unwrap_or_default();