use crate::{Hippox, HippoxConfig, IdentityInformation};
use langhub::{
AudioLLMClient, AudioLLMConfig, ImageLLMClient, ImageLLMConfig, VideoLLMClient, VideoLLMConfig, audio::AudioModelProvider,
chat::ChatModelProvider, image::ImageModelProvider, video::VideoModelProvider,
};
use std::collections::HashMap;
pub struct HippoxBuilder {
llm_provider: Option<ChatModelProvider>,
llm_api_key: Option<String>,
llm_extra_keys: Option<HashMap<String, String>>,
llm_model: Option<String>,
image_provider: Option<ImageModelProvider>,
image_config: Option<ImageLLMConfig>,
video_provider: Option<VideoModelProvider>,
video_config: Option<VideoLLMConfig>,
audio_provider: Option<AudioModelProvider>,
audio_config: Option<AudioLLMConfig>,
config: HippoxConfig,
}
impl HippoxBuilder {
pub fn new(provider: ChatModelProvider) -> Self {
Self {
llm_provider: Some(provider),
llm_api_key: None,
llm_extra_keys: None,
llm_model: None,
image_provider: None,
image_config: None,
video_provider: None,
video_config: None,
audio_provider: None,
audio_config: None,
config: HippoxConfig::default(),
}
}
pub fn new_image(provider: ImageModelProvider, config: ImageLLMConfig) -> Self {
Self {
llm_provider: None,
llm_api_key: None,
llm_extra_keys: None,
llm_model: None,
image_provider: Some(provider),
image_config: Some(config),
video_provider: None,
video_config: None,
audio_provider: None,
audio_config: None,
config: HippoxConfig::default(),
}
}
pub fn new_video(provider: VideoModelProvider, config: VideoLLMConfig) -> Self {
Self {
llm_provider: None,
llm_api_key: None,
llm_extra_keys: None,
llm_model: None,
image_provider: None,
image_config: None,
video_provider: Some(provider),
video_config: Some(config),
audio_provider: None,
audio_config: None,
config: HippoxConfig::default(),
}
}
pub fn new_audio(provider: AudioModelProvider, config: AudioLLMConfig) -> Self {
Self {
llm_provider: None,
llm_api_key: None,
llm_extra_keys: None,
llm_model: None,
image_provider: None,
image_config: None,
video_provider: None,
video_config: None,
audio_provider: Some(provider),
audio_config: Some(config),
config: HippoxConfig::default(),
}
}
pub fn api_key(mut self, key: impl Into<String>) -> Self {
self.llm_api_key = Some(key.into());
self
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.llm_model = Some(model.into());
self
}
pub fn extra_keys(mut self, keys: HashMap<String, String>) -> Self {
self.llm_extra_keys = Some(keys);
self
}
pub fn lang(mut self, lang: impl Into<String>) -> Self {
self.config.lang = lang.into();
self
}
pub fn identity(mut self, f: impl FnOnce(&mut IdentityInformation)) -> Self {
f(&mut self.config.identity_information);
self
}
pub fn image(mut self, provider: ImageModelProvider, config: ImageLLMConfig) -> Self {
self.image_provider = Some(provider);
self.image_config = Some(config);
self
}
pub fn video(mut self, provider: VideoModelProvider, config: VideoLLMConfig) -> Self {
self.video_provider = Some(provider);
self.video_config = Some(config);
self
}
pub fn audio(mut self, provider: AudioModelProvider, config: AudioLLMConfig) -> Self {
self.audio_provider = Some(provider);
self.audio_config = Some(config);
self
}
pub async fn build(self) -> anyhow::Result<Hippox> {
let provider = self
.llm_provider
.ok_or_else(|| anyhow::anyhow!("LLM provider is required for `build()`; use `build_with_model()` for modality-only builders"))?;
Hippox::with_workflow_mode(provider, self.llm_api_key, self.llm_extra_keys, Some(self.config)).await
}
pub async fn build_with_model(self) -> anyhow::Result<Hippox> {
let mut hippox = match self.llm_provider {
Some(provider) => Hippox::with_workflow_mode(provider, self.llm_api_key, self.llm_extra_keys, Some(self.config)).await?,
None => Hippox::without_llm(Some(self.config)).await?,
};
if let (Some(provider), Some(config)) = (self.image_provider, self.image_config) {
let client = Hippox::new_llm_image_with_config(provider, &config)?;
hippox = hippox.with_image_client(client, provider);
}
if let (Some(provider), Some(config)) = (self.video_provider, self.video_config) {
let client = Hippox::new_llm_video_with_config(provider, &config)?;
hippox = hippox.with_video_client(client, provider);
}
if let (Some(provider), Some(config)) = (self.audio_provider, self.audio_config) {
let client = Hippox::new_llm_audio_with_config(provider, &config)?;
hippox = hippox.with_audio_client(client, provider);
}
Ok(hippox)
}
}
impl Hippox {
pub fn builder(provider: ChatModelProvider) -> HippoxBuilder {
HippoxBuilder::new(provider)
}
pub fn builder_image(provider: ImageModelProvider, config: ImageLLMConfig) -> HippoxBuilder {
HippoxBuilder::new_image(provider, config)
}
pub fn builder_video(provider: VideoModelProvider, config: VideoLLMConfig) -> HippoxBuilder {
HippoxBuilder::new_video(provider, config)
}
pub fn builder_audio(provider: AudioModelProvider, config: AudioLLMConfig) -> HippoxBuilder {
HippoxBuilder::new_audio(provider, config)
}
}