kagi-api 0.1.1

use async_trait::async_trait;
use serde::{Deserialize, Serialize};

use crate::v0::Meta;
use crate::{KagiClient, KagiError, KagiResult};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Summary {
    pub meta: Meta,
    pub data: Data,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Data {
    pub output: String,
    pub tokens: usize,
}

/// Parameters url and text are exclusive. You must pass one or the other.
///
/// Total request size is limited to 1MB.
#[derive(Debug, Clone)]
pub enum Subject {
    /// A [url::Url] to a document to summarize.
    Url(url::Url),

    /// Text to summarize.
    Text(String),
}

impl Subject {
    pub fn url_or_text(input: String) -> Self {
        url::Url::parse(&input)
            .map(Subject::Url)
            .unwrap_or_else(|_| Subject::Text(input))
    }
}

/// Several options are provided to control the output the summarizer produces.
///
/// - Summary Types
/// - Summarization Engines
/// - Target Language Codes
/// - Cache
#[derive(Debug, Default)]
pub struct SummaryOptions {
    /// Different summary types are provided that control the structure of the summary output.
    pub summary_type: Option<SummaryType>,

    /// Different summarization engines give you choices over the "flavor" of the summary.
    pub engine: Option<Engine>,

    /// The summarizer can translate the output into a desired language
    pub target_language: Option<Language>,

    /// Whether to allow cached requests & responses. (default is true)
    pub cache: Option<bool>,
}

/// Different summary types are provided that control the structure of the summary output.
#[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum SummaryType {
    /// Paragraph(s) of summary prose
    #[default]
    Summary,

    /// Bulleted list of key points
    Takeaway,
}

/// Different summarization engines are provided that will give you choices over the "flavor"
/// of the summarization text.
#[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Engine {
    /// Friendly, descriptive, fast summary
    #[default]
    Cecil,

    /// Formal, technical, analytical summary
    Agnes,

    /// Informal, creative, friendly summary
    Daphne,

    /// Best-in-class summary using Kagi's enterprise-grade model
    Muriel,
}

/// The summarizer can translate the output into a desired language, using the table of supported
/// language codes below.
///
/// If no language is specified, the document's original language is allowed to influence the
/// summarizer's output. Specifying a language will add a an explicit translation step, to
/// translate the summary to the desired language.
///
/// For example, if a document is mostly written in Spanish, the summary output may itself be in
/// Spanish or contain Spanish passages. Specifying "EN" will ensure all passages are translated
/// as English.
#[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
pub enum Language {
    /// Bulgarian
    BG,
    /// Czech
    CS,
    /// Danish
    DA,
    /// German
    DE,
    /// Greek
    EL,
    /// English
    #[default]
    EN,
    /// Spanish
    ES,
    /// Estonian
    ET,
    /// Finnish
    FI,
    /// French
    FR,
    /// Hungarian
    HU,
    /// Indonesian
    ID,
    /// Italian
    IT,
    /// Japanese
    JA,
    /// Korean
    KO,
    /// Lithuanian
    LT,
    /// Latvian
    LV,
    /// Norwegian
    NB,
    /// Dutch
    NL,
    /// Polish
    PL,
    /// Portuguese
    PT,
    /// Romanian
    RO,
    /// Russian
    RU,
    /// Slovak
    SK,
    /// Slovenian
    SL,
    /// Swedish
    SV,
    /// Turkish
    TR,
    /// Ukrainian
    UK,
    /// Chinese (simplified)
    ZH,
}

/// The Universal Summarizer is an API using powerful LLMs to summarize content on the web, or
/// your own documents, of any length.
#[async_trait]
pub trait UniversalSummarizer {
    async fn summarize(&self, subject: Subject, options: SummaryOptions) -> KagiResult<Summary>;

    async fn summarize_url(&self, url: url::Url, options: SummaryOptions) -> KagiResult<Summary> {
        self.summarize(Subject::Url(url), options).await
    }

    async fn summarize_text(&self, text: String, options: SummaryOptions) -> KagiResult<Summary> {
        self.summarize(Subject::Text(text), options).await
    }
}

#[async_trait]
impl UniversalSummarizer for KagiClient {
    async fn summarize(&self, subject: Subject, options: SummaryOptions) -> KagiResult<Summary> {
        let request_body = build_request_body(subject, options);
        let response = self.query("/v0/summarize", request_body).await?;
        response
            .json::<Summary>()
            .await
            .map_err(KagiError::ReqwestError)
    }
}

fn build_request_body(subject: Subject, options: SummaryOptions) -> serde_json::Value {
    let mut request_body = match subject {
        Subject::Url(url) => serde_json::json!({ "url": url }),
        Subject::Text(text) => serde_json::json!({ "text": text }),
    };

    if let Some(summary_type) = options.summary_type {
        request_body["summary_type"] = serde_json::to_value(summary_type).unwrap();
    }

    if let Some(engine) = options.engine {
        request_body["engine"] = serde_json::to_value(engine).unwrap();
    }

    if let Some(target_language) = options.target_language {
        request_body["target_language"] = serde_json::to_value(target_language).unwrap();
    }

    if let Some(cache) = options.cache {
        request_body["cache"] = serde_json::to_value(cache).unwrap();
    }

    request_body
}

impl std::fmt::Display for SummaryType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SummaryType::Summary => write!(f, "summary"),
            SummaryType::Takeaway => write!(f, "takeaway"),
        }
    }
}

impl std::str::FromStr for SummaryType {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "summary" => Ok(SummaryType::Summary),
            "takeaway" => Ok(SummaryType::Takeaway),
            _ => Err(format!("Unknown summary type '{s}'")),
        }
    }
}

impl std::fmt::Display for Language {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let language = match self {
            Language::BG => "bg",
            Language::CS => "cs",
            Language::DA => "da",
            Language::DE => "de",
            Language::EL => "el",
            Language::EN => "en",
            Language::ES => "es",
            Language::ET => "et",
            Language::FI => "fi",
            Language::FR => "fr",
            Language::HU => "hu",
            Language::ID => "id",
            Language::IT => "it",
            Language::JA => "ja",
            Language::KO => "ko",
            Language::LT => "lt",
            Language::LV => "lv",
            Language::NB => "nb",
            Language::NL => "nl",
            Language::PL => "pl",
            Language::PT => "pt",
            Language::RO => "ro",
            Language::RU => "ru",
            Language::SK => "sk",
            Language::SL => "sl",
            Language::SV => "sv",
            Language::TR => "tr",
            Language::UK => "uk",
            Language::ZH => "zh",
        };

        write!(f, "{}", language)
    }
}

impl std::str::FromStr for Language {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let language = match s.to_lowercase().as_ref() {
            "bg" => Language::BG,
            "cs" => Language::CS,
            "da" => Language::DA,
            "de" => Language::DE,
            "el" => Language::EL,
            "en" => Language::EN,
            "es" => Language::ES,
            "et" => Language::ET,
            "fi" => Language::FI,
            "fr" => Language::FR,
            "hu" => Language::HU,
            "id" => Language::ID,
            "it" => Language::IT,
            "ja" => Language::JA,
            "ko" => Language::KO,
            "lt" => Language::LT,
            "lv" => Language::LV,
            "nb" => Language::NB,
            "nl" => Language::NL,
            "pl" => Language::PL,
            "pt" => Language::PT,
            "ro" => Language::RO,
            "ru" => Language::RU,
            "sk" => Language::SK,
            "sl" => Language::SL,
            "sv" => Language::SV,
            "tr" => Language::TR,
            "uk" => Language::UK,
            "zh" => Language::ZH,
            _ => return Err(format!("Unknown language '{s}'")),
        };

        Ok(language)
    }
}

impl std::fmt::Display for Engine {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let engine = match self {
            Engine::Cecil => "cecil",
            Engine::Agnes => "agnes",
            Engine::Daphne => "daphne",
            Engine::Muriel => "muriel",
        };

        write!(f, "{}", engine)
    }
}

impl std::str::FromStr for Engine {
    type Err = String;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let engine = match s.to_lowercase().as_ref() {
            "cecil" => Engine::Cecil,
            "agnes" => Engine::Agnes,
            "daphne" => Engine::Daphne,
            "muriel" => Engine::Muriel,
            _ => return Err(format!("Unknown engine '{s}")),
        };

        Ok(engine)
    }
}