fix: update type hints for _default_unbatch and _spec attributes in LitAPI class (#435)

* fix: update type hints for _default_unbatch and _spec attributes in LitAPI class

* fix: improve error message for unimplemented unbatch method in LitAPI class and add corresponding test

Author: bhimrazy

src/litserve/api.py

@@ -15,7 +15,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from queue import Queue
-from typing import Optional
+from typing import Callable, Optional
 
 from pydantic import BaseModel
 
@@ -24,8 +24,8 @@
 
 class LitAPI(ABC):
     _stream: bool = False
-    _default_unbatch: callable = None
-    _spec: LitSpec = None
+    _default_unbatch: Optional[Callable] = None
+    _spec: Optional[LitSpec] = None
     _device: Optional[str] = None
     _logger_queue: Optional[Queue] = None
     request_timeout: Optional[float] = None
@@ -76,6 +76,11 @@ def _unbatch_stream(self, output_stream):
 
     def unbatch(self, output):
         """Convert a batched output to a list of outputs."""
+        if self._default_unbatch is None:
+            raise ValueError(
+                "Default implementation for `LitAPI.unbatch` method was not found. "
+                "Please implement the `LitAPI.unbatch` method."
+            )
         return self._default_unbatch(output)
 
     def encode_response(self, output, **kwargs):

tests/test_litapi.py

@@ -72,6 +72,12 @@ def test_default_batch_unbatch():
     assert api.unbatch(output) == inputs, "Default unbatch should not change input"
 
 
+def test_default_unbatch_not_implemented():
+    api = TestDefaultBatchedAPI()
+    with pytest.raises(ValueError, match="Default implementation for `LitAPI.unbatch` method was not found."):
+        api.unbatch(None)
+
+
 class TestStreamAPIBatched(TestStreamAPI):
     def predict(self, x):
         for i in range(4):