Files
poimen-memory/crates/mem-cli/src/http_server.rs
T

401 lines
13 KiB
Rust
Raw Normal View History

use actix_web::{web, App, HttpServer, HttpResponse, HttpRequest, middleware::Logger};
use anyhow::Result;
use chrono;
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 with database and workers
pub struct AppState {
pub api_key: String,
pub start_time: Instant,
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
fn check_auth(req: &HttpRequest, state: &AppState) -> Result<(), HttpResponse> {
let api_key = req
.headers()
.get("apikey")
.and_then(|h| h.to_str().ok())
.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"})));
}
Ok(())
}
/// 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(),
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/{ingest_id}", web::get().to(ingest_status))
.route("/memory/query", web::get().to(query_handler))
.route("/memory/projects", web::get().to(projects_handler))
.route("/memory/skills", web::get().to(skills_handler))
.route("/memory/vault", web::post().to(vault_handler))
})
.bind(("0.0.0.0", port))?
.run()
.await?;
Ok(())
}
/// 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}))
}
/// POST /memory/ingest — queue an ingest job
pub async fn ingest_handler(
req: HttpRequest,
body: web::Json<IngestRequest>,
state: web::Data<AppState>,
) -> HttpResponse {
if let Err(e) = check_auth(&req, &state) {
return e;
}
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();
// 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/{ingest_id} — check ingest status
pub async fn ingest_status(
req: HttpRequest,
ingest_id: web::Path<String>,
state: web::Data<AppState>,
) -> HttpResponse {
if let Err(e) = check_auth(&req, &state) {
return e;
}
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 — 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;
}
let project = match query.get("project") {
Some(p) => p.clone(),
None => {
return HttpResponse::BadRequest().json(json!({"error": "missing project parameter"}))
}
};
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"}))
}
}
}
/// GET /memory/projects — list projects with memory
pub async fn projects_handler(
req: HttpRequest,
state: web::Data<AppState>,
) -> HttpResponse {
if let Err(e) = check_auth(&req, &state) {
return e;
}
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/skills — list extracted skills
pub async fn skills_handler(
req: HttpRequest,
state: web::Data<AppState>,
) -> HttpResponse {
if let Err(e) = check_auth(&req, &state) {
return e;
}
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"}))
}
}
}
/// POST /memory/vault — generate Obsidian vault from memories
pub async fn vault_handler(
req: HttpRequest,
state: web::Data<AppState>,
) -> HttpResponse {
if let Err(e) = check_auth(&req, &state) {
return e;
}
let vault_dir = std::env::var("MEM_HOME").unwrap_or_else(|_| "/data".to_string());
// Get all projects from database
let projects_result = sqlx::query_as::<_, (String,)>(
"SELECT DISTINCT project FROM memories_l1 ORDER BY project",
)
.fetch_all(&state.pool)
.await;
match projects_result {
Ok(projects) => {
let mut generated = 0;
let mut errors = Vec::new();
for (project,) in projects {
let project_vault_dir = format!("{}/vault/{}", vault_dir, project);
// Create project directory
if let Err(e) = std::fs::create_dir_all(&project_vault_dir) {
errors.push(format!("Failed to create {}: {}", project_vault_dir, e));
continue;
}
// Get all L1 memories for this project
let l1s_result = sqlx::query_as::<_, (String, String, String)>(
"SELECT id, query_id, content FROM memories_l1 WHERE project = $1 ORDER BY updated_at DESC",
)
.bind(&project)
.fetch_all(&state.pool)
.await;
match l1s_result {
Ok(l1s) => {
for (id, query_id, content) in l1s {
let filename = format!("{}/{}.md", project_vault_dir, query_id);
let note = format!(
"---\nproject: {}\nlevel: L1\nquery_id: {}\nid: {}\nupdated: {}\n---\n\n{}",
project,
query_id,
id,
chrono::Utc::now().to_rfc3339(),
content
);
if std::fs::write(&filename, note).is_ok() {
generated += 1;
}
}
}
Err(e) => {
errors.push(format!("Failed to fetch L1s for {}: {}", project, e));
}
}
// Get L2 synthesis
let l2_result = sqlx::query_as::<_, (String,)>(
"SELECT content FROM memories_l2 WHERE project = $1",
)
.bind(&project)
.fetch_optional(&state.pool)
.await;
if let Ok(Some((content,))) = l2_result {
let filename = format!("{}/index.md", project_vault_dir);
let note = format!(
"---\nproject: {}\nlevel: L2\ntitle: {} Synthesis\nupdated: {}\n---\n\n{}",
project,
project,
chrono::Utc::now().to_rfc3339(),
content
);
if std::fs::write(&filename, note).is_ok() {
generated += 1;
}
}
}
HttpResponse::Ok().json(json!({
"status": "generated",
"vault_dir": format!("{}/vault", vault_dir),
"notes_created": generated,
"errors": errors
}))
}
Err(_) => {
HttpResponse::InternalServerError().json(json!({"error": "database_error"}))
}
}
}