Optimize Dependencies, Enhance Performance, and Improve Code Quality (#18)

* fix: handle optional stream value in chat request logging

* refactor: update application state and client initialization to accept direct configuration parameters

* refactor(chat): optimize security client usage by removing unnecessary cloning

Replace multiple clones of SecurityClient with direct mutation using with_user_ip method. This improves performance by avoiding redundant cloning while maintaining the same security checks functionality.

* perf(chat): optimize JSON parsing and serialization

This commit improves performance in the chat handler by optimizing JSON operations:

1. Eliminate double JSON parsing:
   - Parse response body once into serde_json::Value
   - Reuse the same parsed value for metrics logging and ChatResponse
   - Reduces redundant parsing operations per request

2. Optimize response serialization:
   - Replace serde_json::to_vec with serde_json::to_writer
   - Write directly to Vec<u8> buffer to minimize allocations
   - Streamline response construction flow

The changes reduce memory allocations and CPU usage, particularly under high load
with many concurrent requests. No behavioral changes, purely performance focused.

Impact:
- Reduces memory allocations in the response path
- Eliminates redundant JSON parsing
- More efficient serialization for masked content responses
- Cleaner error handling flow

* fix(logging): include version in server startup log message

* Bump version to 0.14.0 and update dependencies

- Updated axum from 0.7.4 to 0.8.4
- Updated tower-http from 0.5.1 to 0.6.6
- Updated reqwest from 0.11.24 to 0.12.23
- Updated serde from 1.0.197 to 1.0.219
- Updated serde_yaml from 0.9.31 to 0.9.34
- Updated uuid from 1.7.0 to 1.18.0
- Updated bytes from 1.5.0 to 1.10.1
- Updated http-body-util from 0.1.0 to 0.1.3
- Updated tracing-subscriber from 0.3.18 to 0.3.18 (no change)
- Updated futures-util from 0.3.30 to 0.3 (no change)
- Updated pin-project from 1.1.3 to 1.1.3 (no change)

* Optimize Rust application's library usage and improve code quality (#17)

* Initial plan

* Fix clippy warnings and optimize dependencies for better performance

Co-authored-by: lenoxys <3996456+lenoxys@users.noreply.github.com>

* Further optimize dependencies and improve memory efficiency patterns

Co-authored-by: lenoxys <3996456+lenoxys@users.noreply.github.com>

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: lenoxys <3996456+lenoxys@users.noreply.github.com>

* refactor(chat): simplify state handling and clone security client for IP configuration

* fix(chat): optimize response serialization and streamline streaming request handling

---------

Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>

Author: lenoxys

Cargo.lock

@@ -1,24 +1,23 @@
 [package]
 name = "panw-api-ollama"
-version = "0.13.0"
+version = "0.14.0"
 edition = "2021"
 
 [dependencies]
-axum = "0.7.4"
-tokio = { version = "1.36.0", features = ["full"] }
-tower-http = { version = "0.5.1", features = ["trace"] }
-reqwest = { version = "0.11.24", features = ["json", "stream"] }
-serde = { version = "1.0.197", features = ["derive"] }
+axum = "0.8.4"
+tokio = { version = "1.36.0", features = ["rt-multi-thread", "net", "signal", "macros", "time"] }
+tower-http = { version = "0.6.6", features = ["trace"] }
+reqwest = { version = "0.12.23", features = ["json", "stream"] }
+serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.114"
-serde_yaml = "0.9.31"
-uuid = { version = "1.7.0", features = ["serde", "v4"] }
+serde_yml = "0.0.12"
+uuid = { version = "1.18.0", features = ["v4", "serde"] }
 futures-util = "0.3.30"
 tracing = "0.1.40"
-tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+tracing-subscriber = { version = "0.3.18", features = ["env-filter", "fmt"] }
 thiserror = "1.0.57"
-bytes = "1.5.0"
+bytes = "1.10.1"
 async-stream = "0.3.5"
-http-body-util = "0.1.0"
-chrono = { version = "0.4", features = ["serde"] }
-futures = "0.3"
+http-body-util = "0.1.3"
+chrono = { version = "0.4", features = ["serde", "clock"], default-features = false }
 pin-project = "1.1.3"

Cargo.toml

@@ -24,14 +24,15 @@ use tracing::{debug, info};
 /// configuration, including file access errors, YAML parsing issues,
 /// and validation of configuration values.
 #[derive(Debug, Error)]
+#[allow(clippy::enum_variant_names)]
 pub enum ConfigError {
     /// File I/O errors when reading the configuration file
     #[error("Failed to read config file: {0}")]
     IoError(#[from] std::io::Error),
 
     /// YAML parsing errors in the configuration file
     #[error("Failed to parse config file: {0}")]
-    ParseError(#[from] serde_yaml::Error),
+    ParseError(#[from] serde_yml::Error),
 
     /// Configuration validation errors
     #[error("Validation error: {0}")]
@@ -177,7 +178,7 @@ pub fn load_config(path: &str) -> Result<Config, ConfigError> {
         debug!("Successfully read configuration file");
 
         // Parse YAML
-        let mut config: Config = serde_yaml::from_str(&content)?;
+        let mut config: Config = serde_yml::from_str(&content)?;
         debug!("Successfully parsed YAML configuration");
 
         // Override with environment variables if present

src/config.rs

@@ -65,12 +65,12 @@ pub async fn handle_chat(
     info!("Received chat request for model: {}", request.model);
     debug!(
         "Chat request details: stream={}, messages={}, client_ip={}",
-        request.stream.unwrap(),
+        request.stream.unwrap_or(false),
         request.messages.len(),
         addr.ip()
     );
 
-    // Configure security client with user's IP
+    // Clone security client and configure with user's IP
     let mut security_client = state.security_client.clone();
     security_client.with_user_ip(addr.ip().to_string());
 
@@ -178,18 +178,22 @@ async fn handle_non_streaming_chat(
         ApiError::InternalError("Failed to read response body".to_string())
     })?;
 
-    // Parse response
-    let mut response_body: ChatResponse = serde_json::from_slice(&body_bytes).map_err(|e| {
+    // Parse response once into Value
+    let json_value: serde_json::Value = serde_json::from_slice(&body_bytes).map_err(|e| {
         error!("Failed to parse response: {}", e);
         ApiError::InternalError("Failed to parse response".to_string())
     })?;
 
     debug!("Received response from Ollama, performing security assessment");
 
-    // Extract and log performance metrics if available
-    if let Ok(json) = serde_json::from_slice::<serde_json::Value>(&body_bytes) {
-        log_llm_metrics(&json, false);
-    }
+    // Extract and log performance metrics
+    log_llm_metrics(&json_value, false);
+
+    // Convert to ChatResponse
+    let mut response_body: ChatResponse = serde_json::from_value(json_value).map_err(|e| {
+        error!("Failed to convert response: {}", e);
+        ApiError::InternalError("Failed to convert response".to_string())
+    })?;
 
     // Security assessment on response content
     let assessment = state
@@ -204,20 +208,20 @@ async fn handle_non_streaming_chat(
     }
 
     // If we have masked content, use it
-    let output_bytes = if assessment.is_masked {
+    let response = if assessment.is_masked {
         response_body.message.content = assessment.final_content;
         info!("Chat response passed security checks (with masked content), returning to client");
-        serde_json::to_vec(&response_body)
-            .map(Bytes::from)
-            .map_err(|e| {
-                error!("Failed to serialize modified response: {}", e);
-                ApiError::InternalError("Failed to serialize response".to_string())
-            })?
+
+        let json_bytes = serde_json::to_vec(&response_body).map_err(|e| {
+            error!("Failed to serialize modified response: {}", e);
+            ApiError::InternalError("Failed to serialize response".to_string())
+        })?;
+        build_json_response(Bytes::from(json_bytes))?
     } else {
         info!("Chat response passed security checks, returning to client");
-        body_bytes
+        build_json_response(body_bytes)?
     };
-    Ok(build_json_response(output_bytes)?)
+    Ok(response)
 }
 
 // Handles streaming chat requests using the generic streaming handler.
@@ -242,12 +246,5 @@ async fn handle_streaming_chat(
 
     let model = request.model.clone();
     // For streaming chat, we're dealing with responses from the LLM, so is_prompt should be false
-    handle_streaming_request::<ChatRequest, ChatResponse>(
-        &state,
-        request,
-        "/api/chat",
-        &model,
-        false,
-    )
-    .await
+    handle_streaming_request::<ChatRequest>(&state, request, "/api/chat", &model, false).await
 }

src/handlers/chat.rs

@@ -39,5 +39,5 @@ pub async fn handle_embeddings(
         .bytes()
         .await
         .map_err(|e| ApiError::InternalError(e.to_string()))?;
-    Ok(build_json_response(body_bytes)?)
+    build_json_response(body_bytes)
 }

src/handlers/embeddings.rs

@@ -149,7 +149,7 @@ async fn handle_non_streaming_generate(
     }
 
     // Return safe response
-    Ok(build_json_response(body_bytes)?)
+    build_json_response(body_bytes)
 }
 
 // Handles streaming generate requests.
@@ -171,7 +171,7 @@ async fn handle_streaming_generate(
 
     let model = request.model.clone();
     // For streaming generate, we're dealing with responses from the LLM, so is_prompt should be false
-    handle_streaming_request::<GenerateRequest, GenerateResponse>(
+    handle_streaming_request::<GenerateRequest>(
         &state,
         request,
         "/api/generate",

src/handlers/generate.rs

@@ -20,6 +20,7 @@ pub mod version;
 // security assessment, and internal server issues into a unified error type
 // that can be converted into appropriate HTTP responses.
 #[derive(Debug, thiserror::Error)]
+#[allow(clippy::enum_variant_names)]
 pub enum ApiError {
     // Errors from the Ollama backend service.
     //

src/handlers/mod.rs

@@ -103,7 +103,7 @@ async fn forward_to_ollama<T: Serialize>(
         .bytes()
         .await
         .map_err(|e| ApiError::InternalError(e.to_string()))?;
-    Ok(build_json_response(body_bytes)?)
+    build_json_response(body_bytes)
 }
 
 // Handler for listing models (GET /api/tags)

src/handlers/models.rs

@@ -4,7 +4,7 @@ use axum::{body::Body, response::Response};
 use bytes::Bytes;
 use futures_util::stream::StreamExt;
 use http_body_util::StreamBody;
-use serde::{de::DeserializeOwned, Serialize};
+use serde::Serialize;
 use tracing::{error, info};
 
 // Builds an HTTP response with JSON content type from the provided bytes.
@@ -21,7 +21,7 @@ fn convert_stream_error(err: reqwest::Error) -> reqwest::Error {
 }
 
 // Handles streaming requests to API endpoints, applying security assessment to the streamed responses.
-pub async fn handle_streaming_request<T, R>(
+pub async fn handle_streaming_request<T>(
     state: &AppState,
     request: T,
     endpoint: &str,
@@ -30,7 +30,6 @@ pub async fn handle_streaming_request<T, R>(
 ) -> Result<Response<Body>, ApiError>
 where
     T: Serialize + Send + 'static,
-    R: DeserializeOwned + Serialize + Send + Sync + Unpin + 'static,
 {
     // Get the original stream from ollama client
     let stream = state.ollama_client.stream(endpoint, &request).await?;
@@ -55,16 +54,14 @@ where
         Err(e) => {
             error!("Error in security assessment stream: {:?}", e);
             // Convert error to a user-friendly message
-            let error_message = match e {
-                _ => "Error processing response",
-            };
+            const ERROR_MESSAGE: &str = "Error processing response";
             let error_json = serde_json::json!({
                 "model": model_string,
-                "error": error_message,
+                "error": ERROR_MESSAGE,
                 "done": true
             });
             let error_bytes = serde_json::to_vec(&error_json)
-                .unwrap_or_else(|_| error_message.as_bytes().to_vec());
+                .unwrap_or_else(|_| ERROR_MESSAGE.as_bytes().to_vec());
             Ok(Bytes::from(error_bytes))
         }
     });

src/handlers/utils.rs

@@ -13,5 +13,5 @@ pub async fn handle_version(State(state): State<AppState>) -> Result<Response, A
         .await
         .map_err(|e| ApiError::InternalError(e.to_string()))?;
 
-    Ok(build_json_response(body_bytes)?)
+    build_json_response(body_bytes)
 }