Skip to main content

ai_lib_rust/client/
endpoint.rs

1//! Endpoint resolution and service calls
2
3use crate::protocol::{EndpointConfig, ProtocolError};
4use crate::{Error, Result};
5use std::future::Future;
6
7use super::core::AiClient;
8
9pub trait EndpointExt {
10    fn resolve_endpoint(&self, name: &str) -> Result<&EndpointConfig>;
11
12    /// Call a generic service by name. The returned future is `Send` and safe to use across threads.
13    fn call_service(
14        &self,
15        service_name: &str,
16    ) -> impl Future<Output = Result<serde_json::Value>> + Send;
17
18    /// List models available from the provider. The returned future is `Send` and safe to use across threads.
19    fn list_remote_models(&self) -> impl Future<Output = Result<Vec<String>>> + Send;
20}
21
22impl EndpointExt for AiClient {
23    fn resolve_endpoint(&self, name: &str) -> Result<&EndpointConfig> {
24        self.manifest
25            .endpoints
26            .as_ref()
27            .and_then(|eps| eps.get(name))
28            .ok_or_else(|| {
29                Error::Protocol(ProtocolError::NotFound {
30                    id: name.to_string(),
31                    hint: None,
32                })
33            })
34    }
35
36    /// Call a generic service by name.
37    async fn call_service(&self, service_name: &str) -> Result<serde_json::Value> {
38        let service = self
39            .manifest
40            .services
41            .as_ref()
42            .and_then(|services| services.get(service_name))
43            .ok_or_else(|| {
44                Error::Protocol(ProtocolError::NotFound {
45                    id: service_name.to_string(),
46                    hint: None,
47                })
48            })?;
49
50        self.transport
51            .execute_service(
52                &service.path,
53                &service.method,
54                service.headers.as_ref(),
55                service.query_params.as_ref(),
56            )
57            .await
58    }
59
60    /// List models available from the provider.
61    async fn list_remote_models(&self) -> Result<Vec<String>> {
62        let response = self.call_service("list_models").await?;
63
64        let models: Vec<String> = if let Some(data) = response.get("data") {
65            data.as_array()
66                .unwrap_or(&vec![])
67                .iter()
68                .filter_map(|m| {
69                    m.get("id")
70                        .and_then(|id| id.as_str().map(|s| s.to_string()))
71                })
72                .collect()
73        } else if let Some(models) = response.get("models") {
74            // Gemini style
75            models
76                .as_array()
77                .unwrap_or(&vec![])
78                .iter()
79                .filter_map(|m| {
80                    m.get("name")
81                        .and_then(|n| n.as_str().map(|s| s.to_string()))
82                })
83                .collect()
84        } else {
85            vec![]
86        };
87
88        Ok(models)
89    }
90}