certifai/src/infrastructure/llm.rs

use dioxus::prelude::*;

#[cfg(feature = "server")]
mod inner {
    use serde::{Deserialize, Serialize};

    /// A single message in the OpenAI-compatible chat format used by LiteLLM.
    #[derive(Serialize)]
    pub(super) struct ChatMessage {
        pub role: String,
        pub content: String,
    }

    /// Request body for the OpenAI-compatible chat completions endpoint.
    #[derive(Serialize)]
    pub(super) struct ChatCompletionRequest {
        pub model: String,
        pub messages: Vec<ChatMessage>,
        /// Disable streaming so we get a single JSON response.
        pub stream: bool,
    }

    /// A single choice in the chat completions response.
    #[derive(Deserialize)]
    pub(super) struct ChatChoice {
        pub message: ChatResponseMessage,
    }

    /// The assistant message returned inside a choice.
    #[derive(Deserialize)]
    pub(super) struct ChatResponseMessage {
        pub content: String,
    }

    /// Top-level response from the `/v1/chat/completions` endpoint.
    #[derive(Deserialize)]
    pub(super) struct ChatCompletionResponse {
        pub choices: Vec<ChatChoice>,
    }

    /// Fetch the full text content of a webpage by downloading its HTML
    /// and extracting the main article body, skipping navigation, headers,
    /// footers, and sidebars.
    ///
    /// Uses a tiered extraction strategy:
    /// 1. Try content within `<article>`, `<main>`, or `[role="main"]`
    /// 2. Fall back to all `<p>` tags outside excluded containers
    ///
    /// # Arguments
    ///
    /// * `url` - The article URL to fetch
    ///
    /// # Returns
    ///
    /// The extracted text, or `None` if the fetch/parse fails.
    /// Text is capped at 8000 characters to stay within LLM context limits.
    pub(super) async fn fetch_article_text(url: &str) -> Option<String> {
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(10))
            .build()
            .ok()?;

        let resp = client
            .get(url)
            .header("User-Agent", "CERTifAI/1.0 (Article Summarizer)")
            .send()
            .await
            .ok()?;

        if !resp.status().is_success() {
            return None;
        }

        let html = resp.text().await.ok()?;
        parse_article_html(&html)
    }

    /// Parse article text from raw HTML without any network I/O.
    ///
    /// Uses a tiered extraction strategy:
    /// 1. Try content within `<article>`, `<main>`, or `[role="main"]`
    /// 2. Fall back to all `<p>` tags outside excluded containers
    ///
    /// # Arguments
    ///
    /// * `html` - Raw HTML string to parse
    ///
    /// # Returns
    ///
    /// The extracted text, or `None` if extraction yields < 100 chars.
    /// Output is capped at 8000 characters.
    pub(crate) fn parse_article_html(html: &str) -> Option<String> {
        let document = scraper::Html::parse_document(html);

        // Strategy 1: Extract from semantic article containers.
        // Most news sites wrap the main content in <article>, <main>,
        // or an element with role="main".
        let article_selector = scraper::Selector::parse("article, main, [role='main']").ok()?;
        let paragraph_sel = scraper::Selector::parse("p, h1, h2, h3, li").ok()?;

        let mut text_parts: Vec<String> = Vec::with_capacity(64);

        for container in document.select(&article_selector) {
            for element in container.select(&paragraph_sel) {
                collect_text_fragment(element, &mut text_parts);
            }
        }

        // Strategy 2: If article containers yielded little text, fall back
        // to all <p> tags that are NOT inside nav/header/footer/aside.
        if joined_len(&text_parts) < 200 {
            text_parts.clear();
            let all_p = scraper::Selector::parse("p").ok()?;

            // Tags whose descendants should be excluded from extraction
            const EXCLUDED_TAGS: &[&str] = &["nav", "header", "footer", "aside", "script", "style"];

            for element in document.select(&all_p) {
                // Walk ancestors and skip if inside an excluded container.
                // Checks tag names directly to avoid ego_tree version issues.
                let inside_excluded = element.ancestors().any(|ancestor| {
                    ancestor
                        .value()
                        .as_element()
                        .is_some_and(|el| EXCLUDED_TAGS.contains(&el.name.local.as_ref()))
                });
                if !inside_excluded {
                    collect_text_fragment(element, &mut text_parts);
                }
            }
        }

        let full_text = text_parts.join("\n\n");
        if full_text.len() < 100 {
            return None;
        }

        // Cap at 8000 chars to stay within reasonable LLM context
        let truncated: String = full_text.chars().take(8000).collect();
        Some(truncated)
    }

    /// Extract text from an HTML element and append it to the parts list
    /// if it meets a minimum length threshold.
    fn collect_text_fragment(element: scraper::ElementRef<'_>, parts: &mut Vec<String>) {
        let text: String = element.text().collect::<Vec<_>>().join(" ");
        let trimmed = text.trim().to_string();
        // Skip very short fragments (nav items, buttons, etc.)
        if trimmed.len() >= 30 {
            parts.push(trimmed);
        }
    }

    /// Sum the total character length of all collected text parts.
    pub(crate) fn joined_len(parts: &[String]) -> usize {
        parts.iter().map(|s| s.len()).sum()
    }
}

/// Summarize an article using a LiteLLM proxy.
///
/// First attempts to fetch the full article text from the provided URL.
/// If that fails (paywall, timeout, etc.), falls back to the search snippet.
/// This mirrors how Perplexity fetches and reads source pages before answering.
///
/// # Arguments
///
/// * `snippet` - The search result snippet (fallback content)
/// * `article_url` - The original article URL to fetch full text from
/// * `litellm_url` - Base URL of the LiteLLM proxy (e.g. "http://localhost:4000")
/// * `model` - The model ID to use (e.g. "qwen3-32b")
///
/// # Returns
///
/// A summary string generated by the LLM, or a `ServerFnError` on failure
///
/// # Errors
///
/// Returns `ServerFnError` if the LiteLLM request fails or response parsing fails
#[post("/api/summarize")]
pub async fn summarize_article(
    snippet: String,
    article_url: String,
    litellm_url: String,
    model: String,
) -> Result<String, ServerFnError> {
    use inner::{fetch_article_text, ChatCompletionRequest, ChatCompletionResponse, ChatMessage};

    let state: crate::infrastructure::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    // Use caller-provided values or fall back to ServerState config
    let base_url = if litellm_url.is_empty() {
        state.services.litellm_url.clone()
    } else {
        litellm_url
    };

    let model = if model.is_empty() {
        state.services.litellm_model.clone()
    } else {
        model
    };

    let api_key = state.services.litellm_api_key.clone();

    // Try to fetch the full article; fall back to the search snippet
    let article_text = fetch_article_text(&article_url).await.unwrap_or(snippet);

    let request_body = ChatCompletionRequest {
        model,
        stream: false,
        messages: vec![ChatMessage {
            role: "user".into(),
            content: format!(
                "You are a news summarizer. Summarize the following article text \
                 in 2-3 concise paragraphs. Focus only on the key points and \
                 implications. Do NOT comment on the source, the date, the URL, \
                 the formatting, or whether the content seems complete or not. \
                 Just summarize whatever content is provided.\n\n\
                 {article_text}"
            ),
        }],
    };

    let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));
    let client = reqwest::Client::new();
    let mut request = client
        .post(&url)
        .header("content-type", "application/json")
        .json(&request_body);

    if !api_key.is_empty() {
        request = request.header("Authorization", format!("Bearer {api_key}"));
    }

    let resp = request
        .send()
        .await
        .map_err(|e| ServerFnError::new(format!("LiteLLM request failed: {e}")))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!(
            "LiteLLM returned {status}: {body}"
        )));
    }

    let body: ChatCompletionResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(format!("Failed to parse LiteLLM response: {e}")))?;

    body.choices
        .first()
        .map(|choice| choice.message.content.clone())
        .ok_or_else(|| ServerFnError::new("Empty response from LiteLLM"))
}

/// A lightweight chat message for the follow-up conversation.
/// Uses simple String role ("system"/"user"/"assistant") for OpenAI compatibility.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct FollowUpMessage {
    pub role: String,
    pub content: String,
}

/// Send a follow-up question about an article using a LiteLLM proxy.
///
/// Accepts the full conversation history (system context + prior turns) and
/// returns the assistant's next response. The system message should contain
/// the article text and summary so the LLM has full context.
///
/// # Arguments
///
/// * `messages` - The conversation history including system context
/// * `litellm_url` - Base URL of the LiteLLM proxy
/// * `model` - The model ID to use
///
/// # Returns
///
/// The assistant's response text, or a `ServerFnError` on failure
///
/// # Errors
///
/// Returns `ServerFnError` if the LiteLLM request fails or response parsing fails
#[post("/api/chat")]
pub async fn chat_followup(
    messages: Vec<FollowUpMessage>,
    litellm_url: String,
    model: String,
) -> Result<String, ServerFnError> {
    use inner::{ChatCompletionRequest, ChatCompletionResponse, ChatMessage};

    let state: crate::infrastructure::ServerState =
        dioxus_fullstack::FullstackContext::extract().await?;

    let base_url = if litellm_url.is_empty() {
        state.services.litellm_url.clone()
    } else {
        litellm_url
    };

    let model = if model.is_empty() {
        state.services.litellm_model.clone()
    } else {
        model
    };

    let api_key = state.services.litellm_api_key.clone();

    // Convert FollowUpMessage to inner ChatMessage for the request
    let chat_messages: Vec<ChatMessage> = messages
        .into_iter()
        .map(|m| ChatMessage {
            role: m.role,
            content: m.content,
        })
        .collect();

    let request_body = ChatCompletionRequest {
        model,
        stream: false,
        messages: chat_messages,
    };

    let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/'));
    let client = reqwest::Client::new();
    let mut request = client
        .post(&url)
        .header("content-type", "application/json")
        .json(&request_body);

    if !api_key.is_empty() {
        request = request.header("Authorization", format!("Bearer {api_key}"));
    }

    let resp = request
        .send()
        .await
        .map_err(|e| ServerFnError::new(format!("LiteLLM request failed: {e}")))?;

    if !resp.status().is_success() {
        let status = resp.status();
        let body = resp.text().await.unwrap_or_default();
        return Err(ServerFnError::new(format!(
            "LiteLLM returned {status}: {body}"
        )));
    }

    let body: ChatCompletionResponse = resp
        .json()
        .await
        .map_err(|e| ServerFnError::new(format!("Failed to parse LiteLLM response: {e}")))?;

    body.choices
        .first()
        .map(|choice| choice.message.content.clone())
        .ok_or_else(|| ServerFnError::new("Empty response from LiteLLM"))
}

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

    // -----------------------------------------------------------------------
    // FollowUpMessage serde tests
    // -----------------------------------------------------------------------

    #[test]
    fn followup_message_serde_round_trip() {
        let msg = FollowUpMessage {
            role: "assistant".into(),
            content: "Here is my answer.".into(),
        };
        let json = serde_json::to_string(&msg).expect("serialize FollowUpMessage");
        let back: FollowUpMessage =
            serde_json::from_str(&json).expect("deserialize FollowUpMessage");
        assert_eq!(msg, back);
    }

    #[test]
    fn followup_message_deserialize_from_json_literal() {
        let json = r#"{"role":"system","content":"You are helpful."}"#;
        let msg: FollowUpMessage = serde_json::from_str(json).expect("deserialize literal");
        assert_eq!(msg.role, "system");
        assert_eq!(msg.content, "You are helpful.");
    }

    // -----------------------------------------------------------------------
    // joined_len and parse_article_html tests (server feature required)
    // -----------------------------------------------------------------------

    #[cfg(feature = "server")]
    mod server_tests {
        use super::super::inner::{joined_len, parse_article_html};
        use pretty_assertions::assert_eq;

        #[test]
        fn joined_len_empty_input() {
            assert_eq!(joined_len(&[]), 0);
        }

        #[test]
        fn joined_len_sums_correctly() {
            let parts = vec!["abc".into(), "de".into(), "fghij".into()];
            assert_eq!(joined_len(&parts), 10);
        }

        // -------------------------------------------------------------------
        // parse_article_html tests
        // -------------------------------------------------------------------

        // Helper: generate a string of given length from a repeated word.
        fn lorem(len: usize) -> String {
            "Lorem ipsum dolor sit amet consectetur adipiscing elit "
                .repeat((len / 55) + 1)
                .chars()
                .take(len)
                .collect()
        }

        #[test]
        fn article_tag_extracts_text() {
            let body = lorem(250);
            let html = format!("<html><body><article><p>{body}</p></article></body></html>");
            let result = parse_article_html(&html);
            assert!(result.is_some(), "expected Some for article tag");
            assert!(result.unwrap().contains("Lorem"));
        }

        #[test]
        fn main_tag_extracts_text() {
            let body = lorem(250);
            let html = format!("<html><body><main><p>{body}</p></main></body></html>");
            let result = parse_article_html(&html);
            assert!(result.is_some(), "expected Some for main tag");
        }

        #[test]
        fn fallback_to_p_tags_when_article_main_yield_little() {
            // No <article>/<main>, so falls back to <p> tags
            let body = lorem(250);
            let html = format!("<html><body><div><p>{body}</p></div></body></html>");
            let result = parse_article_html(&html);
            assert!(result.is_some(), "expected fallback to <p> tags");
        }

        #[test]
        fn excludes_nav_footer_aside_content() {
            // Content only inside excluded containers -- should be excluded
            let body = lorem(250);
            let html = format!(
                "<html><body>\
                 <nav><p>{body}</p></nav>\
                 <footer><p>{body}</p></footer>\
                 <aside><p>{body}</p></aside>\
                 </body></html>"
            );
            let result = parse_article_html(&html);
            assert!(result.is_none(), "expected None for excluded-only content");
        }

        #[test]
        fn returns_none_when_text_too_short() {
            let html = "<html><body><p>Short.</p></body></html>";
            let result = parse_article_html(html);
            assert!(result.is_none(), "expected None for short text");
        }

        #[test]
        fn truncates_at_8000_chars() {
            let body = lorem(10000);
            let html = format!("<html><body><article><p>{body}</p></article></body></html>");
            let result = parse_article_html(&html).expect("expected Some");
            assert!(
                result.len() <= 8000,
                "expected <= 8000 chars, got {}",
                result.len()
            );
        }

        #[test]
        fn skips_fragments_under_30_chars() {
            // Only fragments < 30 chars -- should yield None
            let html = "<html><body><article>\
                         <p>Short frag one</p>\
                         <p>Another tiny bit</p>\
                         </article></body></html>";
            let result = parse_article_html(html);
            assert!(result.is_none(), "expected None for tiny fragments");
        }

        #[test]
        fn extracts_from_role_main_attribute() {
            let body = lorem(250);
            let html = format!(
                "<html><body>\
                 <div role=\"main\"><p>{body}</p></div>\
                 </body></html>"
            );
            let result = parse_article_html(&html);
            assert!(result.is_some(), "expected Some for role=main");
        }
    }
}