Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 39 additions & 3 deletions rust/src/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ pub struct AgentLoopConfig {
pub environment_tokens: SampleSpec,
pub tool_invocations: SampleSpec,
pub tool_call_latency_ms: Option<SampleSpec>,
pub user_prefix: Option<String>,
pub request_timeout: Duration,
pub max_retries: usize,
pub retry_delay: Duration,
Expand Down Expand Up @@ -167,6 +168,7 @@ impl AgentLoopConfig {
environment_tokens,
tool_invocations,
tool_call_latency_ms: None,
user_prefix: None,
request_timeout: Duration::from_secs(60),
max_retries: 2,
retry_delay: Duration::from_millis(250),
Expand Down Expand Up @@ -218,6 +220,11 @@ impl AgentLoopConfig {
self
}

pub fn with_user_prefix(mut self, user_prefix: impl Into<String>) -> Self {
self.user_prefix = Some(user_prefix.into());
self
}
Comment on lines +223 to +226

pub fn with_tool_call_latency_ms(mut self, latency_ms: SampleSpec) -> Result<Self> {
latency_ms.validate("tool call latency")?;
self.tool_call_latency_ms = Some(latency_ms);
Expand Down Expand Up @@ -558,7 +565,7 @@ async fn run_agent(
continue;
}

let body = build_request_body(&config, &messages, turn.output_tokens);
let body = build_request_body(&config, &messages, turn.output_tokens, plan.agent_id);
match request_with_retries(&client, &config, &body).await {
Ok(result) => {
report.successful_requests += 1;
Expand Down Expand Up @@ -606,7 +613,12 @@ async fn run_agent(
Ok(report)
}

fn build_request_body(config: &AgentLoopConfig, messages: &[Value], output_tokens: usize) -> Value {
fn build_request_body(
config: &AgentLoopConfig,
messages: &[Value],
output_tokens: usize,
agent_id: usize,
) -> Value {
let mut body = json!({
"model": config.model,
"messages": messages,
Expand Down Expand Up @@ -638,6 +650,9 @@ fn build_request_body(config: &AgentLoopConfig, messages: &[Value], output_token
if let Some(map) = body.as_object_mut() {
map.insert(max_key.to_string(), json!(output_tokens));
map.insert(min_key.to_string(), json!(output_tokens));
if let Some(prefix) = &config.user_prefix {
map.insert("user".to_string(), json!(format!("{prefix}-{agent_id}")));
}
}
body
}
Expand Down Expand Up @@ -1022,13 +1037,34 @@ mod tests {
json!({"role": "tool", "tool_call_id": "call_1", "content": "result"}),
];

let body = build_request_body(&config, &messages, 17);
let body = build_request_body(&config, &messages, 17, 0);
assert_eq!(body["messages"], json!(messages));
assert_eq!(body["max_tokens"], 17);
assert_eq!(body["min_tokens"], 17);
assert_eq!(
body["tool_choice"]["function"]["name"],
Value::String("environment".to_string())
);
assert!(body.get("user").is_none());
}

#[test]
fn user_prefix_stamps_a_per_agent_user_field() {
let config = AgentLoopConfig::try_new(
"http://localhost:8000/v1/chat/completions",
None,
"test-model",
2,
SampleSpec::fixed(8).unwrap(),
SampleSpec::fixed(4).unwrap(),
SampleSpec::fixed(6).unwrap(),
SampleSpec::fixed(2).unwrap(),
)
.unwrap()
.with_user_prefix("loadtest");
let messages = vec![json!({"role": "user", "content": "start"})];

let body = build_request_body(&config, &messages, 5, 1);
assert_eq!(body["user"], Value::String("loadtest-1".to_string()));
}
}
8 changes: 8 additions & 0 deletions rust/src/agent_cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,11 @@ struct Args {
)]
tool_call_latency_lognorm_max_ms: Option<usize>,

/// Stamp each request with an OpenAI `user` field of "<prefix>-<agent_id>",
/// so gateways and routers can key session-sticky routing per agent
#[arg(long)]
user_prefix: Option<String>,

/// API key; when omitted, read --api-key-env
#[arg(long)]
api_key: Option<String>,
Expand Down Expand Up @@ -375,6 +380,9 @@ async fn run(args: Args) -> Result<()> {
if let Some(tool_call_latency_ms) = tool_call_latency_ms {
config = config.with_tool_call_latency_ms(tool_call_latency_ms)?;
}
if let Some(user_prefix) = &args.user_prefix {
config = config.with_user_prefix(user_prefix.clone());
}
Comment on lines +383 to +385
if let Some(seed) = args.seed {
config = config.with_seed(seed);
}
Expand Down
Loading