feat: implement Mistral BYOK feature for OCR processing with user profile integration
This commit is contained in:
@@ -32,10 +32,9 @@ ai:
|
||||
api_key: "${AI_API_KEY}"
|
||||
|
||||
# Mistral OCR for document parsing
|
||||
# Note: API key is set per-user in Profile settings (BYOK)
|
||||
mistral:
|
||||
api_key: "${MISTRAL_API_KEY}"
|
||||
ocr_model: "mistral-ocr-latest"
|
||||
max_pages_per_request: 8
|
||||
max_retries: 2 # Max retry attempts per chunk
|
||||
timeout_secs: 120 # Request timeout in seconds
|
||||
|
||||
|
||||
@@ -57,6 +57,8 @@ pub struct AiConfig {
|
||||
|
||||
#[derive(Debug, Deserialize, Clone)]
|
||||
pub struct MistralConfig {
|
||||
/// API key - NOT loaded from config, set at runtime from user's profile
|
||||
#[serde(skip, default)]
|
||||
pub api_key: String,
|
||||
pub ocr_model: String,
|
||||
pub max_pages_per_request: u32,
|
||||
|
||||
@@ -52,6 +52,8 @@ pub async fn parse_source(
|
||||
State(state): State<OcrState>,
|
||||
Path(id): Path<i32>,
|
||||
) -> Result<Json<ParseResponse>, (StatusCode, Json<ErrorResponse>)> {
|
||||
use crate::models::user::user;
|
||||
|
||||
// 1. Get source from database
|
||||
let source_entity = source::Entity::find_by_id(id)
|
||||
.one(&state.db)
|
||||
@@ -83,6 +85,7 @@ pub async fn parse_source(
|
||||
}
|
||||
|
||||
let file_path = PathBuf::from(&source_entity.file_path);
|
||||
let user_id = source_entity.user_id;
|
||||
|
||||
// 2. Set status to "processing" immediately
|
||||
let mut active_model: source::ActiveModel = source_entity.into();
|
||||
@@ -95,10 +98,39 @@ pub async fn parse_source(
|
||||
}),
|
||||
)
|
||||
})?;
|
||||
// 3. User must have their own Mistral API key configured
|
||||
let user_api_key = if let Ok(Some(user_entity)) = user::Entity::find_by_id(user_id).one(&state.db).await {
|
||||
user_entity.mistral_api_key
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// 3. Spawn background task for OCR processing
|
||||
let api_key = match user_api_key {
|
||||
Some(key) if !key.is_empty() => {
|
||||
tracing::info!("Using user's Mistral API key for source {}", id);
|
||||
key
|
||||
}
|
||||
_ => {
|
||||
// Revert status back to pending since we can't process
|
||||
if let Ok(Some(entity)) = source::Entity::find_by_id(id).one(&state.db).await {
|
||||
let mut revert_model: source::ActiveModel = entity.into();
|
||||
revert_model.status = Set("pending".to_string());
|
||||
let _ = revert_model.update(&state.db).await;
|
||||
}
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
Json(ErrorResponse {
|
||||
error: "Please configure your Mistral API key in Profile settings".to_string(),
|
||||
}),
|
||||
));
|
||||
}
|
||||
};
|
||||
|
||||
let mut mistral_config = state.mistral.clone();
|
||||
mistral_config.api_key = api_key;
|
||||
|
||||
// 4. Spawn background task for OCR processing
|
||||
let db = state.db.clone();
|
||||
let mistral_config = state.mistral.clone();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = process_ocr_background(db, mistral_config, id, file_path).await {
|
||||
@@ -106,7 +138,7 @@ pub async fn parse_source(
|
||||
}
|
||||
});
|
||||
|
||||
// 4. Return immediately
|
||||
// 5. Return immediately
|
||||
Ok(Json(ParseResponse {
|
||||
success: true,
|
||||
biomarkers_count: 0,
|
||||
|
||||
@@ -36,6 +36,7 @@ pub struct UpdateUserRequest {
|
||||
pub alcohol: Option<bool>,
|
||||
pub diet_id: Option<i32>,
|
||||
pub avatar_url: Option<String>,
|
||||
pub mistral_api_key: Option<String>,
|
||||
}
|
||||
|
||||
/// Response for a user.
|
||||
@@ -52,6 +53,7 @@ pub struct UserResponse {
|
||||
pub alcohol: Option<bool>,
|
||||
pub diet: Option<String>,
|
||||
pub avatar_url: Option<String>,
|
||||
pub has_mistral_key: bool,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
@@ -88,16 +90,17 @@ pub async fn list_users(
|
||||
.into_iter()
|
||||
.map(|u| UserResponse {
|
||||
id: u.id,
|
||||
username: u.username,
|
||||
name: u.name,
|
||||
username: u.username.clone(),
|
||||
name: u.name.clone(),
|
||||
role: role_map.get(&u.role_id).cloned().unwrap_or_default(),
|
||||
height_cm: u.height_cm,
|
||||
blood_type: u.blood_type,
|
||||
blood_type: u.blood_type.clone(),
|
||||
birthdate: u.birthdate.map(|d| d.to_string()),
|
||||
smoking: u.smoking,
|
||||
alcohol: u.alcohol,
|
||||
diet: u.diet_id.and_then(|id| diet_map.get(&id).cloned()),
|
||||
avatar_url: u.avatar_url,
|
||||
avatar_url: u.avatar_url.clone(),
|
||||
has_mistral_key: u.mistral_api_key.is_some(),
|
||||
created_at: u.created_at.to_string(),
|
||||
})
|
||||
.collect();
|
||||
@@ -145,6 +148,7 @@ pub async fn get_user(
|
||||
alcohol: u.alcohol,
|
||||
diet: diet_name,
|
||||
avatar_url: u.avatar_url,
|
||||
has_mistral_key: u.mistral_api_key.is_some(),
|
||||
created_at: u.created_at.to_string(),
|
||||
}))
|
||||
}
|
||||
@@ -229,6 +233,7 @@ pub async fn create_user(
|
||||
alcohol: inserted.alcohol,
|
||||
diet: diet_name,
|
||||
avatar_url: inserted.avatar_url,
|
||||
has_mistral_key: inserted.mistral_api_key.is_some(),
|
||||
created_at: inserted.created_at.to_string(),
|
||||
}))
|
||||
}
|
||||
@@ -279,6 +284,9 @@ pub async fn update_user(
|
||||
if req.avatar_url.is_some() {
|
||||
active.avatar_url = Set(req.avatar_url);
|
||||
}
|
||||
if req.mistral_api_key.is_some() {
|
||||
active.mistral_api_key = Set(req.mistral_api_key);
|
||||
}
|
||||
active.updated_at = Set(now);
|
||||
|
||||
let updated = active
|
||||
@@ -317,6 +325,7 @@ pub async fn update_user(
|
||||
alcohol: updated.alcohol,
|
||||
diet: diet_name,
|
||||
avatar_url: updated.avatar_url,
|
||||
has_mistral_key: updated.mistral_api_key.is_some(),
|
||||
created_at: updated.created_at.to_string(),
|
||||
}))
|
||||
}
|
||||
|
||||
@@ -44,6 +44,9 @@ pub struct Model {
|
||||
/// URL to profile avatar icon
|
||||
pub avatar_url: Option<String>,
|
||||
|
||||
/// User's own Mistral API key (BYOK - Bring Your Own Key)
|
||||
pub mistral_api_key: Option<String>,
|
||||
|
||||
pub created_at: DateTime,
|
||||
pub updated_at: DateTime,
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user