1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use std::collections::HashMap;
5use uuid::Uuid;
6
7#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
8#[serde(rename_all = "lowercase")]
9pub enum RunType {
10 Chain,
11 Llm,
12 Tool,
13 Retriever,
14 Embedding,
15 Prompt,
16 Runnable,
17 Custom(String),
18}
19
20impl RunType {
21 pub fn as_str(&self) -> &str {
22 match self {
23 RunType::Chain => "chain",
24 RunType::Llm => "llm",
25 RunType::Tool => "tool",
26 RunType::Retriever => "retriever",
27 RunType::Embedding => "embedding",
28 RunType::Prompt => "prompt",
29 RunType::Runnable => "runnable",
30 RunType::Custom(s) => s,
31 }
32 }
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct Run {
37 pub id: Uuid,
38 pub name: String,
39 #[serde(rename = "run_type")]
40 pub run_type: RunType,
41 pub inputs: Value,
42 #[serde(skip_serializing_if = "Option::is_none")]
43 pub outputs: Option<Value>,
44 #[serde(rename = "start_time")]
45 pub start_time: DateTime<Utc>,
46 #[serde(rename = "end_time", skip_serializing_if = "Option::is_none")]
47 pub end_time: Option<DateTime<Utc>>,
48 #[serde(rename = "parent_run_id", skip_serializing_if = "Option::is_none")]
49 pub parent_run_id: Option<Uuid>,
50 #[serde(rename = "trace_id", skip_serializing_if = "Option::is_none")]
51 pub trace_id: Option<Uuid>,
52 #[serde(rename = "dotted_order", skip_serializing_if = "Option::is_none")]
53 pub dotted_order: Option<String>,
54 #[serde(rename = "session_id", skip_serializing_if = "Option::is_none")]
55 pub session_id: Option<String>,
56 #[serde(rename = "session_name", skip_serializing_if = "Option::is_none")]
57 pub session_name: Option<String>,
58 #[serde(rename = "thread_id", skip_serializing_if = "Option::is_none")]
59 pub thread_id: Option<String>,
60 #[serde(skip_serializing_if = "Option::is_none")]
61 pub error: Option<String>,
62 #[serde(skip_serializing_if = "Vec::is_empty")]
63 pub tags: Vec<String>,
64 #[serde(skip_serializing_if = "HashMap::is_empty")]
65 pub extra: HashMap<String, Value>,
66 #[serde(rename = "prompt_tokens", skip_serializing_if = "Option::is_none")]
68 pub prompt_tokens: Option<u64>,
69 #[serde(rename = "completion_tokens", skip_serializing_if = "Option::is_none")]
70 pub completion_tokens: Option<u64>,
71 #[serde(rename = "total_tokens", skip_serializing_if = "Option::is_none")]
72 pub total_tokens: Option<u64>,
73 #[serde(rename = "total_cost", skip_serializing_if = "Option::is_none")]
74 pub total_cost: Option<f64>,
75 #[serde(rename = "prompt_cost", skip_serializing_if = "Option::is_none")]
76 pub prompt_cost: Option<f64>,
77 #[serde(rename = "completion_cost", skip_serializing_if = "Option::is_none")]
78 pub completion_cost: Option<f64>,
79}
80
81impl Run {
82 pub fn new(name: String, run_type: RunType, inputs: Value) -> Self {
83 let id = Uuid::new_v4();
84 let start_time = Utc::now();
85
86 Self {
87 id,
88 name,
89 run_type,
90 inputs,
91 outputs: None,
92 start_time,
93 end_time: None,
94 parent_run_id: None,
95 trace_id: None,
96 dotted_order: None,
97 session_id: None,
98 session_name: None,
99 thread_id: None,
100 error: None,
101 tags: Vec::new(),
102 extra: HashMap::new(),
103 prompt_tokens: None,
104 completion_tokens: None,
105 total_tokens: None,
106 total_cost: None,
107 prompt_cost: None,
108 completion_cost: None,
109 }
110 }
111
112 pub fn generate_dotted_order(&self, parent_dotted_order: Option<&str>) -> String {
113 let timestamp = self.start_time.format("%Y%m%dT%H%M%S");
116 let microseconds = self.start_time.timestamp_subsec_micros();
117 let uuid_str = self.id.to_string(); let current_part = format!("{}{:06}Z{}", timestamp, microseconds, uuid_str);
120
121 if let Some(parent) = parent_dotted_order {
122 format!("{}.{}", parent, current_part)
123 } else {
124 current_part
125 }
126 }
127
128 pub fn set_error(&mut self, error: &str) {
129 self.error = Some(error.to_string());
130 }
131
132 pub fn end(&mut self, outputs: Value) {
133 self.outputs = Some(outputs);
134 self.end_time = Some(Utc::now());
135 }
136}
137
138#[derive(Debug, Clone, Serialize, Deserialize)]
139pub struct RunUpdate {
140 #[serde(skip_serializing_if = "Option::is_none")]
141 pub outputs: Option<Value>,
142 #[serde(rename = "end_time", skip_serializing_if = "Option::is_none")]
143 pub end_time: Option<DateTime<Utc>>,
144 #[serde(skip_serializing_if = "Option::is_none")]
145 pub error: Option<String>,
146 #[serde(rename = "prompt_tokens", skip_serializing_if = "Option::is_none")]
147 pub prompt_tokens: Option<u64>,
148 #[serde(rename = "completion_tokens", skip_serializing_if = "Option::is_none")]
149 pub completion_tokens: Option<u64>,
150 #[serde(rename = "total_tokens", skip_serializing_if = "Option::is_none")]
151 pub total_tokens: Option<u64>,
152 #[serde(rename = "total_cost", skip_serializing_if = "Option::is_none")]
153 pub total_cost: Option<f64>,
154}
155
156impl From<&Run> for RunUpdate {
157 fn from(run: &Run) -> Self {
158 Self {
159 outputs: run.outputs.clone(),
160 end_time: run.end_time,
161 error: run.error.clone(),
162 prompt_tokens: run.prompt_tokens,
163 completion_tokens: run.completion_tokens,
164 total_tokens: run.total_tokens,
165 total_cost: run.total_cost,
166 }
167 }
168}
169