stakpak_api/
models.rs

1use std::collections::HashMap;
2
3use chrono::{DateTime, Utc};
4use rmcp::model::Content;
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use stakpak_shared::models::{
8    integrations::openai::{
9        AgentModel, ChatMessage, FunctionCall, MessageContent, Role, Tool, ToolCall,
10    },
11    llm::{LLMInput, LLMMessage, LLMMessageContent, LLMMessageTypedContent, LLMTokenUsage},
12};
13use uuid::Uuid;
14
15#[derive(Debug, Clone, Deserialize, Serialize)]
16pub enum ApiStreamError {
17    AgentInputInvalid(String),
18    AgentStateInvalid,
19    AgentNotSupported,
20    AgentExecutionLimitExceeded,
21    AgentInvalidResponseStream,
22    InvalidGeneratedCode,
23    CopilotError,
24    SaveError,
25    Unknown(String),
26}
27
28impl From<&str> for ApiStreamError {
29    fn from(error_str: &str) -> Self {
30        match error_str {
31            s if s.contains("Agent not supported") => ApiStreamError::AgentNotSupported,
32            s if s.contains("Agent state is not valid") => ApiStreamError::AgentStateInvalid,
33            s if s.contains("Agent thinking limit exceeded") => {
34                ApiStreamError::AgentExecutionLimitExceeded
35            }
36            s if s.contains("Invalid response stream") => {
37                ApiStreamError::AgentInvalidResponseStream
38            }
39            s if s.contains("Invalid generated code") => ApiStreamError::InvalidGeneratedCode,
40            s if s.contains(
41                "Our copilot is handling too many requests at this time, please try again later.",
42            ) =>
43            {
44                ApiStreamError::CopilotError
45            }
46            s if s
47                .contains("An error occurred while saving your data. Please try again later.") =>
48            {
49                ApiStreamError::SaveError
50            }
51            s if s.contains("Agent input is not valid: ") => {
52                ApiStreamError::AgentInputInvalid(s.replace("Agent input is not valid: ", ""))
53            }
54            _ => ApiStreamError::Unknown(error_str.to_string()),
55        }
56    }
57}
58
59impl From<String> for ApiStreamError {
60    fn from(error_str: String) -> Self {
61        ApiStreamError::from(error_str.as_str())
62    }
63}
64
65#[derive(Debug, Deserialize, Serialize, Clone)]
66pub struct AgentSession {
67    pub id: Uuid,
68    pub title: String,
69    pub agent_id: AgentID,
70    pub visibility: AgentSessionVisibility,
71    pub checkpoints: Vec<AgentCheckpointListItem>,
72    pub created_at: DateTime<Utc>,
73    pub updated_at: DateTime<Utc>,
74}
75
76#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
77pub enum AgentID {
78    #[default]
79    #[serde(rename = "pablo:v1")]
80    PabloV1,
81}
82
83impl std::str::FromStr for AgentID {
84    type Err = String;
85
86    fn from_str(s: &str) -> Result<Self, Self::Err> {
87        match s {
88            "pablo:v1" => Ok(AgentID::PabloV1),
89            _ => Err(format!("Invalid agent ID: {}", s)),
90        }
91    }
92}
93
94#[derive(Debug, Deserialize, Serialize, Clone)]
95pub enum AgentSessionVisibility {
96    #[serde(rename = "PRIVATE")]
97    Private,
98    #[serde(rename = "PUBLIC")]
99    Public,
100}
101
102impl std::fmt::Display for AgentSessionVisibility {
103    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
104        match self {
105            AgentSessionVisibility::Private => write!(f, "PRIVATE"),
106            AgentSessionVisibility::Public => write!(f, "PUBLIC"),
107        }
108    }
109}
110
111#[derive(Debug, Deserialize, Serialize, Clone)]
112pub struct AgentCheckpointListItem {
113    pub id: Uuid,
114    pub status: AgentStatus,
115    pub execution_depth: usize,
116    pub parent: Option<AgentParentCheckpoint>,
117    pub created_at: DateTime<Utc>,
118    pub updated_at: DateTime<Utc>,
119}
120
121#[derive(Debug, Deserialize, Serialize, Clone)]
122pub struct AgentSessionListItem {
123    pub id: Uuid,
124    pub agent_id: AgentID,
125    pub visibility: AgentSessionVisibility,
126    pub created_at: DateTime<Utc>,
127    pub updated_at: DateTime<Utc>,
128}
129
130impl From<AgentSession> for AgentSessionListItem {
131    fn from(item: AgentSession) -> Self {
132        Self {
133            id: item.id,
134            agent_id: item.agent_id,
135            visibility: item.visibility,
136            created_at: item.created_at,
137            updated_at: item.updated_at,
138        }
139    }
140}
141
142#[derive(Debug, Deserialize, Serialize, Clone)]
143pub struct AgentParentCheckpoint {
144    pub id: Uuid,
145}
146#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
147pub enum AgentStatus {
148    #[serde(rename = "RUNNING")]
149    Running,
150    #[serde(rename = "COMPLETE")]
151    Complete,
152    #[serde(rename = "BLOCKED")]
153    Blocked,
154    #[serde(rename = "FAILED")]
155    Failed,
156}
157impl std::fmt::Display for AgentStatus {
158    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
159        match self {
160            AgentStatus::Running => write!(f, "RUNNING"),
161            AgentStatus::Complete => write!(f, "COMPLETE"),
162            AgentStatus::Blocked => write!(f, "BLOCKED"),
163            AgentStatus::Failed => write!(f, "FAILED"),
164        }
165    }
166}
167
168#[derive(Debug, Deserialize, Serialize, Clone)]
169pub struct RunAgentInput {
170    pub checkpoint_id: Uuid,
171    pub input: AgentInput,
172}
173
174impl PartialEq for RunAgentInput {
175    fn eq(&self, other: &Self) -> bool {
176        self.input == other.input
177    }
178}
179
180#[derive(Debug, Deserialize, Serialize, Clone)]
181pub struct RunAgentOutput {
182    pub checkpoint: AgentCheckpointListItem,
183    pub session: AgentSessionListItem,
184    pub output: AgentOutput,
185}
186
187#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
188#[serde(tag = "agent_id")]
189pub enum AgentInput {
190    #[serde(rename = "pablo:v1")]
191    PabloV1 {
192        messages: Option<Vec<ChatMessage>>,
193        node_states: Option<serde_json::Value>,
194    },
195}
196
197impl AgentInput {
198    pub fn new(agent_id: &AgentID) -> Self {
199        match agent_id {
200            AgentID::PabloV1 => AgentInput::PabloV1 {
201                messages: None,
202                node_states: None,
203            },
204        }
205    }
206    pub fn set_user_prompt(&mut self, prompt: Option<String>) {
207        match self {
208            AgentInput::PabloV1 { messages, .. } => {
209                if let Some(prompt) = prompt {
210                    *messages = Some(vec![ChatMessage {
211                        role: Role::User,
212                        content: Some(MessageContent::String(prompt)),
213                        name: None,
214                        tool_calls: None,
215                        tool_call_id: None,
216                        usage: None,
217                    }]);
218                }
219            }
220        }
221    }
222    pub fn get_agent_id(&self) -> AgentID {
223        match self {
224            AgentInput::PabloV1 { .. } => AgentID::PabloV1,
225        }
226    }
227}
228#[derive(Debug, Deserialize, Serialize, Clone)]
229#[serde(tag = "agent_id")]
230pub enum AgentOutput {
231    #[serde(rename = "pablo:v1")]
232    PabloV1 {
233        messages: Vec<ChatMessage>,
234        node_states: serde_json::Value,
235    },
236}
237
238impl AgentOutput {
239    pub fn get_agent_id(&self) -> AgentID {
240        match self {
241            AgentOutput::PabloV1 { .. } => AgentID::PabloV1,
242        }
243    }
244    pub fn get_messages(&self) -> Vec<ChatMessage> {
245        match self {
246            AgentOutput::PabloV1 { messages, .. } => messages.clone(),
247        }
248    }
249    pub fn set_messages(&mut self, new_messages: Vec<ChatMessage>) {
250        match self {
251            AgentOutput::PabloV1 { messages, .. } => *messages = new_messages,
252        }
253    }
254}
255
256#[derive(Deserialize, Serialize, Debug)]
257pub struct Document {
258    pub content: String,
259    pub uri: String,
260    pub provisioner: ProvisionerType,
261}
262
263#[derive(Deserialize, Serialize, Debug)]
264pub struct SimpleDocument {
265    pub uri: String,
266    pub content: String,
267}
268
269#[derive(Deserialize, Serialize, Debug, Clone)]
270pub struct Block {
271    pub id: Uuid,
272    pub provider: String,
273    pub provisioner: ProvisionerType,
274    pub language: String,
275    pub key: String,
276    pub digest: u64,
277    pub references: Vec<Vec<Segment>>,
278    pub kind: String,
279    pub r#type: Option<String>,
280    pub name: Option<String>,
281    pub config: serde_json::Value,
282    pub document_uri: String,
283    pub code: String,
284    pub start_byte: usize,
285    pub end_byte: usize,
286    pub start_point: Point,
287    pub end_point: Point,
288    pub state: Option<serde_json::Value>,
289    pub updated_at: Option<DateTime<Utc>>,
290    pub created_at: Option<DateTime<Utc>>,
291    pub dependents: Vec<DependentBlock>,
292    pub dependencies: Vec<Dependency>,
293    pub api_group_version: Option<ApiGroupVersion>,
294
295    pub generated_summary: Option<String>,
296}
297
298impl Block {
299    pub fn get_uri(&self) -> String {
300        format!(
301            "{}#L{}-L{}",
302            self.document_uri, self.start_point.row, self.end_point.row
303        )
304    }
305}
306
307#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Hash, Clone)]
308pub enum ProvisionerType {
309    #[serde(rename = "Terraform")]
310    Terraform,
311    #[serde(rename = "Kubernetes")]
312    Kubernetes,
313    #[serde(rename = "Dockerfile")]
314    Dockerfile,
315    #[serde(rename = "GithubActions")]
316    GithubActions,
317    #[serde(rename = "None")]
318    None,
319}
320impl std::str::FromStr for ProvisionerType {
321    type Err = String;
322
323    fn from_str(s: &str) -> Result<Self, Self::Err> {
324        match s.to_lowercase().as_str() {
325            "terraform" => Ok(Self::Terraform),
326            "kubernetes" => Ok(Self::Kubernetes),
327            "dockerfile" => Ok(Self::Dockerfile),
328            "github-actions" => Ok(Self::GithubActions),
329            _ => Ok(Self::None),
330        }
331    }
332}
333impl std::fmt::Display for ProvisionerType {
334    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
335        match self {
336            ProvisionerType::Terraform => write!(f, "terraform"),
337            ProvisionerType::Kubernetes => write!(f, "kubernetes"),
338            ProvisionerType::Dockerfile => write!(f, "dockerfile"),
339            ProvisionerType::GithubActions => write!(f, "github-actions"),
340            ProvisionerType::None => write!(f, "none"),
341        }
342    }
343}
344
345#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
346#[serde(untagged)]
347pub enum Segment {
348    Key(String),
349    Index(usize),
350}
351
352impl std::fmt::Display for Segment {
353    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
354        match self {
355            Segment::Key(key) => write!(f, "{}", key),
356            Segment::Index(index) => write!(f, "{}", index),
357        }
358    }
359}
360impl std::fmt::Debug for Segment {
361    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
362        match self {
363            Segment::Key(key) => write!(f, "{}", key),
364            Segment::Index(index) => write!(f, "{}", index),
365        }
366    }
367}
368
369#[derive(Serialize, Deserialize, Debug, Default, Clone, PartialEq)]
370pub struct Point {
371    pub row: usize,
372    pub column: usize,
373}
374
375#[derive(Deserialize, Serialize, Debug, Clone)]
376pub struct DependentBlock {
377    pub key: String,
378}
379
380#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
381pub struct Dependency {
382    pub id: Option<Uuid>,
383    pub expression: Option<String>,
384    pub from_path: Option<Vec<Segment>>,
385    pub to_path: Option<Vec<Segment>>,
386    #[serde(default = "Vec::new")]
387    pub selectors: Vec<DependencySelector>,
388    #[serde(skip_serializing)]
389    pub key: Option<String>,
390    pub digest: Option<u64>,
391    #[serde(default = "Vec::new")]
392    pub from: Vec<Segment>,
393    pub from_field: Option<Vec<Segment>>,
394    pub to_field: Option<Vec<Segment>>,
395    pub start_byte: Option<usize>,
396    pub end_byte: Option<usize>,
397    pub start_point: Option<Point>,
398    pub end_point: Option<Point>,
399    pub satisfied: bool,
400}
401
402#[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
403pub struct DependencySelector {
404    pub references: Vec<Vec<Segment>>,
405    pub operator: DependencySelectorOperator,
406}
407
408#[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
409pub enum DependencySelectorOperator {
410    Equals,
411    NotEquals,
412    In,
413    NotIn,
414    Exists,
415    DoesNotExist,
416}
417
418#[derive(Serialize, Deserialize, Debug, Clone)]
419pub struct ApiGroupVersion {
420    pub alias: String,
421    pub group: String,
422    pub version: String,
423    pub provisioner: ProvisionerType,
424    pub status: APIGroupVersionStatus,
425}
426
427#[derive(Serialize, Deserialize, Debug, Clone)]
428pub enum APIGroupVersionStatus {
429    #[serde(rename = "UNAVAILABLE")]
430    Unavailable,
431    #[serde(rename = "PENDING")]
432    Pending,
433    #[serde(rename = "AVAILABLE")]
434    Available,
435}
436
437#[derive(Serialize, Deserialize, Debug)]
438pub struct BuildCodeIndexInput {
439    pub documents: Vec<SimpleDocument>,
440}
441
442#[derive(Serialize, Deserialize, Debug, Clone)]
443pub struct IndexError {
444    pub uri: String,
445    pub message: String,
446    pub details: Option<serde_json::Value>,
447}
448
449#[derive(Serialize, Deserialize, Debug, Clone)]
450pub struct BuildCodeIndexOutput {
451    pub blocks: Vec<Block>,
452    pub errors: Vec<IndexError>,
453    pub warnings: Vec<IndexError>,
454}
455
456#[derive(Serialize, Deserialize, Debug, Clone)]
457pub struct CodeIndex {
458    pub last_updated: DateTime<Utc>,
459    pub index: BuildCodeIndexOutput,
460}
461
462#[derive(Debug, Deserialize, Serialize, Clone, Default)]
463pub struct AgentSessionStats {
464    pub aborted_tool_calls: u32,
465    pub analysis_period: Option<String>,
466    pub failed_tool_calls: u32,
467    pub from_date: Option<String>,
468    pub sessions_with_activity: u32,
469    pub successful_tool_calls: u32,
470    pub to_date: Option<String>,
471    pub tools_usage: Vec<ToolUsageStats>,
472    pub total_sessions: u32,
473    pub total_time_saved_seconds: Option<u32>,
474    pub total_tool_calls: u32,
475}
476
477#[derive(Debug, Deserialize, Serialize, Clone)]
478pub struct ToolUsageStats {
479    pub display_name: String,
480    pub time_saved_per_call: Option<f64>,
481    pub time_saved_seconds: Option<u32>,
482    pub tool_name: String,
483    pub usage_counts: ToolUsageCounts,
484}
485
486#[derive(Debug, Deserialize, Serialize, Clone)]
487pub struct ToolUsageCounts {
488    pub aborted: u32,
489    pub failed: u32,
490    pub successful: u32,
491    pub total: u32,
492}
493
494#[derive(Serialize, Deserialize, Debug, Clone, Copy, Default)]
495#[serde(rename_all = "UPPERCASE")]
496pub enum RuleBookVisibility {
497    #[default]
498    Public,
499    Private,
500}
501
502#[derive(Serialize, Deserialize, Debug, Clone)]
503pub struct RuleBook {
504    pub id: String,
505    pub uri: String,
506    pub description: String,
507    pub content: String,
508    pub visibility: RuleBookVisibility,
509    pub tags: Vec<String>,
510    pub created_at: Option<DateTime<Utc>>,
511    pub updated_at: Option<DateTime<Utc>>,
512}
513
514#[derive(Serialize, Deserialize, Debug)]
515pub struct ToolsCallParams {
516    pub name: String,
517    pub arguments: Value,
518}
519
520#[derive(Serialize, Deserialize, Debug)]
521pub struct ToolsCallResponse {
522    pub content: Vec<Content>,
523}
524
525#[derive(Serialize, Deserialize, Clone, Debug)]
526pub struct GetMyAccountResponse {
527    pub username: String,
528    pub id: String,
529    pub first_name: String,
530    pub last_name: String,
531}
532
533impl GetMyAccountResponse {
534    pub fn to_text(&self) -> String {
535        format!(
536            "ID: {}\nUsername: {}\nName: {} {}",
537            self.id, self.username, self.first_name, self.last_name
538        )
539    }
540}
541
542#[derive(Serialize, Deserialize, Debug, Clone)]
543pub struct ListRuleBook {
544    pub id: String,
545    pub uri: String,
546    pub description: String,
547    pub visibility: RuleBookVisibility,
548    pub tags: Vec<String>,
549    pub created_at: Option<DateTime<Utc>>,
550    pub updated_at: Option<DateTime<Utc>>,
551}
552
553#[derive(Serialize, Deserialize, Debug)]
554pub struct ListRulebooksResponse {
555    pub results: Vec<ListRuleBook>,
556}
557
558#[derive(Serialize, Deserialize, Debug)]
559pub struct CreateRuleBookInput {
560    pub uri: String,
561    pub description: String,
562    pub content: String,
563    pub tags: Vec<String>,
564    #[serde(skip_serializing_if = "Option::is_none")]
565    pub visibility: Option<RuleBookVisibility>,
566}
567
568#[derive(Serialize, Deserialize, Debug)]
569pub struct CreateRuleBookResponse {
570    pub id: String,
571}
572
573impl ListRuleBook {
574    pub fn to_text(&self) -> String {
575        format!(
576            "URI: {}\nDescription: {}\nTags: {}\n",
577            self.uri,
578            self.description,
579            self.tags.join(", ")
580        )
581    }
582}
583
584#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
585pub struct SimpleLLMMessage {
586    #[serde(rename = "role")]
587    pub role: SimpleLLMRole,
588    pub content: String,
589}
590
591#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
592#[serde(rename_all = "lowercase")]
593pub enum SimpleLLMRole {
594    User,
595    Assistant,
596}
597
598impl std::fmt::Display for SimpleLLMRole {
599    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
600        match self {
601            SimpleLLMRole::User => write!(f, "user"),
602            SimpleLLMRole::Assistant => write!(f, "assistant"),
603        }
604    }
605}
606
607#[derive(Debug, Deserialize, Serialize)]
608pub struct SearchDocsRequest {
609    pub keywords: String,
610    pub exclude_keywords: Option<String>,
611    pub limit: Option<u32>,
612}
613
614#[derive(Debug, Deserialize, Serialize)]
615pub struct SearchMemoryRequest {
616    pub keywords: Vec<String>,
617    pub start_time: Option<DateTime<Utc>>,
618    pub end_time: Option<DateTime<Utc>>,
619}
620
621#[derive(Debug, Deserialize, Serialize)]
622pub struct SlackReadMessagesRequest {
623    pub channel: String,
624    pub limit: Option<u32>,
625}
626
627#[derive(Debug, Deserialize, Serialize)]
628pub struct SlackReadRepliesRequest {
629    pub channel: String,
630    pub ts: String,
631}
632
633#[derive(Debug, Deserialize, Serialize)]
634pub struct SlackSendMessageRequest {
635    pub channel: String,
636    pub mrkdwn_text: String,
637    pub thread_ts: Option<String>,
638}
639
640#[derive(Debug, Clone, Default, Serialize)]
641pub struct AgentState {
642    pub agent_model: AgentModel,
643    pub messages: Vec<ChatMessage>,
644    pub tools: Option<Vec<Tool>>,
645
646    pub llm_input: Option<LLMInput>,
647    pub llm_output: Option<LLMOutput>,
648
649    pub metadata: Option<HashMap<String, Value>>,
650}
651
652#[derive(Debug, Clone, Default, Serialize)]
653pub struct LLMOutput {
654    pub new_message: LLMMessage,
655    pub usage: LLMTokenUsage,
656}
657
658impl From<&LLMOutput> for ChatMessage {
659    fn from(value: &LLMOutput) -> Self {
660        let message_content = match &value.new_message.content {
661            LLMMessageContent::String(s) => s.clone(),
662            LLMMessageContent::List(l) => l
663                .iter()
664                .map(|c| match c {
665                    LLMMessageTypedContent::Text { text } => text.clone(),
666                    LLMMessageTypedContent::ToolCall { .. } => String::new(),
667                    LLMMessageTypedContent::ToolResult { content, .. } => content.clone(),
668                    LLMMessageTypedContent::Image { .. } => String::new(),
669                })
670                .collect::<Vec<_>>()
671                .join("\n"),
672        };
673        let tool_calls = if let LLMMessageContent::List(items) = &value.new_message.content {
674            let calls: Vec<ToolCall> = items
675                .iter()
676                .filter_map(|item| {
677                    if let LLMMessageTypedContent::ToolCall { id, name, args } = item {
678                        Some(ToolCall {
679                            id: id.clone(),
680                            r#type: "function".to_string(),
681                            function: FunctionCall {
682                                name: name.clone(),
683                                arguments: args.to_string(),
684                            },
685                        })
686                    } else {
687                        None
688                    }
689                })
690                .collect();
691
692            if calls.is_empty() { None } else { Some(calls) }
693        } else {
694            None
695        };
696        ChatMessage {
697            role: Role::Assistant,
698            content: Some(MessageContent::String(message_content)),
699            name: None,
700            tool_calls,
701            tool_call_id: None,
702            usage: None,
703        }
704    }
705}
706
707impl AgentState {
708    pub fn new(
709        agent_model: AgentModel,
710        messages: Vec<ChatMessage>,
711        tools: Option<Vec<Tool>>,
712    ) -> Self {
713        Self {
714            agent_model,
715            messages,
716            tools,
717            llm_input: None,
718            llm_output: None,
719            metadata: None,
720        }
721    }
722
723    pub fn set_messages(&mut self, messages: Vec<ChatMessage>) {
724        self.messages = messages;
725    }
726
727    pub fn set_tools(&mut self, tools: Option<Vec<Tool>>) {
728        self.tools = tools;
729    }
730
731    pub fn set_agent_model(&mut self, agent_model: AgentModel) {
732        self.agent_model = agent_model;
733    }
734
735    pub fn set_llm_input(&mut self, llm_input: Option<LLMInput>) {
736        self.llm_input = llm_input;
737    }
738
739    pub fn set_llm_output(&mut self, new_message: LLMMessage, new_usage: Option<LLMTokenUsage>) {
740        self.llm_output = Some(LLMOutput {
741            new_message,
742            usage: new_usage.unwrap_or_default(),
743        });
744    }
745
746    pub fn append_new_message(&mut self, new_message: ChatMessage) {
747        self.messages.push(new_message);
748    }
749}
stakpak_api/models.rs

stakpak_api/
models.rs
