tiny-agent 0.1.0

use super::ResponseItem;
use crate::shared::{ContentBlock, Message, Role};

pub fn transform_messages(messages: Vec<Message>) -> Vec<ResponseItem> {
    messages.into_iter().flat_map(transform_message).collect()
}

fn transform_message(message: Message) -> Vec<ResponseItem> {
    match message.role {
        Role::Assistant => vec![assistant_message(message.blocks)],
        Role::Tool => tool_result_messages(message.blocks),
        Role::System => vec![plain_message("system", message.blocks)],
        Role::User => vec![plain_message("user", message.blocks)],
    }
}

fn assistant_message(blocks: Vec<ContentBlock>) -> ResponseItem {
    let mut text = String::new();
    let mut tool_calls = Vec::new();

    for block in blocks {
        match block {
            ContentBlock::Text { text: block_text } => {
                text.push_str(&block_text);
            }
            ContentBlock::Thinking { .. } => {}
            ContentBlock::ToolUse { id, name, input } => {
                tool_calls.push(super::OpenAiToolCall {
                    index: tool_calls.len(),
                    id: Some(id),
                    r#type: Some(String::from("function")),
                    function: Some(super::OpenAiFunc {
                        name: Some(name),
                        arguments: Some(input.to_string()),
                    }),
                });
            }
            ContentBlock::ToolResult { .. } => {}
        }
    }

    ResponseItem {
        role: String::from("assistant"),
        content: if text.is_empty() { None } else { Some(text) },
        tool_calls: if tool_calls.is_empty() {
            None
        } else {
            Some(tool_calls)
        },
        tool_call_id: None,
    }
}

fn tool_result_messages(blocks: Vec<ContentBlock>) -> Vec<ResponseItem> {
    blocks
        .into_iter()
        .filter_map(|block| {
            let ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error,
            } = block
            else {
                return None;
            };

            let mut text = flatten_text(content);
            if is_error {
                text = format!("Tool execution failed: {text}");
            }

            Some(ResponseItem {
                role: String::from("tool"),
                content: Some(text),
                tool_calls: None,
                tool_call_id: Some(tool_use_id),
            })
        })
        .collect()
}

fn plain_message(role: &str, blocks: Vec<ContentBlock>) -> ResponseItem {
    ResponseItem {
        role: role.to_string(),
        content: Some(flatten_text(blocks)),
        tool_calls: None,
        tool_call_id: None,
    }
}

fn flatten_text(blocks: Vec<ContentBlock>) -> String {
    let mut out = String::new();
    for block in blocks {
        match block {
            ContentBlock::Text { text } => {
                out.push_str(&text);
            }
            ContentBlock::Thinking { .. } => {}
            ContentBlock::ToolUse { id, name, input } => {
                out.push_str(&format!("[tool_use id={id} name={name} input={}]", input));
            }
            ContentBlock::ToolResult {
                tool_use_id,
                content,
                is_error,
            } => {
                if is_error {
                    out.push_str("[tool_result error ");
                } else {
                    out.push_str("[tool_result ");
                }
                out.push_str(&format!(
                    "tool_use_id={tool_use_id}]{}",
                    flatten_text(content)
                ));
            }
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn openai_projection_does_not_emit_thinking_as_content() {
        let messages = vec![Message {
            role: Role::Assistant,
            blocks: vec![
                ContentBlock::Thinking {
                    text: "private reasoning".to_string(),
                },
                ContentBlock::Text {
                    text: "visible answer".to_string(),
                },
            ],
        }];

        let projected = transform_messages(messages);

        assert_eq!(projected.len(), 1);
        assert_eq!(projected[0].content.as_deref(), Some("visible answer"));
    }
}