//! HTTP server for RAG module - Node.js API compatibility
use axum::{
extract::{Path, Query, State},
http::StatusCode,
response::Json,
routing::{get, post, put, delete},
Router,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tower::ServiceBuilder;
use tower_http::cors::CorsLayer;
use tracing::{info, error};
use rag_module::{RagModule, types::*};
#[derive(Clone)]
struct AppState {
rag: Arc<RagModule>,
}
#[derive(Deserialize)]
struct AddDocumentRequest {
collection_type: String,
document: Document,
}
#[derive(Deserialize)]
struct SearchRequest {
collection_type: String,
query: String,
options: Option<SearchOptions>,
}
#[derive(Deserialize)]
struct ChatMessageRequest {
context_id: String,
role: String,
content: String,
}
#[derive(Serialize)]
struct ApiResponse<T> {
success: bool,
data: Option<T>,
error: Option<String>,
}
impl<T> ApiResponse<T> {
fn success(data: T) -> Self {
Self {\n success: true,\n data: Some(data),\n error: None,\n }\n }\n \n fn error(error: String) -> Self {\n Self {\n success: false,\n data: None,\n error: Some(error),\n }\n }\n}\n\n#[tokio::main]\nasync fn main() -> anyhow::Result<()> {\n // Initialize tracing\n tracing_subscriber::fmt::init();\n \n // Initialize RAG module\n let base_path = std::env::current_dir()?.join(\"rag-data\");\n let rag = rag_module::create_rag_module(base_path).await?;\n rag.initialize().await?;\n \n let state = AppState {\n rag: Arc::new(rag),\n };\n \n // Build the router\n let app = Router::new()\n .route(\"/health\", get(health_check))\n .route(\"/api/documents\", post(add_document))\n .route(\"/api/documents/:collection/:id\", get(get_document))\n .route(\"/api/documents/:collection/:id\", put(update_document))\n .route(\"/api/documents/:collection/:id\", delete(delete_document))\n .route(\"/api/search\", post(search_documents))\n .route(\"/api/chat/message\", post(add_chat_message))\n .route(\"/api/chat/:context_id\", get(get_chat_history))\n .route(\"/api/aws/estate\", post(process_aws_estate))\n .route(\"/api/collections\", get(list_collections))\n .route(\"/api/collections/:name\", get(get_collection_info))\n .layer(\n ServiceBuilder::new()\n .layer(CorsLayer::permissive())\n )\n .with_state(state);\n \n let listener = tokio::net::TcpListener::bind(\"127.0.0.1:3000\").await?;\n info!(\"RAG Module server listening on http://127.0.0.1:3000\");\n \n axum::serve(listener, app).await?;\n \n Ok(())\n}\n\nasync fn health_check() -> Json<ApiResponse<String>> {\n Json(ApiResponse::success(\"OK\".to_string()))\n}\n\nasync fn add_document(\n State(state): State<AppState>,\n Json(request): Json<AddDocumentRequest>,\n) -> Result<Json<ApiResponse<String>>, StatusCode> {\n match state.rag.add_document(&request.collection_type, request.document).await {\n Ok(id) => Ok(Json(ApiResponse::success(id))),\n Err(e) => {\n error!(\"Failed to add document: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn get_document(\n State(state): State<AppState>,\n Path((collection, id)): Path<(String, String)>,\n) -> Result<Json<ApiResponse<Option<Document>>>, StatusCode> {\n match state.rag.get_document(&collection, &id).await {\n Ok(doc) => Ok(Json(ApiResponse::success(doc))),\n Err(e) => {\n error!(\"Failed to get document: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn update_document(\n State(state): State<AppState>,\n Path((collection, id)): Path<(String, String)>,\n Json(document): Json<Document>,\n) -> Result<Json<ApiResponse<()>>, StatusCode> {\n match state.rag.vector_store.update_document(&collection, &id, document).await {\n Ok(_) => Ok(Json(ApiResponse::success(()))),\n Err(e) => {\n error!(\"Failed to update document: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn delete_document(\n State(state): State<AppState>,\n Path((collection, id)): Path<(String, String)>,\n) -> Result<Json<ApiResponse<bool>>, StatusCode> {\n match state.rag.vector_store.delete_document(&collection, &id).await {\n Ok(deleted) => Ok(Json(ApiResponse::success(deleted))),\n Err(e) => {\n error!(\"Failed to delete document: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn search_documents(\n State(state): State<AppState>,\n Json(request): Json<SearchRequest>,\n) -> Result<Json<ApiResponse<Vec<SearchResult>>>, StatusCode> {\n let options = request.options.unwrap_or_default();\n \n match state.rag.search(&request.collection_type, &request.query, options).await {\n Ok(results) => Ok(Json(ApiResponse::success(results))),\n Err(e) => {\n error!(\"Failed to search documents: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn add_chat_message(\n State(state): State<AppState>,\n Json(request): Json<ChatMessageRequest>,\n) -> Result<Json<ApiResponse<String>>, StatusCode> {\n match state.rag.add_chat_message(&request.context_id, &request.role, &request.content).await {\n Ok(id) => Ok(Json(ApiResponse::success(id))),\n Err(e) => {\n error!(\"Failed to add chat message: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn get_chat_history(\n State(state): State<AppState>,\n Path(context_id): Path<String>,\n Query(params): Query<HashMap<String, String>>,\n) -> Result<Json<ApiResponse<Vec<ChatMessage>>>, StatusCode> {\n let limit = params.get(\"limit\")\n .and_then(|s| s.parse().ok());\n \n match state.rag.get_chat_history(&context_id, limit).await {\n Ok(messages) => Ok(Json(ApiResponse::success(messages))),\n Err(e) => {\n error!(\"Failed to get chat history: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn process_aws_estate(\n State(state): State<AppState>,\n Json(data): Json<serde_json::Value>,\n) -> Result<Json<ApiResponse<Vec<String>>>, StatusCode> {\n match state.rag.process_aws_estate(data).await {\n Ok(ids) => Ok(Json(ApiResponse::success(ids))),\n Err(e) => {\n error!(\"Failed to process AWS estate: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn list_collections(\n State(state): State<AppState>,\n) -> Result<Json<ApiResponse<Vec<String>>>, StatusCode> {\n match state.rag.vector_store.list_collections().await {\n Ok(collections) => Ok(Json(ApiResponse::success(collections))),\n Err(e) => {\n error!(\"Failed to list collections: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}\n\nasync fn get_collection_info(\n State(state): State<AppState>,\n Path(name): Path<String>,\n) -> Result<Json<ApiResponse<Option<crate::db::CollectionInfo>>>, StatusCode> {\n match state.rag.vector_store.get_collection_info(&name).await {\n Ok(info) => Ok(Json(ApiResponse::success(info))),\n Err(e) => {\n error!(\"Failed to get collection info: {}\", e);\n Ok(Json(ApiResponse::error(e.to_string())))\n }\n }\n}"