kb-skills/kb_pipeline/platforms/zhihu.py

"""Zhihu Platform Adapter

Provides content extraction for Zhihu platform:
- Answer extraction from question/answer URLs
- Zhuanlan (article) extraction from zhuanlan URLs
- Cookie-based authentication
- Login page detection for expired cookies
"""

from typing import Dict, Any, Optional
from datetime import datetime
import re

from kb_pipeline.platforms.base import PlatformAdapter
from kb_pipeline.errors import AuthError
from kb_pipeline.config import get_settings


class ZhihuAdapter(PlatformAdapter):
    """Zhihu platform adapter for answer and zhuanlan content extraction.

    Uses ZhihuPlaywrightExtractor (from kb_zhihu_metadata) for content extraction.
    Requires KB_ZHIHU_COOKIE environment variable for authentication.

    Example:
        adapter = ZhihuAdapter()
        content_type = adapter.get_content_type("https://www.zhihu.com/question/123/answer/456")
        # -> "answer"

        metadata = await adapter.fetch_metadata("456", "answer")
        content = await adapter.fetch_content("456", "answer")
    """

    PLATFORM = "zhihu"

    # URL patterns for content type detection
    ANSWER_PATTERN = re.compile(r'zhihu\.com/question/\d+/answer/(\d+)')
    ZHUANLAN_PATTERN = re.compile(r'zhuanlan\.zhihu\.com/p/([\w\d]+)')
    QUESTION_PATTERN = re.compile(r'zhihu\.com/question/(\d+)(?!/answer)')

    def __init__(self, headless: bool = True, timeout: int = 30):
        """Initialize ZhihuAdapter.

        Args:
            headless: Run browser in headless mode (default: True)
            timeout: Page load timeout in seconds (default: 30)
        """
        self.headless = headless
        self.timeout = timeout
        self._cookie: Optional[str] = None
        self._extractor = None

    @property
    def cookie(self) -> str:
        """Get Zhihu cookie from settings.

        Returns:
            Cookie string from KB_ZHIHU_COOKIE environment variable.
            Returns empty string if not configured.
        """
        if self._cookie is None:
            settings = get_settings()
            self._cookie = settings.zhihu_cookie.get_secret_value()
        return self._cookie

    def get_content_type(self, url: str) -> str:
        """Identify content type from URL.

        Args:
            url: The content URL to analyze

        Returns:
            Content type: 'answer', 'zhuanlan', or 'unknown'
        """
        if self.ANSWER_PATTERN.search(url):
            return "answer"
        elif self.ZHUANLAN_PATTERN.search(url):
            return "zhuanlan"
        return "unknown"

    def _get_extractor(self):
        """Get or create ZhihuPlaywrightExtractor instance.

        Lazy import to avoid dependency issues when kb_zhihu_metadata is not installed.

        Returns:
            ZhihuPlaywrightExtractor instance
        """
        if self._extractor is None:
            try:
                from kb_zhihu_metadata.zhihu_extractor import ZhihuPlaywrightExtractor
                self._extractor = ZhihuPlaywrightExtractor(
                    headless=self.headless,
                    timeout=self.timeout,
                    cookie=self.cookie
                )
            except ImportError as e:
                raise ImportError(
                    "kb_zhihu_metadata package not found. "
                    "Please install it to use ZhihuAdapter."
                ) from e
        return self._extractor

    def _extract_with_playwright(self, url: str) -> Optional[Dict[str, Any]]:
        """Extract content using Playwright-based extractor.

        Args:
            url: Full URL to extract content from

        Returns:
            Extraction result dictionary or None if extraction failed
        """
        extractor = self._get_extractor()
        return extractor.extract(url)

    def _detect_login_page(self, result: Dict[str, Any]) -> bool:
        """Detect if the result indicates a login page redirect.

        This happens when the cookie has expired or is invalid.

        Args:
            result: Extraction result dictionary

        Returns:
            True if login page is detected, False otherwise
        """
        if not result:
            return False

        title = result.get("title", "")
        content = result.get("content", "")
        content_markdown = result.get("content_markdown", "")

        # Check for login-related indicators
        login_indicators = ["登录", "登陆", "sign in", "login"]
        title_lower = title.lower()

        # If title contains login keywords
        if any(indicator in title_lower for indicator in login_indicators):
            return True

        # If content is suspiciously short (login page has minimal content)
        total_content = content + content_markdown
        if len(total_content.strip()) < 100:
            return True

        return False

    async def fetch_metadata(self, content_id: str, content_type: str) -> Dict[str, Any]:
        """Fetch metadata for the content.

        Args:
            content_id: Answer ID or article ID
            content_type: 'answer' or 'zhuanlan'

        Returns:
            Dictionary containing metadata (title, author, publish_time, etc.)

        Raises:
            AuthError: If cookie is missing or expired
        """
        # Check for missing cookie
        if not self.cookie:
            raise AuthError(
                step_name="zhihu_metadata",
                input_summary=f"content_id={content_id}, type={content_type}",
                auth_type="zhihu_cookie",
            )

        # Build URL based on content type
        if content_type == "answer":
            # For answers, we need the question ID too, but extractor handles this
            url = f"https://www.zhihu.com/question/0/answer/{content_id}"
        elif content_type == "zhuanlan":
            url = f"https://zhuanlan.zhihu.com/p/{content_id}"
        else:
            raise ValueError(f"Unsupported content type: {content_type}")

        # Extract content
        result = self._extract_with_playwright(url)

        # Check for login page
        if self._detect_login_page(result):
            raise AuthError(
                step_name="zhihu_metadata",
                input_summary=f"content_id={content_id}, type={content_type}",
                auth_type="zhihu_cookie",
                cause=Exception("Cookie may have expired, please update KB_ZHIHU_COOKIE")
            )

        if not result:
            raise AuthError(
                step_name="zhihu_metadata",
                input_summary=f"content_id={content_id}, type={content_type}",
                auth_type="zhihu_cookie",
                cause=Exception("Failed to extract content, check if content exists")
            )

        return result

    async def fetch_content(self, content_id: str, content_type: str) -> str:
        """Fetch content in Markdown format.

        Args:
            content_id: Answer ID or article ID
            content_type: 'answer' or 'zhuanlan'

        Returns:
            Markdown-formatted content with YAML frontmatter

        Raises:
            AuthError: If cookie is missing or expired
        """
        # Get metadata (which also validates auth)
        metadata = await self.fetch_metadata(content_id, content_type)

        # Get content markdown
        content_markdown = metadata.get("content_markdown", "") or metadata.get("content", "")

        # Format with YAML frontmatter
        return self._format_with_frontmatter(metadata, content_markdown, content_type)

    def _format_with_frontmatter(
        self,
        metadata: Dict[str, Any],
        content: str,
        content_type: str
    ) -> str:
        """Format content with YAML frontmatter.

        Args:
            metadata: Content metadata dictionary
            content: Markdown content
            content_type: 'answer' or 'zhuanlan'

        Returns:
            Formatted string with YAML frontmatter and content
        """
        frontmatter_lines = [
            "---",
            f"title: {metadata.get('title', '')}",
            f"author: {metadata.get('author', '')}",
            f"source_url: {metadata.get('url', '')}",
            f"platform: {self.PLATFORM}",
            f"content_type: {content_type}",
            f"fetched_at: {datetime.now().isoformat()}",
            "---",
            "",
        ]

        return "\n".join(frontmatter_lines) + content

    def get_output_filename(self, content_id: str, content_type: str, suffix: str) -> str:
        """Generate output filename for the content.

        Args:
            content_id: Answer ID or article ID
            content_type: 'answer' or 'zhuanlan'
            suffix: File suffix (e.g., 'metadata.json', 'content.md')

        Returns:
            Filename string (e.g., 'zhihu_789012_metadata.json')
        """
        return f"zhihu_{content_id}_{suffix}"