Merge pull request #12 from SocAIty/dev

Dev

Author: w4hns1nn

apipod/__init__.py

@@ -1,6 +1,7 @@
 from apipod.api import APIPod
+from apipod.engine.jobs.base_job import BaseJob, LocalJob
 from apipod.engine.jobs.job_progress import JobProgress
-from apipod.engine.jobs.job_result import FileModel, JobResult
+from apipod.engine.jobs.job_result import FileModel, JobLinks, JobMetrics, JobResult
 from media_toolkit import MediaFile, ImageFile, AudioFile, VideoFile, MediaList, MediaDict
 from apipod.common import constants
 
@@ -15,4 +16,20 @@
 except Exception:
     __version__ = "0.0.0"
 
-__all__ = ["APIPod", "JobProgress", "FileModel", "JobResult", "MediaFile", "ImageFile", "AudioFile", "VideoFile", "MediaList", "MediaDict", "constants"]
+__all__ = [
+    "APIPod",
+    "BaseJob",
+    "LocalJob",
+    "JobProgress",
+    "FileModel",
+    "JobLinks",
+    "JobMetrics",
+    "JobResult",
+    "MediaFile",
+    "ImageFile",
+    "AudioFile",
+    "VideoFile",
+    "MediaList",
+    "MediaDict",
+    "constants",
+]

apipod/api.py

@@ -23,6 +23,7 @@ def APIPod(
     |------------- |----------- |---------- |--------------------------- |
     | socaity      | dedicated  | auto      | FastAPI                    |
     | socaity      | dedicated  | localhost | FastAPI + job queue (test) |
+    | socaity      | dedicated  | socaity   | FastAPI + redis (prod)     |
     | socaity      | dedicated  | runpod    | Celery (planned)           |
     | socaity      | dedicated  | scaleway  | Celery (planned)           |
     | socaity      | dedicated  | azure     | Celery (planned)           |
@@ -40,15 +41,20 @@ def APIPod(
     Args:
         orchestrator: "socaity" or "local" (default from env / local).
         compute: "dedicated" or "serverless" (default from env / dedicated).
-        provider: "auto", "localhost", "runpod", "scaleway", "azure" (default from env / localhost).
+        provider: "auto", "localhost", "socaity", "runpod", "scaleway", "azure" (default from env / localhost).
     """
     orchestrator = _resolve_enum(orchestrator, constants.ORCHESTRATOR, APIPOD_ORCHESTRATOR, constants.ORCHESTRATOR.LOCAL)
     compute = _resolve_enum(compute, constants.COMPUTE, APIPOD_COMPUTE, constants.COMPUTE.DEDICATED)
     provider = _resolve_enum(provider, constants.PROVIDER, APIPOD_PROVIDER, constants.PROVIDER.LOCALHOST)
 
     backend_class, use_job_queue = _resolve_backend(orchestrator, compute, provider)
 
-    job_queue = _create_job_queue() if use_job_queue else None
+    custom_job_queue = kwargs.pop("job_queue", None)
+    if custom_job_queue:
+        use_job_queue = True
+        job_queue = custom_job_queue
+    else:
+        job_queue = _create_job_queue() if use_job_queue else None
 
     if backend_class == SocaityFastAPIRouter:
         return backend_class(job_queue=job_queue, *args, **kwargs)
@@ -101,6 +107,9 @@ def _raise_if_unsupported(compute: constants.COMPUTE, provider: constants.PROVID
 
 def _resolve_socaity(compute: constants.COMPUTE, provider: constants.PROVIDER) -> tuple:
     if compute == constants.COMPUTE.DEDICATED:
+        if provider == constants.PROVIDER.SOCAITY:
+            return SocaityFastAPIRouter, True
+
         if provider in (constants.PROVIDER.RUNPOD, constants.PROVIDER.SCALEWAY, constants.PROVIDER.AZURE):
             raise NotImplementedError(
                 f"Celery backend for socaity + dedicated + {provider.value} is planned but not yet available."
@@ -135,4 +144,5 @@ def _resolve_local(compute: constants.COMPUTE, provider: constants.PROVIDER) ->
 
 def _create_job_queue() -> JobQueueInterface:
     from apipod.engine.queue.job_queue import JobQueue
+
     return JobQueue()

apipod/common/constants.py

@@ -14,6 +14,7 @@ class COMPUTE(Enum):
 class PROVIDER(Enum):
     AUTO = "auto"
     LOCALHOST = "localhost"
+    SOCAITY = "socaity"
     RUNPOD = "runpod"
     SCALEWAY = "scaleway"
     AZURE = "azure"

apipod/engine/backend/fastapi/file_handling_mixin.py

@@ -5,7 +5,7 @@
 from apipod.engine.signatures.upload import is_param_media_toolkit_file
 from apipod.engine.jobs.job_result import FileModel, ImageFileModel, AudioFileModel, VideoFileModel
 from apipod.engine.signatures.policies import FastAPISignaturePolicies
-from apipod.engine.files.base_mixin import _BaseFileHandlingMixin
+from apipod.engine.files.base_file_mixin import _BaseFileHandlingMixin
 from apipod.engine.utils import replace_func_signature
 from media_toolkit import MediaList, MediaDict, ImageFile, AudioFile, VideoFile, MediaFile
 import functools

apipod/engine/backend/fastapi/llm_mixin.py

@@ -1,7 +1,6 @@
 from fastapi.responses import StreamingResponse
 
-from apipod.engine.llm.base_mixin import _BaseLLMMixin
-from apipod.common.settings import SERVER_DOMAIN
+from apipod.engine.llm.base_llm_mixin import _BaseLLMMixin
 
 
 class _FastApiLlmMixin(_BaseLLMMixin):
@@ -24,7 +23,6 @@ async def handle_llm_request(self, func, openai_req, should_use_queue, res_model
                 job_params={"payload": openai_req.dict()}
             )
             ret_job = JobResultFactory.from_base_job(job)
-            ret_job.refresh_job_url = f"{SERVER_DOMAIN}/status?job_id={ret_job.id}"
             return ret_job
 
         raw_res = await self._execute_func(func, payload=openai_req, **kwargs)

apipod/engine/backend/fastapi/router.py

@@ -4,10 +4,11 @@
 import logging
 from contextlib import asynccontextmanager
 from typing import Union, Callable, get_type_hints, Generator, AsyncGenerator, Iterator, AsyncIterator
-from fastapi import APIRouter, FastAPI, Response
-from fastapi.responses import JSONResponse
+from fastapi import APIRouter, FastAPI, Request, Response, status
+from fastapi.exceptions import HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
 
-from apipod.common.settings import APIPOD_PORT, APIPOD_HOST, SERVER_DOMAIN
+from apipod.common.settings import APIPOD_PORT, APIPOD_HOST
 from apipod.common.constants import SERVER_HEALTH
 from apipod.engine.jobs.job_result import JobResultFactory, JobResult
 from apipod.engine.endpoint_config import FastApiEndpointConfigurator, EndpointExecutionPlan
@@ -50,10 +51,14 @@ def __init__(
             job_queue: Optional custom JobQueue implementation
             lifespan: Optional async context manager for custom startup/shutdown logic
             args: Additional arguments
-            kwargs: Additional keyword arguments
+            kwargs: May include ``stream_store`` (SSE backend for GET /stream/{job_id}) and
+                ``gateway_stream_url_prefix`` for absolute stream URLs in JobResult, plus
+                additional keyword arguments for parent classes.
         """
         # Extract user-provided lifespan (explicit param or kwarg) before parent init
         user_lifespan = lifespan or kwargs.pop('lifespan', None)
+        stream_store = kwargs.pop("stream_store", None)
+        gateway_stream_url_prefix = kwargs.pop("gateway_stream_url_prefix", "")
 
         # Initialize parent classes
         api_router_params = inspect.signature(APIRouter.__init__).parameters
@@ -92,6 +97,8 @@ def __init__(
 
         self.app: FastAPI = app
         self.prefix = prefix
+        self.stream_store = stream_store
+        self.gateway_stream_url_prefix = gateway_stream_url_prefix
         self.add_standard_routes()
 
         self._endpoint_configurator = FastApiEndpointConfigurator(self)
@@ -173,7 +180,11 @@ def _stop_background_worker(self):
     def add_standard_routes(self):
         """Add standard API routes for status and health checks."""
         if self.job_queue is not None:
-            self.api_route(path="/status", methods=["POST"])(self.get_job)
+            self.api_route(path="/status/{job_id}", methods=["GET"], response_model_exclude_none=True)(self.get_job)
+            self.api_route(path="/status", methods=["POST"], response_model_exclude_none=True)(self.get_job)
+            self.api_route(path="/cancel/{job_id}", methods=["POST"])(self.post_cancel_job)
+            if self.stream_store is not None:
+                self.api_route(path="/stream/{job_id}", methods=["GET"])(self.stream_job_sse)
         self.api_route(path="/health", methods=["GET"])(self.get_health)
 
     def get_health(self) -> Response:
@@ -216,19 +227,76 @@ def get_job(self, job_id: str, return_format: str = 'json') -> JobResult:
         if self.job_queue is None:
             return JobResultFactory.job_not_found(job_id)
 
-        base_job = self.job_queue.get_job(job_id)
-        if base_job is None:
+        ret_job = self.job_queue.get_job_result(job_id)
+        if ret_job is None:
             return JobResultFactory.job_not_found(job_id)
 
-        ret_job = JobResultFactory.from_base_job(base_job)
-        ret_job.refresh_job_url = f"{SERVER_DOMAIN}/status?job_id={ret_job.id}"
-        ret_job.cancel_job_url = f"{SERVER_DOMAIN}/cancel?job_id={ret_job.id}"
-
         if return_format != 'json':
             ret_job = JobResultFactory.gzip_job_result(ret_job)
 
         return ret_job
 
+    def post_cancel_job(self, job_id: str) -> dict:
+        """Cancel a background job (gateway / orchestrator integration)."""
+        job_id = job_id.strip().strip('"').strip("'").strip("?").strip("#")
+        if self.job_queue is None:
+            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Job queue not configured.")
+
+        cancel_fn = getattr(self.job_queue, "cancel_gateway_job", None)
+        if callable(cancel_fn):
+            return cancel_fn(job_id)
+
+        try:
+            self.job_queue.cancel_job(job_id)
+        except NotImplementedError:
+            raise HTTPException(
+                status_code=status.HTTP_501_NOT_IMPLEMENTED,
+                detail="Cancellation is not supported for this job queue.",
+            ) from None
+        return {"id": job_id, "status": "cancelled", "message": "Job cancelled."}
+
+    async def stream_job_sse(self, job_id: str, request: Request):
+        """Server-Sent Events for streaming job output (requires stream_store)."""
+        job_id = job_id.strip().strip('"').strip("'").strip("?").strip("#")
+        if self.stream_store is None:
+            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Streaming not configured.")
+        if self.job_queue is None:
+            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Job queue not configured.")
+
+        jq = self.job_queue
+        job_data = jq.get_job_status(job_id) if hasattr(jq, "get_job_status") else None
+
+        if job_data is None:
+            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Job '{job_id}' not found.")
+
+        st = (job_data.get("status") or "").lower()
+        if st != "streaming" and not self.stream_store.stream_exists(job_id):
+            raise HTTPException(
+                status_code=status.HTTP_409_CONFLICT,
+                detail=f"Job '{job_id}' is not streaming (status: {job_data.get('status')}).",
+            )
+
+        async def _event_generator():
+            try:
+                async for chunk in self.stream_store.read_chunks(job_id):
+                    if await request.is_disconnected():
+                        break
+                    yield chunk
+            except Exception:
+                self._logger.exception("Error during stream delivery | job_id=%s", job_id)
+                yield 'data: {"error": "Internal stream error"}\n\n'
+
+        return StreamingResponse(
+            _event_generator(),
+            media_type="text/event-stream",
+            headers={
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Accel-Buffering": "no",
+            },
+        )
+
+
     def endpoint(self, path: str, methods: list[str] | None = None, max_upload_file_size_mb: int = None, queue_size: int = 500, use_queue: bool = None, *args, **kwargs):
         """
         Unified endpoint decorator.
@@ -356,12 +424,16 @@ async def _unified_worker(*w_args, **w_kwargs):
             plan.max_upload_file_size_mb
         )
 
+        route_kwargs = dict(plan.route_kwargs)
+        if plan.should_use_queue:
+            route_kwargs["response_model_exclude_none"] = True
+
         self.api_route(
             path=plan.path,
             methods=plan.active_methods,
             response_model=JobResult if plan.should_use_queue else res_model,
             *plan.route_args,
-            **plan.route_kwargs
+            **route_kwargs
         )(final_handler)
 
         return final_handler
@@ -472,12 +544,14 @@ def _determine_generator_fun(self, func: Callable) -> bool:
     def _create_task_endpoint_decorator(self, path: str, methods: list[str] | None, max_upload_file_size_mb: int, queue_size: int, args, kwargs):
         """Create a decorator for task endpoints (background job execution)."""
         # FastAPI route decorator (returning JobResult)
+        task_kwargs = dict(kwargs)
+        task_kwargs["response_model_exclude_none"] = True
         fastapi_route_decorator = self.api_route(
             path=path,
             methods=["POST"] if methods is None else methods,
             response_model=JobResult,
             *args,
-            **kwargs
+            **task_kwargs
         )
 
         # Queue decorator

apipod/engine/backend/runpod/llm_mixin.py

@@ -1,4 +1,4 @@
-from apipod.engine.llm.base_mixin import _BaseLLMMixin
+from apipod.engine.llm.base_llm_mixin import _BaseLLMMixin
 
 
 class _RunPodLLMMixin(_BaseLLMMixin):

apipod/engine/backend/runpod/router.py

@@ -10,7 +10,7 @@
 from apipod.engine.jobs.job_progress import JobProgressRunpod, JobProgress
 from apipod.engine.jobs.job_result import JobResultFactory, JobResult
 from apipod.engine.base_backend import _BaseBackend
-from apipod.engine.files.base_mixin import _BaseFileHandlingMixin
+from apipod.engine.files.base_file_mixin import _BaseFileHandlingMixin
 from apipod.engine.backend.runpod.llm_mixin import _RunPodLLMMixin
 
 from apipod.engine.utils import normalize_name
@@ -195,7 +195,10 @@ def _router(self, path, job, **kwargs):
 
         # Prepare result tracking
         start_time = datetime.now(timezone.utc)
-        result = JobResult(id=job['id'], execution_started_at=start_time.strftime(DEFAULT_DATE_TIME_FORMAT))
+        result = JobResult(
+            job_id=job["id"],
+            created_at=start_time.strftime(DEFAULT_DATE_TIME_FORMAT),
+        )
 
         try:
             # Execute the function (Sync or Async Handling)
@@ -218,7 +221,7 @@ def _router(self, path, job, **kwargs):
             print(f"Job {job['id']} failed: {str(e)}")
             traceback.print_exc()
         finally:
-            result.execution_finished_at = datetime.now(timezone.utc).strftime(DEFAULT_DATE_TIME_FORMAT)
+            result.updated_at = datetime.now(timezone.utc).strftime(DEFAULT_DATE_TIME_FORMAT)
 
         result = result.model_dump_json()
         return result