From a2bb86f5a91d4ef813f2e02302c9eb3dc4991b44 Mon Sep 17 00:00:00 2001
From: Aisuko <urakiny@gmail.com>
Date: Fri, 20 Oct 2023 11:25:37 +1100
Subject: [PATCH] Remove services.rs

Signed-off-by: Aisuko <urakiny@gmail.com>
---
 backend/rust/bunker/src/lib.rs     |  7 ++++-
 backend/rust/bunker/src/service.rs | 43 ------------------------------
 backend/rust/burn/src/main.rs      |  8 ++++--
 backend/rust/codegen/src/lib.rs    |  3 ++-
 4 files changed, 14 insertions(+), 47 deletions(-)
 delete mode 100644 backend/rust/bunker/src/service.rs

diff --git a/backend/rust/bunker/src/lib.rs b/backend/rust/bunker/src/lib.rs
index b1a3c3508b1..9e0f781e044 100644
--- a/backend/rust/bunker/src/lib.rs
+++ b/backend/rust/bunker/src/lib.rs
@@ -3,13 +3,17 @@ pub mod pb {
     include!("../generated/backend.rs");
 }
 
+use std::net::SocketAddr;
 use tonic::transport::Server;
 
 pub use crate::pb::backend_server::Backend as BackendService;
 use crate::pb::backend_server::BackendServer;
 
 // Run the backend with the default behavior
-pub async fn run(backend: impl BackendService, addr: impl Into<SocketAddr>) -> anyhow::Result<()> {
+pub async fn run(
+    backend: impl BackendService,
+    addr: impl Into<SocketAddr>,
+) -> Result<(), Box<dyn std::error::Error>> {
     let svc = BackendServer::new(backend);
 
     let r = Server::builder()
@@ -18,3 +22,4 @@ pub async fn run(backend: impl BackendService, addr: impl Into<SocketAddr>) -> a
         .await?;
 
     Ok(r)
+}
diff --git a/backend/rust/bunker/src/service.rs b/backend/rust/bunker/src/service.rs
deleted file mode 100644
index 81cf3132cea..00000000000
--- a/backend/rust/bunker/src/service.rs
+++ /dev/null
@@ -1,43 +0,0 @@
-//! Contains the service trait for the bunker service.
-
-use crate::pb::Result as PbResult;
-use crate::pb::{
-    EmbeddingResult, GenerateImageRequest, HealthMessage, ModelOptions, PredictOptions, Reply,
-    StatusResponse, TokenizationResponse, TranscriptRequest, TranscriptResult, TtsRequest,
-};
-use async_trait::async_trait;
-use tokio_stream::wrappers::ReceiverStream;
-use tonic::{Request, Response, Status};
-
-#[async_trait]
-pub trait BackendService<T = ReceiverStream<Result<Reply, Status>>> {
-    async fn health(&self, request: Request<HealthMessage>) -> Result<Response<Reply>, Status>;
-    async fn predict(&self, request: Request<PredictOptions>) -> Result<Response<Reply>, Status>;
-    async fn load_model(
-        &self,
-        request: Request<ModelOptions>,
-    ) -> Result<Response<PbResult>, Status>;
-    async fn predict_stream(&self, request: Request<PredictOptions>)
-        -> Result<Response<T>, Status>; // https://github.com/rust-lang/rust/issues/29661
-    async fn embedding(
-        &self,
-        request: Request<PredictOptions>,
-    ) -> Result<Response<EmbeddingResult>, Status>;
-    async fn generate_image(
-        &self,
-        request: Request<GenerateImageRequest>,
-    ) -> Result<Response<PbResult>, Status>;
-    async fn audio_transcription(
-        &self,
-        request: Request<TranscriptRequest>,
-    ) -> Result<Response<TranscriptResult>, Status>;
-    async fn tts(&self, request: Request<TtsRequest>) -> Result<Response<PbResult>, Status>;
-    async fn tokenize_string(
-        &self,
-        request: Request<PredictOptions>,
-    ) -> Result<Response<TokenizationResponse>, Status>;
-    async fn status(
-        &self,
-        request: Request<HealthMessage>,
-    ) -> Result<Response<StatusResponse>, Status>;
-}
diff --git a/backend/rust/burn/src/main.rs b/backend/rust/burn/src/main.rs
index 11035a14c58..409363eb568 100644
--- a/backend/rust/burn/src/main.rs
+++ b/backend/rust/burn/src/main.rs
@@ -3,7 +3,8 @@ use bunker::pb::{
     EmbeddingResult, GenerateImageRequest, HealthMessage, ModelOptions, PredictOptions, Reply,
     StatusResponse, TokenizationResponse, TranscriptRequest, TranscriptResult, TtsRequest,
 };
-use bunker::service::BackendService;
+
+use bunker::BackendService;
 use tokio_stream::wrappers::ReceiverStream;
 use tonic::{Request, Response, Status};
 
@@ -14,7 +15,9 @@ use async_trait::async_trait;
 struct BurnBackend;
 
 #[async_trait]
-impl BackendService<ReceiverStream<Result<Reply, Status>>> for BurnBackend {
+impl BackendService for BurnBackend {
+    type PredictStreamStream = ReceiverStream<Result<Reply, Status>>;
+
     async fn health(&self, request: Request<HealthMessage>) -> Result<Response<Reply>, Status> {
         // return a Result<Response<Reply>,Status>
         let reply = Reply {
@@ -84,5 +87,6 @@ impl BackendService<ReceiverStream<Result<Reply, Status>>> for BurnBackend {
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // call bunker::run with BurnBackend
     todo!()
 }
diff --git a/backend/rust/codegen/src/lib.rs b/backend/rust/codegen/src/lib.rs
index 8b137891791..208bc595e4c 100644
--- a/backend/rust/codegen/src/lib.rs
+++ b/backend/rust/codegen/src/lib.rs
@@ -1 +1,2 @@
-
+//! Here is for the more complex situation of code generation. For example, defind the constant for 
+//! the different build target.
\ No newline at end of file