"""Response models for NotebookLM Agent API. This module contains Pydantic models for API response serialization. """ from datetime import datetime from typing import Any, Generic, TypeVar from uuid import UUID from pydantic import BaseModel, ConfigDict, Field T = TypeVar("T") class ErrorDetail(BaseModel): """Error detail information. Attributes: code: Error code for programmatic handling. message: Human-readable error message. details: Additional error details (field errors, etc.). """ model_config = ConfigDict( json_schema_extra={ "example": { "code": "VALIDATION_ERROR", "message": "Invalid input data", "details": [{"field": "title", "error": "Title is too short"}], } } ) code: str = Field( ..., description="Error code for programmatic handling", examples=["VALIDATION_ERROR", "NOT_FOUND", "AUTH_ERROR"], ) message: str = Field( ..., description="Human-readable error message", examples=["Invalid input data", "Notebook not found"], ) details: list[dict[str, Any]] | None = Field( None, description="Additional error details", examples=[[{"field": "title", "error": "Title is too short"}]], ) class ResponseMeta(BaseModel): """Metadata for API responses. Attributes: timestamp: ISO 8601 timestamp of the response. request_id: Unique identifier for the request. """ model_config = ConfigDict( json_schema_extra={ "example": { "timestamp": "2026-04-06T10:30:00Z", "request_id": "550e8400-e29b-41d4-a716-446655440000", } } ) timestamp: datetime = Field( ..., description="ISO 8601 timestamp of the response", examples=["2026-04-06T10:30:00Z"], ) request_id: UUID = Field( ..., description="Unique identifier for the request", examples=["550e8400-e29b-41d4-a716-446655440000"], ) class ApiResponse(BaseModel, Generic[T]): """Standard API response wrapper. This wrapper is used for all API responses to ensure consistency. Attributes: success: Whether the request was successful. data: Response data (None if error). error: Error details (None if success). meta: Response metadata. """ model_config = ConfigDict( json_schema_extra={ "example": { "success": True, "data": {"id": "550e8400-e29b-41d4-a716-446655440000"}, "error": None, "meta": { "timestamp": "2026-04-06T10:30:00Z", "request_id": "550e8400-e29b-41d4-a716-446655440000", }, } } ) success: bool = Field( ..., description="Whether the request was successful", examples=[True, False], ) data: T | None = Field( None, description="Response data (None if error)", ) error: ErrorDetail | None = Field( None, description="Error details (None if success)", ) meta: ResponseMeta = Field( ..., description="Response metadata", ) class Notebook(BaseModel): """Notebook response model. Attributes: id: Unique identifier (UUID). title: Notebook title. description: Optional notebook description. created_at: Creation timestamp. updated_at: Last update timestamp. """ model_config = ConfigDict( json_schema_extra={ "example": { "id": "550e8400-e29b-41d4-a716-446655440000", "title": "My Research Notebook", "description": "A collection of AI research papers", "created_at": "2026-04-06T10:00:00Z", "updated_at": "2026-04-06T10:30:00Z", } } ) id: UUID = Field( ..., description="Unique identifier (UUID)", examples=["550e8400-e29b-41d4-a716-446655440000"], ) title: str = Field( ..., description="Notebook title", examples=["My Research Notebook"], ) description: str | None = Field( None, description="Optional notebook description", examples=["A collection of AI research papers"], ) created_at: datetime = Field( ..., description="Creation timestamp (ISO 8601)", examples=["2026-04-06T10:00:00Z"], ) updated_at: datetime = Field( ..., description="Last update timestamp (ISO 8601)", examples=["2026-04-06T10:30:00Z"], ) class NotebookDetail(Notebook): """Detailed notebook response with counts. Extends Notebook with additional statistics. Attributes: sources_count: Number of sources in the notebook. artifacts_count: Number of artifacts in the notebook. """ model_config = ConfigDict( json_schema_extra={ "example": { "id": "550e8400-e29b-41d4-a716-446655440000", "title": "My Research Notebook", "description": "A collection of AI research papers", "created_at": "2026-04-06T10:00:00Z", "updated_at": "2026-04-06T10:30:00Z", "sources_count": 5, "artifacts_count": 2, } } ) sources_count: int = Field( 0, ge=0, description="Number of sources in the notebook", examples=[5, 10, 0], ) artifacts_count: int = Field( 0, ge=0, description="Number of artifacts in the notebook", examples=[2, 5, 0], ) class PaginationMeta(BaseModel): """Pagination metadata. Attributes: total: Total number of items available. limit: Maximum items per page. offset: Current offset. has_more: Whether more items are available. """ model_config = ConfigDict( json_schema_extra={ "example": { "total": 100, "limit": 20, "offset": 0, "has_more": True, } } ) total: int = Field( ..., ge=0, description="Total number of items available", examples=[100, 50, 0], ) limit: int = Field( ..., ge=1, description="Maximum items per page", examples=[20, 50], ) offset: int = Field( ..., ge=0, description="Current offset", examples=[0, 20, 40], ) has_more: bool = Field( ..., description="Whether more items are available", examples=[True, False], ) class PaginatedNotebooks(BaseModel): """Paginated list of notebooks. Attributes: items: List of notebooks. pagination: Pagination metadata. """ model_config = ConfigDict( json_schema_extra={ "example": { "items": [ { "id": "550e8400-e29b-41d4-a716-446655440000", "title": "Notebook 1", "created_at": "2026-04-06T10:00:00Z", "updated_at": "2026-04-06T10:30:00Z", } ], "pagination": { "total": 100, "limit": 20, "offset": 0, "has_more": True, }, } } ) items: list[Notebook] = Field( ..., description="List of notebooks", ) pagination: PaginationMeta = Field( ..., description="Pagination metadata", ) class Source(BaseModel): """Source response model. Attributes: id: Unique identifier (UUID). notebook_id: Parent notebook ID. type: Source type (url, file, youtube, drive). title: Source title. url: Source URL (if applicable). status: Processing status. created_at: Creation timestamp. """ model_config = ConfigDict( json_schema_extra={ "example": { "id": "550e8400-e29b-41d4-a716-446655440001", "notebook_id": "550e8400-e29b-41d4-a716-446655440000", "type": "url", "title": "Example Article", "url": "https://example.com/article", "status": "ready", "created_at": "2026-04-06T10:00:00Z", } } ) id: UUID = Field( ..., description="Unique identifier (UUID)", examples=["550e8400-e29b-41d4-a716-446655440001"], ) notebook_id: UUID = Field( ..., description="Parent notebook ID", examples=["550e8400-e29b-41d4-a716-446655440000"], ) type: str = Field( ..., description="Source type", examples=["url", "file", "youtube", "drive"], ) title: str = Field( ..., description="Source title", examples=["Example Article"], ) url: str | None = Field( None, description="Source URL (if applicable)", examples=["https://example.com/article"], ) status: str = Field( ..., description="Processing status", examples=["processing", "ready", "error"], ) created_at: datetime = Field( ..., description="Creation timestamp", examples=["2026-04-06T10:00:00Z"], ) class PaginatedSources(BaseModel): """Paginated list of sources. Attributes: items: List of sources. pagination: Pagination metadata. """ items: list[Source] = Field( ..., description="List of sources", ) pagination: PaginationMeta = Field( ..., description="Pagination metadata", ) class HealthStatus(BaseModel): """Health check response. Attributes: status: Health status (healthy, degraded, unhealthy). timestamp: Check timestamp. version: API version. service: Service name. """ model_config = ConfigDict( json_schema_extra={ "example": { "status": "healthy", "timestamp": "2026-04-06T10:30:00Z", "version": "0.1.0", "service": "notebooklm-agent-api", } } ) status: str = Field( ..., description="Health status", examples=["healthy", "degraded", "unhealthy"], ) timestamp: datetime = Field( ..., description="Check timestamp", examples=["2026-04-06T10:30:00Z"], ) version: str = Field( ..., description="API version", examples=["0.1.0"], ) service: str = Field( ..., description="Service name", examples=["notebooklm-agent-api"], ) class SourceReference(BaseModel): """Source reference in chat response. Attributes: source_id: The source ID. title: The source title. snippet: Relevant text snippet from the source. """ model_config = ConfigDict( json_schema_extra={ "example": { "source_id": "550e8400-e29b-41d4-a716-446655440001", "title": "Example Article", "snippet": "Key information from the source...", } } ) source_id: str = Field( ..., description="The source ID", examples=["550e8400-e29b-41d4-a716-446655440001"], ) title: str = Field( ..., description="The source title", examples=["Example Article"], ) snippet: str | None = Field( None, description="Relevant text snippet from the source", examples=["Key information from the source..."], ) class ChatResponse(BaseModel): """Chat response model. Attributes: message: The assistant's response message. sources: List of source references used in the response. timestamp: Response timestamp. """ model_config = ConfigDict( json_schema_extra={ "example": { "message": "Based on the sources, here are the key points...", "sources": [ { "source_id": "550e8400-e29b-41d4-a716-446655440001", "title": "Example Article", "snippet": "Key information...", } ], "timestamp": "2026-04-06T10:30:00Z", } } ) message: str = Field( ..., description="The assistant's response message", examples=["Based on the sources, here are the key points..."], ) sources: list[SourceReference] = Field( default_factory=list, description="List of source references used in the response", ) timestamp: datetime = Field( ..., description="Response timestamp", examples=["2026-04-06T10:30:00Z"], ) class ChatMessage(BaseModel): """Chat message model. Attributes: id: Unique message identifier. role: Message role (user or assistant). content: Message content. timestamp: Message timestamp. sources: Source references (for assistant messages). """ model_config = ConfigDict( json_schema_extra={ "example": { "id": "550e8400-e29b-41d4-a716-446655440002", "role": "assistant", "content": "Based on the sources...", "timestamp": "2026-04-06T10:30:00Z", "sources": [], } } ) id: str = Field( ..., description="Unique message identifier", examples=["550e8400-e29b-41d4-a716-446655440002"], ) role: str = Field( ..., description="Message role (user or assistant)", examples=["user", "assistant"], ) content: str = Field( ..., description="Message content", examples=["What are the key points?"], ) timestamp: datetime = Field( ..., description="Message timestamp", examples=["2026-04-06T10:30:00Z"], ) sources: list[SourceReference] | None = Field( None, description="Source references (for assistant messages)", )