use crate::types::ModelOpsError;
pub struct VllmProvider {
pub endpoint: String,
timeout: std::time::Duration,
}
impl VllmProvider {
pub fn new(endpoint: &str) -> Self {
Self {
endpoint: endpoint.to_string(),
timeout: std::time::Duration::from_secs(30),
}
}
pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = timeout;
self
}
pub async fn complete(&self, prompt: &str) -> Result<String, ModelOpsError> {
let client = reqwest::Client::builder()
.timeout(self.timeout)
.build()
.map_err(|e| ModelOpsError::InferenceFailed(format!("客户端构建失败: {e}")))?;
let url = format!("{}/v1/completions", self.endpoint);
let body = serde_json::json!({
"prompt": prompt,
"max_tokens": 256,
"temperature": 0.7,
});
let response = client
.post(&url)
.json(&body)
.send()
.await
.map_err(|e| ModelOpsError::InferenceFailed(format!("连接 {url} 失败: {e}")))?;
let json: serde_json::Value = response
.json()
.await
.map_err(|e| ModelOpsError::InferenceFailed(format!("解析响应失败: {e}")))?;
json["choices"][0]["text"]
.as_str()
.map(|s| s.to_string())
.ok_or_else(|| {
ModelOpsError::InferenceFailed("响应格式无效: 缺少 choices[0].text".into())
})
}
}