Skip to main content

reagent_rs/services/llm/
client_config.rs

1use std::collections::HashMap;
2
3use crate::{
4    services::llm::{InferenceClient, InferenceClientError},
5    Provider,
6};
7
8#[derive(Debug, Clone, Default)]
9pub struct ClientConfig {
10    pub provider: Option<Provider>,
11    pub base_url: Option<String>,
12    pub api_key: Option<String>,
13    pub organization: Option<String>,
14    pub extra_headers: Option<std::collections::HashMap<String, String>>,
15}
16
17pub trait ClientBuilder {
18    fn provider(self, provider: Option<Provider>) -> Self;
19    fn base_url(self, base_url: Option<impl Into<String>>) -> Self;
20    fn api_key(self, api_key: Option<impl Into<String>>) -> Self;
21    fn organization(self, organization: Option<impl Into<String>>) -> Self;
22    fn extra_headers(self, extra_headers: Option<HashMap<String, String>>) -> Self;
23    fn build(self) -> Result<InferenceClient, InferenceClientError>;
24}
25
26impl ClientBuilder for ClientConfig {
27    fn provider(mut self, provider: Option<Provider>) -> Self {
28        self.provider = provider;
29        self
30    }
31
32    fn base_url(mut self, base_url: Option<impl Into<String>>) -> Self {
33        self.base_url = base_url.map(|s| s.into());
34        self
35    }
36
37    fn api_key(mut self, api_key: Option<impl Into<String>>) -> Self {
38        self.api_key = api_key.map(|s| s.into());
39        self
40    }
41
42    fn organization(mut self, organization: Option<impl Into<String>>) -> Self {
43        self.organization = organization.map(|s| s.into());
44        self
45    }
46
47    fn extra_headers(mut self, extra_headers: Option<HashMap<String, String>>) -> Self {
48        self.extra_headers = extra_headers;
49        self
50    }
51
52    fn build(self) -> Result<InferenceClient, InferenceClientError> {
53        InferenceClient::try_from(ClientConfig {
54            provider: self.provider.or(Some(Provider::Ollama)),
55            base_url: self.base_url,
56            api_key: self.api_key,
57            organization: self.organization,
58            extra_headers: self.extra_headers,
59        })
60    }
61}