84 lines
2.3 KiB
Rust
84 lines
2.3 KiB
Rust
//! Simple agent system for the Gemini SDK
|
|
|
|
use serde::{Deserialize, Serialize};
|
|
|
|
/// A simple AI agent with a system prompt and configuration
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct Agent {
|
|
/// Unique identifier for the agent
|
|
pub id: String,
|
|
/// Human-readable name for the agent
|
|
pub name: String,
|
|
/// System prompt that defines the agent's behavior
|
|
pub system_prompt: String,
|
|
/// Temperature for text generation (0.0 to 2.0)
|
|
pub temperature: f32,
|
|
/// Maximum number of tokens to generate
|
|
pub max_tokens: u32,
|
|
}
|
|
|
|
impl Agent {
|
|
/// Create a new agent
|
|
pub fn new(
|
|
id: impl Into<String>,
|
|
name: impl Into<String>,
|
|
system_prompt: impl Into<String>,
|
|
) -> Self {
|
|
Self {
|
|
id: id.into(),
|
|
name: name.into(),
|
|
system_prompt: system_prompt.into(),
|
|
temperature: 0.7,
|
|
max_tokens: 8192,
|
|
}
|
|
}
|
|
|
|
/// Set the temperature
|
|
pub fn with_temperature(mut self, temperature: f32) -> Self {
|
|
self.temperature = temperature.clamp(0.0, 2.0);
|
|
self
|
|
}
|
|
|
|
/// Set the max tokens
|
|
pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
|
|
self.max_tokens = max_tokens;
|
|
self
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_agent_creation() {
|
|
let agent = Agent::new("test", "Test Agent", "You are a test agent");
|
|
assert_eq!(agent.id, "test");
|
|
assert_eq!(agent.name, "Test Agent");
|
|
assert_eq!(agent.system_prompt, "You are a test agent");
|
|
assert_eq!(agent.temperature, 0.7);
|
|
assert_eq!(agent.max_tokens, 8192);
|
|
}
|
|
|
|
#[test]
|
|
fn test_agent_configuration() {
|
|
let agent = Agent::new("test", "Test", "Test")
|
|
.with_temperature(1.0)
|
|
.with_max_tokens(4096);
|
|
|
|
assert_eq!(agent.temperature, 1.0);
|
|
assert_eq!(agent.max_tokens, 4096);
|
|
}
|
|
|
|
#[test]
|
|
fn test_temperature_clamping() {
|
|
let agent = Agent::new("test", "Test", "Test")
|
|
.with_temperature(3.0); // Should be clamped to 2.0
|
|
assert_eq!(agent.temperature, 2.0);
|
|
|
|
let agent = Agent::new("test", "Test", "Test")
|
|
.with_temperature(-1.0); // Should be clamped to 0.0
|
|
assert_eq!(agent.temperature, 0.0);
|
|
}
|
|
}
|