1use crate::context::{Context, ContextKey, Fact};
24use std::collections::HashSet;
25use std::fmt::Write;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
29pub enum PromptFormat {
30 Plain,
32 #[default]
34 Edn,
35}
36
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum AgentRole {
40 Proposer,
42 Validator,
44 Synthesizer,
46 Analyzer,
48}
49
50impl AgentRole {
51 fn to_keyword(self) -> &'static str {
53 match self {
54 Self::Proposer => ":proposer",
55 Self::Validator => ":validator",
56 Self::Synthesizer => ":synthesizer",
57 Self::Analyzer => ":analyzer",
58 }
59 }
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
64pub enum Constraint {
65 NoInvent,
67 NoContradict,
69 NoHallucinate,
71 CiteSources,
73}
74
75impl Constraint {
76 fn to_keyword(self) -> &'static str {
78 match self {
79 Self::NoInvent => ":no-invent",
80 Self::NoContradict => ":no-contradict",
81 Self::NoHallucinate => ":no-hallucinate",
82 Self::CiteSources => ":cite-sources",
83 }
84 }
85}
86
87#[derive(Debug, Clone, PartialEq, Eq)]
89pub struct OutputContract {
90 pub emit: String,
92 pub key: ContextKey,
94 pub format: Option<String>,
96}
97
98impl OutputContract {
99 #[must_use]
101 pub fn new(emit: impl Into<String>, key: ContextKey) -> Self {
102 Self {
103 emit: emit.into(),
104 key,
105 format: None,
106 }
107 }
108
109 #[must_use]
111 pub fn with_format(mut self, format: impl Into<String>) -> Self {
112 self.format = Some(format.into());
113 self
114 }
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
122pub struct AgentPrompt {
123 pub role: AgentRole,
125 pub objective: String,
127 pub context: PromptContext,
129 pub constraints: HashSet<Constraint>,
131 pub output_contract: OutputContract,
133}
134
135#[derive(Debug, Clone, PartialEq, Eq, Default)]
137pub struct PromptContext {
138 pub facts: Vec<(ContextKey, Vec<Fact>)>,
140}
141
142impl PromptContext {
143 #[must_use]
145 pub fn new() -> Self {
146 Self::default()
147 }
148
149 pub fn add_facts(&mut self, key: ContextKey, facts: Vec<Fact>) {
151 if !facts.is_empty() {
152 self.facts.push((key, facts));
153 }
154 }
155
156 #[must_use]
158 pub fn from_context(ctx: &Context, dependencies: &[ContextKey]) -> Self {
159 let mut prompt_ctx = Self::new();
160 for &key in dependencies {
161 let facts = ctx.get(key).to_vec();
162 prompt_ctx.add_facts(key, facts);
163 }
164 prompt_ctx
165 }
166}
167
168fn context_key_to_keyword(key: ContextKey) -> &'static str {
170 match key {
171 ContextKey::Seeds => ":seeds",
172 ContextKey::Hypotheses => ":hypotheses",
173 ContextKey::Strategies => ":strategies",
174 ContextKey::Constraints => ":constraints",
175 ContextKey::Signals => ":signals",
176 ContextKey::Competitors => ":competitors",
177 ContextKey::Evaluations => ":evaluations",
178 ContextKey::Proposals => ":proposals",
179 ContextKey::Diagnostic => ":diagnostic",
180 }
181}
182
183impl AgentPrompt {
184 #[must_use]
186 pub fn new(
187 role: AgentRole,
188 objective: impl Into<String>,
189 context: PromptContext,
190 output_contract: OutputContract,
191 ) -> Self {
192 Self {
193 role,
194 objective: objective.into(),
195 context,
196 constraints: HashSet::new(),
197 output_contract,
198 }
199 }
200
201 #[must_use]
203 pub fn with_constraint(mut self, constraint: Constraint) -> Self {
204 self.constraints.insert(constraint);
205 self
206 }
207
208 #[must_use]
210 pub fn with_constraints(mut self, constraints: impl IntoIterator<Item = Constraint>) -> Self {
211 self.constraints.extend(constraints);
212 self
213 }
214
215 #[must_use]
226 pub fn to_edn(&self) -> String {
227 let mut s = String::new();
228 s.push_str("{:r ");
229 s.push_str(self.role.to_keyword());
230 s.push_str("\n :o :");
231 s.push_str(&self.objective.replace(' ', "-"));
233 s.push_str("\n :c {");
234
235 let mut first_key = true;
237 for (key, facts) in &self.context.facts {
238 if !first_key {
239 s.push(' ');
240 }
241 first_key = false;
242 s.push_str(context_key_to_keyword(*key));
243 s.push_str(" [{");
244 for (i, fact) in facts.iter().enumerate() {
245 if i > 0 {
246 s.push_str("} {");
247 }
248 s.push_str(":id \"");
249 s.push_str(&escape_string(&fact.id));
250 s.push_str("\" :c \"");
251 s.push_str(&escape_string(&fact.content));
252 s.push('"');
253 }
254 s.push_str("}]");
255 }
256
257 s.push_str("}\n :k #{");
258
259 let mut constraints: Vec<_> = self.constraints.iter().collect();
261 constraints.sort(); for (i, constraint) in constraints.iter().enumerate() {
263 if i > 0 {
264 s.push(' ');
265 }
266 s.push_str(constraint.to_keyword());
267 }
268
269 s.push_str("}\n :out {:emit :");
270 s.push_str(&self.output_contract.emit);
271 s.push_str(" :key ");
272 s.push_str(context_key_to_keyword(self.output_contract.key));
273 if let Some(ref format) = self.output_contract.format {
274 s.push_str(" :format :");
275 s.push_str(format);
276 }
277 s.push_str("}}");
278
279 s
280 }
281
282 #[must_use]
284 pub fn to_plain(&self) -> String {
285 let mut s = String::new();
286 writeln!(s, "Role: {:?}", self.role).unwrap();
287 writeln!(s, "Objective: {}", self.objective).unwrap();
288 writeln!(s, "\nContext:").unwrap();
289
290 for (key, facts) in &self.context.facts {
291 writeln!(s, "\n## {key:?}").unwrap();
292 for fact in facts {
293 writeln!(s, "- {}: {}", fact.id, fact.content).unwrap();
294 }
295 }
296
297 if !self.constraints.is_empty() {
298 writeln!(s, "\nConstraints:").unwrap();
299 for constraint in &self.constraints {
300 writeln!(s, "- {constraint:?}").unwrap();
301 }
302 }
303
304 writeln!(
305 s,
306 "\nOutput: {:?} -> {:?}",
307 self.output_contract.emit, self.output_contract.key
308 )
309 .unwrap();
310
311 s
312 }
313
314 #[must_use]
316 pub fn serialize(&self, format: PromptFormat) -> String {
317 match format {
318 PromptFormat::Edn => self.to_edn(),
319 PromptFormat::Plain => self.to_plain(),
320 }
321 }
322}
323
324fn escape_string(s: &str) -> String {
326 s.replace('\\', "\\\\")
327 .replace('"', "\\\"")
328 .replace('\n', "\\n")
329 .replace('\r', "\\r")
330 .replace('\t', "\\t")
331}
332
333#[cfg(test)]
334mod tests {
335 use super::*;
336 use crate::context::{Context, Fact};
337
338 #[test]
339 fn test_edn_serialization() {
340 let mut ctx = PromptContext::new();
341 ctx.add_facts(
342 ContextKey::Signals,
343 vec![
344 Fact {
345 key: ContextKey::Signals,
346 id: "s1".to_string(),
347 content: "Revenue +15% Q3".to_string(),
348 },
349 Fact {
350 key: ContextKey::Signals,
351 id: "s2".to_string(),
352 content: "Market $2.3B".to_string(),
353 },
354 ],
355 );
356
357 let prompt = AgentPrompt::new(
358 AgentRole::Proposer,
359 "extract-competitors",
360 ctx,
361 OutputContract::new("proposed-fact", ContextKey::Competitors),
362 )
363 .with_constraint(Constraint::NoInvent)
364 .with_constraint(Constraint::NoContradict);
365
366 let edn = prompt.to_edn();
367 assert!(edn.contains(":r :proposer"));
368 assert!(edn.contains(":o :extract-competitors"));
369 assert!(edn.contains(":signals"));
370 assert!(edn.contains(":no-invent"));
371 assert!(edn.contains(":no-contradict"));
372 assert!(edn.contains(":competitors"));
373 }
374
375 #[test]
376 fn test_context_building() {
377 let mut context = Context::new();
378 context
379 .add_fact(Fact {
380 key: ContextKey::Seeds,
381 id: "seed1".to_string(),
382 content: "Test seed".to_string(),
383 })
384 .unwrap();
385
386 let prompt_ctx = PromptContext::from_context(&context, &[ContextKey::Seeds]);
387 assert_eq!(prompt_ctx.facts.len(), 1);
388 assert_eq!(prompt_ctx.facts[0].0, ContextKey::Seeds);
389 assert_eq!(prompt_ctx.facts[0].1.len(), 1);
390 }
391
392 #[test]
393 fn test_escape_string() {
394 assert_eq!(escape_string("hello"), "hello");
395 assert_eq!(escape_string("hello\"world"), "hello\\\"world");
396 assert_eq!(escape_string("hello\nworld"), "hello\\nworld");
397 }
398
399 #[test]
400 fn test_token_efficiency() {
401 let mut ctx = PromptContext::new();
402 ctx.add_facts(
403 ContextKey::Signals,
404 vec![Fact {
405 key: ContextKey::Signals,
406 id: "s1".to_string(),
407 content: "Revenue +15% Q3".to_string(),
408 }],
409 );
410
411 let prompt = AgentPrompt::new(
412 AgentRole::Proposer,
413 "analyze",
414 ctx,
415 OutputContract::new("proposed-fact", ContextKey::Strategies),
416 );
417
418 let edn = prompt.to_edn();
419 let plain = prompt.to_plain();
420
421 println!("EDN length: {}", edn.len());
422 println!("Plain length: {}", plain.len());
423 println!("EDN:\n{edn}");
424 println!("Plain:\n{plain}");
425
426 assert!(!edn.is_empty());
431 assert!(!plain.is_empty());
432 }
433}