use crate::errors::CoreError;
use chrono::Utc;
use serde_json::{Map, Value};
use url::Url;
use uuid::Uuid;
fn parse_url(inference_url: &str) -> Option<Url> {
if let Ok(parsed) = Url::parse(inference_url) {
return Some(parsed);
}
if inference_url.starts_with('/') {
let candidate = format!("http://localhost{}", inference_url);
return Url::parse(&candidate).ok();
}
if !inference_url.contains("://") {
let candidate = format!("http://{}", inference_url);
return Url::parse(&candidate).ok();
}
None
}
pub fn extract_trace_correlation_id(inference_url: Option<&str>) -> Option<String> {
let raw = inference_url?.trim();
if raw.is_empty() {
return None;
}
let parsed = parse_url(raw)?;
let path_segments: Vec<&str> = parsed.path().split('/').filter(|s| !s.is_empty()).collect();
if path_segments.len() >= 3 {
let tail = &path_segments[path_segments.len() - 2..];
if tail == ["chat", "completions"] {
if let Some(segment) = path_segments.get(path_segments.len() - 3) {
if segment.starts_with("trace_")
|| segment.starts_with("cid_")
|| segment.starts_with("eval_")
{
return Some(segment.to_string());
}
}
}
}
for segment in path_segments.iter().rev() {
if segment.starts_with("trace_")
|| segment.starts_with("cid_")
|| segment.starts_with("eval_")
{
return Some(segment.to_string());
}
}
let query = parsed.query().unwrap_or("");
for (key, value) in url::form_urlencoded::parse(query.as_bytes()) {
let key = key.to_string();
if key == "cid" || key == "trace_correlation_id" || key == "trace" {
let val = value.to_string();
if !val.trim().is_empty() {
return Some(val);
}
}
}
None
}
pub fn validate_trace_correlation_id(
trace_correlation_id: Option<&str>,
inference_url: Option<&str>,
fatal: bool,
) -> Result<Option<String>, CoreError> {
if let Some(cid) = trace_correlation_id {
if !cid.trim().is_empty() {
return Ok(Some(cid.to_string()));
}
}
let inference_url = inference_url.unwrap_or("NOT_SET");
let error_msg = format!(
"CRITICAL: Cannot extract trace_correlation_id!\n\nID: UNKNOWN\nInference URL: {}\n\nChecked:\n1. inference_url path segments\n2. inference_url query params\n\nTask app CANNOT proceed without trace_correlation_id.\nThis indicates the RL trainer is not sending it correctly.\n\nSee monorepo/trace_creation_and_judgement.txt 'Fatal Guards' section.\n",
inference_url
);
if fatal {
return Err(CoreError::InvalidInput(error_msg));
}
Ok(None)
}
pub fn include_trace_correlation_id_in_response(
response_data: &Value,
trace_correlation_id: Option<&str>,
) -> Value {
let cid = match trace_correlation_id {
Some(value) if !value.trim().is_empty() => value.trim(),
_ => return response_data.clone(),
};
let mut output = response_data.clone();
let obj = match output.as_object_mut() {
Some(obj) => obj,
None => return output,
};
obj.entry("trace_correlation_id".to_string())
.or_insert_with(|| Value::String(cid.to_string()));
if let Some(trace_val) = obj.get_mut("trace") {
if let Some(trace_obj) = trace_val.as_object_mut() {
let meta = trace_obj
.entry("metadata".to_string())
.or_insert_with(|| Value::Object(Map::new()));
if let Some(meta_obj) = meta.as_object_mut() {
meta_obj
.entry("trace_correlation_id".to_string())
.or_insert_with(|| Value::String(cid.to_string()));
let corr_ids = meta_obj
.entry("correlation_ids".to_string())
.or_insert_with(|| Value::Object(Map::new()));
if let Some(corr_map) = corr_ids.as_object_mut() {
corr_map
.entry("trace_correlation_id".to_string())
.or_insert_with(|| Value::String(cid.to_string()));
}
}
if let Some(session_trace) = trace_obj.get_mut("session_trace") {
if let Some(session_obj) = session_trace.as_object_mut() {
let session_meta = session_obj
.entry("metadata".to_string())
.or_insert_with(|| Value::Object(Map::new()));
if let Some(session_meta_obj) = session_meta.as_object_mut() {
session_meta_obj
.entry("trace_correlation_id".to_string())
.or_insert_with(|| Value::String(cid.to_string()));
}
}
}
}
}
output
}
pub fn build_trace_payload(
messages: &Value,
response: Option<&Value>,
correlation_id: Option<&str>,
session_id: Option<&str>,
metadata: Option<&Value>,
) -> Value {
let messages_value = match messages {
Value::Array(list) => Value::Array(list.clone()),
_ => Value::Array(Vec::new()),
};
let llm_response = match response {
Some(Value::Object(obj)) => {
if obj.contains_key("message") {
Value::Object(obj.clone())
} else if let Some(Value::Array(choices)) = obj.get("choices") {
if let Some(Value::Object(first)) = choices.get(0) {
let mut map = Map::new();
map.insert(
"message".to_string(),
first.get("message").cloned().unwrap_or(Value::Null),
);
map.insert(
"usage".to_string(),
obj.get("usage")
.cloned()
.unwrap_or(Value::Object(Map::new())),
);
map.insert(
"finish_reason".to_string(),
first.get("finish_reason").cloned().unwrap_or(Value::Null),
);
Value::Object(map)
} else {
Value::Object(obj.clone())
}
} else {
Value::Object(obj.clone())
}
}
Some(other) => other.clone(),
None => Value::Object(Map::new()),
};
let mut llm_event = Map::new();
llm_event.insert("type".to_string(), Value::String("lm_call".to_string()));
llm_event.insert(
"event_type".to_string(),
Value::String("lm_call".to_string()),
);
llm_event.insert(
"timestamp".to_string(),
Value::String(Utc::now().to_rfc3339()),
);
let mut llm_request = Map::new();
llm_request.insert("messages".to_string(), messages_value);
llm_event.insert("llm_request".to_string(), Value::Object(llm_request));
llm_event.insert("llm_response".to_string(), llm_response);
llm_event.insert("api_format".to_string(), Value::String("chat".to_string()));
if let Some(cid) = correlation_id {
if !cid.trim().is_empty() {
llm_event.insert("correlation_id".to_string(), Value::String(cid.to_string()));
}
}
let mut event_history = Vec::new();
event_history.push(Value::Object(llm_event));
let mut trace_meta = match metadata {
Some(Value::Object(obj)) => obj.clone(),
_ => Map::new(),
};
let session_id = session_id
.map(|s| s.to_string())
.unwrap_or_else(|| Uuid::new_v4().to_string());
trace_meta
.entry("session_id".to_string())
.or_insert(Value::String(session_id));
if let Some(cid) = correlation_id {
if !cid.trim().is_empty() {
trace_meta
.entry("trace_correlation_id".to_string())
.or_insert(Value::String(cid.to_string()));
let corr_ids = trace_meta
.entry("correlation_ids".to_string())
.or_insert_with(|| Value::Object(Map::new()));
if let Some(corr_map) = corr_ids.as_object_mut() {
corr_map
.entry("trace_correlation_id".to_string())
.or_insert(Value::String(cid.to_string()));
}
}
}
let mut trace = Map::new();
trace.insert(
"schema_version".to_string(),
Value::String("4.0".to_string()),
);
trace.insert("event_history".to_string(), Value::Array(event_history));
trace.insert(
"markov_blanket_message_history".to_string(),
Value::Array(Vec::new()),
);
trace.insert("metadata".to_string(), Value::Object(trace_meta));
Value::Object(trace)
}
pub fn build_trajectory_trace(
messages: &Value,
response: Option<&Value>,
correlation_id: Option<&str>,
session_id: Option<&str>,
metadata: Option<&Value>,
) -> Value {
build_trace_payload(messages, response, correlation_id, session_id, metadata)
}
pub fn include_event_history_in_response(
response_data: &Value,
messages: Option<&Value>,
response: Option<&Value>,
run_id: &str,
correlation_id: Option<&str>,
) -> Value {
let mut output = response_data.clone();
let obj = match output.as_object_mut() {
Some(obj) => obj,
None => return output,
};
let trace_val = obj
.entry("trace".to_string())
.or_insert_with(|| Value::Object(Map::new()));
let trace_obj = match trace_val.as_object_mut() {
Some(obj) => obj,
None => return output,
};
let mut event_history = trace_obj.get("event_history");
if event_history.is_none() {
if let Some(Value::Object(session_trace)) = trace_obj.get("session_trace") {
event_history = session_trace.get("event_history");
}
}
if matches!(event_history, Some(Value::Array(list)) if !list.is_empty()) {
return output;
}
let meta_map = {
let mut map = Map::new();
map.insert("run_id".to_string(), Value::String(run_id.to_string()));
Value::Object(map)
};
let new_trace = build_trace_payload(
messages.unwrap_or(&Value::Array(Vec::new())),
response,
correlation_id,
None,
Some(&meta_map),
);
if let Some(new_trace_obj) = new_trace.as_object() {
let incoming_meta = trace_obj.get("metadata");
let mut merged_meta = match new_trace_obj.get("metadata") {
Some(Value::Object(meta)) => meta.clone(),
_ => Map::new(),
};
if let Some(Value::Object(existing)) = incoming_meta {
for (k, v) in existing.iter() {
merged_meta.insert(k.clone(), v.clone());
}
}
trace_obj.insert("metadata".to_string(), Value::Object(merged_meta));
if let Some(schema) = new_trace_obj.get("schema_version") {
trace_obj
.entry("schema_version".to_string())
.or_insert_with(|| schema.clone());
}
if let Some(events) = new_trace_obj.get("event_history") {
trace_obj.insert("event_history".to_string(), events.clone());
}
if let Some(blanket) = new_trace_obj.get("markov_blanket_message_history") {
trace_obj
.entry("markov_blanket_message_history".to_string())
.or_insert_with(|| blanket.clone());
}
let events_clone = trace_obj.get("event_history").cloned();
if let Some(Value::Object(session_trace)) = trace_obj.get_mut("session_trace") {
if !session_trace.contains_key("event_history") {
if let Some(events) = events_clone {
session_trace.insert("event_history".to_string(), events);
}
}
}
}
output
}
pub fn include_event_history_in_trajectories(
response_data: &Value,
messages_by_trajectory: Option<&Value>,
responses_by_trajectory: Option<&Value>,
run_id: &str,
correlation_id: Option<&str>,
) -> Value {
let messages = messages_by_trajectory
.and_then(|value| value.as_array())
.and_then(|arr| arr.get(0));
let response = responses_by_trajectory
.and_then(|value| value.as_array())
.and_then(|arr| arr.get(0));
include_event_history_in_response(response_data, messages, response, run_id, correlation_id)
}
pub fn verify_trace_correlation_id_in_response(
response_data: &Value,
expected_correlation_id: Option<&str>,
) -> bool {
let expected = match expected_correlation_id {
Some(value) if !value.trim().is_empty() => value,
_ => return false,
};
let obj = match response_data.as_object() {
Some(obj) => obj,
None => return false,
};
if obj.get("trace_correlation_id").and_then(|v| v.as_str()) != Some(expected) {
return false;
}
let trace_val = match obj.get("trace") {
Some(Value::Object(trace)) => trace,
_ => return false,
};
let mut trace_meta_id = trace_val
.get("metadata")
.and_then(|v| v.as_object())
.and_then(|meta| meta.get("trace_correlation_id"))
.and_then(|v| v.as_str());
if trace_meta_id != Some(expected) {
if let Some(Value::Object(session_trace)) = trace_val.get("session_trace") {
trace_meta_id = session_trace
.get("metadata")
.and_then(|v| v.as_object())
.and_then(|meta| meta.get("trace_correlation_id"))
.and_then(|v| v.as_str());
}
}
trace_meta_id == Some(expected)
}