Skip to main content

sim_codec_chat/providers/openai/
encode.rs

1use serde_json::{Value, json};
2use sim_kernel::{CodecId, Error, Expr, Result};
3
4use crate::output_grammar::{OutputGrammarDialect, output_grammar_required, output_grammar_text};
5use crate::{is_model_request_expr, validate_chat_transcript};
6
7use super::OpenAiRequestOptions;
8use super::common::{
9    codec_error, codec_eval_to_codec, flatten_expr, list_field, map_field, marker_is_true,
10    optional_u64_field, string_field, symbol_field,
11};
12
13/// Encodes a model-request transcript into OpenAI chat-completion JSON.
14pub fn encode_openai_request(expr: &Expr, options: &OpenAiRequestOptions) -> Result<Vec<u8>> {
15    if !is_model_request_expr(expr) {
16        return Err(Error::Eval(
17            "openai codec expects a model-request transcript".to_owned(),
18        ));
19    }
20    validate_chat_transcript(expr)?;
21    let entries = request_entries(expr)?;
22    let mut payload = json!({
23        "model": options.model,
24        "stream": options.stream,
25        "messages": transcript_messages(expr)?,
26        "tools": if options.tools { Value::Array(Vec::new()) } else { Value::Null },
27    });
28    attach_output_grammar(entries, &mut payload)?;
29    if options.stream
30        && let Some(object) = payload.as_object_mut()
31    {
32        object.insert("stream_options".to_owned(), json!({"include_usage": true}));
33    }
34    serde_json::to_vec(&payload)
35        .map_err(|err| Error::Eval(format!("openai codec failed to encode request: {err}")))
36}
37
38fn attach_output_grammar(entries: &[(Expr, Expr)], payload: &mut Value) -> Result<()> {
39    let Some(grammar) = output_grammar_text(entries, OutputGrammarDialect::JsonSchema)? else {
40        return Ok(());
41    };
42    let schema = serde_json::from_str::<Value>(&grammar)
43        .map_err(|err| Error::Eval(format!("openai output grammar is not json schema: {err}")))?;
44    let Some(object) = payload.as_object_mut() else {
45        return Err(Error::Eval(
46            "openai request payload must be a json object".to_owned(),
47        ));
48    };
49    object.insert(
50        "response_format".to_owned(),
51        json!({
52            "type": "json_schema",
53            "json_schema": {
54                "name": "sim_output",
55                "strict": output_grammar_required(entries)?,
56                "schema": schema,
57            }
58        }),
59    );
60    Ok(())
61}
62
63fn request_entries(expr: &Expr) -> Result<&[(Expr, Expr)]> {
64    let Expr::Map(entries) = expr else {
65        return Err(Error::Eval(
66            "openai codec expects request transcript as a map".to_owned(),
67        ));
68    };
69    Ok(entries)
70}
71
72/// Encodes a model-response transcript into OpenAI chat-completion JSON.
73pub fn encode_openai_response(expr: &Expr) -> Result<Vec<u8>> {
74    let value = response_json(expr)?;
75    serde_json::to_vec(&value)
76        .map_err(|err| Error::Eval(format!("openai codec failed to encode response: {err}")))
77}
78
79pub(in crate::providers) fn encode_openai_response_for_codec(
80    codec: CodecId,
81    expr: &Expr,
82) -> Result<String> {
83    if !marker_is_true(expr, "model-response") {
84        return Err(codec_error(
85            codec,
86            "openai codec expects a model-response transcript",
87        ));
88    }
89    validate_chat_transcript(expr).map_err(|err| codec_eval_to_codec(codec, err))?;
90    let value = response_json(expr).map_err(|err| codec_eval_to_codec(codec, err))?;
91    serde_json::to_string(&value).map_err(|err| codec_error(codec, err))
92}
93
94fn transcript_messages(expr: &Expr) -> Result<Vec<Value>> {
95    let Expr::Map(entries) = expr else {
96        return Err(Error::Eval(
97            "openai codec expects request transcript as a map".to_owned(),
98        ));
99    };
100    let mut messages = list_field(map_field(entries, "messages")?)?
101        .iter()
102        .map(message_to_json)
103        .collect::<Result<Vec<_>>>()?;
104    messages.push(json!({
105        "role": "user",
106        "content": [{
107            "type": "text",
108            "text": flatten_expr(map_field(entries, "task")?),
109        }],
110    }));
111    Ok(messages)
112}
113
114fn message_to_json(expr: &Expr) -> Result<Value> {
115    let Expr::Map(entries) = expr else {
116        return Err(Error::Eval("openai codec message must be a map".to_owned()));
117    };
118    Ok(json!({
119        "role": symbol_field(entries, "role")?,
120        "content": list_field(map_field(entries, "content")?)?
121            .iter()
122            .map(content_part_to_json)
123            .collect::<Result<Vec<_>>>()?,
124    }))
125}
126
127fn content_part_to_json(expr: &Expr) -> Result<Value> {
128    let Expr::Map(entries) = expr else {
129        return Err(Error::Eval(
130            "openai codec content part must be a map".to_owned(),
131        ));
132    };
133    match symbol_field(entries, "type")?.as_str() {
134        "text" => Ok(json!({
135            "type": "text",
136            "text": string_field(entries, "text")?,
137        })),
138        other => Err(Error::Eval(format!(
139            "openai codec does not support content part type {other}"
140        ))),
141    }
142}
143
144fn response_json(expr: &Expr) -> Result<Value> {
145    let Expr::Map(entries) = expr else {
146        return Err(Error::Eval(
147            "openai codec expects response transcript as a map".to_owned(),
148        ));
149    };
150    let model = string_field(entries, "model")?;
151    let finish_reason = symbol_field(entries, "stop-reason")?;
152    Ok(json!({
153        "id": "chatcmpl-sim",
154        "object": "chat.completion",
155        "created": 0,
156        "model": model,
157        "choices": [{
158            "index": 0,
159            "message": {
160                "role": "assistant",
161                "content": response_text(entries)?,
162            },
163            "finish_reason": finish_reason,
164        }],
165        "usage": response_usage(entries)?,
166    }))
167}
168
169fn response_text(entries: &[(Expr, Expr)]) -> Result<String> {
170    list_field(map_field(entries, "content")?)?
171        .iter()
172        .map(text_content)
173        .collect::<Result<Vec<_>>>()
174        .map(|parts| parts.join(""))
175}
176
177fn text_content(expr: &Expr) -> Result<String> {
178    let Expr::Map(entries) = expr else {
179        return Err(Error::Eval(
180            "openai codec content part must be a map".to_owned(),
181        ));
182    };
183    match symbol_field(entries, "type")?.as_str() {
184        "text" => string_field(entries, "text"),
185        other => Err(Error::Eval(format!(
186            "openai codec does not support content part type {other}"
187        ))),
188    }
189}
190
191fn response_usage(entries: &[(Expr, Expr)]) -> Result<Value> {
192    let Some(usage) = entries.iter().find_map(|(field, value)| match field {
193        Expr::Symbol(symbol) if symbol.name.as_ref() == "usage" => Some(value),
194        _ => None,
195    }) else {
196        return Ok(Value::Null);
197    };
198    let Expr::Map(fields) = usage else {
199        return Err(Error::Eval(
200            "openai codec usage field must be a map".to_owned(),
201        ));
202    };
203    let prompt = optional_u64_field(fields, "input-tokens")?;
204    let completion = optional_u64_field(fields, "output-tokens")?;
205    let total = optional_u64_field(fields, "total-tokens")?.or_else(|| {
206        prompt
207            .zip(completion)
208            .map(|(left, right)| left.saturating_add(right))
209    });
210    Ok(json!({
211        "prompt_tokens": prompt.unwrap_or(0),
212        "completion_tokens": completion.unwrap_or(0),
213        "total_tokens": total.unwrap_or(0),
214    }))
215}