Move specific models from Crawlee to SDK

Author: vdusek

src/apify/apify_storage_client/_key_value_store_client.py

@@ -9,12 +9,8 @@
 
 from apify_client import ApifyClientAsync
 from crawlee.storage_clients._base import KeyValueStoreClient
-from crawlee.storage_clients.models import (
-    KeyValueStoreListKeysPage,
-    KeyValueStoreMetadata,
-    KeyValueStoreRecord,
-    KeyValueStoreRecordMetadata,
-)
+from crawlee.storage_clients.models import KeyValueStoreMetadata, KeyValueStoreRecord, KeyValueStoreRecordMetadata
+from ._models import KeyValueStoreListKeysPage
 
 from apify._crypto import create_hmac_signature
 

src/apify/apify_storage_client/_models.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from typing import Annotated
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from crawlee import Request
+from crawlee._utils.docs import docs_group
+
+
+@docs_group('Data structures')
+class ProlongRequestLockResponse(BaseModel):
+    """Response to prolong request lock calls."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    lock_expires_at: Annotated[datetime, Field(alias='lockExpiresAt')]
+
+
+@docs_group('Data structures')
+class RequestQueueHead(BaseModel):
+    """Model for request queue head.
+
+    Represents a collection of requests retrieved from the beginning of a queue,
+    including metadata about the queue's state and lock information for the requests.
+    """
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    limit: Annotated[int | None, Field(alias='limit', default=None)]
+    """The maximum number of requests that were requested from the queue."""
+
+    had_multiple_clients: Annotated[bool, Field(alias='hadMultipleClients', default=False)]
+    """Indicates whether the queue has been accessed by multiple clients (consumers)."""
+
+    queue_modified_at: Annotated[datetime, Field(alias='queueModifiedAt')]
+    """The timestamp when the queue was last modified."""
+
+    lock_time: Annotated[timedelta | None, Field(alias='lockSecs', default=None)]
+    """The duration for which the returned requests are locked and cannot be processed by other clients."""
+
+    queue_has_locked_requests: Annotated[bool | None, Field(alias='queueHasLockedRequests', default=False)]
+    """Indicates whether the queue contains any locked requests."""
+
+    items: Annotated[list[Request], Field(alias='items', default_factory=list[Request])]
+    """The list of request objects retrieved from the beginning of the queue."""
+
+
+class KeyValueStoreKeyInfo(BaseModel):
+    """Model for a key-value store key info."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    key: Annotated[str, Field(alias='key')]
+    size: Annotated[int, Field(alias='size')]
+
+
+class KeyValueStoreListKeysPage(BaseModel):
+    """Model for listing keys in the key-value store."""
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    count: Annotated[int, Field(alias='count')]
+    limit: Annotated[int, Field(alias='limit')]
+    is_truncated: Annotated[bool, Field(alias='isTruncated')]
+    items: Annotated[list[KeyValueStoreKeyInfo], Field(alias='items', default_factory=list)]
+    exclusive_start_key: Annotated[str | None, Field(alias='exclusiveStartKey', default=None)]
+    next_exclusive_start_key: Annotated[str | None, Field(alias='nextExclusiveStartKey', default=None)]
+
+
+class CachedRequest(BaseModel):
+    """Pydantic model for cached request information."""
+
+    id: str
+    """The ID of the request."""
+
+    was_already_handled: bool
+    """Whether the request was already handled."""
+
+    hydrated: Request | None = None
+    """The hydrated request object (the original one)."""
+
+    lock_expires_at: datetime | None = None
+    """The expiration time of the lock on the request."""
+
+    forefront: bool = False
+    """Whether the request was added to the forefront of the queue."""

src/apify/apify_storage_client/_request_queue_client.py

@@ -13,14 +13,9 @@
 from crawlee import Request
 from crawlee._utils.requests import unique_key_to_request_id
 from crawlee.storage_clients._base import RequestQueueClient
-from crawlee.storage_clients.models import (
-    AddRequestsResponse,
-    CachedRequest,
-    ProcessedRequest,
-    ProlongRequestLockResponse,
-    RequestQueueHead,
-    RequestQueueMetadata,
-)
+from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata
+
+from ._models import CachedRequest, ProlongRequestLockResponse, RequestQueueHead
 
 if TYPE_CHECKING:
     from collections.abc import Sequence