crabtalk_model/
provider.rs

1//! Provider implementation backed by crabllm-provider.
2//!
3//! Wraps `crabllm_provider::Provider` behind wcore's `Model` trait with
4//! type conversion and retry logic.
5
6use crate::{config::ProviderDef, convert};
7use anyhow::Result;
8use async_stream::try_stream;
9use crabllm_core::ApiError;
10use crabllm_provider::Provider as CtProvider;
11use futures_core::Stream;
12use futures_util::StreamExt;
13use rand::Rng;
14use std::time::Duration;
15use wcore::model::{Model, Response, StreamChunk};
16
17/// Unified LLM provider wrapping a crabtalk provider instance.
18#[derive(Clone)]
19pub struct Provider {
20    inner: CtProvider,
21    client: reqwest::Client,
22    model: String,
23    max_retries: u32,
24    timeout: Duration,
25}
26
27/// Strip known endpoint suffixes so both bare origins and full paths work.
28fn normalize_base_url(url: &str) -> String {
29    let url = url.trim_end_matches('/');
30    for suffix in ["/chat/completions", "/messages", "/embeddings"] {
31        if let Some(stripped) = url.strip_suffix(suffix) {
32            return stripped.to_string();
33        }
34    }
35    url.to_string()
36}
37
38/// Construct a `Provider` from a provider definition and model name.
39pub fn build_provider(def: &ProviderDef, model: &str, client: reqwest::Client) -> Result<Provider> {
40    let mut config = def.clone();
41    config.kind = config.effective_kind();
42    let mut inner = CtProvider::from(&config);
43
44    // Apply crabtalk-specific base_url normalization (strip endpoint suffixes).
45    if let CtProvider::Openai {
46        ref mut base_url, ..
47    } = inner
48    {
49        *base_url = normalize_base_url(base_url);
50    }
51
52    Ok(Provider {
53        inner,
54        client,
55        model: model.to_owned(),
56        max_retries: def.max_retries.unwrap_or(2),
57        timeout: Duration::from_secs(def.timeout.unwrap_or(30)),
58    })
59}
60
61impl Model for Provider {
62    async fn send(&self, request: &wcore::model::Request) -> Result<Response> {
63        let mut ct_req = convert::to_ct_request(request);
64        ct_req.stream = Some(false);
65        send_with_retry(
66            &self.inner,
67            &self.client,
68            &ct_req,
69            self.max_retries,
70            self.timeout,
71        )
72        .await
73    }
74
75    fn stream(
76        &self,
77        request: wcore::model::Request,
78    ) -> impl Stream<Item = Result<StreamChunk>> + Send {
79        let inner = self.inner.clone();
80        let client = self.client.clone();
81        let timeout = self.timeout;
82        try_stream! {
83            let mut ct_req = convert::to_ct_request(&request);
84            ct_req.stream = Some(true);
85
86            let boxed = tokio::time::timeout(timeout, inner.chat_completion_stream(&client, &ct_req))
87                .await
88                .map_err(|_| anyhow::anyhow!("stream connection timed out"))?
89                .map_err(format_provider_error)?;
90
91            let mut stream = std::pin::pin!(boxed);
92            while let Some(chunk) = stream.next().await {
93                let ct_chunk = chunk.map_err(format_provider_error)?;
94                yield convert::from_ct_chunk(ct_chunk);
95            }
96        }
97    }
98
99    fn context_limit(&self, model: &str) -> usize {
100        wcore::model::default_context_limit(model)
101    }
102
103    fn active_model(&self) -> String {
104        self.model.clone()
105    }
106}
107
108/// Send a non-streaming request with exponential backoff retry on transient errors.
109async fn send_with_retry(
110    provider: &CtProvider,
111    client: &reqwest::Client,
112    request: &crabllm_core::ChatCompletionRequest,
113    max_retries: u32,
114    timeout: Duration,
115) -> Result<Response> {
116    let mut backoff = Duration::from_millis(100);
117    let mut last_err = None;
118
119    for _ in 0..=max_retries {
120        let result = if timeout.is_zero() {
121            provider.chat_completion(client, request).await
122        } else {
123            tokio::time::timeout(timeout, provider.chat_completion(client, request))
124                .await
125                .map_err(|_| crabllm_core::Error::Timeout)?
126        };
127
128        match result {
129            Ok(resp) => return Ok(convert::from_ct_response(resp)),
130            Err(e) if e.is_transient() => {
131                last_err = Some(e);
132                let jitter = jittered(backoff);
133                tokio::time::sleep(jitter).await;
134                backoff *= 2;
135            }
136            Err(e) => return Err(format_provider_error(e)),
137        }
138    }
139
140    Err(format_provider_error(last_err.unwrap()))
141}
142
143/// Full jitter: random duration in [backoff/2, backoff].
144fn jittered(backoff: Duration) -> Duration {
145    let lo = backoff.as_millis() as u64 / 2;
146    let hi = backoff.as_millis() as u64;
147    if lo >= hi {
148        return backoff;
149    }
150    Duration::from_millis(rand::rng().random_range(lo..=hi))
151}
152
153/// Convert a crabllm error into an anyhow error with a human-readable message.
154///
155/// For provider HTTP errors, attempts to parse the response body as an
156/// OpenAI-compatible API error and extract the `message` field.
157fn format_provider_error(e: crabllm_core::Error) -> anyhow::Error {
158    match e {
159        crabllm_core::Error::Provider { status, body } => {
160            let msg = serde_json::from_str::<ApiError>(&body)
161                .map(|api_err| api_err.error.message)
162                .unwrap_or_else(|_| truncate(&body, 200));
163            anyhow::anyhow!("provider error (HTTP {status}): {msg}")
164        }
165        other => anyhow::anyhow!("{other}"),
166    }
167}
168
169fn truncate(s: &str, max: usize) -> String {
170    match s.char_indices().nth(max) {
171        Some((i, _)) => format!("{}...", &s[..i]),
172        None => s.to_string(),
173    }
174}
crabtalk_model/provider.rs

crabtalk_model/
provider.rs
