adk-server 0.5.0

use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    User,
    Agent,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FileContent {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "mimeType")]
    pub mime_type: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bytes: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub uri: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Part {
    Text {
        text: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        metadata: Option<Map<String, Value>>,
    },
    File {
        file: FileContent,
        #[serde(skip_serializing_if = "Option::is_none")]
        metadata: Option<Map<String, Value>>,
    },
    Data {
        data: Map<String, Value>,
        #[serde(skip_serializing_if = "Option::is_none")]
        metadata: Option<Map<String, Value>>,
    },
}

impl Part {
    pub fn text(text: String) -> Self {
        Part::Text { text, metadata: None }
    }

    pub fn file(file: FileContent) -> Self {
        Part::File { file, metadata: None }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    pub role: Role,
    pub parts: Vec<Part>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<Map<String, Value>>,
    #[serde(rename = "messageId")]
    pub message_id: String,
    #[serde(skip_serializing_if = "Option::is_none", rename = "taskId")]
    pub task_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none", rename = "contextId")]
    pub context_id: Option<String>,
}

impl Message {
    pub fn builder() -> MessageBuilder {
        MessageBuilder::default()
    }
}

#[derive(Default)]
pub struct MessageBuilder {
    role: Option<Role>,
    parts: Vec<Part>,
    metadata: Option<Map<String, Value>>,
    message_id: Option<String>,
    task_id: Option<String>,
    context_id: Option<String>,
}

impl MessageBuilder {
    pub fn role(mut self, role: Role) -> Self {
        self.role = Some(role);
        self
    }

    pub fn parts(mut self, parts: Vec<Part>) -> Self {
        self.parts = parts;
        self
    }

    pub fn metadata(mut self, metadata: Option<Map<String, Value>>) -> Self {
        self.metadata = metadata;
        self
    }

    pub fn message_id(mut self, id: String) -> Self {
        self.message_id = Some(id);
        self
    }

    pub fn build(self) -> Message {
        Message {
            role: self.role.unwrap_or(Role::User),
            parts: self.parts,
            metadata: self.metadata,
            message_id: self.message_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
            task_id: self.task_id,
            context_id: self.context_id,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Artifact {
    #[serde(rename = "artifactId")]
    pub artifact_id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    pub parts: Vec<Part>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub metadata: Option<Map<String, Value>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub extensions: Option<Vec<String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub enum TaskState {
    Submitted,
    Working,
    InputRequired,
    Completed,
    Failed,
    Canceled,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskStatus {
    pub state: TaskState,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskStatusUpdateEvent {
    #[serde(rename = "taskId")]
    pub task_id: String,
    #[serde(skip_serializing_if = "Option::is_none", rename = "contextId")]
    pub context_id: Option<String>,
    pub status: TaskStatus,
    #[serde(rename = "finalUpdate")]
    pub final_update: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskArtifactUpdateEvent {
    #[serde(rename = "taskId")]
    pub task_id: String,
    #[serde(skip_serializing_if = "Option::is_none", rename = "contextId")]
    pub context_id: Option<String>,
    pub artifact: Artifact,
    pub append: bool,
    #[serde(rename = "lastChunk")]
    pub last_chunk: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum UpdateEvent {
    TaskStatusUpdate(TaskStatusUpdateEvent),
    TaskArtifactUpdate(TaskArtifactUpdateEvent),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentSkill {
    pub id: String,
    pub name: String,
    pub description: String,
    pub tags: Vec<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub examples: Option<Vec<String>>,
}

impl AgentSkill {
    pub fn new(id: String, name: String, description: String, tags: Vec<String>) -> Self {
        Self { id, name, description, tags, examples: None }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCapabilities {
    pub streaming: bool,
    #[serde(rename = "pushNotifications")]
    pub push_notifications: bool,
    #[serde(rename = "stateTransitionHistory")]
    pub state_transition_history: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub extensions: Option<Vec<String>>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentCard {
    pub name: String,
    pub description: String,
    pub url: String,
    pub version: String,
    #[serde(rename = "protocolVersion")]
    pub protocol_version: String,
    pub capabilities: AgentCapabilities,
    pub skills: Vec<AgentSkill>,
}

impl AgentCard {
    pub fn builder() -> AgentCardBuilder {
        AgentCardBuilder::default()
    }
}

#[derive(Default)]
pub struct AgentCardBuilder {
    name: Option<String>,
    description: Option<String>,
    url: Option<String>,
    version: Option<String>,
    capabilities: Option<AgentCapabilities>,
    skills: Vec<AgentSkill>,
}

impl AgentCardBuilder {
    pub fn name(mut self, name: String) -> Self {
        self.name = Some(name);
        self
    }

    pub fn description(mut self, description: String) -> Self {
        self.description = Some(description);
        self
    }

    pub fn url(mut self, url: String) -> Self {
        self.url = Some(url);
        self
    }

    pub fn version(mut self, version: String) -> Self {
        self.version = Some(version);
        self
    }

    pub fn capabilities(mut self, capabilities: AgentCapabilities) -> Self {
        self.capabilities = Some(capabilities);
        self
    }

    pub fn skills(mut self, skills: Vec<AgentSkill>) -> Self {
        self.skills = skills;
        self
    }

    /// Build the `AgentCard`, returning an error if required fields are missing.
    ///
    /// # Errors
    ///
    /// Returns a descriptive error string when `name`, `description`, or `url`
    /// have not been set (or were set to empty strings).
    pub fn build(self) -> std::result::Result<AgentCard, String> {
        let name =
            self.name.filter(|s| !s.is_empty()).ok_or("AgentCard requires a non-empty 'name'")?;
        let description = self
            .description
            .filter(|s| !s.is_empty())
            .ok_or("AgentCard requires a non-empty 'description'")?;
        let url =
            self.url.filter(|s| !s.is_empty()).ok_or("AgentCard requires a non-empty 'url'")?;

        Ok(AgentCard {
            name,
            description,
            url,
            version: self.version.unwrap_or_else(|| "1.0.0".to_string()),
            protocol_version: "0.3.0".to_string(),
            capabilities: self.capabilities.unwrap_or(AgentCapabilities {
                streaming: false,
                push_notifications: false,
                state_transition_history: false,
                extensions: None,
            }),
            skills: self.skills,
        })
    }
}