use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::{json, Value};
use crate::distill::{DistilledMemory, Distiller};
use crate::error::Error;
use crate::traits::Result;
pub const DEFAULT_TIMEOUT_SECS: u64 = 300;
pub const DEFAULT_MAX_TRANSCRIPT_CHARS: usize = 100_000;
const SYSTEM_PROMPT: &str = "Extract durable memory from this conversation transcript. \
Produce: (1) a summary under 150 words; (2) facts worth remembering across future \
sessions — decisions made and why, corrections the user gave, user preferences, and \
state changes (X is now Y). Each fact must be one self-contained sentence \
understandable without the transcript. Omit pleasantries, transient debugging \
detail, and anything already obvious from a codebase.";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum SchemaStyle {
#[default]
OpenAi,
LlamaServer,
}
fn distilled_memory_schema() -> Value {
json!({
"type": "object",
"properties": {
"summary": { "type": "string" },
"facts": {
"type": "array",
"items": {
"type": "object",
"properties": {
"kind": {
"type": "string",
"enum": ["decision", "correction", "preference", "state"]
},
"text": { "type": "string" }
},
"required": ["kind", "text"],
"additionalProperties": false
}
}
},
"required": ["summary", "facts"],
"additionalProperties": false
})
}
#[must_use]
pub fn response_format_payload(style: SchemaStyle) -> Value {
let schema = distilled_memory_schema();
match style {
SchemaStyle::OpenAi => json!({
"type": "json_schema",
"json_schema": {
"name": "distilled_memory",
"schema": schema
}
}),
SchemaStyle::LlamaServer => json!({
"type": "json_object",
"schema": schema
}),
}
}
#[must_use]
pub fn truncate_keep_end(text: &str, max_chars: usize) -> &str {
let char_count = text.chars().count();
if char_count <= max_chars {
return text;
}
let skip = char_count - max_chars;
match text.char_indices().nth(skip) {
Some((byte_idx, _)) => &text[byte_idx..],
None => "",
}
}
fn strip_code_fences(content: &str) -> &str {
let trimmed = content.trim();
let Some(after_open) = trimmed.strip_prefix("```") else {
return trimmed;
};
let after_lang = match after_open.find('\n') {
Some(idx) => &after_open[idx + 1..],
None => after_open,
};
match after_lang.rfind("```") {
Some(idx) => after_lang[..idx].trim(),
None => after_lang.trim(),
}
}
#[derive(Debug, Serialize)]
struct ChatRequest<'a> {
model: &'a str,
messages: Vec<ChatMessage<'a>>,
#[serde(skip_serializing_if = "Option::is_none")]
response_format: Option<Value>,
}
#[derive(Debug, Serialize)]
struct ChatMessage<'a> {
role: &'a str,
content: String,
}
#[derive(Debug, Deserialize)]
struct ChatResponse {
#[serde(default)]
choices: Vec<ChatChoice>,
}
#[derive(Debug, Deserialize)]
struct ChatChoice {
message: ChatResponseMessage,
}
#[derive(Debug, Deserialize)]
struct ChatResponseMessage {
#[serde(default)]
content: Option<String>,
}
#[derive(Debug, Clone)]
pub struct OpenAiCompatDistiller {
base_url: String,
model: String,
schema_style: SchemaStyle,
max_transcript_chars: usize,
timeout: Duration,
api_key: Option<String>,
}
impl OpenAiCompatDistiller {
pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Result<Self> {
Ok(Self {
base_url: base_url.into(),
model: model.into(),
schema_style: SchemaStyle::default(),
max_transcript_chars: DEFAULT_MAX_TRANSCRIPT_CHARS,
timeout: Duration::from_secs(DEFAULT_TIMEOUT_SECS),
api_key: None,
})
}
#[must_use]
pub fn with_schema_style(mut self, style: SchemaStyle) -> Self {
self.schema_style = style;
self
}
#[must_use]
pub fn with_timeout_secs(mut self, timeout_secs: u64) -> Self {
self.timeout = Duration::from_secs(timeout_secs);
self
}
#[must_use]
pub fn with_max_transcript_chars(mut self, max_transcript_chars: usize) -> Self {
self.max_transcript_chars = max_transcript_chars;
self
}
#[must_use]
pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
self.api_key = Some(key.into());
self
}
fn build_client(timeout: Duration) -> Result<reqwest::blocking::Client> {
reqwest::blocking::Client::builder()
.timeout(timeout)
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| Error::Distill(format!("failed to build HTTP client: {e}")))
}
fn endpoint(&self) -> String {
format!("{}/chat/completions", self.base_url.trim_end_matches('/'))
}
fn send_raw(
&self,
client: &reqwest::blocking::Client,
body: &ChatRequest<'_>,
) -> std::result::Result<reqwest::blocking::Response, reqwest::Error> {
let mut req = client.post(self.endpoint()).json(body);
if let Some(key) = &self.api_key {
req = req.bearer_auth(key);
}
req.send()
}
fn parse_response(response: reqwest::blocking::Response) -> Result<ChatResponse> {
if !response.status().is_success() {
return Err(Error::Distill(format!(
"non-success status: {}",
response.status()
)));
}
response
.json::<ChatResponse>()
.map_err(|e| Error::Distill(format!("invalid response envelope: {e}")))
}
fn message_content(response: ChatResponse) -> Result<String> {
let content = response
.choices
.into_iter()
.next()
.and_then(|choice| choice.message.content);
match content {
Some(c) if !c.is_empty() => Ok(c),
_ => Err(Error::Distill("response had no message content".into())),
}
}
fn attempt_structured(
&self,
client: &reqwest::blocking::Client,
transcript: &str,
) -> std::result::Result<DistilledMemory, AttemptError> {
let body = ChatRequest {
model: &self.model,
messages: vec![
ChatMessage {
role: "system",
content: SYSTEM_PROMPT.to_owned(),
},
ChatMessage {
role: "user",
content: transcript.to_owned(),
},
],
response_format: Some(response_format_payload(self.schema_style)),
};
let raw = self.send_raw(client, &body).map_err(AttemptError::from)?;
let response = Self::parse_response(raw)?;
let content = Self::message_content(response)?;
serde_json::from_str(&content).map_err(|_| AttemptError::Rejected)
}
fn attempt_prompt_embedded(
&self,
client: &reqwest::blocking::Client,
transcript: &str,
) -> Result<DistilledMemory> {
let schema = distilled_memory_schema();
let prompt = format!(
"{transcript}\n\n---\nRespond with ONLY a JSON object matching this schema, \
and nothing else:\n{schema}"
);
let body = ChatRequest {
model: &self.model,
messages: vec![
ChatMessage {
role: "system",
content: SYSTEM_PROMPT.to_owned(),
},
ChatMessage {
role: "user",
content: prompt,
},
],
response_format: None,
};
let raw = self
.send_raw(client, &body)
.map_err(|e| Error::Distill(format!("request failed: {e}")))?;
let response = Self::parse_response(raw)?;
let content = Self::message_content(response)?;
let stripped = strip_code_fences(&content);
serde_json::from_str(stripped)
.map_err(|e| Error::Distill(format!("failed to parse fallback response: {e}")))
}
}
enum AttemptError {
Transport(reqwest::Error),
Rejected,
}
impl From<reqwest::Error> for AttemptError {
fn from(e: reqwest::Error) -> Self {
if e.is_timeout() || e.is_connect() || e.is_request() {
AttemptError::Transport(e)
} else {
AttemptError::Rejected
}
}
}
impl From<Error> for AttemptError {
fn from(_: Error) -> Self {
AttemptError::Rejected
}
}
impl Distiller for OpenAiCompatDistiller {
fn distill(&self, transcript: &str) -> Result<DistilledMemory> {
let truncated = truncate_keep_end(transcript, self.max_transcript_chars).to_owned();
let this = self.clone();
std::thread::spawn(move || {
let client = Self::build_client(this.timeout)?;
match this.attempt_structured(&client, &truncated) {
Ok(memory) => Ok(memory),
Err(AttemptError::Transport(e)) => {
Err(Error::Distill(format!("request failed: {e}")))
}
Err(AttemptError::Rejected) => this.attempt_prompt_embedded(&client, &truncated),
}
})
.join()
.map_err(|_| Error::Distill("distill thread panicked".into()))?
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::distill::FactKind;
use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpListener;
use std::sync::mpsc;
use std::thread;
struct MockResponse {
status_line: &'static str,
body: String,
}
fn spawn_mock_server(responses: Vec<MockResponse>) -> (String, mpsc::Receiver<String>) {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock listener");
let addr = listener.local_addr().expect("local addr");
let (tx, rx) = mpsc::channel();
thread::spawn(move || {
for response in responses {
let Ok((mut stream, _)) = listener.accept() else {
return;
};
let body = read_request_body(&mut stream);
let _ = tx.send(body);
let payload = format!(
"{}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
response.status_line,
response.body.len(),
response.body
);
let _ = stream.write_all(payload.as_bytes());
let _ = stream.flush();
}
});
(format!("http://{addr}"), rx)
}
fn read_request_body(stream: &mut std::net::TcpStream) -> String {
let mut reader = BufReader::new(stream.try_clone().expect("clone stream"));
let mut captured = String::new();
let mut content_length: usize = 0;
loop {
let mut line = String::new();
if reader.read_line(&mut line).unwrap_or(0) == 0 {
break;
}
let trimmed = line.trim_end();
if let Some(value) = trimmed.to_ascii_lowercase().strip_prefix("content-length:") {
content_length = value.trim().parse().unwrap_or(0);
}
captured.push_str(trimmed);
captured.push('\n');
if trimmed.is_empty() {
break;
}
}
let mut body = vec![0u8; content_length];
if content_length > 0 {
let _ = reader.read_exact(&mut body);
}
captured.push_str(&String::from_utf8(body).unwrap_or_default());
captured
}
fn distiller_for(base_url: &str) -> OpenAiCompatDistiller {
OpenAiCompatDistiller::new(base_url, "test-model")
.expect("construct distiller")
.with_timeout_secs(5)
}
#[test]
fn payload_shapes_for_both_styles() {
let openai = response_format_payload(SchemaStyle::OpenAi);
assert_eq!(openai["type"], "json_schema");
assert_eq!(openai["json_schema"]["name"], "distilled_memory");
assert_eq!(openai["json_schema"]["schema"]["type"], "object");
assert_eq!(
openai["json_schema"]["schema"]["required"],
json!(["summary", "facts"])
);
let llama = response_format_payload(SchemaStyle::LlamaServer);
assert_eq!(llama["type"], "json_object");
assert_eq!(llama["schema"]["type"], "object");
assert_eq!(llama["schema"]["required"], json!(["summary", "facts"]));
}
#[test]
fn schema_style_default_is_openai() {
assert_eq!(SchemaStyle::default(), SchemaStyle::OpenAi);
}
#[test]
fn truncate_keeps_end_and_respects_char_boundaries() {
let text = "ééééé hello world";
let truncated = truncate_keep_end(text, 11);
assert_eq!(truncated.chars().count(), 11);
assert!(truncated.ends_with("hello world"));
assert!(truncated.is_char_boundary(0));
}
#[test]
fn truncate_noop_when_short_enough() {
let text = "short";
assert_eq!(truncate_keep_end(text, 100), text);
assert_eq!(truncate_keep_end(text, 5), text);
}
#[test]
fn strip_fences_removes_json_fence() {
let fenced = "```json\n{\"summary\":\"x\",\"facts\":[]}\n```";
assert_eq!(
strip_code_fences(fenced),
"{\"summary\":\"x\",\"facts\":[]}"
);
}
#[test]
fn strip_fences_removes_plain_fence() {
let fenced = "```\n{\"summary\":\"x\",\"facts\":[]}\n```";
assert_eq!(
strip_code_fences(fenced),
"{\"summary\":\"x\",\"facts\":[]}"
);
}
#[test]
fn strip_fences_passthrough_when_unfenced() {
let plain = "{\"summary\":\"x\",\"facts\":[]}";
assert_eq!(strip_code_fences(plain), plain);
}
#[test]
fn attempt_one_success_with_valid_envelope() {
let body = json!({
"summary": "Discussed deploy fix.",
"facts": [
{"kind": "decision", "text": "We decided to roll back the deploy."}
]
})
.to_string();
let envelope = json!({
"choices": [
{"message": {"content": body}}
]
})
.to_string();
let (url, rx) = spawn_mock_server(vec![MockResponse {
status_line: "HTTP/1.1 200 OK",
body: envelope,
}]);
let distiller = distiller_for(&url);
let result = distiller.distill("hello transcript").expect("distill ok");
assert_eq!(result.summary, "Discussed deploy fix.");
assert_eq!(result.facts.len(), 1);
assert_eq!(result.facts[0].kind, FactKind::Decision);
let request_body = rx.recv().expect("captured request");
assert!(request_body.contains("response_format"));
}
#[test]
fn non_2xx_falls_back_to_prompt_embedded() {
let fallback_body = json!({
"summary": "Fallback summary.",
"facts": []
})
.to_string();
let fallback_envelope = json!({
"choices": [
{"message": {"content": fallback_body}}
]
})
.to_string();
let (url, rx) = spawn_mock_server(vec![
MockResponse {
status_line: "HTTP/1.1 500 Internal Server Error",
body: "{}".to_owned(),
},
MockResponse {
status_line: "HTTP/1.1 200 OK",
body: fallback_envelope,
},
]);
let distiller = distiller_for(&url);
let result = distiller.distill("hello transcript").expect("distill ok");
assert_eq!(result.summary, "Fallback summary.");
let first = rx.recv().expect("first request");
assert!(first.contains("response_format"));
let second = rx.recv().expect("second request");
assert!(!second.contains("response_format"));
assert!(second.contains("additionalProperties"));
}
#[test]
fn fails_open_garbage_content_falls_back() {
let attempt1_envelope = json!({
"choices": [
{"message": {"content": "Sure! Here's some unstructured text about the chat."}}
]
})
.to_string();
let fallback_body = json!({
"summary": "Recovered via fallback.",
"facts": []
})
.to_string();
let fallback_envelope = json!({
"choices": [
{"message": {"content": format!("```json\n{fallback_body}\n```")}}
]
})
.to_string();
let (url, rx) = spawn_mock_server(vec![
MockResponse {
status_line: "HTTP/1.1 200 OK",
body: attempt1_envelope,
},
MockResponse {
status_line: "HTTP/1.1 200 OK",
body: fallback_envelope,
},
]);
let distiller = distiller_for(&url);
let result = distiller.distill("hello transcript").expect("distill ok");
assert_eq!(result.summary, "Recovered via fallback.");
let first = rx.recv().expect("first request");
assert!(first.contains("response_format"));
let second = rx.recv().expect("second request");
assert!(!second.contains("response_format"));
}
#[test]
fn both_attempts_garbage_returns_distill_error() {
let garbage_envelope = json!({
"choices": [
{"message": {"content": "not json at all"}}
]
})
.to_string();
let (url, _rx) = spawn_mock_server(vec![
MockResponse {
status_line: "HTTP/1.1 200 OK",
body: garbage_envelope.clone(),
},
MockResponse {
status_line: "HTTP/1.1 200 OK",
body: garbage_envelope,
},
]);
let distiller = distiller_for(&url);
let err = distiller.distill("hello transcript").unwrap_err();
assert!(matches!(err, Error::Distill(_)));
}
#[test]
fn transport_error_does_not_trigger_fallback() {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind mock listener");
let addr = listener.local_addr().expect("local addr");
let (tx, rx) = mpsc::channel::<()>();
thread::spawn(move || {
for stream in listener.incoming() {
let Ok(stream) = stream else {
return;
};
let _ = tx.send(());
thread::sleep(Duration::from_secs(10));
drop(stream);
}
});
let url = format!("http://{addr}");
let distiller = OpenAiCompatDistiller::new(&url, "test-model")
.expect("construct distiller")
.with_timeout_secs(1);
let err = distiller.distill("hello transcript").unwrap_err();
assert!(matches!(err, Error::Distill(_)));
assert!(rx.recv_timeout(Duration::from_secs(5)).is_ok());
assert!(rx.recv_timeout(Duration::from_millis(100)).is_err());
}
const ALLOWED_TOP_LEVEL_KEYS: &[&str] = &[
"model",
"messages",
"response_format",
"temperature",
"top_p",
"n",
"stop",
"max_tokens",
"max_completion_tokens",
"seed",
"frequency_penalty",
"presence_penalty",
"logit_bias",
"user",
"stream",
"tools",
"tool_choice",
];
const FORBIDDEN_VENDOR_KEYS: &[&str] = &[
"options",
"num_ctx",
"num_batch",
"keep_alive",
"num_predict",
"mirostat",
"repeat_penalty",
"tfs_z",
"num_gpu",
"main_gpu",
];
const ALLOWED_MESSAGE_KEYS: &[&str] =
&["role", "content", "name", "tool_calls", "tool_call_id"];
fn assert_no_vendor_keys(value: &Value) {
match value {
Value::Object(map) => {
for (key, child) in map {
assert!(
!FORBIDDEN_VENDOR_KEYS.contains(&key.as_str()),
"vendor-specific field {key:?} leaked into the request body"
);
assert_no_vendor_keys(child);
}
}
Value::Array(items) => items.iter().for_each(assert_no_vendor_keys),
_ => {}
}
}
fn assert_openai_portable(body: &Value, style: Option<SchemaStyle>) {
let obj = body
.as_object()
.expect("request body must be a JSON object");
for key in obj.keys() {
assert!(
ALLOWED_TOP_LEVEL_KEYS.contains(&key.as_str()),
"non-standard top-level request key {key:?} (not in the OpenAI allowlist)"
);
}
assert_no_vendor_keys(body);
let messages = obj
.get("messages")
.and_then(Value::as_array)
.expect("request body must have a messages array");
for message in messages {
let msg_obj = message.as_object().expect("each message must be an object");
for key in msg_obj.keys() {
assert!(
ALLOWED_MESSAGE_KEYS.contains(&key.as_str()),
"non-standard message key {key:?}"
);
}
}
match style {
None => assert!(
!obj.contains_key("response_format"),
"prompt-embedded path must not emit a response_format"
),
Some(SchemaStyle::OpenAi) => {
let rf = obj
.get("response_format")
.and_then(Value::as_object)
.expect("OpenAi style must emit a response_format object");
assert_eq!(
rf.get("type").and_then(Value::as_str),
Some("json_schema"),
"OpenAi response_format type must be json_schema"
);
assert!(
rf.contains_key("json_schema"),
"OpenAi style must nest the schema under json_schema"
);
assert!(
!rf.contains_key("schema"),
"bare `schema` on response_format is a llama.cpp extension \
and must not appear in OpenAi (strictly portable) style"
);
}
Some(SchemaStyle::LlamaServer) => {
let rf = obj
.get("response_format")
.and_then(Value::as_object)
.expect("LlamaServer style must emit a response_format object");
assert_eq!(
rf.get("type").and_then(Value::as_str),
Some("json_object"),
"LlamaServer response_format type must be json_object"
);
assert!(
rf.contains_key("schema"),
"LlamaServer style is expected to carry the schema deviation"
);
}
}
}
fn conformance_messages() -> Vec<ChatMessage<'static>> {
vec![
ChatMessage {
role: "system",
content: SYSTEM_PROMPT.to_owned(),
},
ChatMessage {
role: "user",
content: "transcript".to_owned(),
},
]
}
#[test]
fn request_body_is_openai_portable_struct_level() {
let openai = ChatRequest {
model: "test-model",
messages: conformance_messages(),
response_format: Some(response_format_payload(SchemaStyle::OpenAi)),
};
assert_openai_portable(
&serde_json::to_value(&openai).expect("serialize"),
Some(SchemaStyle::OpenAi),
);
let llama = ChatRequest {
model: "test-model",
messages: conformance_messages(),
response_format: Some(response_format_payload(SchemaStyle::LlamaServer)),
};
assert_openai_portable(
&serde_json::to_value(&llama).expect("serialize"),
Some(SchemaStyle::LlamaServer),
);
let prompt_embedded = ChatRequest {
model: "test-model",
messages: conformance_messages(),
response_format: None,
};
assert_openai_portable(
&serde_json::to_value(&prompt_embedded).expect("serialize"),
None,
);
}
#[test]
fn with_api_key_sends_authorization_header() {
let body = json!({"summary": "x", "facts": []}).to_string();
let envelope = json!({"choices": [{"message": {"content": body}}]}).to_string();
let (url, rx) = spawn_mock_server(vec![MockResponse {
status_line: "HTTP/1.1 200 OK",
body: envelope,
}]);
let distiller = OpenAiCompatDistiller::new(&url, "test-model")
.expect("construct distiller")
.with_timeout_secs(5)
.with_api_key("my-secret-key");
distiller.distill("transcript").expect("distill ok");
let request = rx.recv().expect("captured request");
assert!(
request
.to_ascii_lowercase()
.contains("authorization: bearer my-secret-key"),
"expected Authorization header in request:\n{request}"
);
}
#[test]
fn without_api_key_no_authorization_header() {
let body = json!({"summary": "x", "facts": []}).to_string();
let envelope = json!({"choices": [{"message": {"content": body}}]}).to_string();
let (url, rx) = spawn_mock_server(vec![MockResponse {
status_line: "HTTP/1.1 200 OK",
body: envelope,
}]);
let distiller = distiller_for(&url);
distiller.distill("transcript").expect("distill ok");
let request = rx.recv().expect("captured request");
assert!(
!request.to_ascii_lowercase().contains("authorization"),
"expected no Authorization header in request:\n{request}"
);
}
#[test]
fn request_body_is_openai_portable_wire_level() {
let body = json!({
"summary": "x",
"facts": []
})
.to_string();
let envelope = json!({
"choices": [
{"message": {"content": body}}
]
})
.to_string();
let (url, rx) = spawn_mock_server(vec![MockResponse {
status_line: "HTTP/1.1 200 OK",
body: envelope,
}]);
let distiller = distiller_for(&url);
distiller.distill("hello transcript").expect("distill ok");
let request = rx.recv().expect("captured request");
let body_start = request.find("\n\n").map_or(0, |i| i + 2);
let parsed: Value =
serde_json::from_str(&request[body_start..]).expect("request body is valid JSON");
assert_openai_portable(&parsed, Some(SchemaStyle::OpenAi));
}
}