1use std::collections::HashMap;
7use std::sync::Arc;
8use std::time::SystemTime;
9
10use async_trait::async_trait;
11use axum::http::HeaderMap;
12use serde::{Deserialize, Serialize};
13use serde_json::Value;
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct ChatMessage {
22 pub role: String,
24 pub content: Value,
26}
27
28impl ChatMessage {
29 #[must_use]
31 pub fn text(role: impl Into<String>, content: impl Into<String>) -> Self {
32 Self {
33 role: role.into(),
34 content: Value::String(content.into()),
35 }
36 }
37
38 #[must_use]
41 pub fn text_view(&self) -> String {
42 match &self.content {
43 Value::String(s) => s.clone(),
44 Value::Array(blocks) => blocks
45 .iter()
46 .filter_map(|b| b.get("text").and_then(Value::as_str))
47 .collect::<Vec<_>>()
48 .join("\n"),
49 _ => String::new(),
50 }
51 }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
60pub struct ModelRequest {
61 pub model: String,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
65 pub system: Option<String>,
66 pub messages: Vec<ChatMessage>,
68 pub max_tokens: u32,
70 #[serde(default)]
72 pub tools: Value,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct ModelResponse {
78 pub model: String,
80 pub text: String,
82 pub in_tokens: u64,
84 pub out_tokens: u64,
86 pub raw: Value,
88}
89
90#[derive(Debug, Clone, thiserror::Error)]
92pub enum ProviderError {
93 #[error("transport error: {0}")]
95 Transport(String),
96 #[error("http {status}: {body}")]
98 Http {
99 status: u16,
101 body: String,
103 },
104 #[error("decode error: {0}")]
106 Decode(String),
107}
108
109impl ProviderError {
110 #[must_use]
113 pub fn is_failover_eligible(&self) -> bool {
114 match self {
115 ProviderError::Transport(_) => true,
116 ProviderError::Http { status, .. } => *status >= 500,
117 ProviderError::Decode(_) => false,
118 }
119 }
120}
121
122#[derive(Clone, Default)]
126pub struct Auth {
127 pub anthropic_key: Option<String>,
129 pub openai_key: Option<String>,
131}
132
133impl std::fmt::Debug for Auth {
134 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
135 f.debug_struct("Auth")
136 .field("anthropic_key", &self.anthropic_key.as_ref().map(|_| "***"))
137 .field("openai_key", &self.openai_key.as_ref().map(|_| "***"))
138 .finish()
139 }
140}
141
142impl Auth {
143 #[must_use]
146 pub fn from_headers(headers: &HeaderMap) -> Self {
147 let anthropic_key = headers
148 .get("x-api-key")
149 .and_then(|v| v.to_str().ok())
150 .map(str::to_owned)
151 .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok());
152 let openai_key = headers
153 .get(axum::http::header::AUTHORIZATION)
154 .and_then(|v| v.to_str().ok())
155 .and_then(|v| v.strip_prefix("Bearer "))
156 .map(str::to_owned)
157 .or_else(|| std::env::var("OPENAI_API_KEY").ok());
158 Self {
159 anthropic_key,
160 openai_key,
161 }
162 }
163}
164
165#[async_trait]
168pub trait Provider: Send + Sync + std::fmt::Debug {
169 async fn complete(
175 &self,
176 req: &ModelRequest,
177 auth: &Auth,
178 ) -> Result<ModelResponse, ProviderError>;
179
180 fn id(&self) -> &str;
182}
183
184#[derive(Serialize)]
185struct AnthropicWireMessage<'a> {
186 role: &'a str,
187 content: &'a Value,
188}
189
190fn wire_model(model: &str) -> &str {
194 model.split_once('/').map_or(model, |(_, m)| m)
195}
196
197fn resolve_api_key(api_key_env: Option<&str>, byok_override: Option<&str>) -> String {
202 api_key_env
203 .and_then(|e| std::env::var(e).ok())
204 .or_else(|| byok_override.map(str::to_owned))
205 .unwrap_or_default()
206}
207
208#[derive(Serialize)]
209struct AnthropicWireRequest<'a> {
210 model: &'a str,
211 #[serde(skip_serializing_if = "Option::is_none")]
212 system: Option<&'a str>,
213 max_tokens: u32,
214 messages: Vec<AnthropicWireMessage<'a>>,
215}
216
217#[derive(Debug, Clone)]
223pub struct AnthropicProvider {
224 pub id: String,
226 pub base_url: String,
228 pub api_key_env: Option<String>,
231 pub http: reqwest::Client,
233}
234
235#[async_trait]
236impl Provider for AnthropicProvider {
237 fn id(&self) -> &str {
238 &self.id
239 }
240
241 async fn complete(
242 &self,
243 req: &ModelRequest,
244 auth: &Auth,
245 ) -> Result<ModelResponse, ProviderError> {
246 let key = resolve_api_key(self.api_key_env.as_deref(), auth.anthropic_key.as_deref());
247 let body = AnthropicWireRequest {
248 model: wire_model(&req.model),
249 system: req.system.as_deref(),
250 max_tokens: req.max_tokens,
251 messages: req
252 .messages
253 .iter()
254 .map(|m| AnthropicWireMessage {
255 role: &m.role,
256 content: &m.content,
257 })
258 .collect(),
259 };
260
261 let url = format!("{}/v1/messages", self.base_url.trim_end_matches('/'));
262 let resp = self
263 .http
264 .post(url)
265 .header("x-api-key", key)
266 .header("anthropic-version", "2023-06-01")
267 .json(&body)
268 .send()
269 .await
270 .map_err(|e| ProviderError::Transport(e.to_string()))?;
271
272 let status = resp.status();
273 let bytes = resp
274 .bytes()
275 .await
276 .map_err(|e| ProviderError::Transport(e.to_string()))?;
277
278 if !status.is_success() {
279 return Err(ProviderError::Http {
280 status: status.as_u16(),
281 body: String::from_utf8_lossy(&bytes).into_owned(),
282 });
283 }
284
285 let json: Value =
286 serde_json::from_slice(&bytes).map_err(|e| ProviderError::Decode(e.to_string()))?;
287 let (text, in_tokens, out_tokens) = anthropic_parse_response(&json)?;
288
289 Ok(ModelResponse {
290 model: req.model.clone(),
291 text,
292 in_tokens,
293 out_tokens,
294 raw: json,
295 })
296 }
297}
298
299fn anthropic_parse_response(json: &Value) -> Result<(String, u64, u64), ProviderError> {
303 let text = json
304 .get("content")
305 .and_then(Value::as_array)
306 .map(|blocks| {
307 blocks
308 .iter()
309 .filter_map(|b| b.get("text").and_then(Value::as_str))
310 .collect::<Vec<_>>()
311 .join("")
312 })
313 .ok_or_else(|| ProviderError::Decode("missing content[].text".to_owned()))?;
314
315 let in_tokens = json
316 .pointer("/usage/input_tokens")
317 .and_then(Value::as_u64)
318 .unwrap_or(0);
319 let out_tokens = json
320 .pointer("/usage/output_tokens")
321 .and_then(Value::as_u64)
322 .unwrap_or(0);
323
324 Ok((text, in_tokens, out_tokens))
325}
326
327fn anthropic_messages_body(req: &ModelRequest, anthropic_version: &str) -> Value {
332 let messages: Vec<Value> = req
333 .messages
334 .iter()
335 .map(|m| serde_json::json!({ "role": m.role, "content": m.content }))
336 .collect();
337 let mut body = serde_json::json!({
338 "anthropic_version": anthropic_version,
339 "max_tokens": req.max_tokens,
340 "messages": messages,
341 });
342 if let Some(system) = req.system.as_deref() {
343 body["system"] = serde_json::json!(system);
344 }
345 body
346}
347
348#[derive(Serialize)]
349struct OpenAiWireMessage<'a> {
350 role: &'a str,
351 content: Value,
352}
353
354#[derive(Serialize)]
355struct OpenAiWireRequest<'a> {
356 model: &'a str,
357 max_tokens: u32,
358 messages: Vec<OpenAiWireMessage<'a>>,
359}
360
361#[derive(Debug, Clone)]
367pub struct OpenAiProvider {
368 pub id: String,
370 pub base_url: String,
372 pub api_key_env: Option<String>,
376 pub http: reqwest::Client,
378}
379
380#[async_trait]
381impl Provider for OpenAiProvider {
382 fn id(&self) -> &str {
383 &self.id
384 }
385
386 async fn complete(
387 &self,
388 req: &ModelRequest,
389 auth: &Auth,
390 ) -> Result<ModelResponse, ProviderError> {
391 let key = resolve_api_key(self.api_key_env.as_deref(), auth.openai_key.as_deref());
392 let mut messages = Vec::with_capacity(req.messages.len() + 1);
393 if let Some(system) = req.system.as_deref() {
394 messages.push(OpenAiWireMessage {
395 role: "system",
396 content: Value::String(system.to_owned()),
397 });
398 }
399 messages.extend(req.messages.iter().map(|m| OpenAiWireMessage {
400 role: &m.role,
401 content: m.content.clone(),
402 }));
403 let body = OpenAiWireRequest {
404 model: wire_model(&req.model),
405 max_tokens: req.max_tokens,
406 messages,
407 };
408
409 let url = format!(
410 "{}/v1/chat/completions",
411 self.base_url.trim_end_matches('/')
412 );
413 let resp = self
414 .http
415 .post(url)
416 .header(axum::http::header::AUTHORIZATION, format!("Bearer {key}"))
417 .json(&body)
418 .send()
419 .await
420 .map_err(|e| ProviderError::Transport(e.to_string()))?;
421
422 let status = resp.status();
423 let bytes = resp
424 .bytes()
425 .await
426 .map_err(|e| ProviderError::Transport(e.to_string()))?;
427
428 if !status.is_success() {
429 return Err(ProviderError::Http {
430 status: status.as_u16(),
431 body: String::from_utf8_lossy(&bytes).into_owned(),
432 });
433 }
434
435 let json: Value =
436 serde_json::from_slice(&bytes).map_err(|e| ProviderError::Decode(e.to_string()))?;
437
438 let text = json
439 .pointer("/choices/0/message/content")
440 .and_then(Value::as_str)
441 .ok_or_else(|| ProviderError::Decode("missing choices[0].message.content".to_owned()))?
442 .to_owned();
443
444 let in_tokens = json
445 .pointer("/usage/prompt_tokens")
446 .and_then(Value::as_u64)
447 .unwrap_or(0);
448 let out_tokens = json
449 .pointer("/usage/completion_tokens")
450 .and_then(Value::as_u64)
451 .unwrap_or(0);
452
453 Ok(ModelResponse {
454 model: req.model.clone(),
455 text,
456 in_tokens,
457 out_tokens,
458 raw: json,
459 })
460 }
461}
462
463fn gemini_request_body(req: &ModelRequest) -> Value {
470 let contents: Vec<Value> = req
471 .messages
472 .iter()
473 .map(|m| {
474 let role = if m.role == "assistant" {
475 "model"
476 } else {
477 "user"
478 };
479 serde_json::json!({ "role": role, "parts": [{ "text": m.text_view() }] })
480 })
481 .collect();
482 let mut body = serde_json::json!({
483 "contents": contents,
484 "generationConfig": { "maxOutputTokens": req.max_tokens },
485 });
486 if let Some(system) = req.system.as_deref() {
487 body["system_instruction"] = serde_json::json!({ "parts": [{ "text": system }] });
488 }
489 body
490}
491
492fn gemini_parse_response(json: &Value) -> Result<(String, u64, u64), ProviderError> {
495 let parts = json
496 .pointer("/candidates/0/content/parts")
497 .and_then(Value::as_array)
498 .ok_or_else(|| ProviderError::Decode("missing candidates[0].content.parts".to_owned()))?;
499 let text = parts
500 .iter()
501 .filter_map(|p| p.get("text").and_then(Value::as_str))
502 .collect::<Vec<_>>()
503 .join("");
504 let in_tokens = json
505 .pointer("/usageMetadata/promptTokenCount")
506 .and_then(Value::as_u64)
507 .unwrap_or(0);
508 let out_tokens = json
509 .pointer("/usageMetadata/candidatesTokenCount")
510 .and_then(Value::as_u64)
511 .unwrap_or(0);
512 Ok((text, in_tokens, out_tokens))
513}
514
515#[derive(Debug, Clone)]
522pub struct GeminiProvider {
523 pub id: String,
525 pub base_url: String,
527 pub api_key_env: Option<String>,
529 pub http: reqwest::Client,
531}
532
533#[async_trait]
534impl Provider for GeminiProvider {
535 fn id(&self) -> &str {
536 &self.id
537 }
538
539 async fn complete(
540 &self,
541 req: &ModelRequest,
542 auth: &Auth,
543 ) -> Result<ModelResponse, ProviderError> {
544 let key = resolve_api_key(self.api_key_env.as_deref(), auth.openai_key.as_deref());
545 let body = gemini_request_body(req);
546 let url = format!(
547 "{}/v1beta/models/{}:generateContent",
548 self.base_url.trim_end_matches('/'),
549 wire_model(&req.model),
550 );
551 let resp = self
552 .http
553 .post(url)
554 .header("x-goog-api-key", key)
555 .json(&body)
556 .send()
557 .await
558 .map_err(|e| ProviderError::Transport(e.to_string()))?;
559
560 let status = resp.status();
561 let bytes = resp
562 .bytes()
563 .await
564 .map_err(|e| ProviderError::Transport(e.to_string()))?;
565 if !status.is_success() {
566 return Err(ProviderError::Http {
567 status: status.as_u16(),
568 body: String::from_utf8_lossy(&bytes).into_owned(),
569 });
570 }
571
572 let json: Value =
573 serde_json::from_slice(&bytes).map_err(|e| ProviderError::Decode(e.to_string()))?;
574 let (text, in_tokens, out_tokens) = gemini_parse_response(&json)?;
575 Ok(ModelResponse {
576 model: req.model.clone(),
577 text,
578 in_tokens,
579 out_tokens,
580 raw: json,
581 })
582 }
583}
584
585struct AwsEnvCredentials {
589 access_key_id: String,
590 secret_access_key: String,
591 session_token: Option<String>,
592}
593
594impl AwsEnvCredentials {
595 fn from_env() -> Result<Self, ProviderError> {
596 let access_key_id = std::env::var("AWS_ACCESS_KEY_ID")
597 .map_err(|_| ProviderError::Transport("AWS_ACCESS_KEY_ID is not set".to_owned()))?;
598 let secret_access_key = std::env::var("AWS_SECRET_ACCESS_KEY")
599 .map_err(|_| ProviderError::Transport("AWS_SECRET_ACCESS_KEY is not set".to_owned()))?;
600 let session_token = std::env::var("AWS_SESSION_TOKEN").ok();
601 Ok(Self {
602 access_key_id,
603 secret_access_key,
604 session_token,
605 })
606 }
607}
608
609fn bedrock_url(region: &str, model: &str) -> String {
611 format!("https://bedrock-runtime.{region}.amazonaws.com/model/{model}/invoke")
612}
613
614fn sign_bedrock(
625 url: &str,
626 region: &str,
627 body: &[u8],
628 creds: &AwsEnvCredentials,
629) -> Result<http::Request<Vec<u8>>, ProviderError> {
630 let host = url
631 .parse::<http::Uri>()
632 .ok()
633 .and_then(|u| u.host().map(str::to_owned))
634 .ok_or_else(|| ProviderError::Transport(format!("invalid bedrock URL: {url}")))?;
635
636 let identity: aws_smithy_runtime_api::client::identity::Identity =
637 aws_credential_types::Credentials::new(
638 creds.access_key_id.clone(),
639 creds.secret_access_key.clone(),
640 creds.session_token.clone(),
641 None,
642 "firstpass",
643 )
644 .into();
645
646 let signing_params: aws_sigv4::http_request::SigningParams<'_> =
647 aws_sigv4::sign::v4::SigningParams::builder()
648 .identity(&identity)
649 .region(region)
650 .name("bedrock")
651 .time(SystemTime::now())
652 .settings(aws_sigv4::http_request::SigningSettings::default())
653 .build()
654 .map_err(|e| ProviderError::Transport(format!("sigv4 signing params: {e}")))?
655 .into();
656
657 let headers = [
658 ("host", host.as_str()),
659 ("content-type", "application/json"),
660 ];
661 let signable = aws_sigv4::http_request::SignableRequest::new(
662 "POST",
663 url,
664 headers.into_iter(),
665 aws_sigv4::http_request::SignableBody::Bytes(body),
666 )
667 .map_err(|e| ProviderError::Transport(format!("sigv4 signable request: {e}")))?;
668
669 let (instructions, _signature) = aws_sigv4::http_request::sign(signable, &signing_params)
670 .map_err(|e| ProviderError::Transport(format!("sigv4 sign: {e}")))?
671 .into_parts();
672
673 let mut req = http::Request::builder()
674 .method("POST")
675 .uri(url)
676 .header("host", host)
677 .header("content-type", "application/json")
678 .body(body.to_vec())
679 .map_err(|e| ProviderError::Transport(format!("build bedrock request: {e}")))?;
680 instructions.apply_to_request_http1x(&mut req);
681 Ok(req)
682}
683
684#[derive(Debug, Clone)]
694pub struct BedrockProvider {
695 pub id: String,
697 pub region: Option<String>,
699 pub http: reqwest::Client,
701}
702
703#[async_trait]
704impl Provider for BedrockProvider {
705 fn id(&self) -> &str {
706 &self.id
707 }
708
709 async fn complete(
710 &self,
711 req: &ModelRequest,
712 _auth: &Auth,
713 ) -> Result<ModelResponse, ProviderError> {
714 let region = self.region.as_deref().ok_or_else(|| {
715 ProviderError::Transport("bedrock provider requires a region".to_owned())
716 })?;
717 let model = wire_model(&req.model);
718 let url = bedrock_url(region, model);
719 let body = anthropic_messages_body(req, "bedrock-2023-05-31");
720 let body_bytes =
721 serde_json::to_vec(&body).map_err(|e| ProviderError::Decode(e.to_string()))?;
722
723 let creds = AwsEnvCredentials::from_env()?;
724 let signed = sign_bedrock(&url, region, &body_bytes, &creds)?;
725 let http_req = reqwest::Request::try_from(signed)
726 .map_err(|e| ProviderError::Transport(format!("build reqwest request: {e}")))?;
727
728 let resp = self
729 .http
730 .execute(http_req)
731 .await
732 .map_err(|e| ProviderError::Transport(e.to_string()))?;
733
734 let status = resp.status();
735 let bytes = resp
736 .bytes()
737 .await
738 .map_err(|e| ProviderError::Transport(e.to_string()))?;
739 if !status.is_success() {
740 return Err(ProviderError::Http {
741 status: status.as_u16(),
742 body: String::from_utf8_lossy(&bytes).into_owned(),
743 });
744 }
745
746 let json: Value =
747 serde_json::from_slice(&bytes).map_err(|e| ProviderError::Decode(e.to_string()))?;
748 let (text, in_tokens, out_tokens) = anthropic_parse_response(&json)?;
749 Ok(ModelResponse {
750 model: req.model.clone(),
751 text,
752 in_tokens,
753 out_tokens,
754 raw: json,
755 })
756 }
757}
758
759fn vertex_url(region: &str, project: &str, model: &str) -> String {
761 format!(
762 "https://{region}-aiplatform.googleapis.com/v1/projects/{project}/locations/{region}/publishers/anthropic/models/{model}:rawPredict"
763 )
764}
765
766#[derive(Debug, Clone)]
776pub struct VertexProvider {
777 pub id: String,
779 pub region: Option<String>,
781 pub project: Option<String>,
783 pub http: reqwest::Client,
785}
786
787#[async_trait]
788impl Provider for VertexProvider {
789 fn id(&self) -> &str {
790 &self.id
791 }
792
793 async fn complete(
794 &self,
795 req: &ModelRequest,
796 _auth: &Auth,
797 ) -> Result<ModelResponse, ProviderError> {
798 let region = self.region.as_deref().ok_or_else(|| {
799 ProviderError::Transport("vertex provider requires a region".to_owned())
800 })?;
801 let project = self.project.as_deref().ok_or_else(|| {
802 ProviderError::Transport("vertex provider requires a project".to_owned())
803 })?;
804 let model = wire_model(&req.model);
805 let url = vertex_url(region, project, model);
806 let body = anthropic_messages_body(req, "vertex-2023-10-16");
807
808 let provider = gcp_auth::provider()
811 .await
812 .map_err(|e| ProviderError::Transport(format!("gcp_auth provider: {e}")))?;
813 let token = provider
814 .token(&["https://www.googleapis.com/auth/cloud-platform"])
815 .await
816 .map_err(|e| ProviderError::Transport(format!("gcp_auth token: {e}")))?;
817
818 let resp = self
819 .http
820 .post(url)
821 .header(
822 axum::http::header::AUTHORIZATION,
823 format!("Bearer {}", token.as_str()),
824 )
825 .json(&body)
826 .send()
827 .await
828 .map_err(|e| ProviderError::Transport(e.to_string()))?;
829
830 let status = resp.status();
831 let bytes = resp
832 .bytes()
833 .await
834 .map_err(|e| ProviderError::Transport(e.to_string()))?;
835 if !status.is_success() {
836 return Err(ProviderError::Http {
837 status: status.as_u16(),
838 body: String::from_utf8_lossy(&bytes).into_owned(),
839 });
840 }
841
842 let json: Value =
843 serde_json::from_slice(&bytes).map_err(|e| ProviderError::Decode(e.to_string()))?;
844 let (text, in_tokens, out_tokens) = anthropic_parse_response(&json)?;
845 Ok(ModelResponse {
846 model: req.model.clone(),
847 text,
848 in_tokens,
849 out_tokens,
850 raw: json,
851 })
852 }
853}
854
855#[derive(Clone)]
857pub struct ProviderRegistry {
858 providers: HashMap<String, Arc<dyn Provider>>,
859}
860
861impl std::fmt::Debug for ProviderRegistry {
862 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
863 f.debug_struct("ProviderRegistry")
864 .field("providers", &self.providers.keys().collect::<Vec<_>>())
865 .finish()
866 }
867}
868
869impl ProviderRegistry {
870 fn build_http_client() -> reqwest::Client {
875 reqwest::Client::builder()
876 .connect_timeout(std::time::Duration::from_secs(10))
877 .timeout(std::time::Duration::from_secs(120))
878 .build()
879 .unwrap_or_else(|_| reqwest::Client::new())
880 }
881
882 #[must_use]
884 pub fn new(anthropic_base: impl Into<String>, openai_base: impl Into<String>) -> Self {
885 Self::from_config(&[], anthropic_base, openai_base)
886 }
887
888 #[must_use]
894 pub fn from_config(
895 defs: &[firstpass_core::ProviderDef],
896 anthropic_base: impl Into<String>,
897 openai_base: impl Into<String>,
898 ) -> Self {
899 let http = Self::build_http_client();
900 let mut providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
901 providers.insert(
902 "anthropic".to_owned(),
903 Arc::new(AnthropicProvider {
904 id: "anthropic".to_owned(),
905 base_url: anthropic_base.into(),
906 api_key_env: None,
907 http: http.clone(),
908 }),
909 );
910 providers.insert(
911 "openai".to_owned(),
912 Arc::new(OpenAiProvider {
913 id: "openai".to_owned(),
914 base_url: openai_base.into(),
915 api_key_env: None,
916 http: http.clone(),
917 }),
918 );
919 for def in defs {
920 let provider: Arc<dyn Provider> = match def.auth {
924 firstpass_core::AuthScheme::AwsSigv4 => Arc::new(BedrockProvider {
925 id: def.id.clone(),
926 region: def.region.clone(),
927 http: http.clone(),
928 }),
929 firstpass_core::AuthScheme::GcpOauth => Arc::new(VertexProvider {
930 id: def.id.clone(),
931 region: def.region.clone(),
932 project: def.project.clone(),
933 http: http.clone(),
934 }),
935 firstpass_core::AuthScheme::ApiKey => match def.dialect {
936 firstpass_core::Dialect::Anthropic => Arc::new(AnthropicProvider {
937 id: def.id.clone(),
938 base_url: def.base_url.clone(),
939 api_key_env: def.api_key_env.clone(),
940 http: http.clone(),
941 }),
942 firstpass_core::Dialect::Openai => Arc::new(OpenAiProvider {
943 id: def.id.clone(),
944 base_url: def.base_url.clone(),
945 api_key_env: def.api_key_env.clone(),
946 http: http.clone(),
947 }),
948 firstpass_core::Dialect::Gemini => Arc::new(GeminiProvider {
949 id: def.id.clone(),
950 base_url: def.base_url.clone(),
951 api_key_env: def.api_key_env.clone(),
952 http: http.clone(),
953 }),
954 },
955 };
956 providers.insert(def.id.clone(), provider);
957 }
958 Self { providers }
959 }
960
961 #[must_use]
963 pub fn from_map(providers: HashMap<String, Arc<dyn Provider>>) -> Self {
964 Self { providers }
965 }
966
967 #[must_use]
969 pub fn get(&self, provider_id: &str) -> Option<Arc<dyn Provider>> {
970 self.providers.get(provider_id).cloned()
971 }
972}
973
974#[cfg(test)]
976#[derive(Debug, Clone, Default)]
977pub struct MockProvider {
978 id: String,
979 outcomes: HashMap<String, Result<ModelResponse, ProviderError>>,
980 calls: Arc<std::sync::Mutex<Vec<String>>>,
983 delay_ms: u64,
986}
987
988#[cfg(test)]
989impl MockProvider {
990 #[must_use]
992 pub fn new(
993 id: impl Into<String>,
994 outcomes: HashMap<String, Result<ModelResponse, ProviderError>>,
995 ) -> Self {
996 Self {
997 id: id.into(),
998 outcomes,
999 calls: Arc::default(),
1000 delay_ms: 0,
1001 }
1002 }
1003
1004 #[must_use]
1006 pub fn with_delay(mut self, ms: u64) -> Self {
1007 self.delay_ms = ms;
1008 self
1009 }
1010
1011 #[must_use]
1014 pub fn call_log(&self) -> Arc<std::sync::Mutex<Vec<String>>> {
1015 Arc::clone(&self.calls)
1016 }
1017}
1018
1019#[cfg(test)]
1020#[async_trait]
1021impl Provider for MockProvider {
1022 fn id(&self) -> &str {
1023 &self.id
1024 }
1025
1026 async fn complete(
1027 &self,
1028 req: &ModelRequest,
1029 _auth: &Auth,
1030 ) -> Result<ModelResponse, ProviderError> {
1031 self.calls.lock().unwrap().push(req.model.clone());
1032 if self.delay_ms > 0 {
1033 tokio::time::sleep(std::time::Duration::from_millis(self.delay_ms)).await;
1034 }
1035 self.outcomes.get(&req.model).cloned().unwrap_or_else(|| {
1036 Err(ProviderError::Decode(format!(
1037 "no mock outcome configured for {}",
1038 req.model
1039 )))
1040 })
1041 }
1042}
1043
1044#[cfg(test)]
1045mod tests {
1046 use super::*;
1047
1048 #[test]
1049 fn wire_model_strips_the_provider_prefix() {
1050 assert_eq!(wire_model("anthropic/claude-haiku-4-5"), "claude-haiku-4-5");
1052 assert_eq!(wire_model("openai/gpt-5.5"), "gpt-5.5");
1053 assert_eq!(wire_model("claude-opus-4-8"), "claude-opus-4-8"); }
1055
1056 #[test]
1057 fn gemini_request_maps_roles_and_system_instruction() {
1058 let req = ModelRequest {
1059 model: "gemini/gemini-2.0-flash".to_owned(),
1060 system: Some("be terse".to_owned()),
1061 messages: vec![
1062 ChatMessage::text("user", "hi"),
1063 ChatMessage::text("assistant", "hello"),
1064 ],
1065 max_tokens: 256,
1066 tools: Value::Null,
1067 };
1068 let body = gemini_request_body(&req);
1069 assert_eq!(body["system_instruction"]["parts"][0]["text"], "be terse");
1071 assert_eq!(body["generationConfig"]["maxOutputTokens"], 256);
1072 assert_eq!(body["contents"][0]["role"], "user");
1074 assert_eq!(body["contents"][0]["parts"][0]["text"], "hi");
1075 assert_eq!(body["contents"][1]["role"], "model");
1076 assert_eq!(body["contents"][1]["parts"][0]["text"], "hello");
1077 }
1078
1079 #[test]
1080 fn gemini_response_parses_text_and_usage() {
1081 let json = serde_json::json!({
1082 "candidates": [{ "content": { "role": "model", "parts": [
1083 { "text": "the answer " }, { "text": "is 42" }
1084 ] } }],
1085 "usageMetadata": { "promptTokenCount": 11, "candidatesTokenCount": 4 }
1086 });
1087 let (text, in_tok, out_tok) = gemini_parse_response(&json).unwrap();
1088 assert_eq!(text, "the answer is 42");
1089 assert_eq!(in_tok, 11);
1090 assert_eq!(out_tok, 4);
1091 assert!(gemini_parse_response(&serde_json::json!({ "candidates": [] })).is_err());
1093 }
1094
1095 #[test]
1096 fn from_config_wires_the_gemini_dialect() {
1097 let defs = vec![firstpass_core::ProviderDef {
1098 id: "gemini".to_owned(),
1099 dialect: firstpass_core::Dialect::Gemini,
1100 base_url: "https://generativelanguage.googleapis.com".to_owned(),
1101 api_key_env: Some("GEMINI_API_KEY".to_owned()),
1102 auth: firstpass_core::AuthScheme::ApiKey,
1103 region: None,
1104 project: None,
1105 }];
1106 let reg = ProviderRegistry::from_config(
1107 &defs,
1108 "https://api.anthropic.com",
1109 "https://api.openai.com",
1110 );
1111 assert_eq!(reg.get("gemini").unwrap().id(), "gemini");
1112 }
1113
1114 #[test]
1115 fn from_config_registers_custom_providers_alongside_builtins() {
1116 let defs = vec![
1117 firstpass_core::ProviderDef {
1118 id: "groq".to_owned(),
1119 dialect: firstpass_core::Dialect::Openai,
1120 base_url: "https://api.groq.com/openai".to_owned(),
1121 api_key_env: Some("GROQ_API_KEY".to_owned()),
1122 auth: firstpass_core::AuthScheme::ApiKey,
1123 region: None,
1124 project: None,
1125 },
1126 firstpass_core::ProviderDef {
1128 id: "openai".to_owned(),
1129 dialect: firstpass_core::Dialect::Openai,
1130 base_url: "https://my-azure.openai.azure.com".to_owned(),
1131 api_key_env: Some("AZURE_OPENAI_KEY".to_owned()),
1132 auth: firstpass_core::AuthScheme::ApiKey,
1133 region: None,
1134 project: None,
1135 },
1136 ];
1137 let reg = ProviderRegistry::from_config(
1138 &defs,
1139 "https://api.anthropic.com",
1140 "https://api.openai.com",
1141 );
1142 assert_eq!(reg.get("anthropic").unwrap().id(), "anthropic");
1144 assert_eq!(reg.get("groq").unwrap().id(), "groq");
1145 assert!(reg.get("does-not-exist").is_none());
1147 }
1148
1149 #[test]
1150 fn anthropic_messages_body_omits_model_and_includes_system_only_when_set() {
1151 let req = ModelRequest {
1152 model: "bedrock/anthropic.claude-3-5-haiku".to_owned(),
1153 system: Some("be terse".to_owned()),
1154 messages: vec![
1155 ChatMessage::text("user", "hi"),
1156 ChatMessage {
1157 role: "assistant".to_owned(),
1158 content: serde_json::json!([{ "type": "text", "text": "hello" }]),
1159 },
1160 ],
1161 max_tokens: 128,
1162 tools: Value::Null,
1163 };
1164 let body = anthropic_messages_body(&req, "bedrock-2023-05-31");
1165 assert!(body.get("model").is_none());
1167 assert_eq!(body["anthropic_version"], "bedrock-2023-05-31");
1168 assert_eq!(body["max_tokens"], 128);
1169 assert_eq!(body["system"], "be terse");
1170 assert_eq!(body["messages"][0]["content"], serde_json::json!("hi"));
1171 assert_eq!(
1173 body["messages"][1]["content"],
1174 serde_json::json!([{ "type": "text", "text": "hello" }])
1175 );
1176
1177 let req_no_system = ModelRequest {
1179 system: None,
1180 ..req
1181 };
1182 let body2 = anthropic_messages_body(&req_no_system, "vertex-2023-10-16");
1183 assert!(body2.get("system").is_none());
1184 }
1185
1186 #[test]
1187 fn bedrock_url_construction_and_missing_region() {
1188 assert_eq!(
1189 bedrock_url("us-east-1", "anthropic.claude-3-5-haiku-20241022-v1:0"),
1190 "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-5-haiku-20241022-v1:0/invoke"
1191 );
1192 }
1193
1194 #[tokio::test]
1195 async fn bedrock_complete_errors_without_a_region() {
1196 let provider = BedrockProvider {
1197 id: "bedrock".to_owned(),
1198 region: None,
1199 http: reqwest::Client::new(),
1200 };
1201 let req = ModelRequest {
1202 model: "bedrock/anthropic.claude-3-5-haiku".to_owned(),
1203 system: None,
1204 messages: vec![],
1205 max_tokens: 16,
1206 tools: Value::Null,
1207 };
1208 let err = provider.complete(&req, &Auth::default()).await.unwrap_err();
1209 assert!(matches!(err, ProviderError::Transport(_)));
1210 }
1211
1212 #[test]
1213 fn bedrock_signing_produces_a_sigv4_authorization_header() {
1214 let creds = AwsEnvCredentials {
1217 access_key_id: "AKIDEXAMPLE".to_owned(),
1218 secret_access_key: "wJalrXUtnFEMI/K7MDENG+bPxRfiCYEXAMPLEKEY".to_owned(),
1219 session_token: None,
1220 };
1221 let url = bedrock_url("us-east-1", "anthropic.claude-3-5-haiku");
1222 let body = br#"{"anthropic_version":"bedrock-2023-05-31"}"#;
1223 let signed = sign_bedrock(&url, "us-east-1", body, &creds).unwrap();
1224
1225 let auth_header = signed
1226 .headers()
1227 .get("authorization")
1228 .and_then(|v| v.to_str().ok())
1229 .expect("authorization header present");
1230 assert!(auth_header.starts_with("AWS4-HMAC-SHA256"));
1231
1232 let host_header = signed
1233 .headers()
1234 .get("host")
1235 .and_then(|v| v.to_str().ok())
1236 .expect("host header present");
1237 assert_eq!(host_header, "bedrock-runtime.us-east-1.amazonaws.com");
1238 }
1239
1240 #[test]
1241 fn vertex_url_construction_and_missing_project() {
1242 assert_eq!(
1243 vertex_url("us-central1", "my-project", "claude-3-5-sonnet"),
1244 "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project/locations/us-central1/publishers/anthropic/models/claude-3-5-sonnet:rawPredict"
1245 );
1246 }
1247
1248 #[tokio::test]
1249 async fn vertex_complete_errors_without_a_project() {
1250 let provider = VertexProvider {
1251 id: "vertex".to_owned(),
1252 region: Some("us-central1".to_owned()),
1253 project: None,
1254 http: reqwest::Client::new(),
1255 };
1256 let req = ModelRequest {
1257 model: "vertex/claude-3-5-sonnet".to_owned(),
1258 system: None,
1259 messages: vec![],
1260 max_tokens: 16,
1261 tools: Value::Null,
1262 };
1263 let err = provider.complete(&req, &Auth::default()).await.unwrap_err();
1264 assert!(matches!(err, ProviderError::Transport(_)));
1265 }
1266
1267 #[test]
1268 fn from_config_wires_bedrock_and_vertex_auth_schemes() {
1269 let defs = vec![
1270 firstpass_core::ProviderDef {
1271 id: "bedrock".to_owned(),
1272 dialect: firstpass_core::Dialect::Anthropic,
1273 base_url: String::new(),
1274 api_key_env: None,
1275 auth: firstpass_core::AuthScheme::AwsSigv4,
1276 region: Some("us-east-1".to_owned()),
1277 project: None,
1278 },
1279 firstpass_core::ProviderDef {
1280 id: "vertex".to_owned(),
1281 dialect: firstpass_core::Dialect::Anthropic,
1282 base_url: String::new(),
1283 api_key_env: None,
1284 auth: firstpass_core::AuthScheme::GcpOauth,
1285 region: Some("us-central1".to_owned()),
1286 project: Some("my-project".to_owned()),
1287 },
1288 ];
1289 let reg = ProviderRegistry::from_config(
1290 &defs,
1291 "https://api.anthropic.com",
1292 "https://api.openai.com",
1293 );
1294 assert_eq!(reg.get("bedrock").unwrap().id(), "bedrock");
1295 assert_eq!(reg.get("vertex").unwrap().id(), "vertex");
1296 }
1297
1298 #[test]
1299 fn resolve_api_key_prefers_configured_env_then_byok() {
1300 let path = std::env::var("PATH").expect("PATH is set");
1302 assert_eq!(resolve_api_key(Some("PATH"), Some("byok")), path);
1304 assert_eq!(
1306 resolve_api_key(Some("FIRSTPASS_DEFINITELY_UNSET_KEY"), Some("byok")),
1307 "byok"
1308 );
1309 assert_eq!(resolve_api_key(None, Some("byok")), "byok");
1311 assert_eq!(resolve_api_key(None, None), "");
1313 }
1314
1315 #[test]
1316 fn anthropic_wire_forwards_tool_and_image_content_verbatim() {
1317 let messages = [
1321 ChatMessage::text("user", "hi"),
1322 ChatMessage {
1323 role: "assistant".to_owned(),
1324 content: serde_json::json!([
1325 { "type": "tool_use", "id": "t1", "name": "calc", "input": { "x": 1 } }
1326 ]),
1327 },
1328 ChatMessage {
1329 role: "user".to_owned(),
1330 content: serde_json::json!([
1331 { "type": "tool_result", "tool_use_id": "t1", "content": "2" },
1332 { "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": "AA==" } }
1333 ]),
1334 },
1335 ];
1336 let body = AnthropicWireRequest {
1337 model: "claude-haiku-4-5",
1338 system: None,
1339 max_tokens: 64,
1340 messages: messages
1341 .iter()
1342 .map(|m| AnthropicWireMessage {
1343 role: &m.role,
1344 content: &m.content,
1345 })
1346 .collect(),
1347 };
1348 let wire = serde_json::to_value(&body).unwrap();
1349 assert_eq!(wire["messages"][0]["content"], serde_json::json!("hi"));
1350 assert_eq!(
1351 wire["messages"][1]["content"],
1352 serde_json::json!([{ "type": "tool_use", "id": "t1", "name": "calc", "input": { "x": 1 } }])
1353 );
1354 assert_eq!(
1355 wire["messages"][2]["content"],
1356 serde_json::json!([
1357 { "type": "tool_result", "tool_use_id": "t1", "content": "2" },
1358 { "type": "image", "source": { "type": "base64", "media_type": "image/png", "data": "AA==" } }
1359 ])
1360 );
1361 }
1362
1363 fn resp(model: &str, text: &str) -> ModelResponse {
1364 ModelResponse {
1365 model: model.to_owned(),
1366 text: text.to_owned(),
1367 in_tokens: 10,
1368 out_tokens: 5,
1369 raw: Value::Null,
1370 }
1371 }
1372
1373 #[test]
1374 fn transport_and_5xx_are_failover_eligible() {
1375 assert!(ProviderError::Transport("boom".into()).is_failover_eligible());
1376 assert!(
1377 ProviderError::Http {
1378 status: 503,
1379 body: String::new()
1380 }
1381 .is_failover_eligible()
1382 );
1383 }
1384
1385 #[test]
1386 fn client_errors_and_decode_failures_are_hard() {
1387 assert!(
1388 !ProviderError::Http {
1389 status: 400,
1390 body: String::new()
1391 }
1392 .is_failover_eligible()
1393 );
1394 assert!(!ProviderError::Decode("bad json".into()).is_failover_eligible());
1395 }
1396
1397 #[test]
1398 fn auth_debug_never_prints_key_material() {
1399 let auth = Auth {
1400 anthropic_key: Some("sk-ant-super-secret".to_owned()),
1401 openai_key: Some("sk-oai-super-secret".to_owned()),
1402 };
1403 let debug = format!("{auth:?}");
1404 assert!(!debug.contains("super-secret"));
1405 }
1406
1407 #[tokio::test]
1408 async fn mock_provider_returns_configured_outcome() {
1409 let mut outcomes = HashMap::new();
1410 outcomes.insert(
1411 "anthropic/claude-haiku-4-5".to_owned(),
1412 Ok(resp("anthropic/claude-haiku-4-5", "hello")),
1413 );
1414 let provider = MockProvider::new("anthropic", outcomes);
1415 let req = ModelRequest {
1416 model: "anthropic/claude-haiku-4-5".to_owned(),
1417 system: None,
1418 messages: vec![],
1419 max_tokens: 100,
1420 tools: Value::Null,
1421 };
1422 let out = provider.complete(&req, &Auth::default()).await.unwrap();
1423 assert_eq!(out.text, "hello");
1424 }
1425}