feat(base models): now all models encapsulates response object.

.github/workflows/ci.yml

@@ -32,10 +32,10 @@ jobs:
           pip install -r requirements.txt
           pip install -r requirements-dev.txt
 
-      - name: Run lint check
+      - name: Run type check
         run: |
-          black --check --verbose deepinfra
+          mypy --install-types deepinfra --non-interactive --verbose
 
-      - name: Run type check
+      - name: Run lint check
         run: |
-          mypy --install-types deepinfra --non-interactive --verbose
+          black --check --verbose deepinfra

deepinfra/models/base/automatic_speech_recognition.py

@@ -26,4 +26,4 @@ def generate(self, body) -> AutomaticSpeechRecognitionResponse:
         response = self.client.post(
             form_data, {"headers": {"content-type": form_data.content_type}}
         )
-        return response.json()
+        return AutomaticSpeechRecognitionResponse(**response.json())

deepinfra/models/base/embeddings.py

@@ -8,11 +8,11 @@ class Embeddings(BaseModel):
     @docs Check the available models at https://deepinfra.com/models/embeddings
     """
 
-    def generate(self, body: EmbeddingsRequest) -> EmbeddingsResponse:
+    def generate(self, body) -> EmbeddingsResponse:
         """
         Generates embeddings.
         :param body:
         :return:
         """
         response = self.client.post(body)
-        return response
+        return EmbeddingsResponse(**response.json())

deepinfra/models/base/text_generation.py

@@ -16,11 +16,11 @@ class TextGeneration(BaseModel):
     @docs Check the available models at https://deepinfra.com/models/text-generation
     """
 
-    def generate(self, body: TextGenerationRequest) -> TextGenerationResponse:
+    def generate(self, body: dict) -> TextGenerationResponse:
         """
         Generates text.
         :param body:
         :return:
         """
         response = self.client.post(body)
-        return response
+        return TextGenerationResponse(**response.json())

deepinfra/models/base/text_to_image.py

@@ -1,4 +1,5 @@
 from deepinfra.models.base import BaseModel
+from deepinfra.types.text_to_image import TextToImageResponse
 
 
 class TextToImage(BaseModel):
@@ -15,4 +16,4 @@ def generate(self, input):
         """
         body = {"input": input}
         response = self.client.post(body)
-        return response
+        return TextToImageResponse(**response.json())

deepinfra/types/text_to_image/__init__.py

@@ -0,0 +1,2 @@
+from .request import TextToImageRequest
+from .response import TextToImageResponse

deepinfra/types/image_generation/request.py → deepinfra/types/text_to_image/request.py

@@ -3,7 +3,7 @@
 
 
 @dataclass
-class ImageGenerationRequest:
+class TextToImageRequest:
     prompt: str
     negative_prompt: Optional[str] = None
     image: Optional[str] = None

deepinfra/types/image_generation/response.py → deepinfra/types/text_to_image/response.py

@@ -10,7 +10,7 @@ class Metrics:
 
 
 @dataclass
-class ImageGenerationResponse:
+class TextToImageResponse:
     request_id: str
     inference_status: InferenceStatus
     input: Dict

deepinfra/utils/form_data.py

@@ -12,7 +12,7 @@ class FormDataUtils:
     """
 
     @staticmethod
-    def get_form_data(data, blob_keys: Optional[List[str]] = None):
+    def get_form_data(data: dict, blob_keys: Optional[List[str]] = None):
         """
         Creates a MultipartEncoder object from the data.
         :param data: