mod private
{
use crate::
{
client::Client,
error::Result,
validation::{ validate_input_text, validate_model_identifier },
};
#[ cfg( feature = "env-config" ) ]
use crate::environment::{ HuggingFaceEnvironment, EnvironmentInterface };
pub use crate::components::inference_shared::
{
ChatMessage,
ChatCompletionRequest,
ChatCompletionResponse,
ChatChoice,
ChatUsage as Usage,
ToolDefinition,
};
use crate::components::tools::Tool;
use serde::{ Serialize, Deserialize };
#[ derive( Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize ) ]
#[ serde( rename_all = "lowercase" ) ]
pub enum InferenceProvider
{
OpenAI,
Cohere,
Together,
Groq,
HfInference,
}
impl InferenceProvider
{
#[ inline ]
#[ must_use ]
pub fn as_str( self ) -> &'static str
{
match self
{
Self::OpenAI => "openai",
Self::Cohere => "cohere",
Self::Together => "together",
Self::Groq => "groq",
Self::HfInference => "hf-inference",
}
}
#[ inline ]
#[ must_use ]
pub fn available_models( self ) -> &'static [ &'static str ]
{
match self
{
Self::OpenAI => &[
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Llama-2-13b-chat-hf",
"meta-llama/Meta-Llama-3-8B-Instruct",
"meta-llama/Meta-Llama-3-70B-Instruct",
"mistralai/Mistral-7B-Instruct-v0.2",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"codellama/CodeLlama-7b-Instruct-hf",
"codellama/CodeLlama-13b-Instruct-hf",
],
Self::Cohere => &[
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Meta-Llama-3-8B-Instruct",
"mistralai/Mistral-7B-Instruct-v0.2",
],
Self::Together => &[
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Llama-2-13b-chat-hf",
"meta-llama/Meta-Llama-3-8B-Instruct",
"meta-llama/Meta-Llama-3-70B-Instruct",
"mistralai/Mistral-7B-Instruct-v0.2",
"mistralai/Mixtral-8x7B-Instruct-v0.1",
"codellama/CodeLlama-7b-Instruct-hf",
],
Self::Groq | Self::HfInference => &[
"meta-llama/Llama-2-7b-chat-hf",
"meta-llama/Meta-Llama-3-8B-Instruct",
"mistralai/Mistral-7B-Instruct-v0.2",
"codellama/CodeLlama-7b-Instruct-hf",
],
}
}
#[ inline ]
#[ must_use ]
pub fn for_model( model : &str ) -> Option< Self >
{
[ Self::OpenAI, Self::Together, Self::Cohere, Self::Groq, Self::HfInference ]
.into_iter()
.find( | provider | provider.available_models().contains( &model ) )
}
}
#[ derive( Debug, Default, Clone ) ]
pub struct ChatCompletionOptions
{
pub tool_choice : Option< String >,
pub max_tokens : Option< u32 >,
pub temperature : Option< f32 >,
pub top_p : Option< f32 >,
}
#[ derive( Debug ) ]
pub struct Providers< E >
where
E : Clone,
{
client : Client< E >,
}
#[ cfg( feature = "env-config" ) ]
impl< E > Providers< E >
where
E : HuggingFaceEnvironment + EnvironmentInterface + Send + Sync + 'static + Clone,
{
#[ inline ]
#[ must_use ]
pub fn new( client : &Client< E > ) -> Self
{
Self
{
client : (*client).clone(),
}
}
#[ inline ]
pub async fn chat_completion(
&self,
model : impl AsRef< str >,
messages : Vec< ChatMessage >,
max_tokens : Option< u32 >,
temperature : Option< f32 >,
top_p : Option< f32 >,
) -> Result< ChatCompletionResponse >
{
let model_id = model.as_ref();
validate_model_identifier( model_id )?;
for message in &messages
{
validate_input_text( &message.content )?;
}
let request = ChatCompletionRequest
{
model : model_id.to_string(),
messages,
max_tokens,
temperature,
top_p,
stream : Some( false ), tools : None,
tool_choice : None,
};
let url = self.client.environment.endpoint_url( "chat/completions" )?;
self.client.post( url.as_str(), &request ).await
}
#[ inline ]
pub async fn simple_chat(
&self,
model : impl AsRef< str >,
user_message : impl Into< String >,
) -> Result< ChatCompletionResponse >
{
let messages = vec!
[
ChatMessage
{
role : "user".to_string(),
content : user_message.into(),
tool_calls : None,
tool_call_id : None,
}
];
self.chat_completion( model, messages, None, None, None ).await
}
#[ inline ]
pub async fn math_completion(
&self,
model : impl AsRef< str >,
math_query : impl Into< String >,
) -> Result< ChatCompletionResponse >
{
let messages = vec!
[
ChatMessage
{
role : "system".to_string(),
content : "You are a helpful math assistant. Solve mathematical problems step by step and provide clear, accurate answers.".to_string(),
tool_calls : None,
tool_call_id : None,
},
ChatMessage
{
role : "user".to_string(),
content : math_query.into(),
tool_calls : None,
tool_call_id : None,
}
];
self.chat_completion( model, messages, Some( 150 ), Some( 0.1 ), Some( 0.95 ) ).await
}
#[ inline ]
#[ must_use ]
pub fn default_pro_model() -> &'static str
{
"meta-llama/Meta-Llama-3-8B-Instruct"
}
#[ inline ]
#[ must_use ]
pub fn fallback_model() -> &'static str
{
"facebook/bart-large-cnn"
}
#[ inline ]
#[ must_use ]
pub fn get_provider_for_model( model : &str ) -> Option< InferenceProvider >
{
InferenceProvider::for_model( model )
}
#[ inline ]
pub async fn chat_completion_with_tools(
&self,
model : impl AsRef< str >,
messages : Vec< ChatMessage >,
tools : Vec< Tool >,
options : ChatCompletionOptions,
) -> Result< ChatCompletionResponse >
{
let model_id = model.as_ref();
validate_model_identifier( model_id )?;
for message in &messages
{
let has_tool_calls = message.tool_calls.as_ref().is_some_and( | tc | !tc.is_empty() );
if !( message.role == "assistant" && has_tool_calls )
{
validate_input_text( &message.content )?;
}
}
let tool_definitions : Vec< ToolDefinition > = tools
.into_iter()
.map( | tool |
{
ToolDefinition
{
tool_type : "function".to_string(),
function : tool,
}
} )
.collect();
let request = ChatCompletionRequest
{
model : model_id.to_string(),
messages,
max_tokens : options.max_tokens,
temperature : options.temperature,
top_p : options.top_p,
stream : Some( false ), tools : Some( tool_definitions ),
tool_choice : options.tool_choice,
};
let url = self.client.environment.endpoint_url( "chat/completions" )?;
self.client.post( url.as_str(), &request ).await
}
}
#[ cfg( not( feature = "env-config" ) ) ]
impl< E > Providers< E >
where
E : Clone,
{
#[ inline ]
#[ must_use ]
pub fn new( client : &Client< E > ) -> Self
{
Self
{
client : (*client).clone(),
}
}
}
}
crate::mod_interface!
{
exposed use private::
{
Providers,
ChatCompletionOptions,
ChatCompletionRequest,
ChatCompletionResponse,
ChatMessage,
ChatChoice,
Usage,
ToolDefinition,
InferenceProvider,
};
}