refactor: make BaseProvider a BaseModel with discriminated union

lib/crewai/src/crewai/state/checkpoint_config.py

@@ -2,12 +2,12 @@
 
 from __future__ import annotations
 
-from typing import Any, Literal
+from typing import Annotated, Any, Literal
 
 from pydantic import BaseModel, Field, model_validator
 
-from crewai.state.provider.core import BaseProvider
 from crewai.state.provider.json_provider import JsonProvider
+from crewai.state.provider.sqlite_provider import SqliteProvider
 
 
 CheckpointEventType = Literal[
@@ -189,7 +189,10 @@ class CheckpointConfig(BaseModel):
         description="Event types that trigger a checkpoint write. "
         'Use ["*"] to checkpoint on every event.',
     )
-    provider: BaseProvider = Field(
+    provider: Annotated[
+        JsonProvider | SqliteProvider,
+        Field(discriminator="provider_type"),
+    ] = Field(
         default_factory=JsonProvider,
         description="Storage backend. Defaults to JsonProvider.",
     )

lib/crewai/src/crewai/state/provider/core.py

@@ -1,39 +1,22 @@
-"""Base protocol for state providers."""
+"""Base class for state providers."""
 
 from __future__ import annotations
 
-from typing import Any, Protocol, runtime_checkable
+from abc import ABC, abstractmethod
 
-from pydantic import GetCoreSchemaHandler
-from pydantic_core import CoreSchema, core_schema
+from pydantic import BaseModel
 
 
-@runtime_checkable
-class BaseProvider(Protocol):
-    """Interface for persisting and restoring runtime state checkpoints.
+class BaseProvider(BaseModel, ABC):
+    """Base class for persisting and restoring runtime state checkpoints.
 
     Implementations handle the storage backend — filesystem, cloud, database,
     etc. — while ``RuntimeState`` handles serialization.
     """
 
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: GetCoreSchemaHandler
-    ) -> CoreSchema:
-        """Allow Pydantic to validate any ``BaseProvider`` instance."""
-
-        def _validate(v: Any) -> BaseProvider:
-            if isinstance(v, BaseProvider):
-                return v
-            raise TypeError(f"Expected a BaseProvider instance, got {type(v)}")
-
-        return core_schema.no_info_plain_validator_function(
-            _validate,
-            serialization=core_schema.plain_serializer_function_ser_schema(
-                lambda v: type(v).__name__, info_arg=False
-            ),
-        )
+    provider_type: str = "base"
 
+    @abstractmethod
     def checkpoint(self, data: str, location: str) -> str:
         """Persist a snapshot synchronously.
 
@@ -46,6 +29,7 @@ def checkpoint(self, data: str, location: str) -> str:
         """
         ...
 
+    @abstractmethod
     async def acheckpoint(self, data: str, location: str) -> str:
         """Persist a snapshot asynchronously.
 
@@ -58,6 +42,7 @@ async def acheckpoint(self, data: str, location: str) -> str:
         """
         ...
 
+    @abstractmethod
     def prune(self, location: str, max_keep: int) -> None:
         """Remove old checkpoints, keeping at most *max_keep*.
 
@@ -67,6 +52,7 @@ def prune(self, location: str, max_keep: int) -> None:
         """
         ...
 
+    @abstractmethod
     def from_checkpoint(self, location: str) -> str:
         """Read a snapshot synchronously.
 
@@ -78,6 +64,7 @@ def from_checkpoint(self, location: str) -> str:
         """
         ...
 
+    @abstractmethod
     async def afrom_checkpoint(self, location: str) -> str:
         """Read a snapshot asynchronously.
 

lib/crewai/src/crewai/state/provider/json_provider.py

@@ -7,6 +7,7 @@
 import logging
 import os
 from pathlib import Path
+from typing import Literal
 import uuid
 
 import aiofiles
@@ -21,6 +22,8 @@
 class JsonProvider(BaseProvider):
     """Persists runtime state checkpoints as JSON files on the local filesystem."""
 
+    provider_type: Literal["json"] = "json"
+
     def checkpoint(self, data: str, location: str) -> str:
         """Write a JSON checkpoint file.
 

lib/crewai/src/crewai/state/provider/sqlite_provider.py

@@ -5,6 +5,7 @@
 from datetime import datetime, timezone
 from pathlib import Path
 import sqlite3
+from typing import Literal
 import uuid
 
 import aiosqlite
@@ -47,6 +48,8 @@ class SqliteProvider(BaseProvider):
     used as the database file path.
     """
 
+    provider_type: Literal["sqlite"] = "sqlite"
+
     def checkpoint(self, data: str, location: str) -> str:
         """Write a checkpoint to the SQLite database.