lib.rs 1.22 KB
Newer Older
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
1
//! Text Generation gRPC client library
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
2
3

mod client;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
4
#[allow(clippy::derive_partial_eq_without_eq)]
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
5
6
7
8
mod pb;
mod sharded_client;

pub use client::Client;
Nicolas Patry's avatar
Nicolas Patry committed
9
10
11
pub use pb::generate::v2::HealthResponse;
pub use pb::generate::v2::InfoResponse as ShardInfo;
pub use pb::generate::v2::{
drbh's avatar
drbh committed
12
13
    Batch, CachedBatch, FinishReason, GeneratedText, Generation, GrammarType,
    NextTokenChooserParameters, Request, StoppingCriteriaParameters, Tokens,
14
};
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
15
16
pub use sharded_client::ShardedClient;
use thiserror::Error;
Olivier Dehaene's avatar
v0.1.0  
Olivier Dehaene committed
17
use tonic::transport;
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
18
19
20
use tonic::Status;

#[derive(Error, Debug, Clone)]
Olivier Dehaene's avatar
Olivier Dehaene committed
21
pub enum ClientError {
22
    #[error("Could not connect to Text Generation server: {0}")]
Olivier Dehaene's avatar
Olivier Dehaene committed
23
    Connection(String),
24
    #[error("Server error: {0}")]
Olivier Dehaene's avatar
Olivier Dehaene committed
25
    Generation(String),
26
27
    #[error("Sharded results are empty")]
    EmptyResults,
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
28
29
30
31
}

impl From<Status> for ClientError {
    fn from(err: Status) -> Self {
32
33
34
        let err = Self::Generation(err.message().to_string());
        tracing::error!("{err}");
        err
Olivier Dehaene's avatar
Olivier Dehaene committed
35
36
37
38
39
    }
}

impl From<transport::Error> for ClientError {
    fn from(err: transport::Error) -> Self {
40
41
42
        let err = Self::Connection(err.to_string());
        tracing::error!("{err}");
        err
Olivier Dehaene's avatar
Init  
Olivier Dehaene committed
43
44
45
46
    }
}

pub type Result<T> = std::result::Result<T, ClientError>;