alith_client/
basic_completion.rs1use 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}