"""Pydantic models for request and response objects in the KOI-net API."""
from typing import Annotated, Literal
from pydantic import BaseModel, Field
from rid_lib import RID, RIDType
from rid_lib.ext import Bundle, Manifest
from ..event import Event
from ..errors import ErrorType
# REQUEST MODELS
[docs]
class PollEvents(BaseModel):
type: Literal["poll_events"] = Field("poll_events")
limit: int = 0
[docs]
class FetchRids(BaseModel):
type: Literal["fetch_rids"] = Field("fetch_rids")
rid_types: list[RIDType] = []
[docs]
class FetchManifests(BaseModel):
type: Literal["fetch_manifests"] = Field("fetch_manifests")
rid_types: list[RIDType] = []
rids: list[RID] = []
[docs]
class FetchBundles(BaseModel):
type: Literal["fetch_bundles"] = Field("fetch_bundles")
rids: list[RID]
# RESPONSE/PAYLOAD MODELS
[docs]
class RidsPayload(BaseModel):
type: Literal["rids_payload"] = Field("rids_payload")
rids: list[RID]
[docs]
class ManifestsPayload(BaseModel):
type: Literal["manifests_payload"] = Field("manifests_payload")
manifests: list[Manifest]
not_found: list[RID] = []
[docs]
class BundlesPayload(BaseModel):
type: Literal["bundles_payload"] = Field("bundles_payload")
bundles: list[Bundle]
not_found: list[RID] = []
deferred: list[RID] = []
[docs]
class EventsPayload(BaseModel):
type: Literal["events_payload"] = Field("events_payload")
events: list[Event]
# ERROR MODELS
[docs]
class ErrorResponse(BaseModel):
type: Literal["error_response"] = Field("error_response")
error: ErrorType
# TYPES
type RequestModels = EventsPayload | PollEvents | FetchRids | FetchManifests | FetchBundles
type ResponseModels = RidsPayload | ManifestsPayload | BundlesPayload | EventsPayload | ErrorResponse
type ApiModels = Annotated[
RequestModels | ResponseModels,
Field(discriminator="type")
]