fix: improve deserialization of python-generated model block traces

Signed-off-by: Nick Mitchell <[email protected]>

Author: starpit

pdl-live-react/src-tauri/src/pdl/ast.rs

@@ -57,17 +57,18 @@ pub enum PdlType {
 /// Timing information
 #[derive(Serialize, Deserialize, Debug, Clone, Default)]
 pub struct Timing {
-    start_nanos: u128,
-    end_nanos: u128,
+    // TODO serde_json doesn't support u128, but (below) as_nanos() returns u128...
+    start_nanos: u64,
+    end_nanos: u64,
     timezone: String,
 }
 
 type TimingError = Box<dyn Error + Send + Sync>;
 impl Timing {
-    fn now() -> Result<u128, TimingError> {
+    fn now() -> Result<u64, TimingError> {
         Ok(::std::time::SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)?
-            .as_nanos())
+            .as_nanos() as u64)
     }
 
     pub fn start() -> Result<Timing, TimingError> {
@@ -257,14 +258,16 @@ pub struct FunctionBlock {
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct PdlUsage {
-    // Completion tokens consumed
+    /// Completion tokens consumed
     pub completion_tokens: u64,
-    // Prompt tokens consumed
+    /// Prompt tokens consumed
     pub prompt_tokens: u64,
-    // Completion nanos
-    pub completion_nanos: u64,
-    // Prompt nanos
-    pub prompt_nanos: u64,
+    /// Completion nanos
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub completion_nanos: Option<u64>,
+    /// Prompt nanos
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub prompt_nanos: Option<u64>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, Default, derive_builder::Builder)]
@@ -277,11 +280,17 @@ pub struct ModelBlock {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub input: Option<Box<PdlBlock>>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    pub platform: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub context: Option<Vec<MessageBlock>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "modelResponse")]
     pub model_response: Option<String>,
     #[serde(rename = "pdl__usage")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub pdl_usage: Option<PdlUsage>,
+    #[serde(rename = "pdl__model_input", skip_serializing_if = "Option::is_none")]
+    pub pdl_model_input: Option<Vec<MessageBlock>>,
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

pdl-live-react/src-tauri/src/pdl/interpreter.rs

@@ -841,9 +841,9 @@ impl<'a> Interpreter<'a> {
                     if let Some(usage) = res.final_data {
                         trace.pdl_usage = Some(PdlUsage {
                             prompt_tokens: usage.prompt_eval_count,
-                            prompt_nanos: usage.prompt_eval_duration,
+                            prompt_nanos: Some(usage.prompt_eval_duration),
                             completion_tokens: usage.eval_count,
-                            completion_nanos: usage.eval_duration,
+                            completion_nanos: Some(usage.eval_duration),
                         });
                     }
                     let output_messages = vec![ChatMessage::assistant(response_string)];