alith_client/
basic_completion.rs

1use alith_interface::{
2    llms::LLMBackend,
3    requests::{
4        completion::{CompletionRequest, CompletionResponse},
5        logit_bias::{LogitBias, LogitBiasTrait},
6        req_components::{RequestConfig, RequestConfigTrait},
7    },
8};
9use alith_prompt::LLMPrompt;
10use std::sync::Arc;
11
12#[derive(Clone)]
13pub struct BasicCompletion {
14    pub base_req: CompletionRequest,
15}
16
17impl BasicCompletion {
18    #[inline]
19    pub fn new(backend: Arc<LLMBackend>) -> Self {
20        Self {
21            base_req: CompletionRequest::new(backend),
22        }
23    }
24
25    #[inline]
26    pub fn prompt(&mut self) -> &mut LLMPrompt {
27        &mut self.base_req.prompt
28    }
29
30    #[inline]
31    pub async fn run(&mut self) -> crate::Result<CompletionResponse> {
32        Ok(self.base_req.request().await?)
33    }
34
35    pub fn parse_response(&self, content: &str) -> crate::Result<String> {
36        if content.is_empty() {
37            return Err(anyhow::format_err!(
38                "parse_response error: content.is_empty()"
39            ));
40        }
41
42        let content = content
43            .strip_prefix("assistant\n\n")
44            .or_else(|| content.strip_prefix("assistant\n"))
45            .or_else(|| content.strip_prefix("assistant"))
46            .unwrap_or(content);
47
48        Ok(content.trim().to_owned())
49    }
50}
51
52impl RequestConfigTrait for BasicCompletion {
53    fn config(&mut self) -> &mut RequestConfig {
54        &mut self.base_req.config
55    }
56
57    fn reset_request(&mut self) {
58        self.base_req.reset_completion_request();
59    }
60}
61
62impl LogitBiasTrait for BasicCompletion {
63    fn lb_mut(&mut self) -> &mut Option<LogitBias> {
64        &mut self.base_req.logit_bias
65    }
66}