starred/minne
1
0
Fork 0
mirror of https://github.com/perstarkse/minne.git synced 2026-03-18 23:44:18 +01:00
Code Issues Packages Projects Releases 10 Wiki Activity

refactor: file_info, rabbitmq, queue

Browse Source
This commit is contained in:
Per Stark
2024-12-10 16:02:40 +01:00
parent f788b3065a
commit 42e2aad2b0
10 changed files with 113 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/bin/server.rs
View File

@@ -1,6 +1,6 @@
use axum::{
extract::DefaultBodyLimit,
routing::{delete, get, post, put},
routing::{get, post},
Router,
};
use std::sync::Arc;
@@ -11,11 +11,8 @@ use zettle_db::{
rabbitmq::{consumer::RabbitMQConsumer, publisher::RabbitMQProducer, RabbitMQConfig},
server::{
routes::{
file::{delete_file_handler, get_file_handler, update_file_handler, upload_handler},
index::index_handler,
ingress::ingress_handler,
query::query_handler,
queue_length::queue_length_handler,
file::upload_handler, index::index_handler, ingress::ingress_handler,
query::query_handler, queue_length::queue_length_handler,
},
AppState,
},
@@ -41,7 +38,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let app_state = AppState {
rabbitmq_producer: Arc::new(RabbitMQProducer::new(&config).await?),
rabbitmq_consumer: Arc::new(RabbitMQConsumer::new(&config).await?),
rabbitmq_consumer: Arc::new(RabbitMQConsumer::new(&config, false).await?),
surreal_db_client: Arc::new(SurrealDbClient::new().await?),
tera: Arc::new(Tera::new("src/server/templates/**/*.html").unwrap()),
};
@@ -68,9 +65,6 @@ fn api_routes_v1() -> Router<AppState> {
// File routes
.route("/file", post(upload_handler))
.layer(DefaultBodyLimit::max(1024 * 1024 * 1024))
.route("/file/:uuid", get(get_file_handler))
.route("/file/:uuid", put(update_file_handler))
.route("/file/:uuid", delete(delete_file_handler))
// Query routes
.route("/query", post(query_handler))
}

9
src/bin/worker.rs
View File

@@ -1,6 +1,6 @@
use tracing::info;
use tracing_subscriber::{fmt, prelude::*, EnvFilter};
use zettle_db::rabbitmq::{consumer::RabbitMQConsumer, RabbitMQConfig };
use zettle_db::rabbitmq::{consumer::RabbitMQConsumer, RabbitMQConfig};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
@@ -12,7 +12,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
.ok();
info!("Starting RabbitMQ consumer");
// Set up RabbitMQ config
let config = RabbitMQConfig {
amqp_addr: "amqp://localhost".to_string(),
@@ -21,11 +21,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
routing_key: "my_key".to_string(),
};
// Create a RabbitMQ consumer
let consumer = RabbitMQConsumer::new(&config).await?;
info!("Consumer connected to RabbitMQ");
let consumer = RabbitMQConsumer::new(&config, true).await?;
// Start consuming messages
consumer.process_messages().await?;

8
src/error.rs
View File

@@ -4,7 +4,10 @@ use serde_json::json;
use thiserror::Error;
use tokio::task::JoinError;
use crate::{ingress::types::ingress_input::IngressContentError, rabbitmq::RabbitMQError};
use crate::{
ingress::types::ingress_input::IngressContentError, rabbitmq::RabbitMQError,
storage::types::file_info::FileError,
};
#[derive(Error, Debug)]
pub enum ProcessingError {
@@ -55,6 +58,8 @@ pub enum ApiError {
RabbitMQError(#[from] RabbitMQError),
#[error("LLM processing error: {0}")]
OpenAIerror(#[from] OpenAIError),
#[error("File error: {0}")]
FileError(#[from] FileError),
}
impl IntoResponse for ApiError {
@@ -69,6 +74,7 @@ impl IntoResponse for ApiError {
(StatusCode::INTERNAL_SERVER_ERROR, self.to_string())
}
ApiError::RabbitMQError(_) => (StatusCode::INTERNAL_SERVER_ERROR, self.to_string()),
ApiError::FileError(_) => (StatusCode::INTERNAL_SERVER_ERROR, self.to_string()),
};
(

8
src/ingress/types/ingress_input.rs
View File

@@ -83,16 +83,16 @@ pub async fn create_ingress_objects(
}
// Look up FileInfo objects using the db and the submitted uuids in input.files
if let Some(file_uuids) = input.files {
for uuid in file_uuids {
if let Some(file_info) = get_item::<FileInfo>(&db_client, &uuid).await? {
if let Some(file_ids) = input.files {
for id in file_ids {
if let Some(file_info) = get_item::<FileInfo>(db_client, &id).await? {
object_list.push(IngressObject::File {
file_info,
instructions: input.instructions.clone(),
category: input.category.clone(),
});
} else {
info!("No file with UUID: {}", uuid);
info!("No file with id: {}", id);
}
}
}

45
src/rabbitmq/consumer.rs
View File

@@ -13,7 +13,7 @@ use tracing::{error, info};
pub struct RabbitMQConsumer {
common: RabbitMQCommon,
pub queue: Queue,
consumer: Consumer,
consumer: Option<Consumer>,
}
impl RabbitMQConsumer {
@@ -22,10 +22,14 @@ impl RabbitMQConsumer {
///
/// # Arguments
/// * `config` - A initialized RabbitMQConfig containing required configurations
/// * `start_consuming` - Set to true to start consuming messages
///
/// # Returns
/// * `Result<Self, RabbitMQError>` - The created client or an error.
pub async fn new(config: &RabbitMQConfig) -> Result<Self, RabbitMQError> {
pub async fn new(
config: &RabbitMQConfig,
start_consuming: bool,
) -> Result<Self, RabbitMQError> {
let common = RabbitMQCommon::new(config).await?;
// Passively declare the exchange (it should already exist)
@@ -36,7 +40,11 @@ impl RabbitMQConsumer {
Self::bind_queue(&common.channel, &config.exchange, &queue, config).await?;
// Initialize the consumer
let consumer = Self::initialize_consumer(&common.channel, config).await?;
let consumer = if start_consuming {
Some(Self::initialize_consumer(&common.channel, config).await?)
} else {
None
};
Ok(Self {
common,
@@ -67,6 +75,27 @@ impl RabbitMQConsumer {
.await
.map_err(|e| RabbitMQError::InitializeConsumerError(e.to_string()))
}
/// Operation to get the current queue length
/// Will redeclare queue to get a updated number
pub async fn get_queue_length(&self) -> Result<u32, RabbitMQError> {
let queue_info = self
.common
.channel
.queue_declare(
&self.queue.name().to_string(),
QueueDeclareOptions {
durable: true,
..QueueDeclareOptions::default()
},
FieldTable::default(),
)
.await
.map_err(|e| RabbitMQError::QueueError(e.to_string()))?;
Ok(queue_info.message_count())
}
/// Declares the queue based on the channel and `RabbitMQConfig`.
/// # Arguments
/// * `channel` - Lapin Channel.
@@ -127,9 +156,15 @@ impl RabbitMQConsumer {
/// `IngressObject` - The object containing content and metadata.
/// `Delivery` - A delivery reciept, required to ack or nack the delivery.
pub async fn consume(&self) -> Result<(IngressObject, Delivery), RabbitMQError> {
// Get consumer or return error if not initialized
let consumer: &lapin::Consumer = self.consumer.as_ref().ok_or_else(|| {
RabbitMQError::ConsumeError(
"Consumer not initialized. Call new() with start_consuming=true".to_string(),
)
})?;
// Receive the next message
let delivery = self
.consumer
let delivery = consumer
.clone()
.next()
.await

97
src/server/routes/file.rs
View File

@@ -1,21 +1,13 @@
use crate::{
server::AppState,
storage::types::file_info::{FileError, FileInfo},
};
use axum::{
extract::{Path, State},
response::IntoResponse,
Json,
};
use crate::{error::ApiError, server::AppState, storage::types::file_info::FileInfo};
use axum::{extract::State, response::IntoResponse, Json};
use axum_typed_multipart::{FieldData, TryFromMultipart, TypedMultipart};
use serde_json::json;
use tempfile::NamedTempFile;
use tracing::info;
use uuid::Uuid;
#[derive(Debug, TryFromMultipart)]
pub struct FileUploadRequest {
#[form_data(limit = "100000")] // Example limit: ~100 KB
#[form_data(limit = "1000000")] // Example limit: ~1000 KB
pub file: FieldData<NamedTempFile>,
}
@@ -25,7 +17,7 @@ pub struct FileUploadRequest {
pub async fn upload_handler(
State(state): State<AppState>,
TypedMultipart(input): TypedMultipart<FileUploadRequest>,
) -> Result<impl IntoResponse, FileError> {
) -> Result<impl IntoResponse, ApiError> {
info!("Received an upload request");
// Process the file upload
@@ -33,7 +25,7 @@ pub async fn upload_handler(
// Prepare the response JSON
let response = json!({
"uuid": file_info.uuid,
"id": file_info.id,
"sha256": file_info.sha256,
"path": file_info.path,
"mime_type": file_info.mime_type,
@@ -44,82 +36,3 @@ pub async fn upload_handler(
// Return the response with HTTP 200
Ok((axum::http::StatusCode::OK, Json(response)))
}
/// Handler to retrieve file information by UUID.
///
/// Route: GET /file/:uuid
pub async fn get_file_handler(
State(state): State<AppState>,
Path(uuid_str): Path<String>,
) -> Result<impl IntoResponse, FileError> {
// Parse UUID
let uuid = Uuid::parse_str(&uuid_str).map_err(|_| FileError::InvalidUuid(uuid_str.clone()))?;
// Retrieve FileInfo
let file_info = FileInfo::get_by_uuid(uuid, &state.surreal_db_client).await?;
// Prepare the response JSON
let response = json!({
"uuid": file_info.uuid,
"sha256": file_info.sha256,
"path": file_info.path,
"mime_type": file_info.mime_type,
});
info!("Retrieved FileInfo: {:?}", file_info);
// Return the response with HTTP 200
Ok((axum::http::StatusCode::OK, Json(response)))
}
/// Handler to update an existing file by UUID.
///
/// Route: PUT /file/:uuid
pub async fn update_file_handler(
State(state): State<AppState>,
Path(uuid_str): Path<String>,
TypedMultipart(input): TypedMultipart<FileUploadRequest>,
) -> Result<impl IntoResponse, FileError> {
// Parse UUID
let uuid = Uuid::parse_str(&uuid_str).map_err(|_| FileError::InvalidUuid(uuid_str.clone()))?;
// Update the file
let updated_file_info = FileInfo::update(uuid, input.file, &state.surreal_db_client).await?;
// Prepare the response JSON
let response = json!({
"uuid": updated_file_info.uuid,
"sha256": updated_file_info.sha256,
"path": updated_file_info.path,
"mime_type": updated_file_info.mime_type,
});
info!("File updated successfully: {:?}", updated_file_info);
// Return the response with HTTP 200
Ok((axum::http::StatusCode::OK, Json(response)))
}
/// Handler to delete a file by UUID.
///
/// Route: DELETE /file/:uuid
pub async fn delete_file_handler(
State(state): State<AppState>,
Path(uuid_str): Path<String>,
) -> Result<impl IntoResponse, FileError> {
// Parse UUID
let uuid = Uuid::parse_str(&uuid_str).map_err(|_| FileError::InvalidUuid(uuid_str.clone()))?;
// Delete the file
FileInfo::delete(uuid, &state.surreal_db_client).await?;
info!("Deleted file with UUID: {}", uuid);
// Prepare the response JSON
let response = json!({
"message": "File deleted successfully",
});
// Return the response with HTTP 204 No Content
Ok((axum::http::StatusCode::NO_CONTENT, Json(response)))
}

1
src/server/routes/index.rs
View File

@@ -8,7 +8,6 @@ use crate::{error::ApiError, server::AppState};
pub async fn index_handler(State(state): State<AppState>) -> Result<Html<String>, ApiError> {
info!("Displaying index page");
// Now you can access the consumer directly from the state
let queue_length = state.rabbitmq_consumer.queue.message_count();
let output = state

45
src/server/routes/queue_length.rs
View File

@@ -1,42 +1,17 @@
use axum::{
http::StatusCode,
response::{IntoResponse, Response},
};
use tracing::{error, info};
use axum::{extract::State, http::StatusCode, response::IntoResponse};
use tracing::info;
use crate::rabbitmq::{consumer::RabbitMQConsumer, RabbitMQConfig};
use crate::{error::ApiError, server::AppState};
pub async fn queue_length_handler() -> Response {
pub async fn queue_length_handler(
State(state): State<AppState>,
) -> Result<impl IntoResponse, ApiError> {
info!("Getting queue length");
// Set up RabbitMQ config
let config = RabbitMQConfig {
amqp_addr: "amqp://localhost".to_string(),
exchange: "my_exchange".to_string(),
queue: "my_queue".to_string(),
routing_key: "my_key".to_string(),
};
let queue_length = state.rabbitmq_consumer.get_queue_length().await?;
// Create a new consumer
match RabbitMQConsumer::new(&config).await {
Ok(consumer) => {
info!("Consumer connected to RabbitMQ");
info!("Queue length: {}", queue_length);
// Get the queue length
let queue_length = consumer.queue.message_count();
info!("Queue length: {}", queue_length);
// Return the queue length with a 200 OK status
(StatusCode::OK, queue_length.to_string()).into_response()
}
Err(e) => {
error!("Failed to create consumer: {:?}", e);
(
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to connect to RabbitMQ".to_string(),
)
.into_response()
}
}
// Return the queue length with a 200 OK status
Ok((StatusCode::OK, queue_length.to_string()))
}

2
src/storage/db.rs
View File

@@ -111,5 +111,5 @@ pub async fn get_item<T>(db_client: &Surreal<Client>, id: &str) -> Result<Option
where
T: for<'de> StoredObject,
{
Ok(db_client.select((T::table_name(), id)).await?)
db_client.select((T::table_name(), id)).await
}

45
src/storage/types/mod.rs
View File

@@ -14,24 +14,57 @@ pub trait StoredObject: Serialize + for<'de> Deserialize<'de> {
#[macro_export]
macro_rules! stored_object {
($name:ident, $table:expr, {$($(#[$attr:meta])* $field:ident: $ty:ty),*}) => {
($name:ident, $table:expr, {$($(#[$attr:meta])* $field:ident: $ty:ty),*}) => {
use axum::async_trait;
use serde::{Deserialize, Deserializer, Serialize};
use surrealdb::sql::Thing;
use $crate::storage::types::StoredObject;
use serde::de::{self, Visitor};
use std::fmt;
fn thing_to_string<'de, D>(deserializer: D) -> Result<String, D::Error>
struct FlexibleIdVisitor;
impl<'de> Visitor<'de> for FlexibleIdVisitor {
type Value = String;
fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
formatter.write_str("a string or a Thing")
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(value.to_string())
}
fn visit_string<E>(self, value: String) -> Result<Self::Value, E>
where
E: de::Error,
{
Ok(value)
}
fn visit_map<A>(self, map: A) -> Result<Self::Value, A::Error>
where
A: de::MapAccess<'de>,
{
// Try to deserialize as Thing
let thing = Thing::deserialize(de::value::MapAccessDeserializer::new(map))?;
Ok(thing.id.to_raw())
}
}
fn deserialize_flexible_id<'de, D>(deserializer: D) -> Result<String, D::Error>
where
D: Deserializer<'de>,
{
let thing = Thing::deserialize(deserializer)?;
Ok(thing.id.to_raw())
deserializer.deserialize_any(FlexibleIdVisitor)
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct $name {
#[serde(deserialize_with = "thing_to_string")]
#[serde(deserialize_with = "deserialize_flexible_id")]
pub id: String,
$(pub $field: $ty),*
}