use std::collections::BTreeMap;
use std::io::Read as _;
use anyhow::{Context, Result};
use clap::ValueEnum;
use serde::Serialize;
use crate::github_issues::CacheUsage;
use crate::jev::protocol::{Answer, SystemOneResponse, Usage};
#[derive(ValueEnum, Clone, Copy, Debug, Default, PartialEq, Eq)]
pub(crate) enum JevFormat {
#[default]
Json,
Yaml,
}
pub(super) const SINGLE_QUESTION_KEY: &str = "answer";
pub(super) fn resolve_state(state: Option<String>) -> Result<String> {
let raw = match state {
Some(s) => s,
None => read_stdin()?,
};
if raw.trim().is_empty() {
anyhow::bail!("state must not be empty (pass it as an argument or on stdin)");
}
Ok(raw)
}
fn read_stdin() -> Result<String> {
let mut buf = String::new();
std::io::stdin()
.read_to_string(&mut buf)
.context("Failed to read state from stdin")?;
Ok(buf)
}
pub(super) fn build_state_value(raw: &str, state_json: bool) -> Result<serde_json::Value> {
if !state_json {
return Ok(serde_json::Value::String(raw.to_string()));
}
let value: serde_json::Value =
serde_yaml::from_str(raw).context("Failed to parse --state-json state as JSON/YAML")?;
match value {
serde_json::Value::String(_)
| serde_json::Value::Object(_)
| serde_json::Value::Array(_) => Ok(value),
other => anyhow::bail!(
"--state-json state must be an object, array or string, got `{other}`; \
drop --state-json to send it as a literal string"
),
}
}
#[derive(Serialize)]
struct SingleAnswerOutput<'a> {
model: &'a str,
answer: &'a Answer,
usage: Usage,
}
#[derive(Serialize)]
struct MultiAnswerOutput<'a> {
model: &'a str,
answers: &'a BTreeMap<String, Answer>,
usage: Usage,
}
fn single_answer(response: &SystemOneResponse) -> Result<&Answer> {
response.answers.get(SINGLE_QUESTION_KEY).with_context(|| {
format!(
"Jev response did not include an answer for the expected key {SINGLE_QUESTION_KEY:?}"
)
})
}
pub(super) fn format_single(response: &SystemOneResponse, format: JevFormat) -> Result<String> {
let answer = single_answer(response)?;
let output = SingleAnswerOutput {
model: &response.model,
answer,
usage: response.usage,
};
format_output(&output, format)
}
pub(super) fn format_multi(response: &SystemOneResponse, format: JevFormat) -> Result<String> {
let output = MultiAnswerOutput {
model: &response.model,
answers: &response.answers,
usage: response.usage,
};
format_output(&output, format)
}
pub(super) fn format_output<T: Serialize>(value: &T, format: JevFormat) -> Result<String> {
match format {
JevFormat::Json => {
let json = serde_json::to_string_pretty(value)
.context("Failed to serialize Jev response as JSON")?;
Ok(format!("{json}\n"))
}
JevFormat::Yaml => {
serde_yaml::to_string(value).context("Failed to serialize Jev response as YAML")
}
}
}
#[derive(Serialize)]
struct WithGithubCache<'a, T> {
#[serde(flatten)]
report: &'a T,
github_cache: CacheUsage,
}
pub(crate) fn format_output_with_cache<T: Serialize>(
value: &T,
cache_use: Option<CacheUsage>,
format: JevFormat,
) -> Result<String> {
match cache_use {
Some(github_cache) => format_output(
&WithGithubCache {
report: value,
github_cache,
},
format,
),
None => format_output(value, format),
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
pub(super) mod test_support {
use crate::jev::client::JevClient;
pub(crate) fn dead_client() -> JevClient {
JevClient::new("http://127.0.0.1:1", "test-key").unwrap()
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use std::collections::BTreeMap;
fn sample_response() -> SystemOneResponse {
SystemOneResponse {
model: "jev-1.13.0".to_string(),
answers: BTreeMap::from([(
SINGLE_QUESTION_KEY.to_string(),
Answer::Noul { noul: 0.5 },
)]),
usage: Usage {
input_tokens: 1,
output_tokens: 2,
},
}
}
#[test]
fn github_cache_usage_is_flattened_beside_the_report_and_only_when_present() {
let response = sample_response();
let usage = CacheUsage {
items_reused: 3,
oldest_age_secs: 150,
};
let plain = format_output_with_cache(&response, None, JevFormat::Json).unwrap();
assert_eq!(plain, format_output(&response, JevFormat::Json).unwrap());
assert!(!plain.contains("github_cache"), "{plain}");
let json = format_output_with_cache(&response, Some(usage), JevFormat::Json).unwrap();
let value: serde_json::Value = serde_json::from_str(&json).unwrap();
assert_eq!(value["model"], "jev-1.13.0", "the report's own fields stay");
assert_eq!(
value["github_cache"],
serde_json::json!({"items_reused": 3, "oldest_age_secs": 150})
);
let yaml = format_output_with_cache(&response, Some(usage), JevFormat::Yaml).unwrap();
assert!(yaml.contains("github_cache:"), "{yaml}");
assert!(yaml.contains("items_reused: 3"), "{yaml}");
assert!(yaml.contains("model: jev-1.13.0"), "{yaml}");
}
#[test]
fn resolve_state_uses_explicit_argument() {
let state = resolve_state(Some("hello".to_string())).unwrap();
assert_eq!(state, "hello");
}
#[test]
fn resolve_state_rejects_empty_argument() {
let err = resolve_state(Some(" \n".to_string())).unwrap_err();
assert!(err.to_string().contains("must not be empty"));
}
#[test]
fn build_state_value_default_is_literal_string() {
let value = build_state_value("{not valid json}", false).unwrap();
assert_eq!(
value,
serde_json::Value::String("{not valid json}".to_string())
);
}
#[test]
fn build_state_value_json_flag_parses_json() {
let value = build_state_value(r#"{"a": 1}"#, true).unwrap();
assert_eq!(value, serde_json::json!({"a": 1}));
}
#[test]
fn build_state_value_json_flag_parses_yaml() {
let value = build_state_value("a: 1\nb: 2\n", true).unwrap();
assert_eq!(value, serde_json::json!({"a": 1, "b": 2}));
}
#[test]
fn build_state_value_json_flag_parses_array() {
let value = build_state_value("[1, 2]", true).unwrap();
assert_eq!(value, serde_json::json!([1, 2]));
}
#[test]
fn build_state_value_json_flag_allows_a_bare_string() {
let value = build_state_value("hello world", true).unwrap();
assert_eq!(value, serde_json::json!("hello world"));
}
#[test]
fn build_state_value_json_flag_rejects_non_string_scalars() {
for raw in ["42", "true", "null", "1.5"] {
let err = build_state_value(raw, true).unwrap_err();
assert!(
err.to_string()
.contains("must be an object, array or string"),
"{raw}: {err}"
);
}
}
#[test]
fn build_state_value_json_flag_rejects_invalid_yaml() {
let err = build_state_value("a: [unterminated", true).unwrap_err();
assert!(err.to_string().contains("Failed to parse"));
}
#[test]
fn single_answer_extracts_the_one_entry() {
let response = sample_response();
let answer = single_answer(&response).unwrap();
assert_eq!(*answer, Answer::Noul { noul: 0.5 });
}
#[test]
fn single_answer_missing_key_errors() {
let response = SystemOneResponse {
model: "jev-1.13.0".to_string(),
answers: BTreeMap::new(),
usage: Usage::default(),
};
let err = single_answer(&response).unwrap_err();
assert!(err.to_string().contains(SINGLE_QUESTION_KEY));
}
#[test]
fn format_single_json_unwraps_answer() {
let out = format_single(&sample_response(), JevFormat::Json).unwrap();
assert!(out.contains("\"model\": \"jev-1.13.0\""));
assert!(out.contains("\"noul\": 0.5"));
assert!(!out.contains("\"answer\": {\n \"answer\""));
assert!(out.ends_with('\n'));
}
#[test]
fn format_single_yaml_unwraps_answer() {
let out = format_single(&sample_response(), JevFormat::Yaml).unwrap();
assert!(out.contains("model: jev-1.13.0"));
assert!(out.contains("noul: 0.5"));
}
#[test]
fn format_multi_json_keeps_answers_map() {
let response = sample_response();
let out = format_multi(&response, JevFormat::Json).unwrap();
assert!(out.contains("\"answers\""));
assert!(out.contains(SINGLE_QUESTION_KEY));
}
#[test]
fn format_multi_yaml_keeps_answers_map() {
let response = sample_response();
let out = format_multi(&response, JevFormat::Yaml).unwrap();
assert!(out.contains("answers:"));
assert!(out.contains(SINGLE_QUESTION_KEY));
}
#[test]
fn jev_format_value_enum_names_are_lowercase() {
assert_eq!(
JevFormat::Json.to_possible_value().unwrap().get_name(),
"json"
);
assert_eq!(
JevFormat::Yaml.to_possible_value().unwrap().get_name(),
"yaml"
);
}
}