reagent_rs/services/llm/
client_config.rs1use 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}