1use async_trait::async_trait;
2use futures_util::Stream;
3use models::*;
4use reqwest::header::HeaderMap;
5use rmcp::model::Content;
6use stakpak_shared::models::integrations::openai::{
7 AgentModel, ChatCompletionResponse, ChatCompletionStreamResponse, ChatMessage, Tool,
8};
9use uuid::Uuid;
10
11pub mod client;
12pub mod error;
13pub mod local;
14pub mod models;
15pub mod stakpak;
16pub mod storage;
17
18pub use client::{
20 AgentClient, AgentClientConfig, DEFAULT_STAKPAK_ENDPOINT, ModelOptions, StakpakConfig,
21};
22
23pub use storage::{
25 BoxedSessionStorage, Checkpoint, CheckpointState, CheckpointSummary, CreateCheckpointRequest,
26 CreateSessionRequest as StorageCreateSessionRequest, CreateSessionResult, ListCheckpointsQuery,
27 ListCheckpointsResult, ListSessionsQuery, ListSessionsResult, LocalStorage, Session,
28 SessionStats, SessionStatus, SessionStorage, SessionSummary, SessionVisibility, StakpakStorage,
29 StorageError, UpdateSessionRequest as StorageUpdateSessionRequest,
30};
31
32#[async_trait]
38pub trait AgentProvider: SessionStorage + Send + Sync {
39 async fn get_my_account(&self) -> Result<GetMyAccountResponse, String>;
41 async fn get_billing_info(
42 &self,
43 account_username: &str,
44 ) -> Result<stakpak_shared::models::billing::BillingResponse, String>;
45
46 async fn list_rulebooks(&self) -> Result<Vec<ListRuleBook>, String>;
48 async fn get_rulebook_by_uri(&self, uri: &str) -> Result<RuleBook, String>;
49 async fn create_rulebook(
50 &self,
51 uri: &str,
52 description: &str,
53 content: &str,
54 tags: Vec<String>,
55 visibility: Option<RuleBookVisibility>,
56 ) -> Result<CreateRuleBookResponse, String>;
57 async fn delete_rulebook(&self, uri: &str) -> Result<(), String>;
58
59 async fn chat_completion(
61 &self,
62 model: AgentModel,
63 messages: Vec<ChatMessage>,
64 tools: Option<Vec<Tool>>,
65 session_id: Option<Uuid>,
66 ) -> Result<ChatCompletionResponse, String>;
67 async fn chat_completion_stream(
68 &self,
69 model: AgentModel,
70 messages: Vec<ChatMessage>,
71 tools: Option<Vec<Tool>>,
72 headers: Option<HeaderMap>,
73 session_id: Option<Uuid>,
74 ) -> Result<
75 (
76 std::pin::Pin<
77 Box<dyn Stream<Item = Result<ChatCompletionStreamResponse, ApiStreamError>> + Send>,
78 >,
79 Option<String>,
80 ),
81 String,
82 >;
83 async fn cancel_stream(&self, request_id: String) -> Result<(), String>;
84
85 async fn search_docs(&self, input: &SearchDocsRequest) -> Result<Vec<Content>, String>;
87
88 async fn memorize_session(&self, checkpoint_id: Uuid) -> Result<(), String>;
90 async fn search_memory(&self, input: &SearchMemoryRequest) -> Result<Vec<Content>, String>;
91
92 async fn slack_read_messages(
94 &self,
95 input: &SlackReadMessagesRequest,
96 ) -> Result<Vec<Content>, String>;
97 async fn slack_read_replies(
98 &self,
99 input: &SlackReadRepliesRequest,
100 ) -> Result<Vec<Content>, String>;
101 async fn slack_send_message(
102 &self,
103 input: &SlackSendMessageRequest,
104 ) -> Result<Vec<Content>, String>;
105}