llm_browser_testkit/
endpoints.rs

1//! Endpoint registry — resolves named endpoints and task-type routing.
2
3use std::collections::HashMap;
4
5use crate::scenario::EndpointConfig;
6use crate::scenario::EndpointType;
7
8/// Classification of a task for endpoint routing.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum TaskType {
11    /// LLM-based element targeting (resolving CSS selectors from natural
12    /// language).
13    Targeting,
14    /// LLM-based assertion evaluation.
15    Assertion,
16}
17
18impl TaskType {
19    /// Returns the routing key string for this task type.
20    #[must_use]
21    pub const fn as_str(self) -> &'static str {
22        match self {
23            Self::Targeting => "targeting",
24            Self::Assertion => "assertion",
25        }
26    }
27}
28
29/// Resolved endpoint ready for use in calls.
30#[derive(Debug, Clone)]
31pub struct ResolvedEndpoint {
32    /// Endpoint name.
33    pub name: String,
34    /// Endpoint type.
35    pub endpoint_type: EndpointType,
36    /// Base URL for HTTP-based endpoints.
37    pub url: String,
38    /// Model name (LLM endpoints only).
39    pub model: Option<String>,
40    /// API key / bearer token.
41    pub api_key: Option<String>,
42    /// Custom HTTP headers.
43    pub headers: HashMap<String, String>,
44    /// Command for MCP subprocess endpoints.
45    pub command: Option<String>,
46    /// Arguments for MCP subprocess endpoints.
47    pub args: Vec<String>,
48    /// Input token pricing per 1M tokens.
49    pub input_price_per_1m: f64,
50    /// Output token pricing per 1M tokens.
51    pub output_price_per_1m: f64,
52    /// Flat cost per call.
53    pub per_call_price: f64,
54}
55
56impl ResolvedEndpoint {
57    /// Creates a default LLM endpoint from environment variables.
58    #[must_use]
59    pub fn default_llm() -> Self {
60        Self {
61            name: "default".to_owned(),
62            endpoint_type: EndpointType::Llm,
63            url: crate::llm_base_url(),
64            model: Some(crate::llm_model()),
65            api_key: std::env::var("HARNESS_LLM_API_KEY").ok(),
66            headers: crate::parse_headers_env(),
67            command: None,
68            args: Vec::new(),
69            input_price_per_1m: 0.0,
70            output_price_per_1m: 0.0,
71            per_call_price: 0.0,
72        }
73    }
74}
75
76/// Registry of all configured endpoints with routing logic.
77#[derive(Debug, Clone)]
78pub struct EndpointRegistry {
79    endpoints: HashMap<String, ResolvedEndpoint>,
80    default_for: HashMap<String, String>,
81}
82
83impl EndpointRegistry {
84    /// Builds a registry from the endpoint definitions in scenario config.
85    ///
86    /// Falls back to a default LLM endpoint derived from env vars / flat
87    /// config fields if no `[config.endpoints]` are defined.
88    #[must_use]
89    pub fn from_config(endpoints: &HashMap<String, EndpointConfig>) -> Self {
90        if endpoints.is_empty() {
91            let default_llm = ResolvedEndpoint::default_llm();
92            let mut map = HashMap::new();
93            let mut default_for = HashMap::new();
94            for tt in &[TaskType::Targeting, TaskType::Assertion] {
95                default_for.insert(tt.as_str().to_owned(), "default".to_owned());
96            }
97            map.insert("default".to_owned(), default_llm);
98            return Self {
99                endpoints: map,
100                default_for,
101            };
102        }
103
104        let mut resolved: HashMap<String, ResolvedEndpoint> = HashMap::new();
105        let mut default_for: HashMap<String, String> = HashMap::new();
106
107        for (name, ec) in endpoints {
108            let re = ResolvedEndpoint {
109                name: name.clone(),
110                endpoint_type: ec.endpoint_type.clone(),
111                url: ec
112                    .url
113                    .clone()
114                    .unwrap_or_else(|| match ec.endpoint_type {
115                        EndpointType::Llm => crate::llm_base_url(),
116                        EndpointType::A2a | EndpointType::Mcp => String::new(),
117                    })
118                    .trim_end_matches('/')
119                    .to_owned(),
120                model: ec.model.clone(),
121                api_key: ec.api_key.clone(),
122                headers: ec.headers.clone(),
123                command: ec.command.clone(),
124                args: ec.args.clone(),
125                input_price_per_1m: ec.pricing.as_ref().map_or(0.0, |p| p.input_per_1m_tokens),
126                output_price_per_1m: ec.pricing.as_ref().map_or(0.0, |p| p.output_per_1m_tokens),
127                per_call_price: ec.pricing.as_ref().map_or(0.0, |p| p.per_call),
128            };
129
130            for df in &ec.default_for {
131                default_for.insert(df.clone(), name.clone());
132            }
133
134            resolved.insert(name.clone(), re);
135        }
136
137        Self {
138            endpoints: resolved,
139            default_for,
140        }
141    }
142
143    /// Resolves an endpoint by explicit name.
144    ///
145    /// Returns `None` if no endpoint with the given name exists.
146    #[must_use]
147    pub fn get(&self, name: &str) -> Option<&ResolvedEndpoint> {
148        self.endpoints.get(name)
149    }
150
151    /// Resolves the best endpoint for a given task type.
152    ///
153    /// Checks for a `default_for` mapping first, then falls back to any LLM
154    /// endpoint, then panics (config error).
155    #[must_use]
156    pub fn resolve_for_task(&self, task: TaskType) -> &ResolvedEndpoint {
157        let key = task.as_str();
158        if let Some(name) = self.default_for.get(key) {
159            if let Some(ep) = self.endpoints.get(name) {
160                return ep;
161            }
162        }
163        // Fallback: first LLM endpoint
164        self.endpoints
165            .values()
166            .find(|ep| ep.endpoint_type == EndpointType::Llm)
167            .unwrap_or_else(|| panic!("no LLM endpoint configured for task {key}"))
168    }
169
170    /// Resolves an endpoint: explicit name takes priority, then task-type
171    /// routing, then first LLM endpoint.
172    #[must_use]
173    pub fn resolve(&self, name: Option<&str>, task: TaskType) -> &ResolvedEndpoint {
174        if let Some(n) = name {
175            if let Some(ep) = self.endpoints.get(n) {
176                return ep;
177            }
178        }
179        self.resolve_for_task(task)
180    }
181
182    /// Returns the number of configured endpoints.
183    #[must_use]
184    pub fn len(&self) -> usize {
185        self.endpoints.len()
186    }
187
188    /// Returns true if no endpoints are configured.
189    #[must_use]
190    pub fn is_empty(&self) -> bool {
191        self.endpoints.is_empty()
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use std::collections::HashMap;
199
200    #[test]
201    fn test_task_type_as_str() {
202        assert_eq!(TaskType::Targeting.as_str(), "targeting");
203        assert_eq!(TaskType::Assertion.as_str(), "assertion");
204    }
205
206    #[test]
207    fn test_registry_empty_config() {
208        let endpoints = HashMap::new();
209        let registry = EndpointRegistry::from_config(&endpoints);
210        assert_eq!(registry.len(), 1);
211        let ep = registry.get("default").unwrap();
212        assert_eq!(ep.endpoint_type, EndpointType::Llm);
213    }
214
215    #[test]
216    fn test_registry_resolve_by_name() {
217        let mut endpoints = HashMap::new();
218        endpoints.insert(
219            "vision".to_owned(),
220            EndpointConfig {
221                endpoint_type: EndpointType::Llm,
222                url: Some("https://api.openai.com".into()),
223                model: Some("gpt-4o".into()),
224                ..Default::default()
225            },
226        );
227
228        let registry = EndpointRegistry::from_config(&endpoints);
229        let ep = registry.get("vision");
230        assert!(ep.is_some());
231        assert_eq!(ep.unwrap().model.as_deref(), Some("gpt-4o"));
232    }
233
234    #[test]
235    fn test_resolve_for_task_with_default() {
236        let mut endpoints = HashMap::new();
237        let ec = EndpointConfig {
238            endpoint_type: EndpointType::Llm,
239            url: Some("http://localhost:8080".into()),
240            model: Some("deepseek".into()),
241            default_for: vec!["targeting".to_owned()],
242            ..Default::default()
243        };
244        endpoints.insert("main".to_owned(), ec);
245
246        let registry = EndpointRegistry::from_config(&endpoints);
247        let ep = registry.resolve_for_task(TaskType::Targeting);
248        assert_eq!(ep.name, "main");
249    }
250
251    #[test]
252    fn test_resolve_explicit_overrides_task() {
253        let mut endpoints = HashMap::new();
254        endpoints.insert(
255            "default".to_owned(),
256            EndpointConfig {
257                endpoint_type: EndpointType::Llm,
258                url: Some("http://default".into()),
259                default_for: vec!["targeting".to_owned()],
260                ..Default::default()
261            },
262        );
263        endpoints.insert(
264            "fast".to_owned(),
265            EndpointConfig {
266                endpoint_type: EndpointType::Llm,
267                url: Some("http://fast".into()),
268                ..Default::default()
269            },
270        );
271
272        let registry = EndpointRegistry::from_config(&endpoints);
273        let ep = registry.resolve(Some("fast"), TaskType::Targeting);
274        assert_eq!(ep.name, "fast");
275    }
276
277    #[test]
278    fn test_default_llm_has_env_values() {
279        let ep = ResolvedEndpoint::default_llm();
280        assert_eq!(ep.endpoint_type, EndpointType::Llm);
281        assert!(ep.model.is_some());
282        assert!(!ep.url.is_empty());
283    }
284}
llm_browser_testkit/endpoints.rs

llm_browser_testkit/
endpoints.rs
