Skip to main content

stakpak_api/
lib.rs

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
18// Re-export unified AgentClient as the primary client
19pub use client::{
20    AgentClient, AgentClientConfig, DEFAULT_STAKPAK_ENDPOINT, ModelOptions, StakpakConfig,
21};
22
23// Re-export storage types
24pub 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/// Unified agent provider trait.
33///
34/// Extends `SessionStorage` so that any `AgentProvider` can also manage
35/// sessions and checkpoints.  This avoids passing two separate trait
36/// objects through the CLI call-chain.
37#[async_trait]
38pub trait AgentProvider: SessionStorage + Send + Sync {
39    // Account
40    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    // Rulebooks
47    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    // Chat
60    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    // Search Docs
86    async fn search_docs(&self, input: &SearchDocsRequest) -> Result<Vec<Content>, String>;
87
88    // Memory
89    async fn memorize_session(&self, checkpoint_id: Uuid) -> Result<(), String>;
90    async fn search_memory(&self, input: &SearchMemoryRequest) -> Result<Vec<Content>, String>;
91
92    // Slack
93    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}