use std::collections::BTreeMap;
use salvor_llm::{Client, Config};
use salvor_tools::{DynTool, ToolHandler, ToolSet};
use serde_json::{Value, json};
use thiserror::Error;
use crate::budgets::{Budgets, Pricing};
use crate::hash::hash_value;
pub const DEFAULT_MAX_RESPONSE_TOKENS: u32 = 4096;
pub struct Agent {
client: Client,
model: String,
system_prompt: Option<String>,
tools: ToolSet,
budgets: Budgets,
pricing: Option<Pricing>,
max_response_tokens: u32,
def_hash: String,
record_prompts: bool,
labels: Option<BTreeMap<String, String>>,
name: Option<String>,
}
impl Agent {
#[must_use]
pub fn builder() -> AgentBuilder {
AgentBuilder::new()
}
#[must_use]
pub fn def_hash(&self) -> &str {
&self.def_hash
}
#[must_use]
pub fn client(&self) -> &Client {
&self.client
}
#[must_use]
pub fn model(&self) -> &str {
&self.model
}
#[must_use]
pub fn system_prompt(&self) -> Option<&str> {
self.system_prompt.as_deref()
}
#[must_use]
pub fn tools(&self) -> &ToolSet {
&self.tools
}
#[must_use]
pub fn budgets(&self) -> &Budgets {
&self.budgets
}
#[must_use]
pub fn pricing(&self) -> Option<&Pricing> {
self.pricing.as_ref()
}
#[must_use]
pub fn max_response_tokens(&self) -> u32 {
self.max_response_tokens
}
#[must_use]
pub fn record_prompts(&self) -> bool {
self.record_prompts
}
#[must_use]
pub fn labels(&self) -> Option<&BTreeMap<String, String>> {
self.labels.as_ref()
}
#[must_use]
pub fn name(&self) -> Option<&str> {
self.name.as_deref()
}
}
#[derive(Default)]
pub struct AgentBuilder {
client: Option<Client>,
config: Option<Config>,
model: Option<String>,
system_prompt: Option<String>,
tools: Vec<Box<dyn DynTool>>,
budgets: Budgets,
pricing: Option<Pricing>,
max_response_tokens: Option<u32>,
record_prompts: bool,
labels: Option<BTreeMap<String, String>>,
name: Option<String>,
}
impl AgentBuilder {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn model(mut self, config: Config, model: impl Into<String>) -> Self {
self.config = Some(config);
self.client = None;
self.model = Some(model.into());
self
}
#[must_use]
pub fn client(mut self, client: Client, model: impl Into<String>) -> Self {
self.client = Some(client);
self.config = None;
self.model = Some(model.into());
self
}
#[must_use]
pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
self.system_prompt = Some(prompt.into());
self
}
#[must_use]
pub fn tool<H: ToolHandler + 'static>(mut self, handler: H) -> Self {
self.tools
.push(Box::new(salvor_tools::TypedTool::new(handler)));
self
}
#[must_use]
pub fn tool_dyn(mut self, tool: Box<dyn DynTool>) -> Self {
self.tools.push(tool);
self
}
#[must_use]
pub fn budgets(mut self, budgets: Budgets) -> Self {
self.budgets = budgets;
self
}
#[must_use]
pub fn pricing(mut self, pricing: Pricing) -> Self {
self.pricing = Some(pricing);
self
}
#[must_use]
pub fn max_response_tokens(mut self, max_tokens: u32) -> Self {
self.max_response_tokens = Some(max_tokens);
self
}
#[must_use]
pub fn record_prompts(mut self, record_prompts: bool) -> Self {
self.record_prompts = record_prompts;
self
}
#[must_use]
pub fn labels(mut self, labels: BTreeMap<String, String>) -> Self {
self.labels = Some(labels);
self
}
#[must_use]
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = Some(name.into());
self
}
pub fn build(self) -> Result<Agent, AgentBuildError> {
let model = self.model.ok_or(AgentBuildError::MissingModel)?;
let client = match (self.client, self.config) {
(Some(client), _) => client,
(None, Some(config)) => Client::new(config).map_err(AgentBuildError::Client)?,
(None, None) => return Err(AgentBuildError::MissingModel),
};
if self.budgets.max_cost_usd.is_some() && self.pricing.is_none() {
return Err(AgentBuildError::CostBudgetWithoutPricing);
}
let mut tools = ToolSet::new();
for tool in self.tools {
let name = tool.name().to_owned();
tools
.register_dyn(tool)
.map_err(|_| AgentBuildError::DuplicateTool { name })?;
}
let def_hash = compute_def_hash(
&model,
self.system_prompt.as_deref(),
&tools,
&self.budgets,
self.pricing.as_ref(),
);
Ok(Agent {
client,
model,
system_prompt: self.system_prompt,
tools,
budgets: self.budgets,
pricing: self.pricing,
max_response_tokens: self
.max_response_tokens
.unwrap_or(DEFAULT_MAX_RESPONSE_TOKENS),
def_hash,
record_prompts: self.record_prompts,
labels: self.labels,
name: self.name,
})
}
}
#[derive(Debug, Error)]
pub enum AgentBuildError {
#[error("an agent needs a model: call .model(config, id) or .client(client, id)")]
MissingModel,
#[error("max_cost_usd is declared but no pricing is set; call .pricing(Pricing {{ .. }})")]
CostBudgetWithoutPricing,
#[error("a tool named `{name}` is registered twice")]
DuplicateTool {
name: String,
},
#[error("client construction failed: {0}")]
Client(salvor_llm::Error),
}
fn compute_def_hash(
model: &str,
system_prompt: Option<&str>,
tools: &ToolSet,
budgets: &Budgets,
pricing: Option<&Pricing>,
) -> String {
let tool_values: Vec<Value> = tools
.descriptors()
.into_iter()
.map(|descriptor| {
json!({
"description": descriptor.description,
"effect": descriptor.effect,
"input_schema": descriptor.input_schema,
"name": descriptor.name,
})
})
.collect();
let value = json!({
"budgets": {
"max_cost_usd": budgets.max_cost_usd,
"max_steps": budgets.max_steps,
"max_tokens": budgets.max_tokens,
"max_wall_time_seconds": budgets.max_wall_time.map(|d| d.as_secs_f64()),
},
"model": model,
"pricing": pricing.map(|p| {
json!({"input_per_mtok": p.input_per_mtok, "output_per_mtok": p.output_per_mtok})
}),
"system_prompt": system_prompt,
"tools": tool_values,
});
hash_value(&value)
}
#[cfg(test)]
mod tests {
use super::*;
use salvor_core::Effect;
use salvor_tools::{ToolCtx, ToolError, ToolOutcome};
use serde_json::json;
struct StubTool {
name: &'static str,
description: &'static str,
}
#[async_trait::async_trait]
impl DynTool for StubTool {
fn name(&self) -> &str {
self.name
}
fn description(&self) -> &str {
self.description
}
fn effect(&self) -> Effect {
Effect::Read
}
fn input_schema(&self) -> Value {
json!({"type": "object"})
}
async fn call_json(
&self,
_ctx: &ToolCtx,
input: Value,
) -> Result<ToolOutcome<Value>, ToolError> {
Ok(ToolOutcome::Output(input))
}
}
fn base_builder() -> AgentBuilder {
Agent::builder()
.model(Config::new(), "test-model")
.system_prompt("prompt")
.tool_dyn(Box::new(StubTool {
name: "alpha",
description: "first",
}))
}
#[test]
fn identical_definitions_share_a_hash() {
let a = Agent::builder()
.model(Config::new(), "test-model")
.tool_dyn(Box::new(StubTool {
name: "alpha",
description: "first",
}))
.tool_dyn(Box::new(StubTool {
name: "beta",
description: "second",
}))
.build()
.unwrap();
let b = Agent::builder()
.model(Config::new(), "test-model")
.tool_dyn(Box::new(StubTool {
name: "beta",
description: "second",
}))
.tool_dyn(Box::new(StubTool {
name: "alpha",
description: "first",
}))
.build()
.unwrap();
assert_eq!(a.def_hash(), b.def_hash());
assert!(a.def_hash().starts_with("sha256:"));
}
#[test]
fn any_definition_change_changes_the_hash() {
let base = base_builder().build().unwrap();
let model_changed = base_builder();
let model_changed = AgentBuilder {
model: Some("other-model".to_owned()),
..model_changed
}
.build()
.unwrap();
assert_ne!(base.def_hash(), model_changed.def_hash());
let prompt_changed = base_builder().system_prompt("different").build().unwrap();
assert_ne!(base.def_hash(), prompt_changed.def_hash());
let tool_changed = Agent::builder()
.model(Config::new(), "test-model")
.system_prompt("prompt")
.tool_dyn(Box::new(StubTool {
name: "alpha",
description: "changed description",
}))
.build()
.unwrap();
assert_ne!(base.def_hash(), tool_changed.def_hash());
let budget_changed = base_builder()
.budgets(Budgets {
max_steps: Some(10),
..Budgets::default()
})
.build()
.unwrap();
assert_ne!(base.def_hash(), budget_changed.def_hash());
let pricing_changed = base_builder()
.pricing(Pricing {
input_per_mtok: 3.0,
output_per_mtok: 15.0,
})
.build()
.unwrap();
assert_ne!(base.def_hash(), pricing_changed.def_hash());
}
#[test]
fn labels_never_affect_the_definition_hash() {
let unlabeled = base_builder().build().unwrap();
let labeled_a = base_builder()
.labels(BTreeMap::from([("build".to_owned(), "42".to_owned())]))
.build()
.unwrap();
let labeled_b = base_builder()
.labels(BTreeMap::from([
("build".to_owned(), "43".to_owned()),
("env".to_owned(), "staging".to_owned()),
]))
.build()
.unwrap();
assert_eq!(unlabeled.def_hash(), labeled_a.def_hash());
assert_eq!(unlabeled.def_hash(), labeled_b.def_hash());
assert_eq!(unlabeled.labels(), None);
assert_eq!(
labeled_a.labels(),
Some(&BTreeMap::from([("build".to_owned(), "42".to_owned())]))
);
}
#[test]
fn name_never_affects_the_definition_hash() {
let unnamed = base_builder().build().unwrap();
let named_a = base_builder().name("triage-agent").build().unwrap();
let named_b = base_builder().name("a-different-name").build().unwrap();
assert_eq!(unnamed.def_hash(), named_a.def_hash());
assert_eq!(unnamed.def_hash(), named_b.def_hash());
assert_eq!(unnamed.name(), None);
assert_eq!(named_a.name(), Some("triage-agent"));
}
#[test]
fn cost_budget_without_pricing_fails_to_build() {
let result = base_builder()
.budgets(Budgets {
max_cost_usd: Some(2.0),
..Budgets::default()
})
.build();
assert!(matches!(
result,
Err(AgentBuildError::CostBudgetWithoutPricing)
));
}
#[test]
fn duplicate_tools_and_missing_model_fail_to_build() {
let duplicate = base_builder()
.tool_dyn(Box::new(StubTool {
name: "alpha",
description: "again",
}))
.build();
assert!(matches!(
duplicate,
Err(AgentBuildError::DuplicateTool { name }) if name == "alpha"
));
assert!(matches!(
Agent::builder().build(),
Err(AgentBuildError::MissingModel)
));
}
}