081f3f0d89
Implement Sprint 3: Chat Functionality
- Add ChatService with send_message and get_history methods
- Add POST /api/v1/notebooks/{id}/chat - Send message
- Add GET /api/v1/notebooks/{id}/chat/history - Get chat history
- Add ChatRequest model (message, include_references)
- Add ChatResponse model (message, sources[], timestamp)
- Add ChatMessage model (id, role, content, timestamp, sources)
- Add SourceReference model (source_id, title, snippet)
- Integrate chat router with main app
Features:
- Send messages to notebook chat
- Get AI responses with source references
- Retrieve chat history
- Support for citations in responses
Tests:
- 14 unit tests for ChatService
- 11 integration tests for chat API
- 25/25 tests passing
Related: Sprint 3 - Chat Functionality
547 lines
14 KiB
Python
547 lines
14 KiB
Python
"""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)",
|
|
)
|