Fix ETLJobResponseType

genai_stack/genai_server/models/etl_models.py

@@ -1,24 +1,20 @@
 import enum
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 from pydantic import BaseModel
 
+from genai_stack.genai_server.schemas import ETLJobStatus
 
-class StatusEnum(enum.Enum):
-    Pending = "Pending"
-    Processing = "Processing"
-    Completed = "Completed"
 
-
-class BaseETLJobType(BaseModel):
-    uuid: str
+class BaseETLJobModel(BaseModel):
+    id: int
     session_id: int
-    status: StatusEnum
-    metadata: dict
+    status: ETLJobStatus
+    metadata: Optional[dict]
 
 
 class ETLJobRequestType(BaseModel):
     __root__: Dict[str, Any]
 
 
-class ETLJobResponseType(BaseModel):
+class ETLJobResponseType(BaseETLJobModel):
     pass

genai_stack/genai_server/routers/etl_routes.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
 from typing import Any
 
 from genai_stack.constant import API, ETL

genai_stack/genai_server/services/etl_service.py

@@ -15,9 +15,7 @@ def submit_job(self, data: Any, stack_session_id: Optional[int] = None) -> ETLJo
             stack_session = get_stack_session(session, stack_session_id=stack_session_id)
 
             etl_job = ETLJob(stack_session=stack_session.id)
-            session.add(
-                etl_job,
-            )
+            session.add(etl_job)
             session.commit()
 
             data = ETLUtil(data).save_request(etl_job.id)
@@ -28,9 +26,10 @@ def submit_job(self, data: Any, stack_session_id: Optional[int] = None) -> ETLJo
             etl_job.data = data
             session.commit()
 
-        return ETLJobResponseType(
-            uuid=etl_job.id,
-            session_id=etl_job.stack_session,
-            status=etl_job.status,
-            metadata=etl_job.meta_data,
-        )
+            response = ETLJobResponseType(
+                id=etl_job.id,
+                session_id=etl_job.stack_session,
+                status=etl_job.status.value,
+                metadata=etl_job.meta_data,
+            )
+            return response