Skip to main content

mesh_llm_api_client/
discover.rs

1use crate::{ClientBuilder, InviteToken, MeshApiError, MeshClient, OwnerKeypair};
2
3#[derive(Clone, Debug, Default)]
4pub struct PublicMeshQuery {
5    pub model: Option<String>,
6    pub min_vram_gb: Option<f64>,
7    pub region: Option<String>,
8    pub target_name: Option<String>,
9    pub relays: Vec<String>,
10}
11
12#[derive(Clone, Debug, Eq, PartialEq)]
13pub struct PublicMesh {
14    pub invite_token: String,
15    pub serving: Vec<String>,
16    pub wanted: Vec<String>,
17    pub on_disk: Vec<String>,
18    pub total_vram_bytes: u64,
19    pub node_count: usize,
20    pub client_count: usize,
21    pub max_clients: usize,
22    pub name: Option<String>,
23    pub region: Option<String>,
24    pub mesh_id: Option<String>,
25    pub publisher_npub: String,
26    pub published_at: u64,
27    pub expires_at: Option<u64>,
28}
29
30pub struct AutoConnectResult {
31    pub client: MeshClient,
32    pub selected_mesh: PublicMesh,
33}
34
35impl PublicMesh {
36    pub fn invite_token(&self) -> &str {
37        &self.invite_token
38    }
39
40    pub fn client_builder(
41        &self,
42        owner_keypair: OwnerKeypair,
43    ) -> Result<ClientBuilder, MeshApiError> {
44        ClientBuilder::from_public_mesh(owner_keypair, self)
45    }
46}
47
48impl From<mesh_client::network::nostr::DiscoveredMesh> for PublicMesh {
49    fn from(value: mesh_client::network::nostr::DiscoveredMesh) -> Self {
50        Self {
51            invite_token: value.listing.invite_token,
52            serving: value.listing.serving,
53            wanted: value.listing.wanted,
54            on_disk: value.listing.on_disk,
55            total_vram_bytes: value.listing.total_vram_bytes,
56            node_count: value.listing.node_count,
57            client_count: value.listing.client_count,
58            max_clients: value.listing.max_clients,
59            name: value.listing.name,
60            region: value.listing.region,
61            mesh_id: value.listing.mesh_id,
62            publisher_npub: value.publisher_npub,
63            published_at: value.published_at,
64            expires_at: value.expires_at,
65        }
66    }
67}
68
69pub async fn discover_public_meshes(
70    query: PublicMeshQuery,
71) -> Result<Vec<PublicMesh>, MeshApiError> {
72    let relays = resolve_relays(&query.relays);
73    let filter = mesh_client::network::nostr::MeshFilter {
74        model: query.model.clone(),
75        min_vram_gb: query.min_vram_gb,
76        region: query.region.clone(),
77    };
78
79    let discovered = mesh_client::network::nostr::discover(&relays, &filter, None)
80        .await
81        .map_err(|error| MeshApiError::Discovery {
82            message: error.to_string(),
83        })?;
84
85    Ok(discovered
86        .into_iter()
87        .filter(|mesh| matches_target_name(mesh, query.target_name.as_deref()))
88        .map(Into::into)
89        .collect())
90}
91
92pub async fn create_auto_client(
93    owner_keypair: OwnerKeypair,
94    query: PublicMeshQuery,
95) -> Result<AutoConnectResult, MeshApiError> {
96    let mesh = select_public_mesh(query).await?;
97    let client = mesh.client_builder(owner_keypair)?.build()?;
98    Ok(AutoConnectResult {
99        client,
100        selected_mesh: mesh,
101    })
102}
103
104impl ClientBuilder {
105    pub fn from_public_mesh(
106        owner_keypair: OwnerKeypair,
107        mesh: &PublicMesh,
108    ) -> Result<Self, MeshApiError> {
109        let token = mesh
110            .invite_token
111            .parse::<InviteToken>()
112            .map_err(|message| MeshApiError::InvalidInviteToken { message })?;
113        Ok(Self::new(owner_keypair, token))
114    }
115}
116
117pub async fn select_public_mesh(query: PublicMeshQuery) -> Result<PublicMesh, MeshApiError> {
118    let meshes = discover_public_meshes(query.clone()).await?;
119    let discovered = meshes
120        .into_iter()
121        .map(public_mesh_to_discovered)
122        .collect::<Vec<_>>();
123
124    match mesh_client::network::nostr::smart_auto(
125        &discovered,
126        0.0,
127        query.target_name.as_deref(),
128        None,
129    ) {
130        mesh_client::network::nostr::AutoDecision::Join { mut candidates } => candidates
131            .drain(..)
132            .next()
133            .map(|(_, mesh)| PublicMesh::from(mesh))
134            .ok_or(MeshApiError::NoPublicMeshFound),
135        mesh_client::network::nostr::AutoDecision::StartNew { .. } => {
136            Err(MeshApiError::NoPublicMeshFound)
137        }
138    }
139}
140
141fn resolve_relays(relays: &[String]) -> Vec<String> {
142    if relays.is_empty() {
143        mesh_client::network::nostr::DEFAULT_RELAYS
144            .iter()
145            .map(|relay| (*relay).to_string())
146            .collect()
147    } else {
148        relays.to_vec()
149    }
150}
151
152fn matches_target_name(
153    mesh: &mesh_client::network::nostr::DiscoveredMesh,
154    target_name: Option<&str>,
155) -> bool {
156    let Some(target_name) = target_name else {
157        return true;
158    };
159
160    mesh.listing
161        .name
162        .as_deref()
163        .map(|name| name.eq_ignore_ascii_case(target_name))
164        .unwrap_or(false)
165}
166
167fn public_mesh_to_discovered(mesh: PublicMesh) -> mesh_client::network::nostr::DiscoveredMesh {
168    mesh_client::network::nostr::DiscoveredMesh {
169        listing: mesh_client::network::nostr::MeshListing {
170            invite_token: mesh.invite_token,
171            serving: mesh.serving,
172            wanted: mesh.wanted,
173            on_disk: mesh.on_disk,
174            total_vram_bytes: mesh.total_vram_bytes,
175            node_count: mesh.node_count,
176            client_count: mesh.client_count,
177            max_clients: mesh.max_clients,
178            name: mesh.name,
179            region: mesh.region,
180            mesh_id: mesh.mesh_id,
181        },
182        publisher_npub: mesh.publisher_npub,
183        published_at: mesh.published_at,
184        expires_at: mesh.expires_at,
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use super::{PublicMesh, matches_target_name};
191
192    fn sample_mesh(name: Option<&str>) -> mesh_client::network::nostr::DiscoveredMesh {
193        mesh_client::network::nostr::DiscoveredMesh {
194            listing: mesh_client::network::nostr::MeshListing {
195                invite_token: "mesh-test:abc123".to_string(),
196                serving: vec!["Qwen".to_string()],
197                wanted: vec![],
198                on_disk: vec![],
199                total_vram_bytes: 32_000_000_000,
200                node_count: 2,
201                client_count: 1,
202                max_clients: 0,
203                name: name.map(str::to_string),
204                region: Some("AU".to_string()),
205                mesh_id: Some("mesh-1".to_string()),
206            },
207            publisher_npub: "npub1test".to_string(),
208            published_at: 1,
209            expires_at: Some(2),
210        }
211    }
212
213    #[test]
214    fn target_name_filter_is_case_insensitive() {
215        let mesh = sample_mesh(Some("Mesh-LLM"));
216        assert!(matches_target_name(&mesh, Some("mesh-llm")));
217        assert!(!matches_target_name(&mesh, Some("other")));
218    }
219
220    #[test]
221    fn public_mesh_can_build_client_builder() {
222        let mesh = PublicMesh::from(sample_mesh(Some("mesh-llm")));
223        let owner_keypair = crate::OwnerKeypair::generate();
224        let builder = mesh.client_builder(owner_keypair);
225        assert!(builder.is_ok());
226    }
227}