🚧 WIP

pyproject.toml

@@ -20,6 +20,7 @@ lint.select = ["E", "F", "W"]
 max-line-length = 88
 
 [tool.pytest.ini_options]
+python_classes = "DisableTestClasses"
 markers = [
     "smoke: quick tests to check basic functionality",
     "sanity: detailed tests to ensure major functions work correctly",

setup.py

@@ -29,10 +29,10 @@ def _setup_long_description() -> Tuple[str, str]:
     ],
     extras_require={
         "dev": [
-            "black",
             "flake8",
             "isort",
             "mypy",
+            "polyfactory",
             "pre-commit",
             "pytest",
             "pytest-asyncio",

src/guidellm/backend/openai.py

@@ -1,9 +1,10 @@
 import functools
 import os
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterable, Iterator, List, Optional
 
 from loguru import logger
-from openai import OpenAI
+from openai import OpenAI, Stream
+from openai.types import Completion
 from transformers import AutoTokenizer
 
 from guidellm.backend import Backend, BackendEngine, GenerativeResponse
@@ -90,39 +91,38 @@ def make_request(
         if self.request_args:
             request_args.update(self.request_args)
 
-        response = self.openai_client.completions.create(
+        response: Iterable[Completion] = self.openai_client.completions.create(
             model=self.model,
             prompt=request.prompt,
             stream=True,
             **request_args,
         )
 
         for chunk in response:
-            if chunk.get("choices"):
-                choice = chunk["choices"][0]
-                if choice.get("finish_reason") == "stop":
+            if chunk.choices:
+                # Available choices: [ stop, length, content_filter ]
+                # from: openai/types/completion_choice.py
+                if (choice := chunk.choices[0]).finish_reason == "stop":
                     logger.debug("Received final response from OpenAI backend")
+
                     yield GenerativeResponse(
                         type_="final",
-                        output=choice["text"],
+                        output=choice.text,
                         prompt=request.prompt,
                         prompt_token_count=(
-                            request.token_count
-                            if request.token_count
-                            else self._token_count(request.prompt)
+                            request.prompt_token_count
+                            or self._token_count(request.prompt)
                         ),
                         output_token_count=(
                             num_gen_tokens
                             if num_gen_tokens
-                            else self._token_count(choice["text"])
+                            else self._token_count(choice.text)
                         ),
                     )
                     break
                 else:
                     logger.debug("Received token from OpenAI backend")
-                    yield GenerativeResponse(
-                        type_="token_iter", add_token=choice["text"]
-                    )
+                    yield GenerativeResponse(type_="token_iter", add_token=choice.text)
 
     def available_models(self) -> List[str]:
         """

tests/conftest.py

@@ -1 +1,2 @@
 
+

tests/unit/backend/conftest.py

@@ -1,37 +1,59 @@
+import random
+from typing import List
+
 import pytest
 from openai.pagination import SyncPage
-from openai.types import Model
+from openai.types import Completion, Model
+
+from guidellm.backend import Backend, BackendEngine, OpenAIBackend
+
+from . import factories
 
 
 @pytest.fixture(autouse=True)
-def openai_models_list_patch(mocker):
+def openai_models_list_patch(mocker) -> List[Model]:
     """
     Mock available models function to avoid OpenAI API call.
     """
 
-    return mocker.patch(
+    items = factories.OpenAIModel.batch(3)
+    mocker.patch(
         "openai.resources.models.Models.list",
-        return_value=SyncPage(
-            object="list",
-            data=[
-                Model(
-                    id="model-id-1",
-                    object="model",
-                    created=1686935002,
-                    owned_by="openai",
-                ),
-                Model(
-                    id="model-id-2",
-                    object="model",
-                    created=1686935003,
-                    owned_by="openai",
-                ),
-                Model(
-                    id="model-id-3",
-                    object="model",
-                    created=1686935004,
-                    owned_by="openai",
-                ),
-            ],
-        ),
+        return_value=SyncPage(object="list", data=items),
     )
+
+    return items
+
+
+@pytest.fixture(autouse=True)
+def openai_completion_create_patch(mocker) -> List[Completion]:
+    """
+    Mock available models function to avoid OpenAI API call.
+    """
+
+    items = factories.OpenAICompletion.batch(random.randint(2, 5))
+    mocker.patch("openai.resources.completions.Completions.create", return_value=items)
+
+    return items
+
+
+@pytest.fixture
+def openai_backend_factory():
+    """
+    Create a test openai backend service.
+    Call without provided arguments returns default Backend service.
+    """
+
+    def inner_wrapper(*_, **kwargs) -> OpenAIBackend:
+        static = {"backend_type": BackendEngine.OPENAI_SERVER}
+        defaults = {
+            "openai_api_key": "dummy api key",
+            "internal_callback_url": "http://localhost:8000",
+        }
+
+        defaults.update(kwargs)
+        defaults.update(static)
+
+        return Backend.create(**defaults)
+
+    return inner_wrapper

tests/unit/backend/factories.py

@@ -0,0 +1,31 @@
+import functools
+import random
+
+from openai.types import Completion, CompletionChoice, Model
+from polyfactory.factories.pydantic_factory import ModelFactory
+
+__all__ = ["OpenAIModel", "OpenAICompletionChoice", "OpenAICompletion"]
+
+
+class OpenAIModel(ModelFactory[Model]):
+    """
+    A model factory for Open AI Model representation.
+    """
+
+    pass
+
+
+class OpenAICompletionChoice(ModelFactory[CompletionChoice]):
+    """
+    A model factory for Open AI Completion Choice representation.
+    """
+
+    finish_reason = "stop"
+
+
+class OpenAICompletion(ModelFactory[Completion]):
+    """
+    A model factory for Open AI Completion representation.
+    """
+
+    choices = functools.partial(OpenAICompletionChoice.batch, random.randint(3, 5))

tests/unit/backend/test_base.py

@@ -24,6 +24,7 @@ def make_request(
     def available_models(self) -> List[str]:
         raise NotImplementedError
 
+    @property
     def default_model(self) -> str:
         raise NotImplementedError
 

tests/unit/backend/test_openai.py

@@ -2,64 +2,77 @@
 This module includes unit tests for the OpenAI Backend Service.
 """
 
+from typing import Callable, List
+
 import pytest
+from openai.types import Completion
 
 from guidellm.backend import Backend, BackendEngine, OpenAIBackend
+from guidellm.core import TextGenerationRequest
 from guidellm.request.test import TestRequestGenerator
 
 
 @pytest.mark.sanity
-def test_openai_backend_creation_with_default_model():
+def test_openai_backend_creation_with_default_model(openai_backend_factory: Callable):
     """
     Test whether the OpenAI Backend service is created correctly
     with all default parameters.
 
     Also checks wheather the `default_models` parameter does not abuse the OpenAI API.
     """
 
-    backend_service = Backend.create(
-        backend_type=BackendEngine.OPENAI_SERVER,
-        openai_api_key="dummy api key",
-        internal_callback_url="http://localhost:8000",
-    )
-
+    backend_service = openai_backend_factory()
     assert isinstance(backend_service, OpenAIBackend)
-    assert backend_service.default_model == "model-id-1"
+    assert backend_service.default_model == backend_service.available_models()[0]
 
 
 @pytest.mark.smoke
 @pytest.mark.parametrize(
     "extra_kwargs",
     [
         {"openai_api_key": "dummy"},
-        {"openai_base_url": "dummy"},
+        {"internal_callback_url": "dummy"},
     ],
 )
 def test_openai_backend_creation_required_arguments(extra_kwargs: dict):
+    """
+    Both OpenAI key & internal callback URL are required to work with OpenAI Backend.
+    """
     with pytest.raises(ValueError):
         Backend.create(
             backend_type=BackendEngine.OPENAI_SERVER,
-            target="https://dummy.com",
             **extra_kwargs,
         )
 
 
-def test_model_tokenizer(mocker):
-    backend_service = Backend.create(
-        backend_type=BackendEngine.OPENAI_SERVER,
-        openai_api_key="dummy api key",
-        internal_callback_url="http://localhost:8000",
-    )
-
+def test_model_tokenizer(openai_backend_factory):
+    backend_service = openai_backend_factory()
     assert backend_service.model_tokenizer("bert-base-uncased")
 
 
-def test_model_tokenizer_no_model():
-    backend_service = Backend.create(
-        backend_type=BackendEngine.OPENAI_SERVER,
-        openai_api_key="dummy api key",
-        internal_callback_url="http://localhost:8000",
-    )
+def test_model_tokenizer_no_model(openai_backend_factory):
+    backend_service = openai_backend_factory()
     tokenizer = backend_service.model_tokenizer("invalid")
-
     assert tokenizer is None
+
+
+def test_make_request(
+    openai_backend_factory, openai_completion_create_patch: List[Completion]
+):
+    request: TextGenerationRequest = TestRequestGenerator().create_item()
+    backend_service: OpenAIBackend = openai_backend_factory()
+    total_generative_responses = 0
+
+    for generative_response, patched_completion in zip(
+        backend_service.make_request(request=request),
+        openai_completion_create_patch,
+    ):
+        total_generative_responses += 1
+        expected_output: str = patched_completion.choices[0].text
+
+        assert generative_response.type_ == "final"
+        assert generative_response.output == expected_output
+        assert generative_response.prompt_token_count == len(request.prompt.split())
+        assert generative_response.output_token_count == len(expected_output.split())
+
+    assert total_generative_responses == 1