feat(core): implement full memory pipeline (#11)
This commit was merged in pull request #11.
This commit is contained in:
@@ -1,18 +1,27 @@
|
||||
use actix_web::{web, App, HttpServer, HttpResponse, HttpRequest, middleware::Logger};
|
||||
use serde_json::json;
|
||||
use std::sync::Mutex;
|
||||
use std::time::Instant;
|
||||
use anyhow::Result;
|
||||
use crate::endpoints::{IngestQueue, IngestRequest};
|
||||
use mem_llm::{EmbeddingsClient, RerankClient};
|
||||
use mem_store::{init_schema, VectorStore};
|
||||
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;
|
||||
|
||||
/// Server state.
|
||||
/// Server state with database and workers
|
||||
pub struct AppState {
|
||||
pub api_key: String,
|
||||
pub start_time: Instant,
|
||||
pub queue: Mutex<IngestQueue>,
|
||||
pub pool: PgPool,
|
||||
pub vector_store: Arc<VectorStore>,
|
||||
pub embeddings: Arc<EmbeddingsClient>,
|
||||
pub ingest_worker: Arc<IngestWorker>,
|
||||
pub query_worker: Arc<QueryWorker>,
|
||||
}
|
||||
|
||||
/// Auth extractor — validates apikey header.
|
||||
/// Auth extractor — validates apikey header
|
||||
fn check_auth(req: &HttpRequest, state: &AppState) -> Result<(), HttpResponse> {
|
||||
let api_key = req
|
||||
.headers()
|
||||
@@ -21,32 +30,50 @@ fn check_auth(req: &HttpRequest, state: &AppState) -> Result<(), HttpResponse> {
|
||||
.map(|s| s.to_string());
|
||||
|
||||
if api_key.as_ref() != Some(&state.api_key) {
|
||||
return Err(HttpResponse::Unauthorized()
|
||||
.json(json!({"error": "unauthorized", "reason": "missing apikey header"})));
|
||||
return Err(HttpResponse::Unauthorized().json(json!({"error": "unauthorized", "reason": "missing apikey header"})));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Start HTTP server.
|
||||
pub async fn start_server(port: u16, api_key: String) -> Result<()> {
|
||||
/// Start HTTP server with database initialization
|
||||
pub async fn start_server(port: u16, api_key: String, database_url: &str) -> Result<()> {
|
||||
// Create connection pool
|
||||
let pool = PgPool::connect(database_url).await?;
|
||||
tracing::info!("Connected to database");
|
||||
|
||||
// Initialize schema
|
||||
init_schema(&pool).await?;
|
||||
tracing::info!("Schema initialized");
|
||||
|
||||
// Create workers
|
||||
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));
|
||||
|
||||
let state = web::Data::new(AppState {
|
||||
api_key,
|
||||
start_time: Instant::now(),
|
||||
queue: Mutex::new(IngestQueue::new()),
|
||||
pool,
|
||||
vector_store,
|
||||
embeddings,
|
||||
ingest_worker,
|
||||
query_worker,
|
||||
});
|
||||
|
||||
tracing::info!("Starting HTTP server on port {}", port);
|
||||
|
||||
HttpServer::new(move || {
|
||||
App::new()
|
||||
.app_data(state.clone())
|
||||
.wrap(Logger::default())
|
||||
.route("/health", web::get().to(health_check))
|
||||
.route("/memory/ingest", web::post().to(ingest_handler))
|
||||
.route("/memory/ingest/{job_id}", web::get().to(ingest_status))
|
||||
.route("/memory/ingest/{ingest_id}", web::get().to(ingest_status))
|
||||
.route("/memory/query", web::get().to(query_handler))
|
||||
.route("/memory/skills", web::get().to(skills_handler))
|
||||
.route("/memory/skills/{name}", web::get().to(skill_detail))
|
||||
.route("/memory/projects", web::get().to(projects_handler))
|
||||
.route("/memory/projects/{id}/status", web::get().to(project_status))
|
||||
.route("/memory/skills", web::get().to(skills_handler))
|
||||
})
|
||||
.bind(("0.0.0.0", port))?
|
||||
.run()
|
||||
@@ -55,14 +82,13 @@ pub async fn start_server(port: u16, api_key: String) -> Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Health check endpoint (no auth required).
|
||||
/// Health check (no auth)
|
||||
pub async fn health_check(state: web::Data<AppState>) -> HttpResponse {
|
||||
let uptime = state.start_time.elapsed().as_secs();
|
||||
HttpResponse::Ok()
|
||||
.json(json!({"status": "ok", "uptime_seconds": uptime}))
|
||||
HttpResponse::Ok().json(json!({"status": "ok", "uptime_seconds": uptime}))
|
||||
}
|
||||
|
||||
/// POST /memory/ingest
|
||||
/// POST /memory/ingest — queue an ingest job
|
||||
pub async fn ingest_handler(
|
||||
req: HttpRequest,
|
||||
body: web::Json<IngestRequest>,
|
||||
@@ -72,88 +98,141 @@ pub async fn ingest_handler(
|
||||
return e;
|
||||
}
|
||||
|
||||
let mut q = state.queue.lock().unwrap();
|
||||
let (job_id, _) = q.submit(&body.project, &body.ingest_id);
|
||||
let project = body.project.clone();
|
||||
let ingest_id = body.ingest_id.clone();
|
||||
let records: Vec<(String, String)> = body
|
||||
.records
|
||||
.iter()
|
||||
.map(|r| (r.text.clone(), body.source.clone()))
|
||||
.collect();
|
||||
|
||||
HttpResponse::Accepted().json(json!({
|
||||
"job_id": job_id,
|
||||
"ingest_id": body.ingest_id,
|
||||
"status_url": format!("/memory/ingest/{}", job_id),
|
||||
"estimated_wait_seconds": 15
|
||||
}))
|
||||
// Create ingest job in DB
|
||||
let job_result = sqlx::query(
|
||||
"INSERT INTO ingest_jobs (id, project, ingest_id, status, created_at)
|
||||
VALUES ($1, $2, $3, 'pending', NOW())
|
||||
ON CONFLICT (ingest_id) DO NOTHING
|
||||
RETURNING id",
|
||||
)
|
||||
.bind(uuid::Uuid::new_v4())
|
||||
.bind(&project)
|
||||
.bind(&ingest_id)
|
||||
.fetch_optional(&state.pool)
|
||||
.await;
|
||||
|
||||
match job_result {
|
||||
Ok(Some(_)) => {
|
||||
// Spawn async ingest task
|
||||
let worker = state.ingest_worker.clone();
|
||||
let proj = project.clone();
|
||||
let id = ingest_id.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = worker.process_ingest(&proj, &id, records).await {
|
||||
tracing::error!("Ingest failed: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
HttpResponse::Accepted().json(json!({
|
||||
"ingest_id": ingest_id,
|
||||
"status": "pending",
|
||||
"status_url": format!("/memory/ingest/{}", ingest_id)
|
||||
}))
|
||||
}
|
||||
Ok(None) => {
|
||||
// Already exists
|
||||
HttpResponse::Conflict().json(json!({
|
||||
"error": "already_ingesting",
|
||||
"ingest_id": ingest_id
|
||||
}))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("DB error: {}", e);
|
||||
HttpResponse::InternalServerError().json(json!({
|
||||
"error": "database_error"
|
||||
}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /memory/ingest/{job_id}
|
||||
/// GET /memory/ingest/{ingest_id} — check ingest status
|
||||
pub async fn ingest_status(
|
||||
req: HttpRequest,
|
||||
job_id: web::Path<String>,
|
||||
ingest_id: web::Path<String>,
|
||||
state: web::Data<AppState>,
|
||||
) -> HttpResponse {
|
||||
if let Err(e) = check_auth(&req, &state) {
|
||||
return e;
|
||||
}
|
||||
|
||||
let q = state.queue.lock().unwrap();
|
||||
match q.get_status(&job_id) {
|
||||
Some(status) => HttpResponse::Ok().json(status),
|
||||
None => HttpResponse::NotFound().json(json!({"error": "job not found"})),
|
||||
let id = ingest_id.into_inner();
|
||||
let result = sqlx::query_as::<_, (String, String, Option<String>)>(
|
||||
"SELECT ingest_id, status, error FROM ingest_jobs WHERE ingest_id = $1",
|
||||
)
|
||||
.bind(&id)
|
||||
.fetch_optional(&state.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(Some((ingest_id, status, error))) => {
|
||||
HttpResponse::Ok().json(json!({
|
||||
"ingest_id": ingest_id,
|
||||
"status": status,
|
||||
"error": error
|
||||
}))
|
||||
}
|
||||
Ok(None) => {
|
||||
HttpResponse::NotFound().json(json!({"error": "not_found"}))
|
||||
}
|
||||
Err(_) => {
|
||||
HttpResponse::InternalServerError().json(json!({"error": "database_error"}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /memory/query
|
||||
/// GET /memory/query — semantic search across memories
|
||||
pub async fn query_handler(
|
||||
req: HttpRequest,
|
||||
query: web::Query<std::collections::HashMap<String, String>>,
|
||||
state: web::Data<AppState>,
|
||||
) -> HttpResponse {
|
||||
if let Err(e) = check_auth(&req, &state) {
|
||||
return e;
|
||||
}
|
||||
|
||||
HttpResponse::Ok().json(json!({
|
||||
"results": [{
|
||||
"level": "L1",
|
||||
"score": 0.95,
|
||||
"text": "Infrastructure root causes",
|
||||
"provenance": ["pi-2026-07-21-xyz"]
|
||||
}]
|
||||
}))
|
||||
}
|
||||
let project = match query.get("project") {
|
||||
Some(p) => p.clone(),
|
||||
None => {
|
||||
return HttpResponse::BadRequest().json(json!({"error": "missing project parameter"}))
|
||||
}
|
||||
};
|
||||
|
||||
/// GET /memory/skills
|
||||
pub async fn skills_handler(
|
||||
req: HttpRequest,
|
||||
state: web::Data<AppState>,
|
||||
) -> HttpResponse {
|
||||
if let Err(e) = check_auth(&req, &state) {
|
||||
return e;
|
||||
let question = match query.get("query") {
|
||||
Some(q) => q.clone(),
|
||||
None => {
|
||||
return HttpResponse::BadRequest().json(json!({"error": "missing query parameter"}))
|
||||
}
|
||||
};
|
||||
|
||||
let limit = query
|
||||
.get("limit")
|
||||
.and_then(|l| l.parse::<i64>().ok())
|
||||
.unwrap_or(5);
|
||||
|
||||
match state.query_worker.query(&project, &question, Some(limit)).await {
|
||||
Ok(results) => {
|
||||
HttpResponse::Ok().json(json!({
|
||||
"query": question,
|
||||
"project": project,
|
||||
"results": results
|
||||
}))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::error!("Query failed: {}", e);
|
||||
HttpResponse::InternalServerError().json(json!({"error": "query_failed"}))
|
||||
}
|
||||
}
|
||||
|
||||
HttpResponse::Ok().json(json!({
|
||||
"skills": [
|
||||
{"name": "infrastructure", "queries": 3},
|
||||
{"name": "errors", "queries": 5}
|
||||
]
|
||||
}))
|
||||
}
|
||||
|
||||
/// GET /memory/skills/{name}
|
||||
pub async fn skill_detail(
|
||||
req: HttpRequest,
|
||||
name: web::Path<String>,
|
||||
state: web::Data<AppState>,
|
||||
) -> HttpResponse {
|
||||
if let Err(e) = check_auth(&req, &state) {
|
||||
return e;
|
||||
}
|
||||
|
||||
HttpResponse::Ok().json(json!({
|
||||
"name": name.into_inner(),
|
||||
"description": "Skill details",
|
||||
"related_queries": 3
|
||||
}))
|
||||
}
|
||||
|
||||
/// GET /memory/projects
|
||||
/// GET /memory/projects — list projects with memory
|
||||
pub async fn projects_handler(
|
||||
req: HttpRequest,
|
||||
state: web::Data<AppState>,
|
||||
@@ -162,28 +241,60 @@ pub async fn projects_handler(
|
||||
return e;
|
||||
}
|
||||
|
||||
HttpResponse::Ok().json(json!({
|
||||
"projects": [
|
||||
{"id": "poimen", "status": "healthy", "memories": 147}
|
||||
]
|
||||
}))
|
||||
let result = sqlx::query_as::<_, (String,)>(
|
||||
"SELECT DISTINCT project FROM memories_l2 ORDER BY project",
|
||||
)
|
||||
.fetch_all(&state.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(rows) => {
|
||||
let projects: Vec<String> = rows.into_iter().map(|(p,)| p).collect();
|
||||
HttpResponse::Ok().json(json!({
|
||||
"projects": projects,
|
||||
"count": projects.len()
|
||||
}))
|
||||
}
|
||||
Err(_) => {
|
||||
HttpResponse::InternalServerError().json(json!({"error": "database_error"}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /memory/projects/{id}/status
|
||||
pub async fn project_status(
|
||||
/// GET /memory/skills — list extracted skills
|
||||
pub async fn skills_handler(
|
||||
req: HttpRequest,
|
||||
id: web::Path<String>,
|
||||
state: web::Data<AppState>,
|
||||
) -> HttpResponse {
|
||||
if let Err(e) = check_auth(&req, &state) {
|
||||
return e;
|
||||
}
|
||||
|
||||
HttpResponse::Ok().json(json!({
|
||||
"project": id.into_inner(),
|
||||
"status": "healthy",
|
||||
"l0_chunks": 412,
|
||||
"l1_memories": 17,
|
||||
"l2_synthesis": 1
|
||||
}))
|
||||
let result = sqlx::query_as::<_, (String, String, String)>(
|
||||
"SELECT name, description, when_to_use FROM skills ORDER BY created_at DESC LIMIT 50",
|
||||
)
|
||||
.fetch_all(&state.pool)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(rows) => {
|
||||
let skills: Vec<serde_json::Value> = rows
|
||||
.into_iter()
|
||||
.map(|(name, desc, when_to_use)| {
|
||||
json!({
|
||||
"name": name,
|
||||
"description": desc,
|
||||
"when_to_use": when_to_use
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
HttpResponse::Ok().json(json!({
|
||||
"skills": skills,
|
||||
"count": skills.len()
|
||||
}))
|
||||
}
|
||||
Err(_) => {
|
||||
HttpResponse::InternalServerError().json(json!({"error": "database_error"}))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user