1use std::borrow::Cow;
18use std::fmt;
19use std::hash::{Hash, Hasher};
20
21use serde::de;
22use serde::{Deserialize, Deserializer, Serialize, Serializer};
23
24use crate::catalog::{Catalog, ModelFacts, ModelSpec, NotFound, Reference};
25use crate::client::env::{self, EnvError};
26use crate::driver::DynModel;
27use crate::http_client::{DynHttpClient, HttpClientExt};
28use crate::operation::Completion;
29use crate::providers::{anthropic, gemini, openai};
30use crate::serve::ErasedHandler;
31use crate::serve::adapters::ModelAdapter;
32use crate::wire::Secret;
33
34pub(crate) const OPENAI_DIALECTS: &[&openai::wire::Dialect] = &[
41 &openai::wire::OPENAI,
42 &openai::wire::AZURE,
43 &openai::wire::DEEPSEEK,
44 &openai::wire::GROQ,
45 &openai::wire::HYPERBOLIC,
46 &openai::wire::MIRA,
47 &openai::wire::PERPLEXITY,
48 &openai::wire::TOGETHER,
49 &openai::wire::HUGGINGFACE,
50 &openai::wire::LLAMACPP,
51 &openai::wire::MISTRAL,
52 &openai::wire::OPENROUTER,
53 &openai::wire::VENICE,
54 &openai::wire::DOUBLEWORD,
55 &openai::wire::ZAI,
56 &openai::wire::MINIMAX,
57 &openai::wire::MOONSHOT,
58 &openai::wire::XIAOMIMIMO,
59 &openai::wire::COHERE,
60 &openai::wire::OLLAMA,
61 &crate::providers::xai::DIALECT,
62 &crate::providers::chatgpt::DIALECT,
63 &crate::providers::copilot::wire::DIALECT,
64];
65
66#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
70pub enum Format {
71 OpenAi,
73 Anthropic,
75 Gemini,
77}
78
79impl Format {
80 pub const ALL: [Format; 3] = [Format::OpenAi, Format::Anthropic, Format::Gemini];
82
83 pub fn as_str(self) -> &'static str {
85 match self {
86 Format::OpenAi => "openai",
87 Format::Anthropic => "anthropic",
88 Format::Gemini => "gemini",
89 }
90 }
91
92 pub fn named(name: &str) -> Option<Self> {
94 Self::ALL.into_iter().find(|format| format.as_str() == name)
95 }
96
97 fn names() -> String {
99 Self::ALL
100 .iter()
101 .map(|format| format.as_str())
102 .collect::<Vec<_>>()
103 .join(", ")
104 }
105}
106
107impl fmt::Display for Format {
108 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109 f.write_str(self.as_str())
110 }
111}
112
113impl Serialize for Format {
114 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
115 serializer.serialize_str(self.as_str())
116 }
117}
118
119impl<'de> Deserialize<'de> for Format {
120 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
121 let name = String::deserialize(deserializer)?;
122 Self::named(&name).ok_or_else(|| {
123 de::Error::custom(format!(
124 "`{name}` is not a protocol family ({})",
125 Format::names()
126 ))
127 })
128 }
129}
130
131#[derive(Debug, Clone, Copy)]
133enum Registered {
134 OpenAi(openai::wire::Dialect),
135 Anthropic(anthropic::wire::Dialect),
136 Gemini,
137}
138
139impl PartialEq for Registered {
141 fn eq(&self, other: &Self) -> bool {
142 self.vendor() == other.vendor() && self.format() == other.format()
143 }
144}
145
146impl Registered {
147 fn vendor(&self) -> &'static str {
148 match self {
149 Self::OpenAi(dialect) => dialect.name,
150 Self::Anthropic(dialect) => dialect.name,
151 Self::Gemini => gemini::PROVIDER_NAME,
152 }
153 }
154
155 fn format(&self) -> Format {
156 match self {
157 Self::OpenAi(_) => Format::OpenAi,
158 Self::Anthropic(_) => Format::Anthropic,
159 Self::Gemini => Format::Gemini,
160 }
161 }
162
163 fn config(&self, api_key: impl Into<Secret>) -> ProviderConfig {
164 match self {
165 Self::OpenAi(dialect) => {
166 ProviderConfig::OpenAi(openai::wire::OpenAIConfig::with_key(dialect, api_key))
167 }
168 Self::Anthropic(dialect) => ProviderConfig::Anthropic(
169 anthropic::wire::AnthropicConfig::with_key(dialect, api_key),
170 ),
171 Self::Gemini => ProviderConfig::Gemini(gemini::GeminiConfig::new(api_key)),
172 }
173 }
174
175 fn config_from_env(&self) -> Result<ProviderConfig, ConnectError> {
180 let tried = self.key_envs();
181 let (api_key, name) = key_from_env(self.vendor(), tried, self.requires_credential())?;
182 Ok(match self {
183 Self::OpenAi(dialect) => {
184 let auth = openai_auth(dialect, name);
185 ProviderConfig::OpenAi(openai::wire::OpenAIConfig::from_env_with_credential(
186 dialect, api_key, auth,
187 )?)
188 }
189 Self::Anthropic(dialect) => ProviderConfig::Anthropic(
190 anthropic::wire::AnthropicConfig::from_env_with_credential(dialect, api_key)?,
191 ),
192 Self::Gemini => ProviderConfig::Gemini(gemini::GeminiConfig::new(api_key)),
193 })
194 }
195
196 fn key_envs(&self) -> Vec<&'static str> {
198 match self {
199 Self::OpenAi(dialect) => openai_key_envs(dialect),
200 Self::Anthropic(dialect) => vec![dialect.api_key_env],
201 Self::Gemini => vec![gemini::API_KEY_ENV],
202 }
203 }
204
205 fn requires_credential(&self) -> bool {
207 match self {
208 Self::OpenAi(dialect) => openai_requires_credential(dialect),
209 Self::Anthropic(_) | Self::Gemini => true,
210 }
211 }
212}
213
214fn openai_key_envs(dialect: &openai::wire::Dialect) -> Vec<&'static str> {
217 std::iter::once(dialect.api_key_env)
218 .chain(
219 dialect
220 .alternate_auth
221 .map(|alternative| alternative.api_key_env),
222 )
223 .collect()
224}
225
226fn openai_requires_credential(dialect: &openai::wire::Dialect) -> bool {
229 !matches!(dialect.quirks.auth, openai::wire::Auth::OptionalBearer)
230}
231
232fn openai_auth(dialect: &openai::wire::Dialect, name: Option<&'static str>) -> openai::wire::Auth {
236 match dialect.alternate_auth {
237 Some(alternative) if name == Some(alternative.api_key_env) => alternative.auth,
238 _ => dialect.quirks.auth,
239 }
240}
241
242fn key_from_env(
246 vendor: &str,
247 tried: Vec<&'static str>,
248 required: bool,
249) -> Result<(String, Option<&'static str>), ConnectError> {
250 for name in &tried {
251 if let Some(value) = env::optional(name)?.filter(|value| !value.is_empty()) {
252 return Ok((value, Some(*name)));
253 }
254 }
255 if required {
256 return Err(ConnectError::MissingKey {
257 vendor: vendor.to_owned(),
258 tried: tried.into_iter().map(Cow::Borrowed).collect(),
259 });
260 }
261 Ok((String::new(), None))
262}
263
264#[derive(Debug)]
268struct CatalogOnly {
269 vendor: &'static str,
271 home: &'static str,
273 credential: bool,
275}
276
277static CATALOG_ONLY: [CatalogOnly; 5] = [
279 CatalogOnly {
280 vendor: "aws_bedrock",
281 home: "the `rig-bedrock` crate",
282 credential: true,
283 },
284 CatalogOnly {
285 vendor: "vertexai",
286 home: "the `rig-vertexai` crate",
287 credential: true,
288 },
289 CatalogOnly {
290 vendor: "gemini-grpc",
291 home: "the `rig-gemini-grpc` crate",
292 credential: true,
293 },
294 CatalogOnly {
295 vendor: "candle",
296 home: "the `rig-candle` crate",
297 credential: false,
298 },
299 CatalogOnly {
300 vendor: "voyageai",
301 home: "`rig_core::providers::voyageai`, which serves embeddings and reranking only",
302 credential: true,
303 },
304];
305
306#[derive(Debug, Clone, Copy)]
308enum Kind {
309 Registered(Registered),
311 CatalogOnly(&'static CatalogOnly),
313 Declared(&'static str),
315}
316
317static DECLARED: std::sync::RwLock<Vec<&'static str>> = std::sync::RwLock::new(Vec::new());
321
322fn declared(vendor: &str) -> Option<&'static str> {
324 DECLARED
325 .read()
326 .unwrap_or_else(std::sync::PoisonError::into_inner)
327 .iter()
328 .copied()
329 .find(|name| *name == vendor)
330}
331
332#[derive(Debug, Clone, Copy)]
341pub struct ProviderId(Kind);
342
343impl PartialEq for ProviderId {
346 fn eq(&self, other: &Self) -> bool {
347 self.vendor() == other.vendor() && self.format() == other.format()
348 }
349}
350
351impl Eq for ProviderId {}
352
353impl Hash for ProviderId {
354 fn hash<H: Hasher>(&self, state: &mut H) {
355 self.vendor().hash(state);
356 self.format().hash(state);
357 }
358}
359
360impl ProviderId {
361 pub fn all() -> impl Iterator<Item = ProviderId> {
365 openai::wire::all()
366 .map(|dialect| Registered::OpenAi(*dialect))
367 .chain(anthropic::wire::all().map(|dialect| Registered::Anthropic(*dialect)))
368 .chain(std::iter::once(Registered::Gemini))
369 .map(|registered| ProviderId(Kind::Registered(registered)))
370 }
371
372 pub fn new(vendor: &str, format: Format) -> Option<Self> {
375 let registered = match format {
376 Format::OpenAi => {
377 openai::wire::by_name(vendor).map(|dialect| Registered::OpenAi(*dialect))
378 }
379 Format::Anthropic => {
380 anthropic::wire::Dialect::by_name(vendor).map(Registered::Anthropic)
381 }
382 Format::Gemini => (vendor == gemini::PROVIDER_NAME).then_some(Registered::Gemini),
383 };
384 registered.map(|registered| Self(Kind::Registered(registered)))
385 }
386
387 pub fn catalog(vendor: &str) -> Option<Self> {
398 let primary = PRIMARY
399 .iter()
400 .find(|(name, _)| *name == vendor)
401 .and_then(|(_, format)| Self::new(vendor, *format));
402 primary
403 .or_else(|| Self::vendor_selections(vendor).next())
404 .or_else(|| {
405 CATALOG_ONLY
406 .iter()
407 .find(|provider| provider.vendor == vendor)
408 .map(|provider| Self(Kind::CatalogOnly(provider)))
409 })
410 .or_else(|| declared(vendor).map(|name| Self(Kind::Declared(name))))
411 }
412
413 pub fn declare(vendor: &str) -> Result<Self, SelectionError> {
443 if vendor.is_empty() || vendor.contains('/') || vendor.contains(char::is_whitespace) {
444 return Err(SelectionError::Malformed {
445 selection: vendor.to_owned(),
446 });
447 }
448 if let Some(known) = Self::catalog(vendor) {
449 return Ok(known);
450 }
451 let mut names = DECLARED
452 .write()
453 .unwrap_or_else(std::sync::PoisonError::into_inner);
454 let name = match names.iter().copied().find(|name| *name == vendor) {
455 Some(name) => name,
456 None => {
457 let name: &'static str = Box::leak(vendor.to_owned().into_boxed_str());
458 names.push(name);
459 name
460 }
461 };
462 Ok(Self(Kind::Declared(name)))
463 }
464
465 pub(crate) fn catalog_vendors() -> impl Iterator<Item = &'static str> {
468 let declared: Vec<&'static str> = DECLARED
469 .read()
470 .unwrap_or_else(std::sync::PoisonError::into_inner)
471 .clone();
472 Self::all()
473 .map(|id| id.vendor())
474 .chain(CATALOG_ONLY.iter().map(|provider| provider.vendor))
475 .chain(declared)
476 }
477
478 pub fn vendor(&self) -> &'static str {
480 match &self.0 {
481 Kind::Registered(registered) => registered.vendor(),
482 Kind::CatalogOnly(provider) => provider.vendor,
483 Kind::Declared(vendor) => vendor,
484 }
485 }
486
487 pub fn format(&self) -> Option<Format> {
490 match &self.0 {
491 Kind::Registered(registered) => Some(registered.format()),
492 Kind::CatalogOnly(_) | Kind::Declared(_) => None,
493 }
494 }
495
496 pub fn is_registered(&self) -> bool {
500 matches!(self.0, Kind::Registered(_))
501 }
502
503 pub fn is_declared(&self) -> bool {
505 matches!(self.0, Kind::Declared(_))
506 }
507
508 pub fn vendor_selections(vendor: &str) -> impl Iterator<Item = ProviderId> + '_ {
510 Self::all().filter(move |id| id.vendor() == vendor)
511 }
512
513 pub fn resolve(selection: &str) -> Result<Self, SelectionError> {
520 let malformed = || SelectionError::Malformed {
521 selection: selection.to_owned(),
522 };
523 let (vendor, format) = match selection.split_once('/') {
524 Some((_, rest)) if rest.contains('/') => return Err(malformed()),
525 Some((vendor, format)) => (vendor, Some(format)),
526 None => (selection, None),
527 };
528 if vendor.is_empty() {
529 return Err(malformed());
530 }
531 let Some(format) = format else {
532 let mut registered = Self::vendor_selections(vendor);
533 let first = registered.next().ok_or_else(|| SelectionError::Unknown {
534 vendor: vendor.to_owned(),
535 })?;
536 return match registered.next() {
537 None => Ok(first),
538 Some(_) => Err(SelectionError::Ambiguous {
539 vendor: vendor.to_owned(),
540 alternatives: alternatives(vendor),
541 }),
542 };
543 };
544 let Some(family) = Format::named(format) else {
545 return Err(SelectionError::UnknownFormat {
546 format: format.to_owned(),
547 });
548 };
549 Self::selection(vendor, family)
550 }
551
552 pub fn config(&self, api_key: impl Into<Secret>) -> Option<ProviderConfig> {
556 match &self.0 {
557 Kind::Registered(registered) => Some(registered.config(api_key)),
558 Kind::CatalogOnly(_) | Kind::Declared(_) => None,
559 }
560 }
561
562 pub fn api_key_envs(&self) -> Vec<&'static str> {
568 match &self.0 {
569 Kind::Registered(registered) => registered.key_envs(),
570 Kind::CatalogOnly(_) | Kind::Declared(_) => Vec::new(),
571 }
572 }
573
574 pub fn requires_credential(&self) -> bool {
580 match &self.0 {
581 Kind::Registered(registered) => registered.requires_credential(),
582 Kind::CatalogOnly(provider) => provider.credential,
583 Kind::Declared(_) => false,
584 }
585 }
586
587 fn registered(self) -> Result<Registered, ConnectError> {
589 match self.0 {
590 Kind::Registered(registered) => Ok(registered),
591 Kind::CatalogOnly(provider) => Err(ConnectError::CatalogOnly {
592 vendor: provider.vendor.to_owned(),
593 served_by: provider.home,
594 }),
595 Kind::Declared(vendor) => Err(ConnectError::CatalogOnly {
596 vendor: vendor.to_owned(),
597 served_by: "whoever declared it: provide it with `Connector::provide`",
598 }),
599 }
600 }
601
602 fn selection(vendor: &str, format: Format) -> Result<Self, SelectionError> {
604 Self::new(vendor, format).ok_or_else(|| {
605 let alternatives = alternatives(vendor);
606 match alternatives.is_empty() {
607 true => SelectionError::Unknown {
608 vendor: vendor.to_owned(),
609 },
610 false => SelectionError::Unregistered {
611 vendor: vendor.to_owned(),
612 format,
613 alternatives,
614 },
615 }
616 })
617 }
618}
619
620const PRIMARY: [(&str, Format); 4] = [
623 ("zai", Format::OpenAi),
624 ("minimax", Format::Anthropic),
625 ("moonshot", Format::OpenAi),
626 ("xiaomimimo", Format::OpenAi),
627];
628
629fn alternatives(vendor: &str) -> Vec<String> {
631 ProviderId::vendor_selections(vendor)
632 .map(|id| id.to_string())
633 .collect()
634}
635
636impl fmt::Display for ProviderId {
637 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
638 match self.format() {
639 Some(format) => write!(f, "{}/{format}", self.vendor()),
640 None => f.write_str(self.vendor()),
641 }
642 }
643}
644
645impl Serialize for ProviderId {
646 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
647 serializer.collect_str(self)
648 }
649}
650
651impl<'de> Deserialize<'de> for ProviderId {
655 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
656 let selection = String::deserialize(deserializer)?;
657 Self::resolve(&selection)
658 .or_else(|error| {
659 Self::catalog(&selection)
660 .filter(|id| !id.is_registered())
661 .ok_or(error)
662 })
663 .map_err(de::Error::custom)
664 }
665}
666
667#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
669pub enum SelectionError {
670 #[error("no registered provider is named `{vendor}`")]
672 Unknown {
673 vendor: String,
675 },
676 #[error(
678 "`{vendor}` speaks no {format} endpoint in this build (it speaks {})",
679 alternatives.join(", ")
680 )]
681 Unregistered {
682 vendor: String,
684 format: Format,
686 alternatives: Vec<String>,
688 },
689 #[error(
691 "`{vendor}` names more than one registered selection: name one of {}",
692 alternatives.join(", ")
693 )]
694 Ambiguous {
695 vendor: String,
697 alternatives: Vec<String>,
699 },
700 #[error("`{format}` is not a protocol family ({})", Format::names())]
702 UnknownFormat {
703 format: String,
705 },
706 #[error("`{selection}` is not a provider selection: expected `vendor` or `vendor/format`")]
708 Malformed {
709 selection: String,
711 },
712}
713
714#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
719pub enum ProviderConfig {
720 #[serde(rename = "openai")]
722 OpenAi(openai::wire::OpenAIConfig),
723 #[serde(rename = "anthropic")]
725 Anthropic(anthropic::wire::AnthropicConfig),
726 #[serde(rename = "gemini")]
728 Gemini(gemini::GeminiConfig),
729}
730
731impl ProviderConfig {
732 pub fn id(&self) -> Option<ProviderId> {
736 ProviderId::new(self.vendor(), self.format())
737 }
738
739 pub fn vendor(&self) -> &'static str {
741 match self {
742 Self::OpenAi(provider) => provider.dialect.name,
743 Self::Anthropic(provider) => provider.dialect.name,
744 Self::Gemini(_) => gemini::PROVIDER_NAME,
745 }
746 }
747
748 pub fn format(&self) -> Format {
750 match self {
751 Self::OpenAi(_) => Format::OpenAi,
752 Self::Anthropic(_) => Format::Anthropic,
753 Self::Gemini(_) => Format::Gemini,
754 }
755 }
756
757 pub fn is_unauthenticated(&self) -> bool {
760 match self {
761 Self::OpenAi(provider) => provider.api_key.is_empty(),
762 Self::Anthropic(provider) => provider.api_key.is_empty(),
763 Self::Gemini(provider) => provider.api_key.is_empty(),
764 }
765 }
766
767 pub fn with_credential(mut self, api_key: impl Into<Secret>) -> Self {
770 let api_key = api_key.into();
771 match &mut self {
772 Self::OpenAi(provider) => provider.api_key = api_key,
773 Self::Anthropic(provider) => provider.api_key = api_key,
774 Self::Gemini(provider) => provider.api_key = api_key,
775 }
776 self
777 }
778
779 pub fn base_url(&self) -> &str {
788 match self {
789 Self::OpenAi(provider) => &provider.base_url,
790 Self::Anthropic(provider) => &provider.base_url,
791 Self::Gemini(provider) => &provider.base_url,
792 }
793 }
794
795 pub fn with_base_url(self, base_url: impl Into<String>) -> Self {
799 let base_url = base_url.into();
800 match self {
801 Self::OpenAi(provider) => Self::OpenAi(provider.with_base_url(base_url)),
802 Self::Anthropic(provider) => Self::Anthropic(provider.with_base_url(base_url)),
803 Self::Gemini(provider) => Self::Gemini(provider.with_base_url(base_url)),
804 }
805 }
806
807 pub fn completion_handler(
811 &self,
812 label: &str,
813 model: &str,
814 http: DynHttpClient,
815 facts: ModelFacts,
816 ) -> ErasedHandler {
817 ErasedHandler::new(ModelAdapter::new(
818 label,
819 self.completion_model(model, http, facts),
820 ))
821 }
822
823 pub fn completion_model(
834 &self,
835 model: &str,
836 http: DynHttpClient,
837 facts: ModelFacts,
838 ) -> DynModel<Completion> {
839 match self {
840 Self::OpenAi(provider) => {
841 let mut model = provider.clone().connect(http).completion(model);
842 model.wire = model.wire.with_facts(facts);
843 model.erase()
844 }
845 Self::Anthropic(provider) => {
846 let mut model = provider.clone().connect(http).completion(model);
847 model.wire = model.wire.with_facts(facts);
848 model.erase()
849 }
850 Self::Gemini(provider) => {
851 let mut model = provider.clone().connect(http).completion(model);
852 model.wire = model.wire.with_facts(facts);
853 model.erase()
854 }
855 }
856 }
857}
858
859#[non_exhaustive]
862#[derive(Clone, Copy, Debug)]
863pub enum ModelSelector<'a> {
864 Spec(&'a ModelSpec),
866 Reference(&'a str),
868}
869
870impl<'a> From<&'a ModelSpec> for ModelSelector<'a> {
871 fn from(spec: &'a ModelSpec) -> Self {
872 Self::Spec(spec)
873 }
874}
875
876impl<'a> From<&'a str> for ModelSelector<'a> {
877 fn from(reference: &'a str) -> Self {
878 Self::Reference(reference)
879 }
880}
881
882impl<'a> From<&'a String> for ModelSelector<'a> {
883 fn from(reference: &'a String) -> Self {
884 Self::Reference(reference)
885 }
886}
887
888#[non_exhaustive]
903#[derive(Clone, Default)]
904pub struct ConnectOptions {
905 api_key: Option<Secret>,
906 account_id: Option<String>,
907 pub(crate) http: Option<DynHttpClient>,
908 format: Option<Format>,
909 base_url: Option<String>,
910}
911
912impl ConnectOptions {
913 pub fn new() -> Self {
915 Self::default()
916 }
917
918 pub fn api_key(mut self, api_key: impl Into<Secret>) -> Self {
921 self.api_key = Some(api_key.into());
922 self
923 }
924
925 pub fn account_id(mut self, account_id: impl Into<String>) -> Self {
938 self.account_id = Some(account_id.into());
939 self
940 }
941
942 pub fn http(mut self, http: impl HttpClientExt + 'static) -> Self {
944 self.http = Some(DynHttpClient::new(http));
945 self
946 }
947
948 pub fn format(mut self, format: Format) -> Self {
951 self.format = Some(format);
952 self
953 }
954
955 pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
957 self.base_url = Some(base_url.into());
958 self
959 }
960}
961
962impl fmt::Debug for ConnectOptions {
964 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
965 f.debug_struct("ConnectOptions")
966 .field("api_key", &self.api_key.as_ref().map(|_| "<set>"))
967 .field("account_id", &self.account_id.as_ref().map(|_| "<set>"))
968 .field("http", &self.http.as_ref().map(|_| "<set>"))
969 .field("format", &self.format)
970 .field("base_url", &self.base_url)
971 .finish()
972 }
973}
974
975#[non_exhaustive]
977#[derive(Debug, thiserror::Error)]
978pub enum ConnectError {
979 #[error("`{reference}` names no model: expected `vendor/model`")]
981 Malformed {
982 reference: String,
984 },
985 #[error("the registry cannot connect to `{vendor}`: its models are served by {served_by}")]
988 CatalogOnly {
989 vendor: String,
991 served_by: &'static str,
993 },
994 #[error(transparent)]
997 NotFound(#[from] NotFound),
998 #[error(transparent)]
1000 Selection(#[from] SelectionError),
1001 #[error(
1004 "no key for `{vendor}`: set {} or pass one with `ConnectOptions::api_key`",
1005 tried.iter().map(|name| format!("`{name}`")).collect::<Vec<_>>().join(" or ")
1006 )]
1007 MissingKey {
1008 vendor: String,
1010 tried: Vec<Cow<'static, str>>,
1013 },
1014 #[error("{0}")]
1017 Provider(String),
1018 #[error(transparent)]
1020 Env(#[from] EnvError),
1021 #[error("no HTTP client: pass one with `ConnectOptions::http`")]
1024 NoHttpClient,
1025}
1026
1027impl Catalog {
1028 #[cfg(feature = "reqwest")]
1042 #[cfg_attr(docsrs, doc(cfg(feature = "reqwest")))]
1043 pub fn connect<'a>(
1044 &self,
1045 model: impl Into<ModelSelector<'a>>,
1046 ) -> Result<DynModel<Completion>, ConnectError> {
1047 self.connect_with(model, ConnectOptions::new())
1048 }
1049
1050 pub fn connect_with<'a>(
1078 &self,
1079 model: impl Into<ModelSelector<'a>>,
1080 options: ConnectOptions,
1081 ) -> Result<DynModel<Completion>, ConnectError> {
1082 let (provider, model, facts) = match model.into() {
1083 ModelSelector::Spec(spec) => (
1084 spec.provider,
1085 spec.id.as_str(),
1086 ModelFacts::new(spec.clone()).with_catalog(self),
1087 ),
1088 ModelSelector::Reference(text) => {
1089 let reference = Reference::parse(text).ok_or_else(|| ConnectError::Malformed {
1090 reference: text.to_owned(),
1091 })?;
1092 let vendor =
1093 ProviderId::catalog(reference.vendor).ok_or_else(|| self.not_found(text))?;
1094 let provider = self
1095 .get_vendor(reference.vendor, reference.model)
1096 .map_or(vendor, |resolved| resolved.spec.provider);
1097 let provider = match reference.format {
1098 Some(format) if provider.is_registered() => {
1099 ProviderId::selection(reference.vendor, format)?
1100 }
1101 _ => provider,
1102 };
1103 (
1104 provider,
1105 reference.model,
1106 self.facts(provider, reference.model),
1107 )
1108 }
1109 };
1110 let provider = match options.format {
1111 Some(format) if provider.is_registered() => {
1112 ProviderId::selection(provider.vendor(), format)?
1113 }
1114 _ => provider,
1115 };
1116 let registered = provider.registered()?;
1117 let config = match options.api_key {
1118 Some(api_key) => registered.config(api_key),
1119 None => registered.config_from_env()?,
1120 };
1121 let config = match (config, options.account_id) {
1122 (ProviderConfig::OpenAi(provider), Some(account_id)) => {
1123 ProviderConfig::OpenAi(provider.with_account_id(account_id))
1124 }
1125 (config, _) => config,
1126 };
1127 let config = match options.base_url {
1128 Some(base_url) => config.with_base_url(base_url),
1129 None => config,
1130 };
1131 let http = match options.http {
1132 Some(http) => http,
1133 None => shared_http()?,
1134 };
1135 Ok(config.completion_model(model, http, facts))
1136 }
1137}
1138
1139#[cfg(feature = "reqwest")]
1141fn shared_http() -> Result<DynHttpClient, ConnectError> {
1142 Ok(rig_reqwest::shared())
1143}
1144
1145#[cfg(not(feature = "reqwest"))]
1147fn shared_http() -> Result<DynHttpClient, ConnectError> {
1148 Err(ConnectError::NoHttpClient)
1149}
1150
1151#[cfg(test)]
1152mod tests;