refactor: task handlers

src/args.rs

@@ -140,9 +140,9 @@ impl OtlpProtocol {
 
 #[derive(Debug, Clone, Copy, Default, PartialEq)]
 pub enum LogFormat {
+    Compact,
     #[default]
     Full,
-    Compact,
     Pretty,
     JSON,
 }

src/clients.rs

@@ -32,7 +32,7 @@ use hyper_timeout::TimeoutConnector;
 use hyper_util::rt::TokioExecutor;
 use tonic::{Request, metadata::MetadataMap};
 use tower::{ServiceBuilder, timeout::TimeoutLayer};
-use tracing::{Span, debug, instrument};
+use tracing::Span;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use url::Url;
 
@@ -205,7 +205,6 @@ impl ClientMap {
     }
 }
 
-#[instrument(skip_all, fields(hostname = service_config.hostname))]
 pub async fn create_http_client(
     default_port: u16,
     service_config: &ServiceConfig,
@@ -220,7 +219,6 @@ pub async fn create_http_client(
     base_url
         .set_port(Some(port))
         .unwrap_or_else(|_| panic!("error setting port: {}", port));
-    debug!(%base_url, "creating HTTP client");
 
     let connect_timeout = Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SEC);
     let request_timeout = Duration::from_secs(
@@ -257,7 +255,6 @@ pub async fn create_http_client(
     Ok(HttpClient::new(base_url, client))
 }
 
-#[instrument(skip_all, fields(hostname = service_config.hostname))]
 pub async fn create_grpc_client<C: Debug + Clone>(
     default_port: u16,
     service_config: &ServiceConfig,
@@ -270,7 +267,6 @@ pub async fn create_grpc_client<C: Debug + Clone>(
     };
     let mut base_url = Url::parse(&format!("{}://{}", protocol, &service_config.hostname)).unwrap();
     base_url.set_port(Some(port)).unwrap();
-    debug!(%base_url, "creating gRPC client");
     let connect_timeout = Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SEC);
     let request_timeout = Duration::from_secs(
         service_config

src/clients/chunker.rs

@@ -22,7 +22,7 @@ use axum::http::HeaderMap;
 use futures::{Future, Stream, StreamExt, TryStreamExt};
 use ginepro::LoadBalancedChannel;
 use tonic::{Code, Request, Response, Status, Streaming};
-use tracing::{Span, debug, info, instrument};
+use tracing::{Span, instrument};
 
 use super::{
     BoxStream, Client, Error, create_grpc_client, errors::grpc_to_http_code,
@@ -68,28 +68,24 @@ impl ChunkerClient {
         }
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn tokenization_task_predict(
         &self,
         model_id: &str,
         request: ChunkerTokenizationTaskRequest,
     ) -> Result<TokenizationResults, Error> {
         let mut client = self.client.clone();
         let request = request_with_headers(request, model_id);
-        debug!(?request, "sending client request");
         let response = client.chunker_tokenization_task_predict(request).await?;
         let span = Span::current();
         trace_context_from_grpc_response(&span, &response);
         Ok(response.into_inner())
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn bidi_streaming_tokenization_task_predict(
         &self,
         model_id: &str,
         request_stream: BoxStream<BidiStreamingChunkerTokenizationTaskRequest>,
     ) -> Result<BoxStream<Result<ChunkerTokenizationStreamResult, Error>>, Error> {
-        info!("sending client stream request");
         let mut client = self.client.clone();
         let request = request_with_headers(request_stream, model_id);
         // NOTE: this is an ugly workaround to avoid bogus higher-ranked lifetime errors.

src/clients/detector.rs

@@ -21,7 +21,6 @@ use axum::http::HeaderMap;
 use http::header::CONTENT_TYPE;
 use hyper::StatusCode;
 use serde::Deserialize;
-use tracing::instrument;
 use url::Url;
 
 use super::{
@@ -79,7 +78,6 @@ pub trait DetectorClientExt: HttpClientExt {
 }
 
 impl<C: DetectorClient + HttpClientExt> DetectorClientExt for C {
-    #[instrument(skip_all, fields(model_id, url))]
     async fn post_to_detector<U: ResponseBody>(
         &self,
         model_id: &str,

src/clients/detector/text_chat.rs

@@ -18,7 +18,7 @@
 use async_trait::async_trait;
 use hyper::HeaderMap;
 use serde::Serialize;
-use tracing::{info, instrument};
+use tracing::info;
 
 use super::{DEFAULT_PORT, DetectorClient, DetectorClientExt};
 use crate::{
@@ -63,7 +63,6 @@ impl TextChatDetectorClient {
         &self.client
     }
 
-    #[instrument(skip_all, fields(model_id, ?headers))]
     pub async fn text_chat(
         &self,
         model_id: &str,

src/clients/detector/text_contents.rs

@@ -20,7 +20,7 @@ use std::collections::BTreeMap;
 use async_trait::async_trait;
 use hyper::HeaderMap;
 use serde::{Deserialize, Serialize};
-use tracing::{info, instrument};
+use tracing::info;
 
 use super::{DEFAULT_PORT, DetectorClient, DetectorClientExt};
 use crate::{
@@ -61,7 +61,6 @@ impl TextContentsDetectorClient {
         &self.client
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn text_contents(
         &self,
         model_id: &str,

src/clients/detector/text_context_doc.rs

@@ -18,7 +18,7 @@
 use async_trait::async_trait;
 use hyper::HeaderMap;
 use serde::{Deserialize, Serialize};
-use tracing::{info, instrument};
+use tracing::info;
 
 use super::{DEFAULT_PORT, DetectorClient, DetectorClientExt};
 use crate::{
@@ -59,7 +59,6 @@ impl TextContextDocDetectorClient {
         &self.client
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn text_context_doc(
         &self,
         model_id: &str,

src/clients/detector/text_generation.rs

@@ -18,7 +18,7 @@
 use async_trait::async_trait;
 use hyper::HeaderMap;
 use serde::Serialize;
-use tracing::{info, instrument};
+use tracing::info;
 
 use super::{DEFAULT_PORT, DetectorClient, DetectorClientExt};
 use crate::{
@@ -59,7 +59,6 @@ impl TextGenerationDetectorClient {
         &self.client
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn text_generation(
         &self,
         model_id: &str,

src/clients/generation.rs

@@ -18,7 +18,6 @@
 use async_trait::async_trait;
 use futures::{StreamExt, TryStreamExt};
 use hyper::HeaderMap;
-use tracing::{debug, instrument};
 
 use super::{BoxStream, Client, Error, NlpClient, TgisClient};
 use crate::{
@@ -63,7 +62,6 @@ impl GenerationClient {
         Self(None)
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn tokenize(
         &self,
         model_id: String,
@@ -79,19 +77,15 @@ impl GenerationClient {
                     return_offsets: false,
                     truncate_input_tokens: 0,
                 };
-                debug!(provider = "tgis", ?request, "sending tokenize request");
                 let mut response = client.tokenize(request, headers).await?;
-                debug!(provider = "tgis", ?response, "received tokenize response");
                 let response = response.responses.swap_remove(0);
                 Ok((response.token_count, response.tokens))
             }
             Some(GenerationClientInner::Nlp(client)) => {
                 let request = TokenizationTaskRequest { text };
-                debug!(provider = "nlp", ?request, "sending tokenize request");
                 let response = client
                     .tokenization_task_predict(&model_id, request, headers)
                     .await?;
-                debug!(provider = "nlp", ?response, "received tokenize response");
                 let tokens = response
                     .results
                     .into_iter()
@@ -103,7 +97,6 @@ impl GenerationClient {
         }
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn generate(
         &self,
         model_id: String,
@@ -120,9 +113,7 @@ impl GenerationClient {
                     requests: vec![GenerationRequest { text }],
                     params,
                 };
-                debug!(provider = "tgis", ?request, "sending generate request");
                 let response = client.generate(request, headers).await?;
-                debug!(provider = "tgis", ?response, "received generate response");
                 Ok(response.into())
             }
             Some(GenerationClientInner::Nlp(client)) => {
@@ -157,18 +148,15 @@ impl GenerationClient {
                         ..Default::default()
                     }
                 };
-                debug!(provider = "nlp", ?request, "sending generate request");
                 let response = client
                     .text_generation_task_predict(&model_id, request, headers)
                     .await?;
-                debug!(provider = "nlp", ?response, "received generate response");
                 Ok(response.into())
             }
             None => Err(Error::ModelNotFound { model_id }),
         }
     }
 
-    #[instrument(skip_all, fields(model_id))]
     pub async fn generate_stream(
         &self,
         model_id: String,
@@ -185,11 +173,6 @@ impl GenerationClient {
                     request: Some(GenerationRequest { text }),
                     params,
                 };
-                debug!(
-                    provider = "tgis",
-                    ?request,
-                    "sending generate_stream request"
-                );
                 let response_stream = client
                     .generate_stream(request, headers)
                     .await?
@@ -229,11 +212,6 @@ impl GenerationClient {
                         ..Default::default()
                     }
                 };
-                debug!(
-                    provider = "nlp",
-                    ?request,
-                    "sending generate_stream request"
-                );
                 let response_stream = client
                     .server_streaming_text_generation_task_predict(&model_id, request, headers)
                     .await?

src/clients/http.rs

@@ -20,7 +20,7 @@ use std::{fmt::Debug, ops::Deref, time::Duration};
 use http_body_util::{BodyExt, Full, combinators::BoxBody};
 use hyper::{
     HeaderMap, Method, Request, StatusCode,
-    body::{Body, Bytes, Incoming},
+    body::{Bytes, Incoming},
 };
 use hyper_rustls::HttpsConnector;
 use hyper_timeout::TimeoutConnector;
@@ -36,7 +36,7 @@ use tower_http::{
         Trace, TraceLayer,
     },
 };
-use tracing::{Span, debug, error, info, info_span, instrument};
+use tracing::{Span, error, info, info_span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use url::Url;
 
@@ -137,7 +137,6 @@ impl HttpClient {
         self.base_url.join(path).unwrap()
     }
 
-    #[instrument(skip_all, fields(url))]
     pub async fn get(
         &self,
         url: Url,
@@ -147,7 +146,6 @@ impl HttpClient {
         self.send(url, Method::GET, headers, body).await
     }
 
-    #[instrument(skip_all, fields(url))]
     pub async fn post(
         &self,
         url: Url,
@@ -157,7 +155,6 @@ impl HttpClient {
         self.send(url, Method::POST, headers, body).await
     }
 
-    #[instrument(skip_all, fields(url))]
     pub async fn send(
         &self,
         url: Url,
@@ -172,12 +169,6 @@ impl HttpClient {
             .uri(url.as_uri());
         match builder.headers_mut() {
             Some(headers_mut) => {
-                debug!(
-                    ?url,
-                    ?headers,
-                    ?body,
-                    "sending client request"
-                );
                 headers_mut.extend(headers);
                 let body =
                     Full::new(Bytes::from(serde_json::to_vec(&body).map_err(|e| {
@@ -211,13 +202,6 @@ impl HttpClient {
                             message: format!("client request timeout: {}", e),
                         }),
                 }?;
-
-                debug!(
-                    status = ?response.status(),
-                    headers = ?response.headers(),
-                    size = ?response.size_hint(),
-                    "incoming client response"
-                );
                 let span = Span::current();
                 trace::trace_context_from_http_response(&span, &response);
                 Ok(response.into())

src/clients/openai.rs

@@ -24,7 +24,6 @@ use http_body_util::BodyExt;
 use hyper::{HeaderMap, StatusCode};
 use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
-use tracing::{info, instrument};
 
 use super::{
     Client, Error, HttpClient, create_http_client, detector::ContentAnalysisResponse,
@@ -70,14 +69,12 @@ impl OpenAiClient {
         &self.client
     }
 
-    #[instrument(skip_all, fields(request.model))]
     pub async fn chat_completions(
         &self,
         request: ChatCompletionsRequest,
         headers: HeaderMap,
     ) -> Result<ChatCompletionsResponse, Error> {
         let url = self.inner().endpoint(CHAT_COMPLETIONS_ENDPOINT);
-        info!("sending Open AI chat completion request to {}", url);
         if request.stream {
             let (tx, rx) = mpsc::channel(32);
             let mut event_stream = self
@@ -296,11 +293,11 @@ pub struct ChatCompletionsRequest {
 #[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct DetectorConfig {
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub input: Option<HashMap<String, DetectorParams>>,
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub input: HashMap<String, DetectorParams>,
 
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub output: Option<HashMap<String, DetectorParams>>,
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub output: HashMap<String, DetectorParams>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]