dynamo_llm/preprocessor/prompt/template/
oai.rs1use super::*;
17
18use minijinja::{context, value::Value};
19
20use crate::protocols::openai::{
21 chat_completions::NvCreateChatCompletionRequest, completions::NvCreateCompletionRequest,
22};
23use tracing;
24
25impl OAIChatLikeRequest for NvCreateChatCompletionRequest {
26 fn messages(&self) -> Value {
27 Value::from_serialize(&self.inner.messages)
28 }
29
30 fn tools(&self) -> Option<Value> {
31 if self.inner.tools.is_none() {
32 None
33 } else {
34 Some(Value::from_serialize(&self.inner.tools))
35 }
36 }
37
38 fn tool_choice(&self) -> Option<Value> {
39 if self.inner.tool_choice.is_none() {
40 None
41 } else {
42 Some(Value::from_serialize(&self.inner.tool_choice))
43 }
44 }
45
46 fn should_add_generation_prompt(&self) -> bool {
47 if let Some(last) = self.inner.messages.last() {
48 matches!(
49 last,
50 async_openai::types::ChatCompletionRequestMessage::User(_)
51 )
52 } else {
53 true
54 }
55 }
56}
57
58impl OAIChatLikeRequest for NvCreateCompletionRequest {
59 fn messages(&self) -> minijinja::value::Value {
60 let message = async_openai::types::ChatCompletionRequestMessage::User(
61 async_openai::types::ChatCompletionRequestUserMessage {
62 content: async_openai::types::ChatCompletionRequestUserMessageContent::Text(
63 crate::protocols::openai::completions::prompt_to_string(&self.inner.prompt),
64 ),
65 name: None,
66 },
67 );
68
69 minijinja::value::Value::from_serialize(vec![message])
70 }
71
72 fn should_add_generation_prompt(&self) -> bool {
73 true
74 }
75}
76
77impl OAIPromptFormatter for HfTokenizerConfigJsonFormatter {
78 fn supports_add_generation_prompt(&self) -> bool {
79 self.supports_add_generation_prompt
80 }
81
82 fn render(&self, req: &dyn OAIChatLikeRequest) -> Result<String> {
83 let mixins = Value::from_dyn_object(self.mixins.clone());
84
85 let tools = req.tools();
86 let has_tools = tools.is_some();
87 let add_generation_prompt = req.should_add_generation_prompt();
88
89 tracing::trace!(
90 "Rendering prompt with tools: {:?}, add_generation_prompt: {}",
91 has_tools,
92 add_generation_prompt
93 );
94
95 let ctx = context! {
96 messages => req.messages(),
97 tools => tools,
98 bos_token => self.config.bos_tok(),
99 eos_token => self.config.eos_tok(),
100 unk_token => self.config.unk_tok(),
101 add_generation_prompt => add_generation_prompt,
102 ..mixins
103 };
104
105 let ctx = context! { ..ctx, ..context! {
106
107 }};
108
109 let tmpl = if has_tools {
110 self.env.get_template("tool_use")?
111 } else {
112 self.env.get_template("default")?
113 };
114
115 Ok(tmpl.render(&ctx)?)
116 }
117}