Skip to main content

foundation_models/
transcript.rs

1//! Transcript inspection and restoration.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::time::{SystemTime, UNIX_EPOCH};
5
6use serde_json::{json, Map, Value};
7
8use crate::content::GeneratedContent;
9use crate::error::FMError;
10use crate::generation::GenerationOptions;
11use crate::prompt::{
12    Instructions, ResponseFormat, Segment, StructuredSegment, TextSegment, ToolDefinition,
13};
14
15static NEXT_SYNTHETIC_ID: AtomicU64 = AtomicU64::new(1);
16
17fn synthetic_id(prefix: &str) -> String {
18    let millis = SystemTime::now()
19        .duration_since(UNIX_EPOCH)
20        .unwrap_or_default()
21        .as_millis();
22    let counter = NEXT_SYNTHETIC_ID.fetch_add(1, Ordering::Relaxed);
23    format!("{prefix}-{millis}-{counter}")
24}
25
26/// A session transcript.
27#[derive(Debug, Clone, PartialEq, Default)]
28pub struct Transcript {
29    entries: Vec<Entry>,
30}
31
32impl Transcript {
33    /// Create an empty transcript.
34    #[must_use]
35    pub const fn new() -> Self {
36        Self {
37            entries: Vec::new(),
38        }
39    }
40
41    /// Create a transcript from entries.
42    #[must_use]
43    pub fn from_entries(entries: Vec<Entry>) -> Self {
44        Self { entries }
45    }
46
47    /// Borrow the transcript entries.
48    #[must_use]
49    pub fn entries(&self) -> &[Entry] {
50        &self.entries
51    }
52
53    /// Push a transcript entry.
54    pub fn push(&mut self, entry: Entry) {
55        self.entries.push(entry);
56    }
57
58    /// Parse a FoundationModels transcript JSON string.
59    ///
60    /// # Errors
61    ///
62    /// Returns [`FMError::DecodingFailure`] if `json` does not match the SDK's
63    /// transcript encoding.
64    pub fn from_json_str(json: &str) -> Result<Self, FMError> {
65        let root: Value = serde_json::from_str(json)
66            .map_err(|error| FMError::DecodingFailure(error.to_string()))?;
67        let entries = root
68            .get("transcript")
69            .and_then(|transcript| transcript.get("entries"))
70            .and_then(Value::as_array)
71            .ok_or_else(|| {
72                FMError::DecodingFailure("transcript JSON is missing transcript.entries".into())
73            })?;
74        let entries = entries
75            .iter()
76            .map(Entry::from_json_value)
77            .collect::<Result<Vec<_>, _>>()?;
78        Ok(Self { entries })
79    }
80
81    /// Serialize the transcript back to FoundationModels' native JSON shape.
82    ///
83    /// # Errors
84    ///
85    /// Returns [`FMError::InvalidArgument`] if one of the entries contains an
86    /// invalid JSON payload.
87    pub fn to_json_string(&self) -> Result<String, FMError> {
88        serde_json::to_string(&json!({
89            "version": 1,
90            "type": "FoundationModels.Transcript",
91            "transcript": {
92                "entries": self.entries.iter().map(Entry::to_json_value).collect::<Result<Vec<_>, _>>()?
93            }
94        }))
95        .map_err(|error| FMError::InvalidArgument(format!("failed to encode transcript JSON: {error}")))
96    }
97}
98
99impl From<Vec<Entry>> for Transcript {
100    fn from(entries: Vec<Entry>) -> Self {
101        Self::from_entries(entries)
102    }
103}
104
105/// One transcript entry.
106#[derive(Debug, Clone, PartialEq)]
107pub enum Entry {
108    Instructions(TranscriptInstructions),
109    Prompt(TranscriptPrompt),
110    ToolCalls(ToolCalls),
111    ToolOutput(ToolOutput),
112    Response(TranscriptResponse),
113}
114
115impl Entry {
116    fn from_json_value(value: &Value) -> Result<Self, FMError> {
117        let role = value
118            .get("role")
119            .and_then(Value::as_str)
120            .ok_or_else(|| FMError::DecodingFailure("transcript entry is missing role".into()))?;
121        match role {
122            "instructions" => Ok(Self::Instructions(TranscriptInstructions::from_json_value(
123                value,
124            )?)),
125            "user" => Ok(Self::Prompt(TranscriptPrompt::from_json_value(value)?)),
126            "tool" => Ok(Self::ToolOutput(ToolOutput::from_json_value(value)?)),
127            "response" if value.get("toolCalls").is_some() => {
128                Ok(Self::ToolCalls(ToolCalls::from_json_value(value)?))
129            }
130            "response" => Ok(Self::Response(TranscriptResponse::from_json_value(value)?)),
131            other => Err(FMError::DecodingFailure(format!(
132                "unsupported transcript role `{other}`"
133            ))),
134        }
135    }
136
137    fn to_json_value(&self) -> Result<Value, FMError> {
138        match self {
139            Self::Instructions(entry) => entry.to_json_value(),
140            Self::Prompt(entry) => entry.to_json_value(),
141            Self::ToolCalls(entry) => entry.to_json_value(),
142            Self::ToolOutput(entry) => entry.to_json_value(),
143            Self::Response(entry) => entry.to_json_value(),
144        }
145    }
146}
147
148/// An instructions transcript entry.
149#[derive(Debug, Clone, PartialEq)]
150pub struct TranscriptInstructions {
151    pub id: Option<String>,
152    pub instructions: Instructions,
153    pub tool_definitions: Vec<ToolDefinition>,
154}
155
156impl TranscriptInstructions {
157    fn from_json_value(value: &Value) -> Result<Self, FMError> {
158        Ok(Self {
159            id: value
160                .get("id")
161                .and_then(Value::as_str)
162                .map(ToOwned::to_owned),
163            instructions: Instructions::from(parse_segments(value.get("contents"))?),
164            tool_definitions: parse_tool_definitions(value.get("tools"))?,
165        })
166    }
167
168    fn to_json_value(&self) -> Result<Value, FMError> {
169        let mut object = Map::new();
170        object.insert("role".into(), Value::String("instructions".into()));
171        object.insert(
172            "id".into(),
173            Value::String(
174                self.id
175                    .clone()
176                    .unwrap_or_else(|| synthetic_id("instructions")),
177            ),
178        );
179        object.insert(
180            "contents".into(),
181            segments_to_json(self.instructions.segments())?,
182        );
183        if !self.tool_definitions.is_empty() {
184            object.insert(
185                "tools".into(),
186                Value::Array(
187                    self.tool_definitions
188                        .iter()
189                        .map(ToolDefinition::to_transcript_json_value)
190                        .collect(),
191                ),
192            );
193        }
194        Ok(Value::Object(object))
195    }
196}
197
198/// A user-prompt transcript entry.
199#[derive(Debug, Clone, PartialEq)]
200pub struct TranscriptPrompt {
201    pub id: Option<String>,
202    pub prompt: crate::prompt::Prompt,
203    pub options: GenerationOptions,
204    pub response_format: Option<ResponseFormat>,
205}
206
207impl TranscriptPrompt {
208    fn from_json_value(value: &Value) -> Result<Self, FMError> {
209        Ok(Self {
210            id: value
211                .get("id")
212                .and_then(Value::as_str)
213                .map(ToOwned::to_owned),
214            prompt: crate::prompt::Prompt::from(parse_segments(value.get("contents"))?),
215            options: GenerationOptions::from_transcript_json_value(value.get("options")),
216            response_format: value
217                .get("responseFormat")
218                .map(ResponseFormat::from_transcript_json_value)
219                .transpose()?,
220        })
221    }
222
223    fn to_json_value(&self) -> Result<Value, FMError> {
224        let mut object = Map::new();
225        object.insert("role".into(), Value::String("user".into()));
226        object.insert(
227            "id".into(),
228            Value::String(self.id.clone().unwrap_or_else(|| synthetic_id("prompt"))),
229        );
230        object.insert("contents".into(), segments_to_json(self.prompt.segments())?);
231        object.insert("options".into(), self.options.to_transcript_json_value());
232        if let Some(response_format) = &self.response_format {
233            object.insert(
234                "responseFormat".into(),
235                response_format.to_transcript_json_value(),
236            );
237        }
238        Ok(Value::Object(object))
239    }
240}
241
242/// A transcript entry that records tool calls the model made.
243#[derive(Debug, Clone, PartialEq)]
244pub struct ToolCalls {
245    pub id: Option<String>,
246    pub calls: Vec<ToolCall>,
247}
248
249impl ToolCalls {
250    fn from_json_value(value: &Value) -> Result<Self, FMError> {
251        Ok(Self {
252            id: value
253                .get("id")
254                .and_then(Value::as_str)
255                .map(ToOwned::to_owned),
256            calls: value
257                .get("toolCalls")
258                .and_then(Value::as_array)
259                .map_or(&[] as &[Value], Vec::as_slice)
260                .iter()
261                .map(ToolCall::from_json_value)
262                .collect::<Result<Vec<_>, _>>()?,
263        })
264    }
265
266    fn to_json_value(&self) -> Result<Value, FMError> {
267        Ok(json!({
268            "role": "response",
269            "id": self.id.clone().unwrap_or_else(|| synthetic_id("tool-calls")),
270            "toolCalls": self.calls.iter().map(ToolCall::to_json_value).collect::<Result<Vec<_>, _>>()?,
271        }))
272    }
273}
274
275/// One tool call entry.
276#[derive(Debug, Clone, PartialEq)]
277pub struct ToolCall {
278    pub id: String,
279    pub tool_name: String,
280    pub arguments: GeneratedContent,
281}
282
283impl ToolCall {
284    fn from_json_value(value: &Value) -> Result<Self, FMError> {
285        let arguments = value
286            .get("arguments")
287            .and_then(Value::as_str)
288            .ok_or_else(|| FMError::DecodingFailure("tool call is missing arguments".into()))?;
289        Ok(Self {
290            id: value
291                .get("id")
292                .and_then(Value::as_str)
293                .unwrap_or_default()
294                .to_string(),
295            tool_name: value
296                .get("name")
297                .and_then(Value::as_str)
298                .unwrap_or_default()
299                .to_string(),
300            arguments: GeneratedContent::from_json_str(arguments)?,
301        })
302    }
303
304    fn to_json_value(&self) -> Result<Value, FMError> {
305        Ok(json!({
306            "id": self.id,
307            "name": self.tool_name,
308            "arguments": self.arguments.json_string()?,
309        }))
310    }
311}
312
313/// A tool output transcript entry.
314#[derive(Debug, Clone, PartialEq)]
315pub struct ToolOutput {
316    pub id: String,
317    pub tool_name: String,
318    pub tool_call_id: Option<String>,
319    pub segments: Vec<Segment>,
320}
321
322impl ToolOutput {
323    fn from_json_value(value: &Value) -> Result<Self, FMError> {
324        Ok(Self {
325            id: value
326                .get("id")
327                .and_then(Value::as_str)
328                .unwrap_or_default()
329                .to_string(),
330            tool_name: value
331                .get("toolName")
332                .and_then(Value::as_str)
333                .unwrap_or_default()
334                .to_string(),
335            tool_call_id: value
336                .get("toolCallID")
337                .and_then(Value::as_str)
338                .map(ToOwned::to_owned),
339            segments: parse_segments(value.get("contents"))?,
340        })
341    }
342
343    fn to_json_value(&self) -> Result<Value, FMError> {
344        Ok(json!({
345            "role": "tool",
346            "id": self.id,
347            "toolCallID": self.tool_call_id.clone().unwrap_or_else(|| self.id.clone()),
348            "toolName": self.tool_name,
349            "contents": segments_to_json(&self.segments)?,
350        }))
351    }
352}
353
354/// A model response transcript entry.
355#[derive(Debug, Clone, PartialEq)]
356pub struct TranscriptResponse {
357    pub id: Option<String>,
358    pub asset_ids: Vec<String>,
359    pub segments: Vec<Segment>,
360}
361
362impl TranscriptResponse {
363    fn from_json_value(value: &Value) -> Result<Self, FMError> {
364        Ok(Self {
365            id: value
366                .get("id")
367                .and_then(Value::as_str)
368                .map(ToOwned::to_owned),
369            asset_ids: value
370                .get("assets")
371                .and_then(Value::as_array)
372                .map(|assets| {
373                    assets
374                        .iter()
375                        .filter_map(Value::as_str)
376                        .map(ToOwned::to_owned)
377                        .collect::<Vec<_>>()
378                })
379                .unwrap_or_default(),
380            segments: parse_segments(value.get("contents"))?,
381        })
382    }
383
384    fn to_json_value(&self) -> Result<Value, FMError> {
385        Ok(json!({
386            "role": "response",
387            "id": self.id.clone().unwrap_or_else(|| synthetic_id("response")),
388            "assets": self.asset_ids,
389            "contents": segments_to_json(&self.segments)?,
390        }))
391    }
392}
393
394fn parse_segments(value: Option<&Value>) -> Result<Vec<Segment>, FMError> {
395    value
396        .and_then(Value::as_array)
397        .map_or(&[] as &[Value], Vec::as_slice)
398        .iter()
399        .map(|segment| {
400            let segment_type = segment
401                .get("type")
402                .and_then(Value::as_str)
403                .ok_or_else(|| FMError::DecodingFailure("segment is missing type".into()))?;
404            match segment_type {
405                "text" => Ok(Segment::Text(TextSegment {
406                    id: segment
407                        .get("id")
408                        .and_then(Value::as_str)
409                        .map(ToOwned::to_owned),
410                    text: segment
411                        .get("text")
412                        .and_then(Value::as_str)
413                        .unwrap_or_default()
414                        .to_string(),
415                })),
416                "structure" => {
417                    let structure = segment.get("structure").ok_or_else(|| {
418                        FMError::DecodingFailure("structured segment is missing structure".into())
419                    })?;
420                    let content = structure.get("content").ok_or_else(|| {
421                        FMError::DecodingFailure("structured segment is missing content".into())
422                    })?;
423                    Ok(Segment::Structure(StructuredSegment {
424                        id: segment
425                            .get("id")
426                            .and_then(Value::as_str)
427                            .map(ToOwned::to_owned),
428                        source: structure
429                            .get("source")
430                            .and_then(Value::as_str)
431                            .unwrap_or("GeneratedContent")
432                            .to_string(),
433                        content: GeneratedContent::from_json_str(
434                            &serde_json::to_string(content).map_err(|error| {
435                                FMError::InvalidArgument(format!(
436                                    "structured segment content is not valid JSON: {error}"
437                                ))
438                            })?,
439                        )?,
440                    }))
441                }
442                other => Err(FMError::DecodingFailure(format!(
443                    "unsupported segment type `{other}`"
444                ))),
445            }
446        })
447        .collect()
448}
449
450fn segments_to_json(segments: &[Segment]) -> Result<Value, FMError> {
451    Ok(Value::Array(
452        segments
453            .iter()
454            .map(|segment| match segment {
455                Segment::Text(TextSegment { id, text }) => Ok(json!({
456                    "type": "text",
457                    "id": id.clone().unwrap_or_else(|| synthetic_id("segment-text")),
458                    "text": text,
459                })),
460                Segment::Structure(StructuredSegment {
461                    id,
462                    source,
463                    content,
464                }) => {
465                    let content_value: Value = serde_json::from_str(&content.json_string()?)
466                        .map_err(|error| {
467                            FMError::InvalidArgument(format!(
468                                "structured segment content is not valid JSON: {error}"
469                            ))
470                        })?;
471                    Ok(json!({
472                        "type": "structure",
473                        "id": id.clone().unwrap_or_else(|| synthetic_id("segment-structure")),
474                        "structure": {
475                            "source": source,
476                            "content": content_value,
477                        }
478                    }))
479                }
480            })
481            .collect::<Result<Vec<_>, _>>()?,
482    ))
483}
484
485fn parse_tool_definitions(value: Option<&Value>) -> Result<Vec<ToolDefinition>, FMError> {
486    value
487        .and_then(Value::as_array)
488        .map_or(&[] as &[Value], Vec::as_slice)
489        .iter()
490        .map(|tool| {
491            let function = tool.get("function").ok_or_else(|| {
492                FMError::DecodingFailure("tool definition is missing function body".into())
493            })?;
494            let parameters = function.get("parameters").ok_or_else(|| {
495                FMError::DecodingFailure("tool definition is missing parameters".into())
496            })?;
497            Ok(ToolDefinition::new(
498                function
499                    .get("name")
500                    .and_then(Value::as_str)
501                    .unwrap_or_default(),
502                function
503                    .get("description")
504                    .and_then(Value::as_str)
505                    .unwrap_or_default(),
506                crate::schema::GenerationSchema::from_json_schema_unchecked(
507                    serde_json::to_string(parameters).map_err(|error| {
508                        FMError::InvalidArgument(format!(
509                            "tool parameters are not valid JSON: {error}"
510                        ))
511                    })?,
512                ),
513            ))
514        })
515        .collect()
516}