Fix JobStoreDocument validator

src/jobflow/core/job.py

@@ -10,7 +10,6 @@
 from monty.json import MSONable, jsanitize
 
 from jobflow.core.reference import OnMissing, OutputReference
-from jobflow.schemas.job_output_schema import JobStoreDocument
 from jobflow.utils.uuid import suuid
 
 if typing.TYPE_CHECKING:
@@ -560,6 +559,7 @@ def run(self, store: jobflow.JobStore) -> Response:
 
         from jobflow import CURRENT_JOB
         from jobflow.core.flow import get_flow
+        from jobflow.core.schemas import JobStoreDocument
 
         index_str = f", {self.index}" if self.index != 1 else ""
         logger.info(f"Starting job - {self.name} ({self.uuid}{index_str})")

src/jobflow/core/schemas.py

@@ -0,0 +1,34 @@
+"""A Pydantic model for Jobstore document."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class JobStoreDocument(BaseModel):
+    """A Pydantic model for Jobstore document."""
+
+    uuid: str = Field(
+        None, description="An unique identifier for the job. Generated automatically."
+    )
+    index: int = Field(
+        None,
+        description="The index of the job (number of times the job has been replaced).",
+    )
+    output: Any = Field(
+        None,
+        description="This is a reference to the future job output.",
+    )
+    completed_at: str = Field(None, description="The time the job was completed.")
+    metadata: dict = Field(
+        None,
+        description="Metadata information supplied by the user.",
+    )
+    hosts: list[str] = Field(
+        None,
+        description="The list of UUIDs of the hosts containing the job.",
+    )
+    name: str = Field(
+        None,
+        description="The name of the job.",
+    )

src/jobflow/core/store.py

@@ -18,7 +18,7 @@
 
     from maggma.core import Sort
 
-    from jobflow.schemas.job_output_schema import JobStoreDocument
+    from jobflow.core.schemas import JobStoreDocument
 
     obj_type = Union[str, Enum, type[MSONable], list[Union[Enum, str, type[MSONable]]]]
     save_type = Optional[dict[str, obj_type]]

tests/core/test_job_output_schema.py → tests/core/test_schemas.py

@@ -5,7 +5,7 @@
 
 @pytest.fixture()
 def sample_data():
-    from jobflow.schemas.job_output_schema import JobStoreDocument
+    from jobflow.core.schemas import JobStoreDocument
 
     return JobStoreDocument(
         uuid="abc123",
@@ -33,7 +33,7 @@ def test_job_store_document_model(sample_data):
 
 def test_job_store_update(memory_jobstore, sample_data):
     # Storing document as a JobStoreDocument
-    from jobflow.schemas.job_output_schema import JobStoreDocument
+    from jobflow.core.schemas import JobStoreDocument
 
     d = {
         "index": 1,