autoagents_llm/completion/
mod.rs

1use async_trait::async_trait;
2
3use crate::{chat::ChatResponse, error::LLMError, ToolCall};
4
5/// A request for text completion from an LLM provider.
6#[derive(Debug, Clone)]
7pub struct CompletionRequest {
8    /// The input prompt text to complete
9    pub prompt: String,
10    /// Optional maximum number of tokens to generate
11    pub max_tokens: Option<u32>,
12    /// Optional temperature parameter to control randomness (0.0-1.0)
13    pub temperature: Option<f32>,
14}
15
16/// A response containing generated text from a completion request.
17#[derive(Debug, Clone)]
18pub struct CompletionResponse {
19    /// The generated completion text
20    pub text: String,
21}
22
23impl ChatResponse for CompletionResponse {
24    fn text(&self) -> Option<String> {
25        Some(self.text.clone())
26    }
27
28    fn tool_calls(&self) -> Option<Vec<ToolCall>> {
29        None
30    }
31}
32
33impl CompletionRequest {
34    /// Creates a new completion request with just a prompt.
35    ///
36    /// # Arguments
37    ///
38    /// * `prompt` - The input text to complete
39    pub fn new(prompt: impl Into<String>) -> Self {
40        Self {
41            prompt: prompt.into(),
42            max_tokens: None,
43            temperature: None,
44        }
45    }
46
47    /// Creates a builder for constructing a completion request.
48    ///
49    /// # Arguments
50    ///
51    /// * `prompt` - The input text to complete
52    pub fn builder(prompt: impl Into<String>) -> CompletionRequestBuilder {
53        CompletionRequestBuilder {
54            prompt: prompt.into(),
55            max_tokens: None,
56            temperature: None,
57        }
58    }
59}
60
61/// Builder for constructing completion requests with optional parameters.
62#[derive(Debug, Clone)]
63pub struct CompletionRequestBuilder {
64    /// The input prompt text to complete
65    pub prompt: String,
66    /// Optional maximum number of tokens to generate
67    pub max_tokens: Option<u32>,
68    /// Optional temperature parameter to control randomness (0.0-1.0)
69    pub temperature: Option<f32>,
70}
71
72impl CompletionRequestBuilder {
73    /// Sets the maximum number of tokens to generate.
74    pub fn max_tokens(mut self, val: u32) -> Self {
75        self.max_tokens = Some(val);
76        self
77    }
78
79    /// Sets the temperature parameter for controlling randomness.
80    pub fn temperature(mut self, val: f32) -> Self {
81        self.temperature = Some(val);
82        self
83    }
84
85    /// Builds the completion request with the configured parameters.
86    pub fn build(self) -> CompletionRequest {
87        CompletionRequest {
88            prompt: self.prompt,
89            max_tokens: self.max_tokens,
90            temperature: self.temperature,
91        }
92    }
93}
94
95/// Trait for providers that support text completion requests.
96#[async_trait]
97pub trait CompletionProvider {
98    /// Sends a completion request to generate text.
99    ///
100    /// # Arguments
101    ///
102    /// * `req` - The completion request parameters
103    ///
104    /// # Returns
105    ///
106    /// The generated completion text or an error
107    async fn complete(&self, req: &CompletionRequest) -> Result<CompletionResponse, LLMError>;
108}
109
110impl std::fmt::Display for CompletionResponse {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        write!(f, "{}", self.text)
113    }
114}