Skip to main content

vtcode_core/core/agent/
conversation.rs

1//! Helpers for composing agent conversations and bridging provider-specific message formats.
2
3use crate::core::agent::task::{ContextItem, Task};
4use crate::llm::provider::{ContentPart, Message, MessageContent, MessageRole, ToolCall};
5use crate::llm::providers::gemini::wire::{Content, FunctionCall, FunctionResponse, InlineData, Part};
6use serde_json::{Value, json};
7use std::collections::HashMap;
8use std::fmt::Write;
9
10const GEMINI_PRESERVED_PARTS_PREFIX: &str = "__vtcode_gemini_parts__:";
11
12/// Build the initial conversation payload (without the system instruction message).
13pub fn build_conversation(task: &Task, contexts: &[ContextItem]) -> Vec<Content> {
14    let mut conversation = Vec::with_capacity(3);
15    let mut task_content = String::with_capacity(task.title.len() + task.description.len() + 20);
16    let _ = write!(task_content, "Task: {}\nDescription: {}", task.title, task.description);
17    conversation.push(Content::user_text(task_content));
18
19    if let Some(instructions) = task.instructions.as_ref() {
20        conversation.push(Content::user_text(instructions.clone()));
21    }
22
23    if !contexts.is_empty() {
24        let mut context_content = String::from("Relevant Context:");
25        for ctx in contexts {
26            let _ = write!(context_content, "\nContext [{}]: {}", ctx.id, ctx.content);
27        }
28        conversation.push(Content::user_text(context_content));
29    }
30
31    conversation
32}
33
34/// Convert Gemini `Content` structures into universal provider messages.
35pub fn messages_from_conversation(conversation: &[Content]) -> Vec<Message> {
36    let mut messages = Vec::with_capacity(conversation.len());
37    for content in conversation {
38        let part_count = content.parts.len();
39        let mut content_parts = Vec::with_capacity(part_count);
40        let mut tool_calls = Vec::new();
41        let mut tool_responses = Vec::new();
42
43        for part in &content.parts {
44            match part {
45                Part::Text { text: part_text, .. } => {
46                    if let Some(ContentPart::Text { text }) = content_parts.last_mut() {
47                        if !text.is_empty() {
48                            text.push('\n');
49                        }
50                        text.push_str(part_text);
51                    } else if !part_text.is_empty() {
52                        content_parts.push(ContentPart::text(part_text.clone()));
53                    }
54                }
55                Part::InlineData { inline_data } => {
56                    content_parts.push(ContentPart::image(inline_data.data.clone(), inline_data.mime_type.clone()));
57                }
58                Part::FunctionCall { function_call, thought_signature } => {
59                    let mut tool_call = ToolCall::function(
60                        function_call.id.clone().unwrap_or_default(),
61                        function_call.name.clone(),
62                        function_call.args.to_string(),
63                    );
64                    tool_call.thought_signature = thought_signature.clone();
65                    tool_calls.push(tool_call);
66                }
67                Part::FunctionResponse { function_response, .. } => {
68                    let id = function_response.id.clone().unwrap_or_else(|| "unknown".to_string());
69                    let response_str = function_response.response.to_string();
70                    tool_responses.push(Message::tool_response(id, response_str));
71                }
72                Part::ToolCall { .. }
73                | Part::ToolResponse { .. }
74                | Part::ExecutableCode { .. }
75                | Part::CodeExecutionResult { .. } => {}
76                Part::CacheControl { .. } => {}
77            }
78        }
79
80        if !tool_responses.is_empty() {
81            messages.extend(tool_responses);
82            if !content_parts.is_empty() {
83                messages.push(Message::user_with_parts(content_parts));
84            }
85            continue;
86        }
87
88        let mut message = match content.role.as_str() {
89            "model" => {
90                if content_parts.is_empty() {
91                    Message::assistant(String::new())
92                } else {
93                    Message::assistant_with_parts(content_parts)
94                }
95            }
96            _ => {
97                if content_parts.is_empty() {
98                    Message::user(String::new())
99                } else {
100                    Message::user_with_parts(content_parts)
101                }
102            }
103        };
104
105        if !tool_calls.is_empty() {
106            message.tool_calls = Some(tool_calls);
107        }
108
109        if let Some(raw_parts_detail) = preserved_parts_detail(&content.parts) {
110            message = message.with_reasoning_details(Some(vec![json!(raw_parts_detail)]));
111        }
112
113        messages.push(message);
114    }
115
116    messages
117}
118
119/// Convert Gemini `Content` structures into universal provider messages.
120///
121/// System instructions travel separately via `LLMRequest.system_prompt` on the active request.
122pub fn build_messages_from_conversation(conversation: &[Content]) -> Vec<Message> {
123    messages_from_conversation(conversation)
124}
125
126fn parts_from_message_content(content: &MessageContent) -> Vec<Part> {
127    match content {
128        MessageContent::Text(text) => {
129            if text.is_empty() {
130                Vec::new()
131            } else {
132                vec![Part::Text { text: text.clone(), thought_signature: None }]
133            }
134        }
135        MessageContent::Parts(parts) => {
136            let mut converted = Vec::with_capacity(parts.len());
137            for part in parts {
138                match part {
139                    ContentPart::Text { text } => {
140                        if !text.is_empty() {
141                            converted.push(Part::Text { text: text.clone(), thought_signature: None });
142                        }
143                    }
144                    ContentPart::Image { data, mime_type, .. } => {
145                        converted.push(Part::InlineData {
146                            inline_data: InlineData { mime_type: mime_type.clone(), data: data.clone() },
147                        });
148                    }
149                    ContentPart::File { filename, file_id, file_url, .. } => {
150                        let fallback = filename
151                            .clone()
152                            .or_else(|| file_id.clone())
153                            .or_else(|| file_url.clone())
154                            .unwrap_or_else(|| "attached file".to_string());
155                        converted.push(Part::Text {
156                            text: format!("[File input not directly supported: {fallback}]"),
157                            thought_signature: None,
158                        });
159                    }
160                }
161            }
162            converted
163        }
164    }
165}
166
167fn tool_call_arguments(arguments: &str) -> Value {
168    serde_json::from_str(arguments).unwrap_or_else(|_| Value::String(arguments.to_string()))
169}
170
171fn tool_response_value(content: &MessageContent) -> Value {
172    let text = content.as_text();
173    serde_json::from_str(text.as_ref()).unwrap_or_else(|_| json!({ "result": text.as_ref() }))
174}
175
176/// Rebuild Gemini-style conversation content from archived provider messages.
177///
178/// System messages are skipped because exec regenerates the current system prompt.
179pub fn conversation_from_messages(messages: &[Message]) -> Vec<Content> {
180    let mut conversation = Vec::with_capacity(messages.len());
181    let mut tool_names_by_call_id: HashMap<String, String> = HashMap::with_capacity(messages.len());
182
183    for message in messages {
184        match message.role {
185            MessageRole::System => {}
186            MessageRole::User => {
187                let parts = parts_from_message_content(&message.content);
188                if !parts.is_empty() {
189                    conversation.push(Content { role: "user".to_string(), parts });
190                }
191            }
192            MessageRole::Assistant => {
193                let parts = preserved_parts_from_message(message).unwrap_or_else(|| {
194                    let mut rebuilt_parts = parts_from_message_content(&message.content);
195                    if let Some(tool_calls) = &message.tool_calls {
196                        for tool_call in tool_calls {
197                            let Some(function) = &tool_call.function else {
198                                continue;
199                            };
200
201                            let id = (!tool_call.id.is_empty()).then(|| tool_call.id.clone());
202                            if let Some(call_id) = id.as_ref() {
203                                tool_names_by_call_id.insert(call_id.clone(), function.name.clone());
204                            }
205
206                            rebuilt_parts.push(Part::FunctionCall {
207                                function_call: FunctionCall {
208                                    name: function.name.clone(),
209                                    args: tool_call_arguments(&function.arguments),
210                                    id,
211                                },
212                                thought_signature: tool_call.thought_signature.clone(),
213                            });
214                        }
215                    }
216                    rebuilt_parts
217                });
218
219                for part in &parts {
220                    if let Part::FunctionCall { function_call, .. } = part
221                        && let Some(call_id) = function_call.id.as_ref()
222                    {
223                        tool_names_by_call_id.insert(call_id.clone(), function_call.name.clone());
224                    }
225                }
226
227                if !parts.is_empty() {
228                    conversation.push(Content { role: "model".to_string(), parts });
229                }
230            }
231            MessageRole::Tool => {
232                let Some(call_id) = message.tool_call_id.as_ref().filter(|value| !value.is_empty()).cloned() else {
233                    let parts = parts_from_message_content(&message.content);
234                    if !parts.is_empty() {
235                        conversation.push(Content { role: "user".to_string(), parts });
236                    }
237                    continue;
238                };
239
240                let tool_name = message
241                    .origin_tool
242                    .clone()
243                    .or_else(|| tool_names_by_call_id.get(&call_id).cloned())
244                    .unwrap_or_else(|| "tool".to_string());
245
246                conversation.push(Content {
247                    role: "function".to_string(),
248                    parts: vec![Part::FunctionResponse {
249                        function_response: FunctionResponse {
250                            name: tool_name,
251                            response: tool_response_value(&message.content),
252                            id: Some(call_id),
253                        },
254                        thought_signature: None,
255                    }],
256                });
257            }
258        }
259    }
260
261    conversation
262}
263
264fn preserved_parts_from_message(message: &Message) -> Option<Vec<Part>> {
265    let details = message.reasoning_details.as_ref()?;
266    for detail in details {
267        let Some(text) = detail.as_str() else {
268            continue;
269        };
270        let Some(payload) = text.strip_prefix(GEMINI_PRESERVED_PARTS_PREFIX) else {
271            continue;
272        };
273        if let Ok(parts) = serde_json::from_str::<Vec<Part>>(payload) {
274            return Some(parts);
275        }
276    }
277    None
278}
279
280fn preserved_parts_detail(parts: &[Part]) -> Option<String> {
281    let should_preserve = parts.iter().any(|part| {
282        part.thought_signature().is_some()
283            || matches!(
284                part,
285                Part::ToolCall { .. }
286                    | Part::ToolResponse { .. }
287                    | Part::ExecutableCode { .. }
288                    | Part::CodeExecutionResult { .. }
289                    | Part::FunctionResponse { .. }
290                    | Part::InlineData { .. }
291            )
292    });
293    if !should_preserve {
294        return None;
295    }
296
297    serde_json::to_string(parts)
298        .ok()
299        .map(|serialized| format!("{GEMINI_PRESERVED_PARTS_PREFIX}{serialized}"))
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use crate::llm::provider::{FunctionCall, ToolCall};
306
307    fn sample_task() -> Task {
308        Task {
309            id: "task-1".to_owned(),
310            title: "Example".to_owned(),
311            description: "Do something".to_owned(),
312            instructions: Some("Follow steps".to_owned()),
313        }
314    }
315
316    #[test]
317    fn conversation_builds_expected_steps() {
318        let task = sample_task();
319        let contexts = vec![ContextItem { id: "ctx1".into(), content: "Data".into() }];
320        let conversation = build_conversation(&task, &contexts);
321        assert_eq!(conversation.len(), 3);
322    }
323
324    #[test]
325    fn messages_mirror_conversation_without_system_prompt() {
326        let task = sample_task();
327        let conversation = build_conversation(&task, &[]);
328        let messages = build_messages_from_conversation(&conversation);
329        assert_eq!(messages.len(), conversation.len());
330        assert!(messages.iter().all(|message| message.role != MessageRole::System));
331    }
332
333    #[test]
334    fn archived_messages_rebuild_function_history() {
335        let history = vec![
336            Message::system("Base".to_string()),
337            Message::user("Inspect src/main.rs".to_string()),
338            Message::assistant_with_tools(
339                "Running read_file".to_string(),
340                vec![ToolCall {
341                    id: "call-1".to_string(),
342                    call_type: "function".to_string(),
343                    function: Some(FunctionCall {
344                        namespace: None,
345                        name: "read_file".to_string(),
346                        arguments: "{\"path\":\"src/main.rs\"}".to_string(),
347                    }),
348                    text: None,
349                    thought_signature: None,
350                }],
351            ),
352            Message::tool_response("call-1".to_string(), "{\"content\":\"fn main() {}\"}".to_string()),
353            Message::assistant("Done".to_string()),
354        ];
355
356        let conversation = conversation_from_messages(&history);
357        let rebuilt = build_messages_from_conversation(&conversation);
358
359        assert_eq!(rebuilt[0].role, MessageRole::User);
360        assert_eq!(rebuilt[0].content.as_text().as_ref(), "Inspect src/main.rs");
361        assert_eq!(rebuilt[1].role, MessageRole::Assistant);
362        assert_eq!(
363            rebuilt[1]
364                .tool_calls
365                .as_ref()
366                .and_then(|calls| calls.first())
367                .and_then(|call| call.function.as_ref())
368                .map(|function| function.name.as_str()),
369            Some("read_file")
370        );
371        assert_eq!(rebuilt[2].role, MessageRole::Tool);
372        assert_eq!(rebuilt[2].tool_call_id.as_deref(), Some("call-1"));
373        assert_eq!(rebuilt[3].role, MessageRole::Assistant);
374        assert_eq!(rebuilt[3].content.as_text().as_ref(), "Done");
375    }
376}