ai-lib-core 0.9.6

AI-Protocol execution runtime core (protocol, client, pipeline, transport)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

//! Endpoint resolution and service calls

use crate::protocol::{EndpointConfig, ProtocolError, ServiceConfig};
use crate::{Error, Result};
use std::collections::HashMap;
use std::future::Future;

use super::core::AiClient;

pub trait EndpointExt {
    fn resolve_endpoint(&self, name: &str) -> Result<&EndpointConfig>;

    /// Call a generic service by name. The returned future is `Send` and safe to use across threads.
    fn call_service(
        &self,
        service_name: &str,
    ) -> impl Future<Output = Result<serde_json::Value>> + Send;

    /// List models available from the provider. The returned future is `Send` and safe to use across threads.
    fn list_remote_models(&self) -> impl Future<Output = Result<Vec<String>>> + Send;
}

impl EndpointExt for AiClient {
    fn resolve_endpoint(&self, name: &str) -> Result<&EndpointConfig> {
        self.manifest
            .endpoints
            .as_ref()
            .and_then(|eps: &HashMap<String, EndpointConfig>| eps.get(name))
            .ok_or_else(|| {
                Error::Protocol(ProtocolError::NotFound {
                    id: name.to_string(),
                    hint: None,
                })
            })
    }

    /// Call a generic service by name.
    async fn call_service(&self, service_name: &str) -> Result<serde_json::Value> {
        let service = self
            .manifest
            .services
            .as_ref()
            .and_then(|services: &HashMap<String, ServiceConfig>| services.get(service_name))
            .ok_or_else(|| {
                Error::Protocol(ProtocolError::NotFound {
                    id: service_name.to_string(),
                    hint: None,
                })
            })?;

        self.transport
            .execute_service(
                &service.path,
                &service.method,
                service.headers.as_ref(),
                service.query_params.as_ref(),
            )
            .await
    }

    /// List models available from the provider.
    async fn list_remote_models(&self) -> Result<Vec<String>> {
        let response = self.call_service("list_models").await?;

        let models: Vec<String> = if let Some(data) = response.get("data") {
            data.as_array()
                .unwrap_or(&vec![])
                .iter()
                .filter_map(|m| {
                    m.get("id")
                        .and_then(|id| id.as_str().map(|s| s.to_string()))
                })
                .collect()
        } else if let Some(models) = response.get("models") {
            // Gemini style
            models
                .as_array()
                .unwrap_or(&vec![])
                .iter()
                .filter_map(|m| {
                    m.get("name")
                        .and_then(|n| n.as_str().map(|s| s.to_string()))
                })
                .collect()
        } else {
            vec![]
        };

        Ok(models)
    }
}