[prompt] xml-format EnvironmentContext (#2272)
## Summary Before we land #2243, let's start printing environment_context in our preferred format. This struct will evolve over time with new information, xml gives us a balance of human readable without too much parsing, llm readable, and extensible. Also moves us over to an Option-based struct, so we can easily provide diffs to the model. ## Testing - [x] Updated tests to reflect new format
This commit is contained in:
@@ -508,10 +508,10 @@ impl Session {
|
|||||||
conversation_items.push(Prompt::format_user_instructions_message(user_instructions));
|
conversation_items.push(Prompt::format_user_instructions_message(user_instructions));
|
||||||
}
|
}
|
||||||
conversation_items.push(ResponseItem::from(EnvironmentContext::new(
|
conversation_items.push(ResponseItem::from(EnvironmentContext::new(
|
||||||
turn_context.cwd.to_path_buf(),
|
Some(turn_context.cwd.clone()),
|
||||||
turn_context.approval_policy,
|
Some(turn_context.approval_policy),
|
||||||
turn_context.sandbox_policy.clone(),
|
Some(turn_context.sandbox_policy.clone()),
|
||||||
sess.user_shell.clone(),
|
Some(sess.user_shell.clone()),
|
||||||
)));
|
)));
|
||||||
sess.record_conversation_items(&conversation_items).await;
|
sess.record_conversation_items(&conversation_items).await;
|
||||||
|
|
||||||
@@ -1068,10 +1068,11 @@ async fn submission_loop(
|
|||||||
turn_context = Arc::new(new_turn_context);
|
turn_context = Arc::new(new_turn_context);
|
||||||
if cwd.is_some() || approval_policy.is_some() || sandbox_policy.is_some() {
|
if cwd.is_some() || approval_policy.is_some() || sandbox_policy.is_some() {
|
||||||
sess.record_conversation_items(&[ResponseItem::from(EnvironmentContext::new(
|
sess.record_conversation_items(&[ResponseItem::from(EnvironmentContext::new(
|
||||||
new_cwd,
|
cwd,
|
||||||
new_approval_policy,
|
approval_policy,
|
||||||
new_sandbox_policy,
|
sandbox_policy,
|
||||||
sess.user_shell.clone(),
|
// Shell is not configurable from turn to turn
|
||||||
|
None,
|
||||||
))])
|
))])
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,11 +8,10 @@ use crate::protocol::AskForApproval;
|
|||||||
use crate::protocol::SandboxPolicy;
|
use crate::protocol::SandboxPolicy;
|
||||||
use crate::shell::Shell;
|
use crate::shell::Shell;
|
||||||
use codex_protocol::config_types::SandboxMode;
|
use codex_protocol::config_types::SandboxMode;
|
||||||
use std::fmt::Display;
|
|
||||||
use std::path::PathBuf;
|
use std::path::PathBuf;
|
||||||
|
|
||||||
/// wraps environment context message in a tag for the model to parse more easily.
|
/// wraps environment context message in a tag for the model to parse more easily.
|
||||||
pub(crate) const ENVIRONMENT_CONTEXT_START: &str = "<environment_context>\n";
|
pub(crate) const ENVIRONMENT_CONTEXT_START: &str = "<environment_context>";
|
||||||
pub(crate) const ENVIRONMENT_CONTEXT_END: &str = "</environment_context>";
|
pub(crate) const ENVIRONMENT_CONTEXT_END: &str = "</environment_context>";
|
||||||
|
|
||||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, DeriveDisplay)]
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, DeriveDisplay)]
|
||||||
@@ -25,58 +24,87 @@ pub enum NetworkAccess {
|
|||||||
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
||||||
#[serde(rename = "environment_context", rename_all = "snake_case")]
|
#[serde(rename = "environment_context", rename_all = "snake_case")]
|
||||||
pub(crate) struct EnvironmentContext {
|
pub(crate) struct EnvironmentContext {
|
||||||
pub cwd: PathBuf,
|
pub cwd: Option<PathBuf>,
|
||||||
pub approval_policy: AskForApproval,
|
pub approval_policy: Option<AskForApproval>,
|
||||||
pub sandbox_mode: SandboxMode,
|
pub sandbox_mode: Option<SandboxMode>,
|
||||||
pub network_access: NetworkAccess,
|
pub network_access: Option<NetworkAccess>,
|
||||||
pub shell: Shell,
|
pub shell: Option<Shell>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl EnvironmentContext {
|
impl EnvironmentContext {
|
||||||
pub fn new(
|
pub fn new(
|
||||||
cwd: PathBuf,
|
cwd: Option<PathBuf>,
|
||||||
approval_policy: AskForApproval,
|
approval_policy: Option<AskForApproval>,
|
||||||
sandbox_policy: SandboxPolicy,
|
sandbox_policy: Option<SandboxPolicy>,
|
||||||
shell: Shell,
|
shell: Option<Shell>,
|
||||||
) -> Self {
|
) -> Self {
|
||||||
Self {
|
Self {
|
||||||
cwd,
|
cwd,
|
||||||
approval_policy,
|
approval_policy,
|
||||||
sandbox_mode: match sandbox_policy {
|
sandbox_mode: match sandbox_policy {
|
||||||
SandboxPolicy::DangerFullAccess => SandboxMode::DangerFullAccess,
|
Some(SandboxPolicy::DangerFullAccess) => Some(SandboxMode::DangerFullAccess),
|
||||||
SandboxPolicy::ReadOnly => SandboxMode::ReadOnly,
|
Some(SandboxPolicy::ReadOnly) => Some(SandboxMode::ReadOnly),
|
||||||
SandboxPolicy::WorkspaceWrite { .. } => SandboxMode::WorkspaceWrite,
|
Some(SandboxPolicy::WorkspaceWrite { .. }) => Some(SandboxMode::WorkspaceWrite),
|
||||||
|
None => None,
|
||||||
},
|
},
|
||||||
network_access: match sandbox_policy {
|
network_access: match sandbox_policy {
|
||||||
SandboxPolicy::DangerFullAccess => NetworkAccess::Enabled,
|
Some(SandboxPolicy::DangerFullAccess) => Some(NetworkAccess::Enabled),
|
||||||
SandboxPolicy::ReadOnly => NetworkAccess::Restricted,
|
Some(SandboxPolicy::ReadOnly) => Some(NetworkAccess::Restricted),
|
||||||
SandboxPolicy::WorkspaceWrite { network_access, .. } => {
|
Some(SandboxPolicy::WorkspaceWrite { network_access, .. }) => {
|
||||||
if network_access {
|
if network_access {
|
||||||
NetworkAccess::Enabled
|
Some(NetworkAccess::Enabled)
|
||||||
} else {
|
} else {
|
||||||
NetworkAccess::Restricted
|
Some(NetworkAccess::Restricted)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
None => None,
|
||||||
},
|
},
|
||||||
shell,
|
shell,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
impl Display for EnvironmentContext {
|
impl EnvironmentContext {
|
||||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
/// Serializes the environment context to XML. Libraries like `quick-xml`
|
||||||
writeln!(
|
/// require custom macros to handle Enums with newtypes, so we just do it
|
||||||
f,
|
/// manually, to keep things simple. Output looks like:
|
||||||
"Current working directory: {}",
|
///
|
||||||
self.cwd.to_string_lossy()
|
/// ```xml
|
||||||
)?;
|
/// <environment_context>
|
||||||
writeln!(f, "Approval policy: {}", self.approval_policy)?;
|
/// <cwd>...</cwd>
|
||||||
writeln!(f, "Sandbox mode: {}", self.sandbox_mode)?;
|
/// <approval_policy>...</approval_policy>
|
||||||
writeln!(f, "Network access: {}", self.network_access)?;
|
/// <sandbox_mode>...</sandbox_mode>
|
||||||
if let Some(shell_name) = self.shell.name() {
|
/// <network_access>...</network_access>
|
||||||
writeln!(f, "Shell: {shell_name}")?;
|
/// <shell>...</shell>
|
||||||
|
/// </environment_context>
|
||||||
|
/// ```
|
||||||
|
pub fn serialize_to_xml(self) -> String {
|
||||||
|
let mut lines = vec![ENVIRONMENT_CONTEXT_START.to_string()];
|
||||||
|
if let Some(cwd) = self.cwd {
|
||||||
|
lines.push(format!(" <cwd>{}</cwd>", cwd.to_string_lossy()));
|
||||||
}
|
}
|
||||||
Ok(())
|
if let Some(approval_policy) = self.approval_policy {
|
||||||
|
lines.push(format!(
|
||||||
|
" <approval_policy>{}</approval_policy>",
|
||||||
|
approval_policy
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if let Some(sandbox_mode) = self.sandbox_mode {
|
||||||
|
lines.push(format!(" <sandbox_mode>{}</sandbox_mode>", sandbox_mode));
|
||||||
|
}
|
||||||
|
if let Some(network_access) = self.network_access {
|
||||||
|
lines.push(format!(
|
||||||
|
" <network_access>{}</network_access>",
|
||||||
|
network_access
|
||||||
|
));
|
||||||
|
}
|
||||||
|
if let Some(shell) = self.shell
|
||||||
|
&& let Some(shell_name) = shell.name()
|
||||||
|
{
|
||||||
|
lines.push(format!(" <shell>{}</shell>", shell_name));
|
||||||
|
}
|
||||||
|
lines.push(ENVIRONMENT_CONTEXT_END.to_string());
|
||||||
|
lines.join("\n")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -86,7 +114,7 @@ impl From<EnvironmentContext> for ResponseItem {
|
|||||||
id: None,
|
id: None,
|
||||||
role: "user".to_string(),
|
role: "user".to_string(),
|
||||||
content: vec![ContentItem::InputText {
|
content: vec![ContentItem::InputText {
|
||||||
text: format!("{ENVIRONMENT_CONTEXT_START}{ec}{ENVIRONMENT_CONTEXT_END}"),
|
text: ec.serialize_to_xml(),
|
||||||
}],
|
}],
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -89,10 +89,15 @@ async fn prefixes_context_and_instructions_once_and_consistently_across_requests
|
|||||||
let shell = default_user_shell().await;
|
let shell = default_user_shell().await;
|
||||||
|
|
||||||
let expected_env_text = format!(
|
let expected_env_text = format!(
|
||||||
"<environment_context>\nCurrent working directory: {}\nApproval policy: on-request\nSandbox mode: read-only\nNetwork access: restricted\n{}</environment_context>",
|
r#"<environment_context>
|
||||||
|
<cwd>{}</cwd>
|
||||||
|
<approval_policy>on-request</approval_policy>
|
||||||
|
<sandbox_mode>read-only</sandbox_mode>
|
||||||
|
<network_access>restricted</network_access>
|
||||||
|
{}</environment_context>"#,
|
||||||
cwd.path().to_string_lossy(),
|
cwd.path().to_string_lossy(),
|
||||||
match shell.name() {
|
match shell.name() {
|
||||||
Some(name) => format!("Shell: {name}\n"),
|
Some(name) => format!(" <shell>{}</shell>\n", name),
|
||||||
None => String::new(),
|
None => String::new(),
|
||||||
}
|
}
|
||||||
);
|
);
|
||||||
@@ -190,12 +195,10 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
|
|||||||
.unwrap();
|
.unwrap();
|
||||||
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
|
wait_for_event(&codex, |ev| matches!(ev, EventMsg::TaskComplete(_))).await;
|
||||||
|
|
||||||
// Change everything about the turn context.
|
|
||||||
let new_cwd = TempDir::new().unwrap();
|
|
||||||
let writable = TempDir::new().unwrap();
|
let writable = TempDir::new().unwrap();
|
||||||
codex
|
codex
|
||||||
.submit(Op::OverrideTurnContext {
|
.submit(Op::OverrideTurnContext {
|
||||||
cwd: Some(new_cwd.path().to_path_buf()),
|
cwd: None,
|
||||||
approval_policy: Some(AskForApproval::Never),
|
approval_policy: Some(AskForApproval::Never),
|
||||||
sandbox_policy: Some(SandboxPolicy::WorkspaceWrite {
|
sandbox_policy: Some(SandboxPolicy::WorkspaceWrite {
|
||||||
writable_roots: vec![writable.path().to_path_buf()],
|
writable_roots: vec![writable.path().to_path_buf()],
|
||||||
@@ -227,7 +230,6 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
|
|||||||
|
|
||||||
let body1 = requests[0].body_json::<serde_json::Value>().unwrap();
|
let body1 = requests[0].body_json::<serde_json::Value>().unwrap();
|
||||||
let body2 = requests[1].body_json::<serde_json::Value>().unwrap();
|
let body2 = requests[1].body_json::<serde_json::Value>().unwrap();
|
||||||
|
|
||||||
// prompt_cache_key should remain constant across overrides
|
// prompt_cache_key should remain constant across overrides
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
body1["prompt_cache_key"], body2["prompt_cache_key"],
|
body1["prompt_cache_key"], body2["prompt_cache_key"],
|
||||||
@@ -243,16 +245,13 @@ async fn overrides_turn_context_but_keeps_cached_prefix_and_key_constant() {
|
|||||||
"content": [ { "type": "input_text", "text": "hello 2" } ]
|
"content": [ { "type": "input_text", "text": "hello 2" } ]
|
||||||
});
|
});
|
||||||
// After overriding the turn context, the environment context should be emitted again
|
// After overriding the turn context, the environment context should be emitted again
|
||||||
// reflecting the new cwd, approval policy and sandbox settings.
|
// reflecting the new approval policy and sandbox settings. Omit cwd because it did
|
||||||
let shell = default_user_shell().await;
|
// not change.
|
||||||
let expected_env_text_2 = format!(
|
let expected_env_text_2 = r#"<environment_context>
|
||||||
"<environment_context>\nCurrent working directory: {}\nApproval policy: never\nSandbox mode: workspace-write\nNetwork access: enabled\n{}</environment_context>",
|
<approval_policy>never</approval_policy>
|
||||||
new_cwd.path().to_string_lossy(),
|
<sandbox_mode>workspace-write</sandbox_mode>
|
||||||
match shell.name() {
|
<network_access>enabled</network_access>
|
||||||
Some(name) => format!("Shell: {name}\n"),
|
</environment_context>"#;
|
||||||
None => String::new(),
|
|
||||||
}
|
|
||||||
);
|
|
||||||
let expected_env_msg_2 = serde_json::json!({
|
let expected_env_msg_2 = serde_json::json!({
|
||||||
"type": "message",
|
"type": "message",
|
||||||
"id": serde_json::Value::Null,
|
"id": serde_json::Value::Null,
|
||||||
|
|||||||
Reference in New Issue
Block a user