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, Debug, Clone)]
526pub struct APIKeyScope {
527 pub r#type: String,
528 pub name: String,
529}
530
531impl std::fmt::Display for APIKeyScope {
532 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
533 write!(f, "{} ({})", self.name, self.r#type)
534 }
535}
536
537#[derive(Serialize, Deserialize, Clone, Debug)]
538pub struct GetMyAccountResponse {
539 pub username: String,
540 pub id: String,
541 pub first_name: String,
542 pub last_name: String,
543 pub email: String,
544 pub scope: Option<APIKeyScope>,
545}
546
547impl GetMyAccountResponse {
548 pub fn to_text(&self) -> String {
549 format!(
550 "ID: {}\nUsername: {}\nName: {} {}\nEmail: {}",
551 self.id, self.username, self.first_name, self.last_name, self.email
552 )
553 }
554}
555
556#[derive(Serialize, Deserialize, Debug, Clone)]
557pub struct ListRuleBook {
558 pub id: String,
559 pub uri: String,
560 pub description: String,
561 pub visibility: RuleBookVisibility,
562 pub tags: Vec<String>,
563 pub created_at: Option<DateTime<Utc>>,
564 pub updated_at: Option<DateTime<Utc>>,
565}
566
567#[derive(Serialize, Deserialize, Debug)]
568pub struct ListRulebooksResponse {
569 pub results: Vec<ListRuleBook>,
570}
571
572#[derive(Serialize, Deserialize, Debug)]
573pub struct CreateRuleBookInput {
574 pub uri: String,
575 pub description: String,
576 pub content: String,
577 pub tags: Vec<String>,
578 #[serde(skip_serializing_if = "Option::is_none")]
579 pub visibility: Option<RuleBookVisibility>,
580}
581
582#[derive(Serialize, Deserialize, Debug)]
583pub struct CreateRuleBookResponse {
584 pub id: String,
585}
586
587impl ListRuleBook {
588 pub fn to_text(&self) -> String {
589 format!(
590 "URI: {}\nDescription: {}\nTags: {}\n",
591 self.uri,
592 self.description,
593 self.tags.join(", ")
594 )
595 }
596}
597
598#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
599pub struct SimpleLLMMessage {
600 #[serde(rename = "role")]
601 pub role: SimpleLLMRole,
602 pub content: String,
603}
604
605#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
606#[serde(rename_all = "lowercase")]
607pub enum SimpleLLMRole {
608 User,
609 Assistant,
610}
611
612impl std::fmt::Display for SimpleLLMRole {
613 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
614 match self {
615 SimpleLLMRole::User => write!(f, "user"),
616 SimpleLLMRole::Assistant => write!(f, "assistant"),
617 }
618 }
619}
620
621#[derive(Debug, Deserialize, Serialize)]
622pub struct SearchDocsRequest {
623 pub keywords: String,
624 pub exclude_keywords: Option<String>,
625 pub limit: Option<u32>,
626}
627
628#[derive(Debug, Deserialize, Serialize)]
629pub struct SearchMemoryRequest {
630 pub keywords: Vec<String>,
631 pub start_time: Option<DateTime<Utc>>,
632 pub end_time: Option<DateTime<Utc>>,
633}
634
635#[derive(Debug, Deserialize, Serialize)]
636pub struct SlackReadMessagesRequest {
637 pub channel: String,
638 pub limit: Option<u32>,
639}
640
641#[derive(Debug, Deserialize, Serialize)]
642pub struct SlackReadRepliesRequest {
643 pub channel: String,
644 pub ts: String,
645}
646
647#[derive(Debug, Deserialize, Serialize)]
648pub struct SlackSendMessageRequest {
649 pub channel: String,
650 pub mrkdwn_text: String,
651 pub thread_ts: Option<String>,
652}
653
654#[derive(Debug, Clone, Default, Serialize)]
655pub struct AgentState {
656 pub agent_model: AgentModel,
657 pub messages: Vec<ChatMessage>,
658 pub tools: Option<Vec<Tool>>,
659
660 pub llm_input: Option<LLMInput>,
661 pub llm_output: Option<LLMOutput>,
662
663 pub metadata: Option<HashMap<String, Value>>,
664}
665
666#[derive(Debug, Clone, Default, Serialize)]
667pub struct LLMOutput {
668 pub new_message: LLMMessage,
669 pub usage: LLMTokenUsage,
670}
671
672impl From<&LLMOutput> for ChatMessage {
673 fn from(value: &LLMOutput) -> Self {
674 let message_content = match &value.new_message.content {
675 LLMMessageContent::String(s) => s.clone(),
676 LLMMessageContent::List(l) => l
677 .iter()
678 .map(|c| match c {
679 LLMMessageTypedContent::Text { text } => text.clone(),
680 LLMMessageTypedContent::ToolCall { .. } => String::new(),
681 LLMMessageTypedContent::ToolResult { content, .. } => content.clone(),
682 LLMMessageTypedContent::Image { .. } => String::new(),
683 })
684 .collect::<Vec<_>>()
685 .join("\n"),
686 };
687 let tool_calls = if let LLMMessageContent::List(items) = &value.new_message.content {
688 let calls: Vec<ToolCall> = items
689 .iter()
690 .filter_map(|item| {
691 if let LLMMessageTypedContent::ToolCall { id, name, args } = item {
692 Some(ToolCall {
693 id: id.clone(),
694 r#type: "function".to_string(),
695 function: FunctionCall {
696 name: name.clone(),
697 arguments: args.to_string(),
698 },
699 })
700 } else {
701 None
702 }
703 })
704 .collect();
705
706 if calls.is_empty() { None } else { Some(calls) }
707 } else {
708 None
709 };
710 ChatMessage {
711 role: Role::Assistant,
712 content: Some(MessageContent::String(message_content)),
713 name: None,
714 tool_calls,
715 tool_call_id: None,
716 usage: None,
717 }
718 }
719}
720
721impl AgentState {
722 pub fn new(
723 agent_model: AgentModel,
724 messages: Vec<ChatMessage>,
725 tools: Option<Vec<Tool>>,
726 ) -> Self {
727 Self {
728 agent_model,
729 messages,
730 tools,
731 llm_input: None,
732 llm_output: None,
733 metadata: None,
734 }
735 }
736
737 pub fn set_messages(&mut self, messages: Vec<ChatMessage>) {
738 self.messages = messages;
739 }
740
741 pub fn set_tools(&mut self, tools: Option<Vec<Tool>>) {
742 self.tools = tools;
743 }
744
745 pub fn set_agent_model(&mut self, agent_model: AgentModel) {
746 self.agent_model = agent_model;
747 }
748
749 pub fn set_llm_input(&mut self, llm_input: Option<LLMInput>) {
750 self.llm_input = llm_input;
751 }
752
753 pub fn set_llm_output(&mut self, new_message: LLMMessage, new_usage: Option<LLMTokenUsage>) {
754 self.llm_output = Some(LLMOutput {
755 new_message,
756 usage: new_usage.unwrap_or_default(),
757 });
758 }
759
760 pub fn append_new_message(&mut self, new_message: ChatMessage) {
761 self.messages.push(new_message);
762 }
763}