Skip to content

Commit

Permalink
docs: Enhance InferenceGatewayAPI with health check and detailed docu…
Browse files Browse the repository at this point in the history
…mentation

Signed-off-by: Eden Reich <[email protected]>
  • Loading branch information
edenreich committed Jan 22, 2025
1 parent dc47832 commit c70a9e1
Showing 1 changed file with 43 additions and 0 deletions.
43 changes: 43 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,21 +1,35 @@
//! Inference Gateway SDK for Rust
//!
//! This crate provides a Rust client for the Inference Gateway API, allowing interaction
//! with various LLM providers through a unified interface.
use reqwest::blocking::Client;
use serde::{Deserialize, Serialize};
use std::{error::Error, fmt};

/// Represents a model available through a provider
#[derive(Debug, Serialize, Deserialize)]
pub struct Model {
/// Unique identifier of the model
pub id: String,
/// Type of the object (always "model")
pub object: String,
/// Organization that owns the model
pub owned_by: String,
/// Unix timestamp of when the model was created
pub created: i64,
}

/// Collection of models available from a specific provider
#[derive(Debug, Serialize, Deserialize)]
pub struct ProviderModels {
/// The LLM provider
pub provider: Provider,
/// List of available models
pub models: Vec<Model>,
}

/// Supported LLM providers
#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Provider {
Expand All @@ -40,51 +54,76 @@ impl fmt::Display for Provider {
}
}

/// A message in a conversation with an LLM
#[derive(Debug, Serialize, Deserialize)]
pub struct Message {
/// Role of the message sender ("system", "user" or "assistant")
pub role: String,
/// Content of the message
pub content: String,
}

/// Request payload for generating content
#[derive(Debug, Serialize)]
struct GenerateRequest {
/// Name of the model
model: String,
/// Conversation history and prompt
messages: Vec<Message>,
}

#[derive(Debug, Deserialize)]
pub struct GenerateResponse {
/// Provider that generated the response
pub provider: String,
/// Content of the response
pub response: ResponseContent,
}

#[derive(Debug, Deserialize)]
pub struct ResponseContent {
/// Role of the responder (typically "assistant")
pub role: String,
/// Model that generated the response
pub model: String,
/// Generated content
pub content: String,
}

/// Client for interacting with the Inference Gateway API
pub struct InferenceGatewayClient {
base_url: String,
client: Client,
token: Option<String>,
}

/// Core API interface for the Inference Gateway
pub trait InferenceGatewayAPI {
/// Lists available models from all providers
fn list_models(&self) -> Result<Vec<ProviderModels>, Box<dyn Error>>;

/// Generates content using a specified model
///
/// # Arguments
/// * `provider` - The LLM provider to use
/// * `model` - Name of the model
/// * `messages` - Conversation history and prompt
fn generate_content(
&self,
provider: Provider,
model: &str,
messages: Vec<Message>,
) -> Result<GenerateResponse, Box<dyn Error>>;

/// Checks if the API is available
fn health_check(&self) -> Result<bool, Box<dyn Error>>;
}

impl InferenceGatewayClient {
/// Creates a new client instance
///
/// # Arguments
/// * `base_url` - Base URL of the Inference Gateway API
pub fn new(base_url: &str) -> Self {
Self {
base_url: base_url.to_string(),
Expand All @@ -93,6 +132,10 @@ impl InferenceGatewayClient {
}
}

/// Sets an authentication token for the client
///
/// # Arguments
/// * `token` - JWT token for authentication
pub fn with_token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self
Expand Down

0 comments on commit c70a9e1

Please sign in to comment.