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    /// Iterate over transcript entries.
54    pub fn iter(&self) -> impl Iterator<Item = &Entry> {
55        self.entries.iter()
56    }
57
58    /// Number of transcript entries.
59    #[must_use]
60    pub fn len(&self) -> usize {
61        self.entries.len()
62    }
63
64    /// Whether the transcript is empty.
65    #[must_use]
66    pub fn is_empty(&self) -> bool {
67        self.entries.is_empty()
68    }
69
70    /// Push a transcript entry.
71    pub fn push(&mut self, entry: Entry) {
72        self.entries.push(entry);
73    }
74
75    /// Parse a FoundationModels transcript JSON string.
76    ///
77    /// # Errors
78    ///
79    /// Returns [`FMError::DecodingFailure`] if `json` does not match the SDK's
80    /// transcript encoding.
81    pub fn from_json_str(json: &str) -> Result<Self, FMError> {
82        let root: Value = serde_json::from_str(json)
83            .map_err(|error| FMError::DecodingFailure(error.to_string()))?;
84        let entries = root
85            .get("transcript")
86            .and_then(|transcript| transcript.get("entries"))
87            .and_then(Value::as_array)
88            .ok_or_else(|| {
89                FMError::DecodingFailure("transcript JSON is missing transcript.entries".into())
90            })?;
91        let entries = entries
92            .iter()
93            .map(Entry::from_json_value)
94            .collect::<Result<Vec<_>, _>>()?;
95        Ok(Self { entries })
96    }
97
98    /// Serialize the transcript back to FoundationModels' native JSON shape.
99    ///
100    /// # Errors
101    ///
102    /// Returns [`FMError::InvalidArgument`] if one of the entries contains an
103    /// invalid JSON payload.
104    pub fn to_json_string(&self) -> Result<String, FMError> {
105        serde_json::to_string(&json!({
106            "version": 1,
107            "type": "FoundationModels.Transcript",
108            "transcript": {
109                "entries": self.entries.iter().map(Entry::to_json_value).collect::<Result<Vec<_>, _>>()?
110            }
111        }))
112        .map_err(|error| FMError::InvalidArgument(format!("failed to encode transcript JSON: {error}")))
113    }
114}
115
116impl From<Vec<Entry>> for Transcript {
117    fn from(entries: Vec<Entry>) -> Self {
118        Self::from_entries(entries)
119    }
120}
121
122impl<'a> IntoIterator for &'a Transcript {
123    type Item = &'a Entry;
124    type IntoIter = std::slice::Iter<'a, Entry>;
125
126    fn into_iter(self) -> Self::IntoIter {
127        self.entries.iter()
128    }
129}
130
131impl IntoIterator for Transcript {
132    type Item = Entry;
133    type IntoIter = std::vec::IntoIter<Entry>;
134
135    fn into_iter(self) -> Self::IntoIter {
136        self.entries.into_iter()
137    }
138}
139
140/// One transcript entry.
141#[derive(Debug, Clone, PartialEq)]
142pub enum Entry {
143    Instructions(TranscriptInstructions),
144    Prompt(TranscriptPrompt),
145    ToolCalls(ToolCalls),
146    ToolOutput(ToolOutput),
147    Response(TranscriptResponse),
148}
149
150impl Entry {
151    /// Best-effort identifier for this transcript entry.
152    #[must_use]
153    pub fn id(&self) -> Option<&str> {
154        match self {
155            Self::Instructions(entry) => entry.id.as_deref(),
156            Self::Prompt(entry) => entry.id.as_deref(),
157            Self::ToolCalls(entry) => entry.id.as_deref(),
158            Self::ToolOutput(entry) => Some(entry.id.as_str()),
159            Self::Response(entry) => entry.id.as_deref(),
160        }
161    }
162
163    fn from_json_value(value: &Value) -> Result<Self, FMError> {
164        let role = value
165            .get("role")
166            .and_then(Value::as_str)
167            .ok_or_else(|| FMError::DecodingFailure("transcript entry is missing role".into()))?;
168        match role {
169            "instructions" => Ok(Self::Instructions(TranscriptInstructions::from_json_value(
170                value,
171            )?)),
172            "user" => Ok(Self::Prompt(TranscriptPrompt::from_json_value(value)?)),
173            "tool" => Ok(Self::ToolOutput(ToolOutput::from_json_value(value)?)),
174            "response" if value.get("toolCalls").is_some() => {
175                Ok(Self::ToolCalls(ToolCalls::from_json_value(value)?))
176            }
177            "response" => Ok(Self::Response(TranscriptResponse::from_json_value(value)?)),
178            other => Err(FMError::DecodingFailure(format!(
179                "unsupported transcript role `{other}`"
180            ))),
181        }
182    }
183
184    fn to_json_value(&self) -> Result<Value, FMError> {
185        match self {
186            Self::Instructions(entry) => entry.to_json_value(),
187            Self::Prompt(entry) => entry.to_json_value(),
188            Self::ToolCalls(entry) => entry.to_json_value(),
189            Self::ToolOutput(entry) => entry.to_json_value(),
190            Self::Response(entry) => entry.to_json_value(),
191        }
192    }
193}
194
195/// An instructions transcript entry.
196#[derive(Debug, Clone, PartialEq)]
197pub struct TranscriptInstructions {
198    pub id: Option<String>,
199    pub instructions: Instructions,
200    pub tool_definitions: Vec<ToolDefinition>,
201}
202
203impl TranscriptInstructions {
204    /// Create an instructions entry.
205    #[must_use]
206    pub fn new(instructions: Instructions) -> Self {
207        Self {
208            id: None,
209            instructions,
210            tool_definitions: Vec::new(),
211        }
212    }
213
214    fn from_json_value(value: &Value) -> Result<Self, FMError> {
215        Ok(Self {
216            id: value
217                .get("id")
218                .and_then(Value::as_str)
219                .map(ToOwned::to_owned),
220            instructions: Instructions::from(parse_segments(value.get("contents"))?),
221            tool_definitions: parse_tool_definitions(value.get("tools"))?,
222        })
223    }
224
225    fn to_json_value(&self) -> Result<Value, FMError> {
226        let mut object = Map::new();
227        object.insert("role".into(), Value::String("instructions".into()));
228        object.insert(
229            "id".into(),
230            Value::String(
231                self.id
232                    .clone()
233                    .unwrap_or_else(|| synthetic_id("instructions")),
234            ),
235        );
236        object.insert(
237            "contents".into(),
238            segments_to_json(self.instructions.segments())?,
239        );
240        if !self.tool_definitions.is_empty() {
241            object.insert(
242                "tools".into(),
243                Value::Array(
244                    self.tool_definitions
245                        .iter()
246                        .map(ToolDefinition::to_transcript_json_value)
247                        .collect(),
248                ),
249            );
250        }
251        Ok(Value::Object(object))
252    }
253}
254
255/// A user-prompt transcript entry.
256#[derive(Debug, Clone, PartialEq)]
257pub struct TranscriptPrompt {
258    pub id: Option<String>,
259    pub prompt: crate::prompt::Prompt,
260    pub options: GenerationOptions,
261    pub response_format: Option<ResponseFormat>,
262}
263
264impl TranscriptPrompt {
265    /// Create a prompt entry.
266    #[must_use]
267    pub fn new(prompt: crate::prompt::Prompt) -> Self {
268        Self {
269            id: None,
270            prompt,
271            options: GenerationOptions::new(),
272            response_format: None,
273        }
274    }
275
276    fn from_json_value(value: &Value) -> Result<Self, FMError> {
277        Ok(Self {
278            id: value
279                .get("id")
280                .and_then(Value::as_str)
281                .map(ToOwned::to_owned),
282            prompt: crate::prompt::Prompt::from(parse_segments(value.get("contents"))?),
283            options: GenerationOptions::from_transcript_json_value(value.get("options")),
284            response_format: value
285                .get("responseFormat")
286                .map(ResponseFormat::from_transcript_json_value)
287                .transpose()?,
288        })
289    }
290
291    fn to_json_value(&self) -> Result<Value, FMError> {
292        let mut object = Map::new();
293        object.insert("role".into(), Value::String("user".into()));
294        object.insert(
295            "id".into(),
296            Value::String(self.id.clone().unwrap_or_else(|| synthetic_id("prompt"))),
297        );
298        object.insert("contents".into(), segments_to_json(self.prompt.segments())?);
299        object.insert("options".into(), self.options.to_transcript_json_value());
300        if let Some(response_format) = &self.response_format {
301            object.insert(
302                "responseFormat".into(),
303                response_format.to_transcript_json_value(),
304            );
305        }
306        Ok(Value::Object(object))
307    }
308}
309
310/// A transcript entry that records tool calls the model made.
311#[derive(Debug, Clone, PartialEq)]
312pub struct ToolCalls {
313    pub id: Option<String>,
314    pub calls: Vec<ToolCall>,
315}
316
317impl<'a> IntoIterator for &'a ToolCalls {
318    type Item = &'a ToolCall;
319    type IntoIter = std::slice::Iter<'a, ToolCall>;
320
321    fn into_iter(self) -> Self::IntoIter {
322        self.calls.iter()
323    }
324}
325
326impl IntoIterator for ToolCalls {
327    type Item = ToolCall;
328    type IntoIter = std::vec::IntoIter<ToolCall>;
329
330    fn into_iter(self) -> Self::IntoIter {
331        self.calls.into_iter()
332    }
333}
334
335impl ToolCalls {
336    /// Create a tool-calls entry.
337    #[must_use]
338    pub fn new(calls: Vec<ToolCall>) -> Self {
339        Self { id: None, calls }
340    }
341
342    /// Borrow the tool calls.
343    #[must_use]
344    pub fn calls(&self) -> &[ToolCall] {
345        &self.calls
346    }
347
348    /// Iterate over tool calls.
349    pub fn iter(&self) -> impl Iterator<Item = &ToolCall> {
350        self.calls.iter()
351    }
352
353    /// Number of tool calls in this entry.
354    #[must_use]
355    pub fn len(&self) -> usize {
356        self.calls.len()
357    }
358
359    /// Whether this tool-call entry is empty.
360    #[must_use]
361    pub fn is_empty(&self) -> bool {
362        self.calls.is_empty()
363    }
364
365    fn from_json_value(value: &Value) -> Result<Self, FMError> {
366        Ok(Self {
367            id: value
368                .get("id")
369                .and_then(Value::as_str)
370                .map(ToOwned::to_owned),
371            calls: value
372                .get("toolCalls")
373                .and_then(Value::as_array)
374                .map_or(&[] as &[Value], Vec::as_slice)
375                .iter()
376                .map(ToolCall::from_json_value)
377                .collect::<Result<Vec<_>, _>>()?,
378        })
379    }
380
381    fn to_json_value(&self) -> Result<Value, FMError> {
382        Ok(json!({
383            "role": "response",
384            "id": self.id.clone().unwrap_or_else(|| synthetic_id("tool-calls")),
385            "toolCalls": self.calls.iter().map(ToolCall::to_json_value).collect::<Result<Vec<_>, _>>()?,
386        }))
387    }
388}
389
390/// One tool call entry.
391#[derive(Debug, Clone, PartialEq)]
392pub struct ToolCall {
393    pub id: String,
394    pub tool_name: String,
395    pub arguments: GeneratedContent,
396}
397
398impl ToolCall {
399    /// Create a tool call.
400    #[must_use]
401    pub fn new(
402        id: impl Into<String>,
403        tool_name: impl Into<String>,
404        arguments: GeneratedContent,
405    ) -> Self {
406        Self {
407            id: id.into(),
408            tool_name: tool_name.into(),
409            arguments,
410        }
411    }
412
413    fn from_json_value(value: &Value) -> Result<Self, FMError> {
414        let arguments = value
415            .get("arguments")
416            .and_then(Value::as_str)
417            .ok_or_else(|| FMError::DecodingFailure("tool call is missing arguments".into()))?;
418        Ok(Self {
419            id: value
420                .get("id")
421                .and_then(Value::as_str)
422                .unwrap_or_default()
423                .to_string(),
424            tool_name: value
425                .get("name")
426                .and_then(Value::as_str)
427                .unwrap_or_default()
428                .to_string(),
429            arguments: GeneratedContent::from_json_str(arguments)?,
430        })
431    }
432
433    fn to_json_value(&self) -> Result<Value, FMError> {
434        Ok(json!({
435            "id": self.id,
436            "name": self.tool_name,
437            "arguments": self.arguments.json_string()?,
438        }))
439    }
440}
441
442/// A tool output transcript entry.
443#[derive(Debug, Clone, PartialEq)]
444pub struct ToolOutput {
445    pub id: String,
446    pub tool_name: String,
447    pub tool_call_id: Option<String>,
448    pub segments: Vec<Segment>,
449}
450
451impl ToolOutput {
452    /// Create a tool output entry.
453    #[must_use]
454    pub fn new(
455        id: impl Into<String>,
456        tool_name: impl Into<String>,
457        segments: Vec<Segment>,
458    ) -> Self {
459        let id = id.into();
460        Self {
461            id: id.clone(),
462            tool_name: tool_name.into(),
463            tool_call_id: Some(id),
464            segments,
465        }
466    }
467
468    fn from_json_value(value: &Value) -> Result<Self, FMError> {
469        Ok(Self {
470            id: value
471                .get("id")
472                .and_then(Value::as_str)
473                .unwrap_or_default()
474                .to_string(),
475            tool_name: value
476                .get("toolName")
477                .and_then(Value::as_str)
478                .unwrap_or_default()
479                .to_string(),
480            tool_call_id: value
481                .get("toolCallID")
482                .and_then(Value::as_str)
483                .map(ToOwned::to_owned),
484            segments: parse_segments(value.get("contents"))?,
485        })
486    }
487
488    fn to_json_value(&self) -> Result<Value, FMError> {
489        Ok(json!({
490            "role": "tool",
491            "id": self.id,
492            "toolCallID": self.tool_call_id.clone().unwrap_or_else(|| self.id.clone()),
493            "toolName": self.tool_name,
494            "contents": segments_to_json(&self.segments)?,
495        }))
496    }
497}
498
499/// A model response transcript entry.
500#[derive(Debug, Clone, PartialEq)]
501pub struct TranscriptResponse {
502    pub id: Option<String>,
503    pub asset_ids: Vec<String>,
504    pub segments: Vec<Segment>,
505}
506
507impl TranscriptResponse {
508    /// Create a response entry.
509    #[must_use]
510    pub fn new(segments: Vec<Segment>) -> Self {
511        Self {
512            id: None,
513            asset_ids: Vec::new(),
514            segments,
515        }
516    }
517
518    fn from_json_value(value: &Value) -> Result<Self, FMError> {
519        Ok(Self {
520            id: value
521                .get("id")
522                .and_then(Value::as_str)
523                .map(ToOwned::to_owned),
524            asset_ids: value
525                .get("assets")
526                .and_then(Value::as_array)
527                .map(|assets| {
528                    assets
529                        .iter()
530                        .filter_map(Value::as_str)
531                        .map(ToOwned::to_owned)
532                        .collect::<Vec<_>>()
533                })
534                .unwrap_or_default(),
535            segments: parse_segments(value.get("contents"))?,
536        })
537    }
538
539    fn to_json_value(&self) -> Result<Value, FMError> {
540        Ok(json!({
541            "role": "response",
542            "id": self.id.clone().unwrap_or_else(|| synthetic_id("response")),
543            "assets": self.asset_ids,
544            "contents": segments_to_json(&self.segments)?,
545        }))
546    }
547}
548
549fn parse_segments(value: Option<&Value>) -> Result<Vec<Segment>, FMError> {
550    value
551        .and_then(Value::as_array)
552        .map_or(&[] as &[Value], Vec::as_slice)
553        .iter()
554        .map(|segment| {
555            let segment_type = segment
556                .get("type")
557                .and_then(Value::as_str)
558                .ok_or_else(|| FMError::DecodingFailure("segment is missing type".into()))?;
559            match segment_type {
560                "text" => Ok(Segment::Text(TextSegment {
561                    id: segment
562                        .get("id")
563                        .and_then(Value::as_str)
564                        .map(ToOwned::to_owned),
565                    text: segment
566                        .get("text")
567                        .and_then(Value::as_str)
568                        .unwrap_or_default()
569                        .to_string(),
570                })),
571                "structure" => {
572                    let structure = segment.get("structure").ok_or_else(|| {
573                        FMError::DecodingFailure("structured segment is missing structure".into())
574                    })?;
575                    let content = structure.get("content").ok_or_else(|| {
576                        FMError::DecodingFailure("structured segment is missing content".into())
577                    })?;
578                    Ok(Segment::Structure(StructuredSegment {
579                        id: segment
580                            .get("id")
581                            .and_then(Value::as_str)
582                            .map(ToOwned::to_owned),
583                        source: structure
584                            .get("source")
585                            .and_then(Value::as_str)
586                            .unwrap_or("GeneratedContent")
587                            .to_string(),
588                        content: GeneratedContent::from_json_str(
589                            &serde_json::to_string(content).map_err(|error| {
590                                FMError::InvalidArgument(format!(
591                                    "structured segment content is not valid JSON: {error}"
592                                ))
593                            })?,
594                        )?,
595                    }))
596                }
597                other => Err(FMError::DecodingFailure(format!(
598                    "unsupported segment type `{other}`"
599                ))),
600            }
601        })
602        .collect()
603}
604
605fn segments_to_json(segments: &[Segment]) -> Result<Value, FMError> {
606    Ok(Value::Array(
607        segments
608            .iter()
609            .map(|segment| match segment {
610                Segment::Text(TextSegment { id, text }) => Ok(json!({
611                    "type": "text",
612                    "id": id.clone().unwrap_or_else(|| synthetic_id("segment-text")),
613                    "text": text,
614                })),
615                Segment::Structure(StructuredSegment {
616                    id,
617                    source,
618                    content,
619                }) => {
620                    let content_value: Value = serde_json::from_str(&content.json_string()?)
621                        .map_err(|error| {
622                            FMError::InvalidArgument(format!(
623                                "structured segment content is not valid JSON: {error}"
624                            ))
625                        })?;
626                    Ok(json!({
627                        "type": "structure",
628                        "id": id.clone().unwrap_or_else(|| synthetic_id("segment-structure")),
629                        "structure": {
630                            "source": source,
631                            "content": content_value,
632                        }
633                    }))
634                }
635            })
636            .collect::<Result<Vec<_>, _>>()?,
637    ))
638}
639
640fn parse_tool_definitions(value: Option<&Value>) -> Result<Vec<ToolDefinition>, FMError> {
641    value
642        .and_then(Value::as_array)
643        .map_or(&[] as &[Value], Vec::as_slice)
644        .iter()
645        .map(|tool| {
646            let function = tool.get("function").ok_or_else(|| {
647                FMError::DecodingFailure("tool definition is missing function body".into())
648            })?;
649            let parameters = function.get("parameters").ok_or_else(|| {
650                FMError::DecodingFailure("tool definition is missing parameters".into())
651            })?;
652            Ok(ToolDefinition::new(
653                function
654                    .get("name")
655                    .and_then(Value::as_str)
656                    .unwrap_or_default(),
657                function
658                    .get("description")
659                    .and_then(Value::as_str)
660                    .unwrap_or_default(),
661                crate::schema::GenerationSchema::from_json_schema_unchecked(
662                    serde_json::to_string(parameters).map_err(|error| {
663                        FMError::InvalidArgument(format!(
664                            "tool parameters are not valid JSON: {error}"
665                        ))
666                    })?,
667                ),
668            ))
669        })
670        .collect()
671}