More not allowing org id to be None in some cases (mypy only) (#2620)

skyvern/forge/sdk/artifact/manager.py

@@ -87,14 +87,15 @@ class ArtifactManager:
 
     async def create_log_artifact(
         self,
+        *,
         log_entity_type: LogEntityType,
         log_entity_id: str,
         artifact_type: ArtifactType,
+        organization_id: str,
         step_id: str | None = None,
         task_id: str | None = None,
         workflow_run_id: str | None = None,
         workflow_run_block_id: str | None = None,
-        organization_id: str | None = None,
         data: bytes | None = None,
         path: str | None = None,
     ) -> str:

skyvern/forge/sdk/artifact/models.py

@@ -77,7 +77,7 @@ class Artifact(BaseModel):
     observer_thought_id: str | None = None
     ai_suggestion_id: str | None = None
     signed_url: str | None = None
-    organization_id: str | None = None
+    organization_id: str
 
     def __getitem__(self, key: str) -> Any:
         return getattr(self, key)

skyvern/forge/sdk/artifact/storage/s3.py

@@ -77,7 +77,7 @@ class S3Storage(BaseStorage):
         )
         await self.async_client.upload_file(artifact.uri, data, storage_class=sc)
 
-    async def _get_storage_class_for_org(self, organization_id: str | None) -> S3StorageClass:
+    async def _get_storage_class_for_org(self, organization_id: str) -> S3StorageClass:
         return S3StorageClass.STANDARD
 
     async def retrieve_artifact(self, artifact: Artifact) -> bytes | None:

skyvern/forge/sdk/log_artifacts.py

@@ -105,10 +105,11 @@ async def save_workflow_run_block_logs(workflow_run_block_id: str) -> None:
 
 
 async def _save_log_artifacts(
+    *,
     log: list[dict],
     log_entity_type: LogEntityType,
     log_entity_id: str,
-    organization_id: str | None,
+    organization_id: str | None = None,
     step_id: str | None = None,
     task_id: str | None = None,
     workflow_run_id: str | None = None,
@@ -117,7 +118,13 @@ async def _save_log_artifacts(
     try:
         if not settings.ENABLE_LOG_ARTIFACTS:
             return
-
+        if not organization_id:
+            LOG.error(
+                "Organization ID is required to save log artifacts",
+                log_entity_type=log_entity_type,
+                log_entity_id=log_entity_id,
+            )
+            return
         log_json = json.dumps(log, cls=SkyvernJSONLogEncoder, indent=2)
 
         log_artifact = await app.DATABASE.get_artifact_by_entity_id(

skyvern/forge/sdk/schemas/ai_suggestions.py

@@ -12,7 +12,7 @@ class AISuggestion(AISuggestionBase):
     model_config = ConfigDict(from_attributes=True)
     ai_suggestion_type: str
     ai_suggestion_id: str
-    organization_id: str | None = None
+    organization_id: str
 
     created_at: datetime
     modified_at: datetime

skyvern/forge/sdk/schemas/task_v2.py

@@ -30,7 +30,7 @@ class TaskV2(BaseModel):
 
     observer_cruise_id: str = Field(alias="task_id")
     status: TaskV2Status
-    organization_id: str | None = None
+    organization_id: str
     workflow_run_id: str | None = None
     workflow_id: str | None = None
     workflow_permanent_id: str | None = None
@@ -99,7 +99,7 @@ class Thought(BaseModel):
 
     observer_thought_id: str = Field(alias="thought_id")
     observer_cruise_id: str = Field(alias="task_id")
-    organization_id: str | None = None
+    organization_id: str
     workflow_run_id: str | None = None
     workflow_run_block_id: str | None = None
     workflow_id: str | None = None

skyvern/forge/sdk/schemas/workflow_runs.py

@@ -15,7 +15,7 @@ class WorkflowRunBlock(BaseModel):
     workflow_run_block_id: str
     block_workflow_run_id: str | None = None
     workflow_run_id: str
-    organization_id: str | None = None
+    organization_id: str
     description: str | None = None
     parent_workflow_run_block_id: str | None = None
     block_type: BlockType