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    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 Model types from stakai
24pub use stakai::{Model, ModelCost, ModelLimit};
25
26// Re-export storage types
27pub use storage::{
28    BoxedSessionStorage, Checkpoint, CheckpointState, CheckpointSummary, CreateCheckpointRequest,
29    CreateSessionRequest as StorageCreateSessionRequest, CreateSessionResult, ListCheckpointsQuery,
30    ListCheckpointsResult, ListSessionsQuery, ListSessionsResult, LocalStorage, Session,
31    SessionStats, SessionStatus, SessionStorage, SessionSummary, SessionVisibility, StakpakStorage,
32    StorageError, UpdateSessionRequest as StorageUpdateSessionRequest,
33};
34
35/// Find a model by ID string
36///
37/// Parses the model string and searches the model cache:
38/// - Format "provider/model_id" searches within that specific provider
39/// - Plain "model_id" searches all providers
40///
41/// When `use_stakpak` is true, the model is transformed for Stakpak API routing.
42pub fn find_model(model_str: &str, use_stakpak: bool) -> Option<Model> {
43    const PROVIDERS: &[&str] = &["anthropic", "openai", "google"];
44
45    let (provider_hint, model_id) = parse_model_string(model_str);
46
47    // Search with provider hint first, then fall back to searching all
48    let model = provider_hint
49        .and_then(|p| find_in_provider(p, model_id))
50        .or_else(|| {
51            PROVIDERS
52                .iter()
53                .find_map(|&p| find_in_provider(p, model_id))
54        })?;
55
56    Some(if use_stakpak {
57        transform_for_stakpak(model)
58    } else {
59        model
60    })
61}
62
63/// Parse "provider/model_id" or plain "model_id"
64fn parse_model_string(s: &str) -> (Option<&str>, &str) {
65    match s.find('/') {
66        Some(idx) => {
67            let provider = &s[..idx];
68            let model_id = &s[idx + 1..];
69            let normalized = match provider {
70                "gemini" => "google",
71                p => p,
72            };
73            (Some(normalized), model_id)
74        }
75        None => (None, s),
76    }
77}
78
79/// Find a model by ID within a specific provider
80fn find_in_provider(provider_id: &str, model_id: &str) -> Option<Model> {
81    let models = stakai::load_models_for_provider(provider_id).ok()?;
82
83    // Try exact match first
84    if let Some(model) = models.iter().find(|m| m.id == model_id) {
85        return Some(model.clone());
86    }
87
88    // Try prefix match (e.g., "gpt-5.2-2026-01-15" matches catalog's "gpt-5.2")
89    // Find the longest matching prefix
90    let mut best_match: Option<&Model> = None;
91    let mut best_len = 0;
92
93    for model in &models {
94        if model_id.starts_with(&model.id) && model.id.len() > best_len {
95            best_match = Some(model);
96            best_len = model.id.len();
97        }
98    }
99
100    best_match.cloned()
101}
102
103/// Transform a model for Stakpak API routing
104fn transform_for_stakpak(model: Model) -> Model {
105    Model {
106        id: format!("{}/{}", model.provider, model.id),
107        provider: "stakpak".into(),
108        name: model.name,
109        reasoning: model.reasoning,
110        cost: model.cost,
111        limit: model.limit,
112        release_date: model.release_date,
113    }
114}
115
116/// Unified agent provider trait.
117///
118/// Extends `SessionStorage` so that any `AgentProvider` can also manage
119/// sessions and checkpoints.  This avoids passing two separate trait
120/// objects through the CLI call-chain.
121#[async_trait]
122pub trait AgentProvider: SessionStorage + Send + Sync {
123    // Account
124    async fn get_my_account(&self) -> Result<GetMyAccountResponse, String>;
125    async fn get_billing_info(
126        &self,
127        account_username: &str,
128    ) -> Result<stakpak_shared::models::billing::BillingResponse, String>;
129
130    // Rulebooks
131    async fn list_rulebooks(&self) -> Result<Vec<ListRuleBook>, String>;
132    async fn get_rulebook_by_uri(&self, uri: &str) -> Result<RuleBook, String>;
133    async fn create_rulebook(
134        &self,
135        uri: &str,
136        description: &str,
137        content: &str,
138        tags: Vec<String>,
139        visibility: Option<RuleBookVisibility>,
140    ) -> Result<CreateRuleBookResponse, String>;
141    async fn delete_rulebook(&self, uri: &str) -> Result<(), String>;
142
143    // Chat
144    async fn chat_completion(
145        &self,
146        model: Model,
147        messages: Vec<ChatMessage>,
148        tools: Option<Vec<Tool>>,
149        session_id: Option<Uuid>,
150    ) -> Result<ChatCompletionResponse, String>;
151    async fn chat_completion_stream(
152        &self,
153        model: Model,
154        messages: Vec<ChatMessage>,
155        tools: Option<Vec<Tool>>,
156        headers: Option<HeaderMap>,
157        session_id: Option<Uuid>,
158    ) -> Result<
159        (
160            std::pin::Pin<
161                Box<dyn Stream<Item = Result<ChatCompletionStreamResponse, ApiStreamError>> + Send>,
162            >,
163            Option<String>,
164        ),
165        String,
166    >;
167    async fn cancel_stream(&self, request_id: String) -> Result<(), String>;
168
169    // Search Docs
170    async fn search_docs(&self, input: &SearchDocsRequest) -> Result<Vec<Content>, String>;
171
172    // Memory
173    async fn memorize_session(&self, checkpoint_id: Uuid) -> Result<(), String>;
174    async fn search_memory(&self, input: &SearchMemoryRequest) -> Result<Vec<Content>, String>;
175
176    // Slack
177    async fn slack_read_messages(
178        &self,
179        input: &SlackReadMessagesRequest,
180    ) -> Result<Vec<Content>, String>;
181    async fn slack_read_replies(
182        &self,
183        input: &SlackReadRepliesRequest,
184    ) -> Result<Vec<Content>, String>;
185    async fn slack_send_message(
186        &self,
187        input: &SlackSendMessageRequest,
188    ) -> Result<Vec<Content>, String>;
189
190    // Models
191    async fn list_models(&self) -> Vec<Model>;
192}
stakpak_api/lib.rs

stakpak_api/
lib.rs
