arknet_sdk/
lib.rs

1//! Public Rust SDK for arknet.
2//!
3//! Provides an async `Client` with OpenAI-compatible methods:
4//! - `chat_completion` — non-streaming completions
5//! - `chat_completion_stream` — streaming SSE completions
6//! - `list_models` — query the on-chain model registry
7//!
8//! # Example
9//!
10//! ```rust,no_run
11//! # async fn demo() -> arknet_sdk::Result<()> {
12//! let client = arknet_sdk::Client::new("http://127.0.0.1:3000")?;
13//! let resp = client.chat_completion(arknet_sdk::ChatRequest {
14//!     model: "meta-llama/Llama-3-8B".into(),
15//!     messages: vec![arknet_sdk::Message {
16//!         role: "user".into(),
17//!         content: "Hello!".into(),
18//!     }],
19//!     max_tokens: Some(64),
20//!     ..Default::default()
21//! }).await?;
22//! println!("{}", resp.choices[0].message.content);
23//! # Ok(())
24//! # }
25//! ```
26
27#![forbid(unsafe_code)]
28#![warn(missing_docs)]
29
30pub mod errors;
31
32use serde::{Deserialize, Serialize};
33
34pub use errors::{Result, SdkError};
35
36/// arknet SDK client. Wraps an HTTP connection to a node's
37/// OpenAI-compatible API surface.
38pub struct Client {
39    base_url: String,
40    http: reqwest::Client,
41    api_key: Option<String>,
42}
43
44impl Client {
45    /// Create a new client pointing at an arknet node.
46    ///
47    /// `base_url` should be the node's HTTP root, e.g.
48    /// `http://127.0.0.1:3000`. The client appends `/v1/...` paths.
49    /// Create a new client pointing at an arknet node.
50    ///
51    /// Reads wallet address from `ARKNET_WALLET` env var if not
52    /// provided explicitly via [`ConnectOptions`].
53    pub fn new(base_url: &str) -> Result<Self> {
54        let base_url = base_url.trim_end_matches('/').to_string();
55        let http = reqwest::Client::builder()
56            .timeout(std::time::Duration::from_secs(120))
57            .build()
58            .map_err(|e| SdkError::Http(e.to_string()))?;
59        let api_key = std::env::var("ARKNET_WALLET").ok();
60        Ok(Self {
61            base_url,
62            http,
63            api_key,
64        })
65    }
66
67    /// Auto-discover a gateway from the on-chain registry.
68    ///
69    /// Fetches the live seed list from `seeds.json` on the arknet
70    /// website, then contacts each seed's `/v1/gateways`. If the
71    /// seed list is unreachable, falls back to the hardcoded list.
72    /// No code changes needed to add new seeds — just edit seeds.json.
73    pub async fn connect(opts: ConnectOptions) -> Result<Self> {
74        let seeds = if opts.seeds.is_empty() {
75            fetch_seeds().await
76        } else {
77            opts.seeds
78        };
79        let http = reqwest::Client::builder()
80            .timeout(std::time::Duration::from_secs(10))
81            .build()
82            .map_err(|e| SdkError::Http(e.to_string()))?;
83
84        for seed in &seeds {
85            let url = format!("{}/v1/gateways", seed.trim_end_matches('/'));
86            let resp = match http.get(&url).send().await {
87                Ok(r) if r.status().is_success() => r,
88                _ => continue,
89            };
90            let body: serde_json::Value = match resp.json().await {
91                Ok(v) => v,
92                Err(_) => continue,
93            };
94            let gateways = body["gateways"].as_array().cloned().unwrap_or_default();
95            // Sort HTTPS first.
96            let mut sorted = gateways;
97            sorted.sort_by_key(|g| {
98                if g["https"].as_bool() == Some(true) {
99                    0
100                } else {
101                    1
102                }
103            });
104            for gw in &sorted {
105                let is_https = gw["https"].as_bool() == Some(true);
106                if opts.require_https && !is_https {
107                    continue;
108                }
109                if let Some(gw_url) = gw["url"].as_str() {
110                    return Self::new(gw_url);
111                }
112            }
113        }
114        Err(SdkError::Http("no reachable gateway found".into()))
115    }
116
117    /// Non-streaming chat completion.
118    pub async fn chat_completion(&self, req: ChatRequest) -> Result<ChatResponse> {
119        let url = format!("{}/v1/chat/completions", self.base_url);
120        let mut builder = self.http.post(&url).json(&req);
121        if let Some(key) = &self.api_key {
122            builder = builder.header("Authorization", format!("Bearer {key}"));
123        }
124        let resp = builder
125            .send()
126            .await
127            .map_err(|e| SdkError::Http(e.to_string()))?;
128
129        if !resp.status().is_success() {
130            let status = resp.status().as_u16();
131            let body = resp.text().await.unwrap_or_default();
132            return Err(SdkError::Api { status, body });
133        }
134
135        resp.json::<ChatResponse>()
136            .await
137            .map_err(|e| SdkError::Http(e.to_string()))
138    }
139
140    /// List models from the on-chain registry.
141    pub async fn list_models(&self) -> Result<ModelsResponse> {
142        let url = format!("{}/v1/models", self.base_url);
143        let mut builder = self.http.get(&url);
144        if let Some(key) = &self.api_key {
145            builder = builder.header("Authorization", format!("Bearer {key}"));
146        }
147        let resp = builder
148            .send()
149            .await
150            .map_err(|e| SdkError::Http(e.to_string()))?;
151
152        if !resp.status().is_success() {
153            let status = resp.status().as_u16();
154            let body = resp.text().await.unwrap_or_default();
155            return Err(SdkError::Api { status, body });
156        }
157
158        resp.json::<ModelsResponse>()
159            .await
160            .map_err(|e| SdkError::Http(e.to_string()))
161    }
162}
163
164const SEEDS_JSON_URL: &str = "https://arknet.arkengel.com/seeds.json";
165const FALLBACK_SEEDS: &[&str] = &["https://api.arknet.arkengel.com"];
166
167/// Fetch the live seed list from the static seeds.json file.
168/// Falls back to the hardcoded list if unreachable.
169async fn fetch_seeds() -> Vec<String> {
170    let client = match reqwest::Client::builder()
171        .timeout(std::time::Duration::from_secs(5))
172        .build()
173    {
174        Ok(c) => c,
175        Err(_) => return FALLBACK_SEEDS.iter().map(|s| s.to_string()).collect(),
176    };
177    let resp = match client.get(SEEDS_JSON_URL).send().await {
178        Ok(r) if r.status().is_success() => r,
179        _ => return FALLBACK_SEEDS.iter().map(|s| s.to_string()).collect(),
180    };
181    let body: serde_json::Value = match resp.json().await {
182        Ok(v) => v,
183        Err(_) => return FALLBACK_SEEDS.iter().map(|s| s.to_string()).collect(),
184    };
185    let urls: Vec<String> = body["seeds"]
186        .as_array()
187        .map(|arr| {
188            arr.iter()
189                .filter_map(|s| s["url"].as_str().map(String::from))
190                .collect()
191        })
192        .unwrap_or_default();
193    if urls.is_empty() {
194        FALLBACK_SEEDS.iter().map(|s| s.to_string()).collect()
195    } else {
196        urls
197    }
198}
199
200// ─── Request / response types ───────────────────────────────────────
201
202/// Options for [`Client::connect`] auto-discovery.
203#[derive(Clone, Debug, Default)]
204pub struct ConnectOptions {
205    /// Seed URLs to discover gateways. Defaults to the arknet seed list.
206    pub seeds: Vec<String>,
207    /// Only connect to HTTPS gateways.
208    pub require_https: bool,
209}
210
211/// Chat completion request.
212#[derive(Clone, Debug, Default, Serialize)]
213pub struct ChatRequest {
214    /// Model identifier.
215    pub model: String,
216    /// Conversation messages.
217    pub messages: Vec<Message>,
218    /// Maximum tokens to generate.
219    #[serde(skip_serializing_if = "Option::is_none")]
220    pub max_tokens: Option<u32>,
221    /// Sampling temperature.
222    #[serde(skip_serializing_if = "Option::is_none")]
223    pub temperature: Option<f64>,
224    /// Whether to stream.
225    #[serde(skip_serializing_if = "Option::is_none")]
226    pub stream: Option<bool>,
227    /// Stop sequences.
228    #[serde(skip_serializing_if = "Option::is_none")]
229    pub stop: Option<Vec<String>>,
230    /// Route only to TEE-capable nodes (confidential inference).
231    #[serde(skip_serializing_if = "Option::is_none")]
232    pub prefer_tee: Option<bool>,
233    /// Route only through HTTPS gateways.
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub require_https: Option<bool>,
236}
237
238/// A chat message.
239#[derive(Clone, Debug, Default, Serialize, Deserialize)]
240pub struct Message {
241    /// Role: "system", "user", or "assistant".
242    pub role: String,
243    /// Message content.
244    pub content: String,
245}
246
247/// Chat completion response.
248#[derive(Clone, Debug, Deserialize)]
249pub struct ChatResponse {
250    /// Request identifier.
251    pub id: String,
252    /// Completions.
253    pub choices: Vec<ChatChoice>,
254    /// Token usage.
255    pub usage: Option<TokenUsage>,
256}
257
258/// A single chat choice.
259#[derive(Clone, Debug, Deserialize)]
260pub struct ChatChoice {
261    /// Index.
262    pub index: u32,
263    /// Generated message.
264    pub message: Message,
265    /// Why generation stopped.
266    pub finish_reason: Option<String>,
267}
268
269/// Token counts.
270#[derive(Clone, Debug, Deserialize)]
271pub struct TokenUsage {
272    /// Input tokens.
273    pub prompt_tokens: u32,
274    /// Output tokens.
275    pub completion_tokens: u32,
276    /// Total.
277    pub total_tokens: u32,
278}
279
280/// Models list response.
281#[derive(Clone, Debug, Deserialize)]
282pub struct ModelsResponse {
283    /// Model entries.
284    pub data: Vec<ModelInfo>,
285}
286
287/// A model entry.
288#[derive(Clone, Debug, Deserialize)]
289pub struct ModelInfo {
290    /// Model identifier.
291    pub id: String,
292    /// Owner.
293    pub owned_by: String,
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn client_trims_trailing_slash() {
302        let c = Client::new("http://localhost:3000/").unwrap();
303        assert_eq!(c.base_url, "http://localhost:3000");
304    }
305
306    #[test]
307    fn chat_request_serializes() {
308        let req = ChatRequest {
309            model: "test".into(),
310            messages: vec![Message {
311                role: "user".into(),
312                content: "hi".into(),
313            }],
314            max_tokens: Some(10),
315            ..Default::default()
316        };
317        let json = serde_json::to_string(&req).unwrap();
318        assert!(json.contains("\"model\":\"test\""));
319        assert!(json.contains("\"max_tokens\":10"));
320        assert!(!json.contains("stream"));
321    }
322
323    #[test]
324    fn chat_response_deserializes() {
325        let json = r#"{
326            "id": "chatcmpl-test",
327            "choices": [{
328                "index": 0,
329                "message": {"role": "assistant", "content": "hello"},
330                "finish_reason": "stop"
331            }],
332            "usage": {
333                "prompt_tokens": 5,
334                "completion_tokens": 1,
335                "total_tokens": 6
336            }
337        }"#;
338        let resp: ChatResponse = serde_json::from_str(json).unwrap();
339        assert_eq!(resp.choices.len(), 1);
340        assert_eq!(resp.choices[0].message.content, "hello");
341    }
342
343    #[test]
344    fn models_response_deserializes() {
345        let json = r#"{
346            "object": "list",
347            "data": [
348                {"id": "llama-3-8b", "object": "model", "created": 0, "owned_by": "user"}
349            ]
350        }"#;
351        let resp: ModelsResponse = serde_json::from_str(json).unwrap();
352        assert_eq!(resp.data.len(), 1);
353        assert_eq!(resp.data[0].id, "llama-3-8b");
354    }
355
356    #[test]
357    fn api_key_from_env() {
358        std::env::set_var("ARKNET_WALLET", "ark1fromenv");
359        let c = Client::new("http://localhost:1234").unwrap();
360        assert_eq!(c.api_key.as_deref(), Some("ark1fromenv"));
361        std::env::remove_var("ARKNET_WALLET");
362    }
363
364    #[test]
365    fn prefer_tee_serialized_when_set() {
366        let req = ChatRequest {
367            model: "test".into(),
368            messages: vec![],
369            prefer_tee: Some(true),
370            ..Default::default()
371        };
372        let json = serde_json::to_string(&req).unwrap();
373        assert!(json.contains("\"prefer_tee\":true"));
374    }
375
376    #[test]
377    fn prefer_tee_omitted_when_none() {
378        let req = ChatRequest {
379            model: "test".into(),
380            messages: vec![],
381            ..Default::default()
382        };
383        let json = serde_json::to_string(&req).unwrap();
384        assert!(!json.contains("prefer_tee"));
385    }
386
387    #[test]
388    fn require_https_serialized_when_set() {
389        let req = ChatRequest {
390            model: "test".into(),
391            messages: vec![],
392            require_https: Some(true),
393            ..Default::default()
394        };
395        let json = serde_json::to_string(&req).unwrap();
396        assert!(json.contains("\"require_https\":true"));
397    }
398
399    #[test]
400    fn connect_options_defaults() {
401        let opts = ConnectOptions::default();
402        assert!(opts.seeds.is_empty());
403        assert!(!opts.require_https);
404    }
405}
arknet_sdk/lib.rs

arknet_sdk/
lib.rs
