harn_vm/stdlib/template/
llm_context.rs1use crate::value::VmDictExt;
26use std::cell::RefCell;
27use std::collections::BTreeMap;
28use std::sync::atomic::{AtomicU64, Ordering};
29
30use crate::value::VmValue;
31
32#[derive(Debug, Clone)]
36pub struct LlmRenderContext {
37 pub provider: String,
38 pub model: String,
39 pub family: String,
40 pub capabilities: VmValue,
43}
44
45impl LlmRenderContext {
46 pub fn resolve(provider: &str, model: &str) -> Self {
49 let (provider, model) = crate::llm::managed_supply::logical_route(provider, model)
50 .unwrap_or_else(|_| (provider.to_string(), model.to_string()));
51 let caps = crate::llm::capabilities::lookup(&provider, &model);
52 let capabilities =
53 crate::llm::config_builtins::capabilities_to_vm_value(&provider, &model, &caps);
54 Self {
55 family: crate::llm_config::model_family(&provider, &model),
56 provider,
57 model,
58 capabilities,
59 }
60 }
61
62 pub fn to_vm_value(&self) -> VmValue {
65 let mut dict = BTreeMap::new();
66 dict.put_str("provider", self.provider.as_str());
67 dict.put_str("model", self.model.as_str());
68 dict.put_str("family", self.family.as_str());
69 dict.insert("capabilities".to_string(), self.capabilities.clone());
70 VmValue::dict(dict)
71 }
72}
73
74thread_local! {
75 static LLM_RENDER_STACK: RefCell<Vec<LlmRenderContextFrame>> = const { RefCell::new(Vec::new()) };
76}
77
78static NEXT_LLM_RENDER_FRAME_ID: AtomicU64 = AtomicU64::new(1);
79
80#[derive(Debug, Clone)]
84pub(crate) struct LlmRenderContextFrame {
85 id: u64,
86 context: LlmRenderContext,
87}
88
89fn next_frame_id() -> u64 {
90 NEXT_LLM_RENDER_FRAME_ID.fetch_add(1, Ordering::Relaxed)
91}
92
93fn push_llm_render_context_frame(ctx: LlmRenderContext) -> u64 {
94 let id = next_frame_id();
95 LLM_RENDER_STACK.with(|stack| {
96 stack
97 .borrow_mut()
98 .push(LlmRenderContextFrame { id, context: ctx });
99 });
100 id
101}
102
103pub fn push_llm_render_context(ctx: LlmRenderContext) {
107 push_llm_render_context_frame(ctx);
108}
109
110pub fn pop_llm_render_context() -> Option<LlmRenderContext> {
114 LLM_RENDER_STACK.with(|stack| stack.borrow_mut().pop().map(|frame| frame.context))
115}
116
117pub fn current_llm_render_context() -> Option<LlmRenderContext> {
121 LLM_RENDER_STACK.with(|stack| stack.borrow().last().map(|frame| frame.context.clone()))
122}
123
124pub(crate) fn swap_llm_render_stack(
130 next: Vec<LlmRenderContextFrame>,
131) -> Vec<LlmRenderContextFrame> {
132 LLM_RENDER_STACK.with(|stack| std::mem::replace(&mut *stack.borrow_mut(), next))
133}
134
135pub(crate) fn reset_llm_render_stack() {
138 LLM_RENDER_STACK.with(|stack| stack.borrow_mut().clear());
139}
140
141pub struct LlmRenderContextGuard {
145 frame_id: u64,
146 expected_depth: usize,
150}
151
152impl LlmRenderContextGuard {
153 pub fn enter(ctx: LlmRenderContext) -> Self {
154 let frame_id = push_llm_render_context_frame(ctx);
155 let depth = LLM_RENDER_STACK.with(|stack| stack.borrow().len());
156 Self {
157 frame_id,
158 expected_depth: depth,
159 }
160 }
161}
162
163impl Drop for LlmRenderContextGuard {
164 fn drop(&mut self) {
165 LLM_RENDER_STACK.with(|stack| {
166 let mut stack = stack.borrow_mut();
167 if stack.last().is_some_and(|frame| frame.id == self.frame_id) {
168 stack.pop();
169 return;
170 }
171
172 if let Some(index) = stack.iter().position(|frame| frame.id == self.frame_id) {
173 debug_assert_eq!(
174 stack.len(),
175 self.expected_depth,
176 "LlmRenderContextGuard nested-drop order violated",
177 );
178 debug_assert_eq!(index + 1, stack.len(), "nested-drop order violated");
179 stack.remove(index);
180 }
181 });
182 }
183}
184
185#[cfg(test)]
186mod tests {
187 use super::*;
188
189 fn derive_family(provider: &str, model: &str) -> String {
190 crate::llm_config::model_family(provider, model)
191 }
192
193 #[test]
194 fn family_from_model_id_takes_precedence() {
195 assert_eq!(
196 derive_family("openrouter", "anthropic/claude-3-5-sonnet"),
197 "anthropic-claude"
198 );
199 assert_eq!(derive_family("openrouter", "openai/gpt-4o"), "openai-gpt");
200 assert_eq!(
201 derive_family("openrouter", "google/gemini-1.5-pro"),
202 "google-gemini"
203 );
204 assert_eq!(derive_family("llamacpp", "qwen3.6-35b-a3b"), "qwen");
205 }
206
207 #[test]
208 fn family_falls_back_to_provider_alias() {
209 assert_eq!(
210 derive_family("anthropic", "unknown-future-model"),
211 "anthropic-claude"
212 );
213 assert_eq!(derive_family("azure", "deployment-xyz"), "openai-gpt");
214 assert_eq!(derive_family("vertex", "model-xyz"), "google-gemini");
215 assert_eq!(derive_family("local", "anonymous-snapshot"), "local");
216 assert_eq!(derive_family("", ""), "unknown");
217 }
218
219 #[test]
220 fn push_pop_stack_round_trip() {
221 reset_llm_render_stack();
222 assert!(current_llm_render_context().is_none());
223 push_llm_render_context(LlmRenderContext::resolve("anthropic", "claude-3-5-sonnet"));
224 assert_eq!(
225 current_llm_render_context().map(|c| c.family),
226 Some("anthropic-claude".to_string()),
227 );
228 push_llm_render_context(LlmRenderContext::resolve("openai", "gpt-4o"));
229 assert_eq!(
230 current_llm_render_context().map(|c| c.family),
231 Some("openai-gpt".to_string()),
232 );
233 pop_llm_render_context();
234 assert_eq!(
235 current_llm_render_context().map(|c| c.family),
236 Some("anthropic-claude".to_string()),
237 );
238 pop_llm_render_context();
239 assert!(current_llm_render_context().is_none());
240 }
241
242 #[test]
243 fn guard_pops_on_drop() {
244 reset_llm_render_stack();
245 {
246 let _guard = LlmRenderContextGuard::enter(LlmRenderContext::resolve(
247 "anthropic",
248 "claude-3-5-sonnet",
249 ));
250 assert!(current_llm_render_context().is_some());
251 }
252 assert!(current_llm_render_context().is_none());
253 }
254}