kb-skills/kb_workflow/pipeline/link_parser.py

"""Link Parser module

Provides unified URL parsing for various content platforms.
Supports Bilibili (video, article) and Zhihu (answer, zhuanlan) URLs.
"""

import re
from dataclasses import dataclass, field
from typing import Optional

import requests

from kb_workflow.pipeline.retry import retry_on_network_error


@dataclass
class ParsedLink:
    """Represents a parsed link with platform and content information

    Attributes:
        platform: Platform name (bilibili, zhihu)
        content_type: Content type (video, article, answer, zhuanlan)
        content_id: Unique content identifier
        original_url: Original URL provided by user
        resolved_url: Resolved URL after short link expansion (defaults to original_url)
    """
    platform: str
    content_type: str
    content_id: str
    original_url: str
    resolved_url: str = field(default="")

    def __post_init__(self):
        """Set resolved_url to original_url if not provided"""
        if not self.resolved_url:
            self.resolved_url = self.original_url


class LinkParser:
    """Unified URL parser for multiple platforms

    Parses URLs from Bilibili and Zhihu, extracting platform, content type,
    and content ID. Supports short link resolution for b23.tv URLs.
    """

    # Bilibili patterns
    BILIBILI_VIDEO_PATTERN = re.compile(
        r'https?://(?:www\.)?bilibili\.com/video/(BV[a-zA-Z0-9]+)',
        re.IGNORECASE
    )
    BILIBILI_ARTICLE_PATTERN = re.compile(
        r'https?://(?:www\.)?bilibili\.com/(?:read|article)/(cv[a-zA-Z0-9]+)',
        re.IGNORECASE
    )

    # Zhihu patterns
    ZHIHU_ANSWER_PATTERN = re.compile(
        r'https?://(?:www\.)?zhihu\.com/question/\d+/answer/(\d+)',
        re.IGNORECASE
    )
    ZHIHU_ZHUANLAN_PATTERN = re.compile(
        r'https?://zhuanlan\.zhihu\.com/p/(\d+)',
        re.IGNORECASE
    )

    # Short link patterns
    B23_TV_PATTERN = re.compile(
        r'https?://(?:www\.)?b23\.tv/([a-zA-Z0-9]+)',
        re.IGNORECASE
    )

    # Timeout for short link resolution
    SHORT_LINK_TIMEOUT = 10

    @retry_on_network_error
    def _resolve_shortlink(self, url: str) -> Optional[str]:
        """Resolve short link to its final URL

        Uses HTTP HEAD request with redirects to get the final URL.

        Args:
            url: Short link URL to resolve

        Returns:
            Final URL after redirects, or None if resolution failed
        """
        try:
            response = requests.head(
                url,
                allow_redirects=True,
                timeout=self.SHORT_LINK_TIMEOUT
            )
            if response.status_code == 200:
                return response.url
            return None
        except (requests.RequestException, OSError):
            return None

    def parse(self, url: str) -> Optional[ParsedLink]:
        """Parse a URL and extract platform, content type, and ID

        Args:
            url: URL to parse

        Returns:
            ParsedLink with extracted information, or None if URL is not recognized
        """
        if not url or not isinstance(url, str):
            return None

        url = url.strip()

        # Try b23.tv short link first
        if self.B23_TV_PATTERN.match(url):
            resolved_url = self._resolve_shortlink(url)
            if resolved_url:
                # Recursively parse the resolved URL
                result = self._parse_standard_url(resolved_url)
                if result:
                    result.original_url = url
                    result.resolved_url = resolved_url
                return result
            return None

        return self._parse_standard_url(url)

    def _parse_standard_url(self, url: str) -> Optional[ParsedLink]:
        """Parse a standard (non-short) URL

        Args:
            url: URL to parse

        Returns:
            ParsedLink with extracted information, or None if URL is not recognized
        """
        # Try Bilibili video
        match = self.BILIBILI_VIDEO_PATTERN.match(url)
        if match:
            return ParsedLink(
                platform="bilibili",
                content_type="video",
                content_id=match.group(1),
                original_url=url,
                resolved_url=url
            )

        # Try Bilibili article
        match = self.BILIBILI_ARTICLE_PATTERN.match(url)
        if match:
            return ParsedLink(
                platform="bilibili",
                content_type="article",
                content_id=match.group(1),
                original_url=url,
                resolved_url=url
            )

        # Try Zhihu answer
        match = self.ZHIHU_ANSWER_PATTERN.match(url)
        if match:
            return ParsedLink(
                platform="zhihu",
                content_type="answer",
                content_id=match.group(1),
                original_url=url,
                resolved_url=url
            )

        # Try Zhihu zhuanlan
        match = self.ZHIHU_ZHUANLAN_PATTERN.match(url)
        if match:
            return ParsedLink(
                platform="zhihu",
                content_type="zhuanlan",
                content_id=match.group(1),
                original_url=url,
                resolved_url=url
            )

        # URL not recognized
        return None