1use std::collections::BTreeMap;
2use std::collections::HashMap;
3
4use talos_core::message::{SystemCacheMarker, SystemCacheType};
5use talos_plugin::{HookContext, HookEvent, HookOutcome, HookRegistry};
6use talos_skill::SkillIndex;
7
8use super::assets::{DEFAULT_IDENTITY, TOOL_CALLING_STRICT};
9use super::sections::{PromptSection, PromptSectionKind};
10use super::types::{ActivatedSkillContext, CacheMarker, CacheType, ContextFile, ToolDescription};
11
12#[derive(Debug, Clone)]
37pub struct SystemPromptBuilder {
38 identity: String,
40 tools: Vec<ToolDescription>,
42 skill_index: Vec<SkillIndex>,
44 activated_skill: Option<ActivatedSkillContext>,
46 context_files: Vec<ContextFile>,
48 user_preferences: String,
50 custom_prompt: Option<String>,
52 append_prompt: Option<String>,
54 memory_section: Option<String>,
56 todo_section: Option<String>,
58 tool_call_format: &'static str,
59 template_vars: HashMap<String, String>,
61}
62
63impl SystemPromptBuilder {
64 #[must_use]
69 pub fn new() -> Self {
70 let mut template_vars = HashMap::new();
71 template_vars.insert(
72 "workspace_info".to_string(),
73 "Workspace information unavailable.".to_string(),
74 );
75 template_vars.insert(
76 "model_info".to_string(),
77 "Provider model metadata unavailable.".to_string(),
78 );
79
80 Self {
81 identity: DEFAULT_IDENTITY.to_string(),
82 tools: Vec::new(),
83 skill_index: Vec::new(),
84 activated_skill: None,
85 context_files: Vec::new(),
86 user_preferences: String::new(),
87 custom_prompt: None,
88 append_prompt: None,
89 memory_section: None,
90 todo_section: None,
91 tool_call_format: "",
92 template_vars,
93 }
94 }
95
96 pub fn with_strict_tool_format(mut self) -> Self {
97 self.tool_call_format = TOOL_CALLING_STRICT;
98 self
99 }
100
101 pub fn with_tool_format(mut self, format: &'static str) -> Self {
102 self.tool_call_format = format;
103 self
104 }
105
106 #[must_use]
108 pub fn with_template_var(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
109 self.template_vars.insert(key.into(), value.into());
110 self
111 }
112
113 #[must_use]
115 pub fn with_workspace_info(self, value: impl Into<String>) -> Self {
116 self.with_template_var("workspace_info", value)
117 }
118
119 #[must_use]
121 pub fn with_model_info(self, value: impl Into<String>) -> Self {
122 self.with_template_var("model_info", value)
123 }
124
125 #[must_use]
130 pub fn with_tools(mut self, tools: Vec<ToolDescription>) -> Self {
131 self.tools = tools;
132 self
133 }
134
135 #[must_use]
140 pub fn with_skill_index(mut self, skills: Vec<SkillIndex>) -> Self {
141 self.skill_index = skills;
142 self
143 }
144
145 #[must_use]
150 pub fn with_activated_skill(mut self, skill: Option<ActivatedSkillContext>) -> Self {
151 self.activated_skill = skill;
152 self
153 }
154
155 #[must_use]
161 pub fn with_context_files(mut self, files: Vec<ContextFile>) -> Self {
162 self.context_files = files;
163 self
164 }
165
166 #[must_use]
168 pub fn with_user_preferences(mut self, prefs: String) -> Self {
169 self.user_preferences = prefs;
170 self
171 }
172
173 #[must_use]
178 pub fn with_memory_section(mut self, section: Option<String>) -> Self {
179 self.memory_section = section;
180 self
181 }
182
183 #[must_use]
188 pub fn with_todo_section(mut self, section: Option<String>) -> Self {
189 self.todo_section = section;
190 self
191 }
192
193 #[must_use]
199 pub fn with_custom_prompt(mut self, prompt: String) -> Self {
200 self.custom_prompt = Some(prompt);
201 self
202 }
203
204 #[must_use]
208 pub fn with_append_prompt(mut self, prompt: String) -> Self {
209 self.append_prompt = Some(prompt);
210 self
211 }
212
213 pub fn clear_append_prompt(&mut self) {
215 self.append_prompt = None;
216 }
217
218 pub fn set_append_prompt_opt(&mut self, prompt: Option<String>) {
222 self.append_prompt = prompt;
223 }
224
225 fn render_template(&self, template: &str, extra_vars: &[(&str, String)]) -> String {
226 let mut rendered = template.to_string();
227 let mut vars = self.template_vars.clone();
228 for (key, value) in extra_vars {
229 vars.insert((*key).to_string(), value.clone());
230 }
231
232 for (key, value) in vars {
233 rendered = rendered.replace(&format!("{{{{{key}}}}}"), &value);
234 }
235 rendered
236 }
237
238 fn current_datetime() -> String {
239 let seconds = std::time::SystemTime::now()
240 .duration_since(std::time::UNIX_EPOCH)
241 .map(|duration| duration.as_secs())
242 .unwrap_or(0);
243 format!("unix_seconds={seconds}")
244 }
245
246 fn tool_protocol_hint(&self) -> String {
247 if self.tool_call_format.is_empty() {
248 "Native tool calling is enabled. Use provider-native tool calls; do not emit textual tool-call JSON unless the provider requires a fallback.".to_string()
249 } else {
250 self.tool_call_format.trim().to_string()
251 }
252 }
253
254 fn prompt_sections(&self) -> Vec<PromptSection> {
255 let mut sections: Vec<PromptSection> = Vec::new();
256
257 let stable_vars = [("tool_protocol_hint", self.tool_protocol_hint())];
258
259 let identity = if let Some(ref custom) = self.custom_prompt {
260 self.render_template(custom, &stable_vars)
261 } else {
262 self.render_template(&self.identity, &stable_vars)
263 };
264 sections.push(PromptSection {
265 text: format!("# Identity\n{identity}\n"),
266 kind: PromptSectionKind::Cacheable,
267 });
268
269 if self.tools.is_empty() {
270 sections.push(PromptSection {
271 text: String::from("# Tools\nNo tools available.\n"),
272 kind: PromptSectionKind::Cacheable,
273 });
274 } else {
275 let mut families: BTreeMap<_, Vec<&ToolDescription>> = BTreeMap::new();
276 for tool in &self.tools {
277 families.entry(tool.family).or_default().push(tool);
278 }
279
280 sections.push(PromptSection {
281 text: String::from("# Tools\nTool definitions are grouped by stable family.\n"),
282 kind: PromptSectionKind::Cacheable,
283 });
284
285 for (family, mut sorted_tools) in families {
286 sorted_tools.sort_by(|a, b| a.name.cmp(&b.name));
287 let mut tools_section = format!("# Tool Family: {family:?}\n");
288 for tool in sorted_tools {
289 tools_section.push_str(&format!("## {}\n{}\n", tool.name, tool.description));
290 if let Some(props) = tool.parameters.get("properties")
291 && let Some(required) = tool.parameters.get("required")
292 {
293 let req_list: Vec<&str> = required
294 .as_array()
295 .map(|a| a.iter().filter_map(|v| v.as_str()).collect())
296 .unwrap_or_default();
297 let mut param_parts = Vec::new();
298 for (key, val) in props.as_object().unwrap_or(&serde_json::Map::new()) {
299 let desc = val
300 .get("description")
301 .and_then(|d| d.as_str())
302 .unwrap_or("");
303 let ptype = val.get("type").and_then(|t| t.as_str()).unwrap_or("any");
304 let req = if req_list.contains(&key.as_str()) {
305 "required"
306 } else {
307 "optional"
308 };
309 param_parts
310 .push(format!(" - {} ({}): {} [{}]", key, ptype, desc, req));
311 }
312 if !param_parts.is_empty() {
313 tools_section.push_str("Parameters:\n");
314 tools_section.push_str(¶m_parts.join("\n"));
315 tools_section.push_str("\n\n");
316 }
317 }
318 tools_section.push('\n');
319 }
320 sections.push(PromptSection {
321 text: tools_section,
322 kind: PromptSectionKind::Cacheable,
323 });
324 }
325 }
326
327 if self.skill_index.is_empty() {
328 sections.push(PromptSection {
329 text: String::from("# Skills\nNo skills available.\n"),
330 kind: PromptSectionKind::Cacheable,
331 });
332 } else {
333 let mut skills_section = String::from("# Skills\n");
334 for skill in &self.skill_index {
335 skills_section.push_str(&format!("- **{}**: {}\n", skill.name, skill.description));
336 }
337 skills_section.push('\n');
338 sections.push(PromptSection {
339 text: skills_section,
340 kind: PromptSectionKind::Cacheable,
341 });
342 }
343
344 if let Some(ref skill) = self.activated_skill {
345 sections.push(PromptSection {
346 text: format!(
347 "# Activated Skill: {}\n{}\n",
348 skill.name.trim(),
349 skill.content.trim()
350 ),
351 kind: PromptSectionKind::Cacheable,
352 });
353 }
354
355 if self.context_files.is_empty() {
356 sections.push(PromptSection {
357 text: String::from("# Context\nNo context files loaded.\n"),
358 kind: PromptSectionKind::Dynamic,
359 });
360 } else {
361 let mut context_section = String::from("# Context\n");
362 for file in &self.context_files {
363 context_section.push_str(&format!("--- {} ---\n{}\n\n", file.path, file.content));
364 }
365 sections.push(PromptSection {
366 text: context_section,
367 kind: PromptSectionKind::Dynamic,
368 });
369 }
370
371 if let Some(ref memory) = self.memory_section {
372 sections.push(PromptSection {
373 text: format!("# Memory\n{memory}\n"),
374 kind: PromptSectionKind::Dynamic,
375 });
376 }
377
378 if let Some(ref todos) = self.todo_section {
379 sections.push(PromptSection {
380 text: format!("# Session Todos\n{todos}\n"),
381 kind: PromptSectionKind::Dynamic,
382 });
383 }
384
385 if !self.user_preferences.is_empty() {
386 sections.push(PromptSection {
387 text: format!("# User Preferences\n{}\n", self.user_preferences),
388 kind: PromptSectionKind::Dynamic,
389 });
390 }
391
392 let runtime_section = self.render_template(
393 "# Runtime Context\nCurrent datetime: {{datetime}}\n",
394 &[("datetime", Self::current_datetime())],
395 );
396 sections.push(PromptSection {
397 text: runtime_section,
398 kind: PromptSectionKind::Dynamic,
399 });
400
401 if let Some(ref append) = self.append_prompt {
402 sections.push(PromptSection {
403 text: format!("# Additional Instructions\n{append}\n"),
404 kind: PromptSectionKind::Dynamic,
405 });
406 }
407
408 sections
409 }
410
411 #[must_use]
423 pub fn build(&self) -> String {
424 self.prompt_sections()
425 .into_iter()
426 .map(|section| section.text)
427 .collect::<Vec<_>>()
428 .join("\n")
429 }
430
431 #[allow(dead_code)]
436 pub(crate) async fn build_with_hooks(
437 &self,
438 hook_registry: &HookRegistry,
439 ctx: &HookContext,
440 ) -> Result<(String, Vec<SystemCacheMarker>), String> {
441 let (prompt, markers) = self.build_with_cache_markers();
442 let original_prompt = prompt.clone();
443 let outcome = hook_registry
444 .dispatch(ctx, HookEvent::OnSystemPromptBuilt { prompt: &prompt })
445 .await;
446
447 match outcome {
448 HookOutcome::Continue(HookEvent::OnSystemPromptBuilt { prompt })
449 | HookOutcome::Skip(HookEvent::OnSystemPromptBuilt { prompt }) => {
450 let prompt = prompt.to_string();
451 let markers = if prompt == original_prompt {
452 markers.into_iter().map(Into::into).collect()
453 } else {
454 Vec::new()
455 };
456 Ok((prompt, markers))
457 }
458 HookOutcome::Deny { reason, .. } => Err(reason),
459 HookOutcome::Continue(_) | HookOutcome::Skip(_) => {
460 Ok((prompt, markers.into_iter().map(Into::into).collect()))
461 }
462 }
463 }
464
465 pub(crate) async fn build_with_hooks_from_prompt(
471 &self,
472 hook_registry: &HookRegistry,
473 ctx: &HookContext,
474 prompt: &str,
475 stable_prefix_len: usize,
476 ) -> Result<(String, Vec<SystemCacheMarker>), String> {
477 let original_prompt = prompt.to_string();
478 let outcome = hook_registry
479 .dispatch(ctx, HookEvent::OnSystemPromptBuilt { prompt })
480 .await;
481
482 match outcome {
483 HookOutcome::Continue(HookEvent::OnSystemPromptBuilt { prompt })
484 | HookOutcome::Skip(HookEvent::OnSystemPromptBuilt { prompt }) => {
485 let prompt = prompt.to_string();
486 let markers = if prompt == original_prompt && stable_prefix_len > 0 {
487 vec![SystemCacheMarker {
488 offset: 0,
489 length: stable_prefix_len,
490 cache_type: SystemCacheType::Ephemeral,
491 }]
492 } else {
493 Vec::new()
494 };
495 Ok((prompt, markers))
496 }
497 HookOutcome::Deny { reason, .. } => Err(reason),
498 HookOutcome::Continue(_) | HookOutcome::Skip(_) => {
499 let markers = if stable_prefix_len > 0 {
500 vec![SystemCacheMarker {
501 offset: 0,
502 length: stable_prefix_len,
503 cache_type: SystemCacheType::Ephemeral,
504 }]
505 } else {
506 Vec::new()
507 };
508 Ok((prompt.to_string(), markers))
509 }
510 }
511 }
512
513 #[must_use]
522 pub fn build_stable_prefix(&self) -> String {
523 let sections = self.prompt_sections();
524 let mut prefix = String::new();
525 let mut first = true;
526 for section in §ions {
527 if section.kind != PromptSectionKind::Cacheable {
528 break;
529 }
530 if !first {
531 prefix.push('\n');
532 }
533 prefix.push_str(§ion.text);
534 first = false;
535 }
536 prefix
537 }
538
539 #[must_use]
545 pub fn build_dynamic_suffix(&self) -> String {
546 let sections = self.prompt_sections();
547 let mut suffix = String::new();
548 let mut first = true;
549 for section in §ions {
550 if section.kind == PromptSectionKind::Cacheable {
551 continue;
552 }
553 if !first {
554 suffix.push('\n');
555 }
556 suffix.push_str(§ion.text);
557 first = false;
558 }
559 suffix
560 }
561
562 #[must_use]
575 pub fn build_with_cache_markers(&self) -> (String, Vec<CacheMarker>) {
576 let mut markers: Vec<CacheMarker> = Vec::new();
577 let sections = self.prompt_sections();
578 let mut prompt = String::new();
579
580 for (index, section) in sections.iter().enumerate() {
581 if index > 0 {
582 prompt.push('\n');
583 }
584 let offset = prompt.len();
585 prompt.push_str(§ion.text);
586 if section.kind == PromptSectionKind::Cacheable {
587 markers.push(CacheMarker {
588 offset,
589 length: section.text.len(),
590 cache_type: CacheType::Ephemeral,
591 });
592 }
593 }
594
595 (prompt, markers)
596 }
597
598 #[must_use]
604 pub fn total_tokens(&self) -> usize {
605 let prompt = self.build();
606 prompt.chars().count().div_ceil(4)
608 }
609
610 pub fn log_size(&self) {
615 let prompt = self.build();
616 let char_count = prompt.chars().count();
617 let token_estimate = self.total_tokens();
618 eprintln!("System prompt: {char_count} characters, ~{token_estimate} tokens");
619 }
620}
621
622impl Default for SystemPromptBuilder {
623 fn default() -> Self {
624 Self::new()
625 }
626}