feat: capture RL training logprobs

Author: Lixtt

core/src/agent/llm_turn.rs

@@ -33,17 +33,14 @@ impl AgentLoop {
             "LLM completion started"
         );
 
-        let selected_tool_names =
-            crate::tools::select_tools_for_messages(&self.config.tools, &state.messages)
-                .into_iter()
-                .map(|tool| tool.name)
-                .collect::<Vec<_>>();
+        let selected_tools =
+            crate::tools::select_tools_for_messages(&self.config.tools, &state.messages);
         self.config.rl_trajectory_recorder.record_llm_request(
             session_id.unwrap_or(""),
             turn,
             &state.messages,
             augmented_system.as_deref(),
-            &selected_tool_names,
+            &selected_tools,
             estimate_prompt_tokens(&state.messages, augmented_system.as_deref()),
         );
 

core/src/agent/tests.rs

@@ -263,6 +263,7 @@ impl MockLlmClient {
                 cache_write_tokens: None,
             },
             stop_reason: Some("end_turn".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         }
     }
@@ -291,6 +292,7 @@ impl MockLlmClient {
                 cache_write_tokens: None,
             },
             stop_reason: Some("tool_use".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         }
     }
@@ -1080,6 +1082,7 @@ async fn test_agent_hitl_multiple_tool_calls() {
                 cache_write_tokens: None,
             },
             stop_reason: Some("tool_use".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         },
         MockLlmClient::text_response("Both executed!"),
@@ -1164,6 +1167,7 @@ async fn test_agent_hitl_partial_approval() {
                 cache_write_tokens: None,
             },
             stop_reason: Some("tool_use".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         },
         MockLlmClient::text_response("First worked, second rejected."),
@@ -1893,6 +1897,7 @@ async fn test_agent_multiple_tools_single_turn() {
                 cache_write_tokens: None,
             },
             stop_reason: Some("tool_use".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         },
         MockLlmClient::text_response("Both commands ran"),

core/src/agent_api.rs

@@ -215,6 +215,14 @@ pub struct SessionOptions {
     /// training or service data collection. If unset, the same config can
     /// be enabled by `A3S_CODE_TRAJECTORY_PATH`.
     pub rl_trajectory: Option<crate::rl_trajectory::RlTrajectoryConfig>,
+    /// Request token-level log probabilities from compatible LLM providers.
+    ///
+    /// This is off by default because many public providers reject logprob
+    /// requests with tool calls. Training/evaluation harnesses using compatible
+    /// OpenAI-style backends can enable it explicitly.
+    pub llm_logprobs: Option<bool>,
+    /// Number of alternative token logprobs to request per generated token.
+    pub llm_top_logprobs: Option<usize>,
     /// Auto-save after each completed `send()` or default-history `stream()` call.
     pub auto_save: bool,
     /// Optional artifact retention limits for large tool/program outputs.

core/src/agent_api/session_config.rs

@@ -74,13 +74,44 @@ pub(super) fn resolve_session_llm_client(
         }
     }
 
+    let logprobs = opts
+        .llm_logprobs
+        .or_else(|| env_bool("A3S_CODE_LLM_LOGPROBS"))
+        .or_else(|| env_bool("A3S_CODE_OPENAI_LOGPROBS"));
+    if let Some(enabled) = logprobs {
+        llm_config = llm_config.with_logprobs(enabled);
+    }
+
+    let top_logprobs = opts
+        .llm_top_logprobs
+        .or_else(|| env_usize("A3S_CODE_LLM_TOP_LOGPROBS"))
+        .or_else(|| env_usize("A3S_CODE_OPENAI_TOP_LOGPROBS"));
+    if let Some(top_logprobs) = top_logprobs {
+        llm_config = llm_config.with_top_logprobs(top_logprobs);
+    }
+
     if let Some(session_id) = session_id {
         llm_config = llm_config.with_session_id(session_id);
     }
 
     Ok(crate::llm::create_client_with_config(llm_config))
 }
 
+fn env_bool(name: &str) -> Option<bool> {
+    let value = std::env::var(name).ok()?;
+    match value.trim().to_ascii_lowercase().as_str() {
+        "1" | "true" | "yes" | "on" => Some(true),
+        "0" | "false" | "no" | "off" => Some(false),
+        _ => None,
+    }
+}
+
+fn env_usize(name: &str) -> Option<usize> {
+    std::env::var(name)
+        .ok()
+        .and_then(|value| value.trim().parse::<usize>().ok())
+}
+
 pub(super) struct ResolvedSessionMemory {
     pub(super) memory: Option<Arc<crate::memory::AgentMemory>>,
     pub(super) init_warning: Option<String>,

core/src/agent_api/session_options.rs

@@ -43,6 +43,8 @@ impl std::fmt::Debug for SessionOptions {
             .field("session_store", &self.session_store.is_some())
             .field("session_id", &self.session_id)
             .field("rl_trajectory", &self.rl_trajectory)
+            .field("llm_logprobs", &self.llm_logprobs)
+            .field("llm_top_logprobs", &self.llm_top_logprobs)
             .field("auto_save", &self.auto_save)
             .field("artifact_store_limits", &self.artifact_store_limits)
             .field("max_parse_retries", &self.max_parse_retries)
@@ -376,6 +378,19 @@ impl SessionOptions {
         self
     }
 
+    /// Request token-level log probabilities from compatible LLM providers.
+    pub fn with_llm_logprobs(mut self, enabled: bool) -> Self {
+        self.llm_logprobs = Some(enabled);
+        self
+    }
+
+    /// Request up to `top_logprobs` alternative logprobs per generated token.
+    pub fn with_llm_top_logprobs(mut self, top_logprobs: usize) -> Self {
+        self.llm_logprobs = Some(true);
+        self.llm_top_logprobs = Some(top_logprobs);
+        self
+    }
+
     /// Enable auto-save after each `send()` call
     pub fn with_auto_save(mut self, enabled: bool) -> Self {
         self.auto_save = enabled;

core/src/agent_api/tests.rs

@@ -30,6 +30,7 @@ impl StaticStreamingClient {
                 cache_write_tokens: None,
             },
             stop_reason: Some("end_turn".to_string()),
+            token_logprobs: Vec::new(),
             meta: None,
         }
     }

core/src/llm/anthropic.rs

@@ -245,6 +245,7 @@ impl AnthropicClient {
                     cache_write_tokens: parsed.usage.cache_creation_input_tokens,
                 },
                 stop_reason: Some(parsed.stop_reason),
+                token_logprobs: Vec::new(),
                 meta: Some(LlmResponseMeta {
                     provider: Some(self.provider_name.clone()),
                     request_model: Some(self.model.clone()),
@@ -581,6 +582,7 @@ impl AnthropicClient {
                                                 },
                                                 usage: usage.clone(),
                                                 stop_reason: stop_reason.clone(),
+                                                token_logprobs: Vec::new(),
                                                 meta: Some(LlmResponseMeta {
                                                     provider: Some(provider_name.clone()),
                                                     request_model: Some(request_model.clone()),

core/src/llm/factory.rs

@@ -26,6 +26,10 @@ pub struct LlmConfig {
     pub max_tokens: Option<usize>,
     /// Extended thinking budget in tokens (Anthropic only).
     pub thinking_budget: Option<usize>,
+    /// Request token-level log probabilities from OpenAI-compatible providers.
+    pub logprobs: Option<bool>,
+    /// Number of alternative logprobs per token when logprobs are requested.
+    pub top_logprobs: Option<usize>,
     /// When true, temperature is never sent to the API (e.g., o1 models).
     pub disable_temperature: bool,
 }
@@ -47,6 +51,8 @@ impl std::fmt::Debug for LlmConfig {
             .field("temperature", &self.temperature)
             .field("max_tokens", &self.max_tokens)
             .field("thinking_budget", &self.thinking_budget)
+            .field("logprobs", &self.logprobs)
+            .field("top_logprobs", &self.top_logprobs)
             .field("disable_temperature", &self.disable_temperature)
             .finish()
     }
@@ -70,6 +76,8 @@ impl LlmConfig {
             temperature: None,
             max_tokens: None,
             thinking_budget: None,
+            logprobs: None,
+            top_logprobs: None,
             disable_temperature: false,
         }
     }
@@ -114,6 +122,17 @@ impl LlmConfig {
         self
     }
 
+    pub fn with_logprobs(mut self, enabled: bool) -> Self {
+        self.logprobs = Some(enabled);
+        self
+    }
+
+    pub fn with_top_logprobs(mut self, top_logprobs: usize) -> Self {
+        self.logprobs = Some(true);
+        self.top_logprobs = Some(top_logprobs);
+        self
+    }
+
     pub(crate) fn resolved_headers(&self) -> HashMap<String, String> {
         let mut headers = self.headers.clone();
         if let (Some(header_name), Some(session_id)) = (&self.session_id_header, &self.session_id) {
@@ -168,6 +187,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc<dyn LlmClient> {
             if let Some(max) = config.max_tokens {
                 client = client.with_max_tokens(max);
             }
+            if let Some(enabled) = config.logprobs {
+                client = client.with_logprobs(enabled);
+            }
+            if let Some(top_logprobs) = config.top_logprobs {
+                client = client.with_top_logprobs(top_logprobs);
+            }
             Arc::new(client)
         }
         "glm" | "zhipu" | "bigmodel" => {
@@ -183,6 +208,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc<dyn LlmClient> {
             if let Some(max) = config.max_tokens {
                 client = client.with_max_tokens(max);
             }
+            if let Some(enabled) = config.logprobs {
+                client = client.with_logprobs(enabled);
+            }
+            if let Some(top_logprobs) = config.top_logprobs {
+                client = client.with_top_logprobs(top_logprobs);
+            }
             Arc::new(client)
         }
         // OpenAI-compatible providers (deepseek, groq, together, ollama, etc.)
@@ -208,6 +239,12 @@ pub fn create_client_with_config(config: LlmConfig) -> Arc<dyn LlmClient> {
             if let Some(max) = config.max_tokens {
                 client = client.with_max_tokens(max);
             }
+            if let Some(enabled) = config.logprobs {
+                client = client.with_logprobs(enabled);
+            }
+            if let Some(top_logprobs) = config.top_logprobs {
+                client = client.with_top_logprobs(top_logprobs);
+            }
             Arc::new(client)
         }
     }