Skip to main content

ai_lib_rust/client/
endpoint.rs

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