Skip to main content

telos_agent/model/provider/deepseek/
client.rs

1use async_trait::async_trait;
2use futures_core::stream::Stream;
3use serde_json::Value;
4
5use crate::error::{AgentError, ProviderError};
6use crate::model::provider::{CompletionRequest, CompletionResponse, ModelProvider, ProviderEvent};
7
8use super::config::DeepSeekConfig;
9use super::error::{classify_reqwest_error, map_deepseek_http_error};
10use super::request::{build_chat_request, build_fim_request};
11use super::response::{parse_chat_response, parse_fim_response};
12use super::stream::stream_json_to;
13use super::types::{
14    DeepSeekBalance, DeepSeekChatOptions, DeepSeekFimRequest, DeepSeekFimResponse,
15    DeepSeekModelList,
16};
17use super::url::{normalize_beta_url, normalize_chat_url, normalize_v1_url};
18
19/// [`ModelProvider`] implementation backed by the DeepSeek API.
20pub struct DeepSeekProvider {
21    pub(super) client: reqwest::Client,
22    pub(super) api_key: String,
23    pub(super) model: String,
24    pub(super) chat_url: String,
25    pub(super) beta_chat_url: String,
26    pub(super) fim_url: String,
27    pub(super) models_url: String,
28    pub(super) balance_url: String,
29}
30
31impl DeepSeekProvider {
32    pub fn new(config: DeepSeekConfig) -> Self {
33        Self {
34            client: reqwest::Client::new(),
35            api_key: config.api_key,
36            model: config.model,
37            chat_url: normalize_chat_url(&config.base_url),
38            beta_chat_url: normalize_beta_url(&config.base_url, "chat/completions"),
39            fim_url: normalize_beta_url(&config.base_url, "completions"),
40            models_url: normalize_v1_url(&config.base_url, "models"),
41            balance_url: normalize_v1_url(&config.base_url, "user/balance"),
42        }
43    }
44
45    pub async fn complete_with_options(
46        &self,
47        request: CompletionRequest,
48        options: DeepSeekChatOptions,
49    ) -> Result<CompletionResponse, AgentError> {
50        let body = build_chat_request(&self.model, &request, false, &options);
51        let response = self.send_json_to(self.chat_url_for_options(&options), body).await?;
52        let value: Value = response
53            .json()
54            .await
55            .map_err(|err| AgentError::Provider(ProviderError::InvalidResponse(err.to_string())))?;
56        parse_chat_response(value)
57    }
58
59    pub fn stream_complete_with_options<'a>(
60        &'a self,
61        request: CompletionRequest,
62        options: DeepSeekChatOptions,
63    ) -> std::pin::Pin<Box<dyn Stream<Item = Result<ProviderEvent, AgentError>> + Send + 'a>> {
64        let body = build_chat_request(&self.model, &request, true, &options);
65        let url = self.chat_url_for_options(&options).to_string();
66        Box::pin(stream_json_to(&self.client, &self.api_key, url, body))
67    }
68
69    pub async fn fim_complete(
70        &self,
71        request: DeepSeekFimRequest,
72    ) -> Result<DeepSeekFimResponse, AgentError> {
73        let body = build_fim_request(&self.model, request);
74        let response = self.send_json_to(&self.fim_url, body).await?;
75        let value: Value = response
76            .json()
77            .await
78            .map_err(|err| AgentError::Provider(ProviderError::InvalidResponse(err.to_string())))?;
79        parse_fim_response(value)
80    }
81
82    pub async fn list_models(&self) -> Result<DeepSeekModelList, AgentError> {
83        self.send_get_json(&self.models_url).await
84    }
85
86    pub async fn balance(&self) -> Result<DeepSeekBalance, AgentError> {
87        self.send_get_json(&self.balance_url).await
88    }
89
90    async fn send_json(&self, body: Value) -> Result<reqwest::Response, AgentError> {
91        self.send_json_to(&self.chat_url, body).await
92    }
93
94    async fn send_json_to(&self, url: &str, body: Value) -> Result<reqwest::Response, AgentError> {
95        let response = self
96            .client
97            .post(url)
98            .bearer_auth(&self.api_key)
99            .header("User-Agent", concat!("telos_agent/", env!("CARGO_PKG_VERSION")))
100            .json(&body)
101            .send()
102            .await
103            .map_err(classify_reqwest_error)?;
104
105        if response.status().is_success() {
106            Ok(response)
107        } else {
108            Err(map_deepseek_http_error(response).await)
109        }
110    }
111
112    async fn send_get_json<T>(&self, url: &str) -> Result<T, AgentError>
113    where
114        T: serde::de::DeserializeOwned,
115    {
116        let response = self
117            .client
118            .get(url)
119            .bearer_auth(&self.api_key)
120            .header("User-Agent", concat!("telos_agent/", env!("CARGO_PKG_VERSION")))
121            .send()
122            .await
123            .map_err(classify_reqwest_error)?;
124
125        if !response.status().is_success() {
126            return Err(map_deepseek_http_error(response).await);
127        }
128
129        response
130            .json()
131            .await
132            .map_err(|err| AgentError::Provider(ProviderError::InvalidResponse(err.to_string())))
133    }
134
135    fn chat_url_for_options(&self, options: &DeepSeekChatOptions) -> &str {
136        if options.prefix.is_some() { &self.beta_chat_url } else { &self.chat_url }
137    }
138}
139
140#[async_trait]
141impl ModelProvider for DeepSeekProvider {
142    async fn complete(&self, request: CompletionRequest) -> Result<CompletionResponse, AgentError> {
143        let body =
144            build_chat_request(&self.model, &request, false, &DeepSeekChatOptions::default());
145        let response = self.send_json(body).await?;
146        let value: Value = response
147            .json()
148            .await
149            .map_err(|err| AgentError::Provider(ProviderError::InvalidResponse(err.to_string())))?;
150        parse_chat_response(value)
151    }
152
153    fn stream_complete<'a>(
154        &'a self,
155        request: CompletionRequest,
156    ) -> std::pin::Pin<Box<dyn Stream<Item = Result<ProviderEvent, AgentError>> + Send + 'a>> {
157        let body = build_chat_request(&self.model, &request, true, &DeepSeekChatOptions::default());
158        Box::pin(stream_json_to(&self.client, &self.api_key, self.chat_url.clone(), body))
159    }
160}