kb-skills/kb_workflow/crawler/models.py

"""Zhihu Crawler Data Models

Pydantic models for structured data from Zhihu API responses.
"""

from typing import Optional, Union, List
from pydantic import BaseModel, Field


class AuthorInfo(BaseModel):
    """Author information from Zhihu content.

    Attributes:
        name: Author display name
        url_token: Author URL token (e.g., 'zhang-san-12')
        headline: Author bio/headline
    """
    name: str = Field(..., description="Author display name")
    url_token: Optional[str] = Field(None, description="Author URL token")
    headline: Optional[str] = Field(None, description="Author bio/headline")


class ContentMetadata(BaseModel):
    """Metadata for Zhihu content (answers, articles, etc).

    Attributes:
        author: Author information
        voteup_count: Number of upvotes/likes
        comment_count: Number of comments
        created_time: Unix timestamp of creation
        updated_time: Unix timestamp of last update
    """
    author: Optional[AuthorInfo] = Field(None, description="Author info")
    voteup_count: int = Field(default=0, description="Upvote count")
    comment_count: int = Field(default=0, description="Comment count")
    created_time: Optional[int] = Field(None, description="Creation timestamp")
    updated_time: Optional[int] = Field(None, description="Update timestamp")


class AnswerData(BaseModel):
    """Data model for Zhihu answer.

    Attributes:
        id: Answer ID
        question_id: Question ID this answer belongs to
        question_title: Title of the question
        content: Answer content in HTML format
        excerpt: Plain text excerpt/summary
        metadata: Additional metadata (author, votes, etc.)
    """
    id: int = Field(..., description="Answer ID")
    question_id: int = Field(..., description="Question ID")
    question_title: str = Field(..., description="Question title")
    content: str = Field(default="", description="Answer content (HTML)")
    excerpt: str = Field(default="", description="Plain text excerpt")
    metadata: ContentMetadata = Field(
        default_factory=ContentMetadata,
        description="Content metadata"
    )


class ArticleData(BaseModel):
    """Data model for Zhihu zhuanlan article.

    Attributes:
        id: Article ID
        title: Article title
        content: Article content in HTML format
        excerpt: Plain text excerpt/summary
        metadata: Additional metadata (author, votes, etc.)
    """
    id: int = Field(..., description="Article ID")
    title: str = Field(..., description="Article title")
    content: str = Field(default="", description="Article content (HTML)")
    excerpt: str = Field(default="", description="Plain text excerpt")
    metadata: ContentMetadata = Field(
        default_factory=ContentMetadata,
        description="Content metadata"
    )


class CommentAuthor(BaseModel):
    """Comment author information.

    Attributes:
        name: Author display name
        url_token: Author URL token (e.g., 'zhang-san-12')
    """
    name: str = Field(..., description="Author display name")
    url_token: Optional[str] = Field(None, description="Author URL token")


class CommentData(BaseModel):
    """Data model for Zhihu comment.

    Attributes:
        id: Comment ID
        content: Comment content in plain text
        author: Comment author information
        created_time: Unix timestamp of creation
        like_count: Number of likes/upvotes
        child_count: Number of child/reply comments
        is_author: Whether the commenter is the content author
    """
    id: int = Field(..., description="Comment ID")
    content: str = Field(default="", description="Comment content")
    author: Optional[CommentAuthor] = Field(None, description="Comment author")
    created_time: Optional[int] = Field(None, description="Creation timestamp")
    like_count: int = Field(default=0, description="Like/upvote count")
    child_count: int = Field(default=0, description="Number of child comments")
    is_author: bool = Field(default=False, description="Whether commenter is content author")


class CommentsResult(BaseModel):
    """Result of comment crawling operation.

    Attributes:
        content_type: Type of content ("answers" or "articles")
        content_id: ID of the content
        comments: List of comments
        total_count: Total number of comments (may exceed returned count)
    """
    content_type: str = Field(..., description="Content type: answers or articles")
    content_id: int = Field(..., description="Content ID")
    comments: List[CommentData] = Field(
        default_factory=list,
        description="List of comments"
    )
    total_count: int = Field(default=0, description="Total comment count")


class CrawlResult(BaseModel):
    """Result of a crawl operation.

    Attributes:
        success: Whether the crawl was successful
        content_type: Type of content ("answer" or "article")
        data: Extracted content data if successful
        comments: Optional list of comments (when include_comments=True)
        error: Error message if failed
    """
    success: bool = Field(..., description="Whether crawl succeeded")
    content_type: str = Field(default="", description="Content type: answer or article")
    data: Optional[Union[AnswerData, ArticleData]] = Field(
        None,
        description="Extracted content data"
    )
    comments: Optional[List[CommentData]] = Field(
        None,
        description="Comments if include_comments was True"
    )
    error: Optional[str] = Field(None, description="Error message if failed")