ai_lib_rust/client/
endpoint.rs1use 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 fn call_service(
15 &self,
16 service_name: &str,
17 ) -> impl Future<Output = Result<serde_json::Value>> + Send;
18
19 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 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 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 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}