hehe_agent/
config.rs

1use serde::{Deserialize, Serialize};
2use std::time::Duration;
3
4#[derive(Clone, Debug, Serialize, Deserialize)]
5pub struct AgentConfig {
6    #[serde(default = "default_name")]
7    pub name: String,
8
9    pub system_prompt: String,
10
11    pub model: String,
12
13    #[serde(default = "default_temperature")]
14    pub temperature: f32,
15
16    #[serde(default = "default_max_tokens")]
17    pub max_tokens: Option<usize>,
18
19    #[serde(default = "default_max_iterations")]
20    pub max_iterations: usize,
21
22    #[serde(default = "default_max_context_messages")]
23    pub max_context_messages: usize,
24
25    #[serde(default = "default_tool_timeout_secs")]
26    pub tool_timeout_secs: u64,
27
28    #[serde(default = "default_tools_enabled")]
29    pub tools_enabled: bool,
30}
31
32fn default_name() -> String {
33    "assistant".to_string()
34}
35
36fn default_temperature() -> f32 {
37    0.7
38}
39
40fn default_max_tokens() -> Option<usize> {
41    None
42}
43
44fn default_max_iterations() -> usize {
45    10
46}
47
48fn default_max_context_messages() -> usize {
49    50
50}
51
52fn default_tool_timeout_secs() -> u64 {
53    60
54}
55
56fn default_tools_enabled() -> bool {
57    true
58}
59
60impl Default for AgentConfig {
61    fn default() -> Self {
62        Self {
63            name: default_name(),
64            system_prompt: String::new(),
65            model: "gpt-4o".to_string(),
66            temperature: default_temperature(),
67            max_tokens: default_max_tokens(),
68            max_iterations: default_max_iterations(),
69            max_context_messages: default_max_context_messages(),
70            tool_timeout_secs: default_tool_timeout_secs(),
71            tools_enabled: default_tools_enabled(),
72        }
73    }
74}
75
76impl AgentConfig {
77    pub fn new(model: impl Into<String>, system_prompt: impl Into<String>) -> Self {
78        Self {
79            model: model.into(),
80            system_prompt: system_prompt.into(),
81            ..Default::default()
82        }
83    }
84
85    pub fn with_name(mut self, name: impl Into<String>) -> Self {
86        self.name = name.into();
87        self
88    }
89
90    pub fn with_temperature(mut self, temperature: f32) -> Self {
91        self.temperature = temperature;
92        self
93    }
94
95    pub fn with_max_tokens(mut self, max_tokens: usize) -> Self {
96        self.max_tokens = Some(max_tokens);
97        self
98    }
99
100    pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
101        self.max_iterations = max_iterations;
102        self
103    }
104
105    pub fn with_max_context_messages(mut self, max_messages: usize) -> Self {
106        self.max_context_messages = max_messages;
107        self
108    }
109
110    pub fn with_tool_timeout(mut self, timeout: Duration) -> Self {
111        self.tool_timeout_secs = timeout.as_secs();
112        self
113    }
114
115    pub fn with_tools_enabled(mut self, enabled: bool) -> Self {
116        self.tools_enabled = enabled;
117        self
118    }
119
120    pub fn tool_timeout(&self) -> Duration {
121        Duration::from_secs(self.tool_timeout_secs)
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn test_config_default() {
131        let config = AgentConfig::default();
132        assert_eq!(config.name, "assistant");
133        assert_eq!(config.max_iterations, 10);
134        assert!(config.tools_enabled);
135    }
136
137    #[test]
138    fn test_config_builder() {
139        let config = AgentConfig::new("gpt-4o", "You are helpful.")
140            .with_name("my-agent")
141            .with_temperature(0.5)
142            .with_max_iterations(5);
143
144        assert_eq!(config.name, "my-agent");
145        assert_eq!(config.model, "gpt-4o");
146        assert_eq!(config.temperature, 0.5);
147        assert_eq!(config.max_iterations, 5);
148    }
149}
hehe_agent/config.rs

hehe_agent/
config.rs
