#![allow(unused_variables)]
use axum::{
extract::{Path, Query, State},
http::StatusCode,
Json,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use crate::api::models::{ApiError, ApiResponse};
use crate::api::server::AppState;
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Document {
pub id: String,
pub content: String,
pub metadata: HashMap<String, serde_json::Value>,
pub created_at: String,
pub updated_at: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub chunks: Option<Vec<DocumentChunk>>,
}
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct DocumentChunk {
pub id: String,
pub content: String,
pub index: usize,
pub start_offset: usize,
pub end_offset: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<HashMap<String, serde_json::Value>>,
}
#[derive(Debug, Deserialize)]
pub struct CreateDocumentRequest {
pub id: Option<String>,
pub content: String,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
pub chunking: Option<ChunkingConfig>,
#[serde(default = "default_true")]
pub embed: bool,
pub vector_store: Option<String>,
}
fn default_true() -> bool {
true
}
#[derive(Debug, Deserialize, Clone)]
pub struct ChunkingConfig {
#[serde(default = "default_strategy")]
pub strategy: String,
#[serde(default = "default_chunk_size")]
pub chunk_size: usize,
#[serde(default = "default_overlap")]
pub overlap: usize,
#[serde(default)]
pub split_on_headers: bool,
pub min_chunk_size: Option<usize>,
}
fn default_strategy() -> String {
"sentence".to_string()
}
fn default_chunk_size() -> usize {
512
}
fn default_overlap() -> usize {
50
}
#[derive(Debug, Deserialize)]
pub struct BatchCreateRequest {
pub documents: Vec<CreateDocumentRequest>,
pub chunking: Option<ChunkingConfig>,
pub vector_store: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct UpdateDocumentRequest {
pub content: Option<String>,
pub metadata: Option<HashMap<String, serde_json::Value>>,
#[serde(default)]
pub rechunk: bool,
#[serde(default)]
pub reembed: bool,
}
#[derive(Debug, Deserialize)]
pub struct SearchDocumentsRequest {
pub query: String,
#[serde(default = "default_search_type")]
pub search_type: String,
#[serde(default = "default_limit")]
pub limit: usize,
pub filter: Option<HashMap<String, serde_json::Value>>,
#[serde(default = "default_true")]
pub include_content: bool,
#[serde(default)]
pub include_chunks: bool,
#[serde(default)]
pub highlight: bool,
pub alpha: Option<f32>,
}
fn default_search_type() -> String {
"hybrid".to_string()
}
fn default_limit() -> usize {
10
}
#[derive(Debug, Serialize)]
pub struct DocumentSearchResult {
pub id: String,
pub score: f32,
#[serde(skip_serializing_if = "Option::is_none")]
pub content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub highlights: Option<Vec<String>>,
pub metadata: HashMap<String, serde_json::Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub chunks: Option<Vec<ChunkSearchResult>>,
}
#[derive(Debug, Serialize)]
pub struct ChunkSearchResult {
pub chunk_id: String,
pub score: f32,
pub content: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub highlight: Option<String>,
pub index: usize,
}
#[derive(Debug, Serialize)]
pub struct SearchResponse {
pub results: Vec<DocumentSearchResult>,
pub total: usize,
pub query_time_ms: u64,
}
#[derive(Debug, Deserialize)]
pub struct ListDocumentsQuery {
pub limit: Option<usize>,
pub offset: Option<usize>,
pub sort: Option<String>,
pub order: Option<String>,
pub filter: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct ChunkDocumentRequest {
pub config: ChunkingConfig,
#[serde(default = "default_true")]
pub embed: bool,
pub vector_store: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct SimilarDocumentsRequest {
#[serde(default = "default_limit")]
pub limit: usize,
#[serde(default)]
pub include_content: bool,
}
pub async fn list_documents(
State(state): State<AppState>,
Query(query): Query<ListDocumentsQuery>,
) -> Result<Json<ApiResponse<Vec<Document>>>, ApiError> {
let docs = state.db.list_documents(
"default",
).map_err(|e| ApiError::internal(format!("Failed to list documents: {}", e)))?;
let documents: Vec<Document> = docs
.into_iter()
.map(|d| {
let metadata = match d.metadata {
Some(serde_json::Value::Object(map)) => map.into_iter().collect(),
_ => HashMap::new(),
};
Document {
id: d.id,
content: d.content,
metadata,
created_at: d.created_at,
updated_at: d.updated_at,
chunks: None,
}
})
.collect();
Ok(Json(ApiResponse::success(documents)))
}
pub async fn create_document(
State(state): State<AppState>,
Json(req): Json<CreateDocumentRequest>,
) -> Result<(StatusCode, Json<ApiResponse<Document>>), ApiError> {
let id = req.id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
let metadata_value = if req.metadata.is_empty() {
None
} else {
Some(serde_json::Value::Object(
req.metadata.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
))
};
let _doc_id = state.db.create_document(
"default",
&id,
&req.content,
metadata_value,
).map_err(|e| ApiError::internal(format!("Failed to create document: {}", e)))?;
let document = Document {
id: id.clone(),
content: req.content.clone(),
metadata: req.metadata,
created_at: chrono::Utc::now().to_rfc3339(),
updated_at: chrono::Utc::now().to_rfc3339(),
chunks: None,
};
Ok((StatusCode::CREATED, Json(ApiResponse::success(document))))
}
pub async fn batch_create_documents(
State(state): State<AppState>,
Json(req): Json<BatchCreateRequest>,
) -> Result<(StatusCode, Json<ApiResponse<serde_json::Value>>), ApiError> {
let docs_data: Vec<crate::types::DocumentData> = req.documents.iter().map(|d| {
let metadata = if d.metadata.is_empty() {
None
} else {
Some(serde_json::Value::Object(
d.metadata.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect()
))
};
crate::types::DocumentData {
id: d.id.clone().unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
content: d.content.clone(),
metadata,
created_at: chrono::Utc::now().to_rfc3339(),
updated_at: chrono::Utc::now().to_rfc3339(),
chunks: vec![],
}
}).collect();
let ids = state.db.batch_create_documents(
"default",
docs_data,
).map_err(|e| ApiError::internal(format!("Failed to batch create: {}", e)))?;
let count = ids.len();
Ok((StatusCode::CREATED, Json(ApiResponse::success(serde_json::json!({
"created_count": count,
})))))
}
pub async fn get_document(
State(state): State<AppState>,
Path(doc_id): Path<String>,
Query(params): Query<HashMap<String, String>>,
) -> Result<Json<ApiResponse<Document>>, ApiError> {
let include_chunks = params.get("include_chunks")
.map(|s| s == "true")
.unwrap_or(false);
let doc = state.db.get_document("default", &doc_id)
.map_err(|e| ApiError::not_found(format!("Document not found: {}", e)))?;
let metadata = match doc.metadata {
Some(serde_json::Value::Object(map)) => map.into_iter().collect(),
_ => HashMap::new(),
};
let document = Document {
id: doc.id,
content: doc.content,
metadata,
created_at: doc.created_at,
updated_at: doc.updated_at,
chunks: if include_chunks {
None
} else {
None
},
};
Ok(Json(ApiResponse::success(document)))
}
pub async fn update_document(
State(state): State<AppState>,
Path(doc_id): Path<String>,
Json(req): Json<UpdateDocumentRequest>,
) -> Result<Json<ApiResponse<Document>>, ApiError> {
let metadata_value = req.metadata.clone().map(|map| {
serde_json::Value::Object(
map.into_iter().collect()
)
});
state.db.update_document(
"default",
&doc_id,
req.content.as_deref().unwrap_or(""),
metadata_value,
).map_err(|e| ApiError::internal(format!("Failed to update document: {}", e)))?;
let document = Document {
id: doc_id.clone(),
content: req.content.clone().unwrap_or_default(),
metadata: req.metadata.unwrap_or_default(),
created_at: chrono::Utc::now().to_rfc3339(),
updated_at: chrono::Utc::now().to_rfc3339(),
chunks: None,
};
Ok(Json(ApiResponse::success(document)))
}
pub async fn delete_document(
State(state): State<AppState>,
Path(doc_id): Path<String>,
) -> Result<StatusCode, ApiError> {
state.db.delete_document("default", &doc_id)
.map_err(|e| ApiError::internal(format!("Failed to delete document: {}", e)))?;
Ok(StatusCode::NO_CONTENT)
}
pub async fn search_documents(
State(state): State<AppState>,
Json(req): Json<SearchDocumentsRequest>,
) -> Result<Json<ApiResponse<SearchResponse>>, ApiError> {
let start = std::time::Instant::now();
let raw_results = state.db.search_documents(
"default",
&req.query,
).map_err(|e| ApiError::internal(format!("Search failed: {}", e)))?;
let total = raw_results.len();
let search_results: Vec<DocumentSearchResult> = raw_results
.into_iter()
.map(|doc| {
let metadata = match doc.metadata {
Some(serde_json::Value::Object(map)) => map.into_iter().collect(),
_ => HashMap::new(),
};
DocumentSearchResult {
id: doc.id,
score: 1.0, content: if req.include_content { Some(doc.content) } else { None },
highlights: None, metadata,
chunks: None, }
})
.collect();
Ok(Json(ApiResponse::success(SearchResponse {
results: search_results,
total,
query_time_ms: start.elapsed().as_millis() as u64,
})))
}
pub async fn get_chunks(
State(state): State<AppState>,
Path(doc_id): Path<String>,
) -> Result<Json<ApiResponse<Vec<DocumentChunk>>>, ApiError> {
let chunks = state.db.get_document_chunks("default", &doc_id)
.map_err(|e| ApiError::internal(format!("Failed to get chunks: {}", e)))?;
let result: Vec<DocumentChunk> = chunks
.into_iter()
.enumerate()
.map(|(index, (id, _score))| DocumentChunk {
id,
content: String::new(), index,
start_offset: 0,
end_offset: 0,
metadata: None,
})
.collect();
Ok(Json(ApiResponse::success(result)))
}
pub async fn chunk_document(
State(state): State<AppState>,
Path(doc_id): Path<String>,
Json(req): Json<ChunkDocumentRequest>,
) -> Result<Json<ApiResponse<Vec<DocumentChunk>>>, ApiError> {
let chunk_ids = state.db.rechunk_document(
"default",
&doc_id,
512,
).map_err(|e| ApiError::internal(format!("Failed to chunk document: {}", e)))?;
let result: Vec<DocumentChunk> = chunk_ids
.into_iter()
.enumerate()
.map(|(index, id)| DocumentChunk {
id: id.clone(),
content: String::new(), index,
start_offset: 0,
end_offset: 0,
metadata: None,
})
.collect();
Ok(Json(ApiResponse::success(result)))
}
pub async fn similar_documents(
State(state): State<AppState>,
Path(doc_id): Path<String>,
Json(req): Json<SimilarDocumentsRequest>,
) -> Result<Json<ApiResponse<Vec<DocumentSearchResult>>>, ApiError> {
let results = state.db.find_similar_documents(
"default",
&doc_id,
req.limit,
).map_err(|e| ApiError::internal(format!("Failed to find similar: {}", e)))?;
let search_results: Vec<DocumentSearchResult> = results
.into_iter()
.map(|(doc, score)| {
let metadata = match doc.metadata {
Some(serde_json::Value::Object(map)) => map.into_iter().collect(),
_ => HashMap::new(),
};
DocumentSearchResult {
id: doc.id,
score,
content: if req.include_content { Some(doc.content) } else { None },
highlights: None,
metadata,
chunks: None,
}
})
.collect();
Ok(Json(ApiResponse::success(search_results)))
}