agent_runtime/llm/
types.rs

1use serde::{Deserialize, Serialize};
2
3#[cfg(test)]
4#[path = "types_test.rs"]
5mod types_test;
6
7/// Role of a message in a conversation
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
9#[serde(rename_all = "lowercase")]
10pub enum Role {
11    System,
12    User,
13    Assistant,
14}
15
16/// A single message in a chat conversation
17#[derive(Debug, Clone, Serialize, Deserialize)]
18pub struct ChatMessage {
19    pub role: Role,
20    pub content: String,
21}
22
23impl ChatMessage {
24    pub fn system(content: impl Into<String>) -> Self {
25        Self {
26            role: Role::System,
27            content: content.into(),
28        }
29    }
30
31    pub fn user(content: impl Into<String>) -> Self {
32        Self {
33            role: Role::User,
34            content: content.into(),
35        }
36    }
37
38    pub fn assistant(content: impl Into<String>) -> Self {
39        Self {
40            role: Role::Assistant,
41            content: content.into(),
42        }
43    }
44}
45
46/// Request for chat completion
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct ChatRequest {
49    pub messages: Vec<ChatMessage>,
50
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub temperature: Option<f32>,
53
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub max_tokens: Option<u32>,
56
57    #[serde(skip_serializing_if = "Option::is_none")]
58    pub top_p: Option<f32>,
59}
60
61impl ChatRequest {
62    pub fn new(messages: Vec<ChatMessage>) -> Self {
63        Self {
64            messages,
65            temperature: None,
66            max_tokens: None,
67            top_p: None,
68        }
69    }
70
71    pub fn with_temperature(mut self, temperature: f32) -> Self {
72        self.temperature = Some(temperature);
73        self
74    }
75
76    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
77        self.max_tokens = Some(max_tokens);
78        self
79    }
80
81    pub fn with_top_p(mut self, top_p: f32) -> Self {
82        self.top_p = Some(top_p);
83        self
84    }
85}
86
87/// Response from chat completion
88#[derive(Debug, Clone, Serialize, Deserialize)]
89pub struct ChatResponse {
90    pub content: String,
91    pub model: String,
92
93    #[serde(skip_serializing_if = "Option::is_none")]
94    pub usage: Option<Usage>,
95
96    #[serde(skip_serializing_if = "Option::is_none")]
97    pub finish_reason: Option<String>,
98}
99
100/// Token usage statistics
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct Usage {
103    pub prompt_tokens: u32,
104    pub completion_tokens: u32,
105    pub total_tokens: u32,
106}
agent_runtime/llm/types.rs

agent_runtime/llm/
types.rs
