agentic_llm/
traits.rs

1//! Core traits and types for LLM adapters.
2
3use async_trait::async_trait;
4use futures::Stream;
5use std::pin::Pin;
6
7use crate::error::LLMError;
8
9/// Role in a conversation.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Role {
12    /// System message (instructions)
13    System,
14    /// User message
15    User,
16    /// Assistant response
17    Assistant,
18}
19
20/// A message in a conversation.
21#[derive(Debug, Clone)]
22pub struct LLMMessage {
23    /// Role of the message sender
24    pub role: Role,
25    /// Content of the message
26    pub content: String,
27}
28
29impl LLMMessage {
30    /// Create a system message.
31    #[must_use]
32    pub fn system(content: impl Into<String>) -> Self {
33        Self {
34            role: Role::System,
35            content: content.into(),
36        }
37    }
38
39    /// Create a user message.
40    #[must_use]
41    pub fn user(content: impl Into<String>) -> Self {
42        Self {
43            role: Role::User,
44            content: content.into(),
45        }
46    }
47
48    /// Create an assistant message.
49    #[must_use]
50    pub fn assistant(content: impl Into<String>) -> Self {
51        Self {
52            role: Role::Assistant,
53            content: content.into(),
54        }
55    }
56}
57
58/// Token usage statistics.
59#[derive(Debug, Clone, Default)]
60pub struct TokenUsage {
61    /// Tokens in the prompt
62    pub prompt: u32,
63    /// Tokens in the completion
64    pub completion: u32,
65    /// Total tokens used
66    pub total: u32,
67}
68
69/// Reason for completion finishing.
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub enum FinishReason {
72    /// Normal stop (end of response)
73    Stop,
74    /// Hit max tokens limit
75    Length,
76    /// Error occurred
77    Error,
78}
79
80/// Response from an LLM.
81#[derive(Debug, Clone)]
82pub struct LLMResponse {
83    /// Generated content
84    pub content: String,
85    /// Token usage statistics
86    pub tokens_used: TokenUsage,
87    /// Reason for finishing
88    pub finish_reason: FinishReason,
89    /// Model that generated the response
90    pub model: String,
91}
92
93/// A chunk from streaming response.
94#[derive(Debug, Clone)]
95pub struct StreamChunk {
96    /// Partial content
97    pub content: String,
98    /// Whether this is the final chunk
99    pub done: bool,
100    /// Token usage (only available on final chunk)
101    pub tokens_used: Option<TokenUsage>,
102    /// Finish reason (only available on final chunk)
103    pub finish_reason: Option<FinishReason>,
104}
105
106/// Trait for LLM adapters.
107///
108/// Implement this trait to add support for a new LLM provider.
109#[async_trait]
110pub trait LLMAdapter: Send + Sync {
111    /// Get the provider name (e.g., "openai", "anthropic").
112    fn provider(&self) -> &str;
113
114    /// Get the model name being used.
115    fn model(&self) -> &str;
116
117    /// Generate a completion from messages.
118    ///
119    /// # Errors
120    ///
121    /// Returns an error if the API call fails.
122    async fn generate(&self, messages: &[LLMMessage]) -> Result<LLMResponse, LLMError>;
123
124    /// Generate a streaming completion.
125    ///
126    /// Returns a stream of chunks that can be processed as they arrive.
127    fn generate_stream(
128        &self,
129        messages: &[LLMMessage],
130    ) -> Pin<Box<dyn Stream<Item = Result<StreamChunk, LLMError>> + Send + '_>>;
131
132    /// Check if the LLM provider is accessible.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error if the health check fails.
137    async fn health_check(&self) -> Result<bool, LLMError>;
138}
agentic_llm/traits.rs

agentic_llm/
traits.rs
