🚧 WIP

.gitignore

@@ -160,3 +160,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# MacOS files
+.DS_Store

Makefile

@@ -15,8 +15,9 @@ style:
 	isort src tests
 	flake8 src tests --max-line-length 88
 
-# test:
-#     pytest tests
+
+test:
+	python -m pytest --cache-clear -vvv -x -s ./tests
 
 build:
 	python setup.py sdist bdist_wheel
@@ -31,4 +32,4 @@ clean:
 	rm -rf .mypy_cache
 	rm -rf .pytest_cache
 
-.PHONY: install install-dev quality style test test-unit test-integration test-e2e test-smoke test-sanity test-regression build clean
+.PHONY: install install-dev quality style test build clean

setup.py

@@ -1,46 +1,49 @@
-from setuptools import setup, find_packages
 from typing import Tuple
 
+from setuptools import find_packages, setup
+
 
 def _setup_long_description() -> Tuple[str, str]:
     return open("README.md", "r", encoding="utf-8").read(), "text/markdown"
 
 
 setup(
-    name='guidellm',
-    version='0.1.0',
-    author='Neuralmagic, Inc.',
-    description='Guidance platform for deploying and managing large language models.',
+    name="guidellm",
+    version="0.1.0",
+    author="Neuralmagic, Inc.",
+    description="Guidance platform for deploying and managing large language models.",
     long_description=_setup_long_description()[0],
     long_description_content_type=_setup_long_description()[1],
     license="Apache",
     url="https://github.com/neuralmagic/guidellm",
-    packages=find_packages(where='src'),
-    package_dir={'': 'src'},
+    packages=find_packages(where="src"),
+    package_dir={"": "src"},
     include_package_data=True,
     install_requires=[
-        'datasets',
-        'loguru',
-        'numpy',
-        'openai',
-        'requests',
-        'transformers',
+        "datasets",
+        "loguru",
+        "numpy",
+        "openai",
+        "requests",
+        "transformers",
     ],
     extras_require={
-        'dev': [
-            'pytest',
-            'sphinx',
-            'ruff',
-            'mypy',
-            'black',
-            'isort',
-            'flake8',
-            'pre-commit',
+        "dev": [
+            "black",
+            "flake8",
+            "isort",
+            "mypy",
+            "pre-commit",
+            "pytest",
+            "pytest-asyncio",
+            "pytest-mock",
+            "ruff",
+            "sphinx",
         ],
     },
     entry_points={
-        'console_scripts': [
-            'guidellm=guidellm.main:main',
+        "console_scripts": [
+            "guidellm=guidellm.main:main",
         ],
     },
     python_requires=">=3.8.0",

src/guidellm/backend/__init__.py

@@ -1,9 +1,9 @@
-from .base import Backend, BackendTypes, GenerativeResponse
+from .base import Backend, BackendEngine, GenerativeResponse
 from .openai import OpenAIBackend
 
 __all__ = [
     "Backend",
-    "BackendTypes",
+    "BackendEngine",
     "GenerativeResponse",
     "OpenAIBackend",
 ]

src/guidellm/backend/base.py

@@ -1,18 +1,26 @@
+import functools
 import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
-from typing import Iterator, List, Optional, Type, Union
+from typing import Generic, Iterator, List, Optional, Type, Union
 
 from loguru import logger
 
 from guidellm.core.request import TextGenerationRequest
 from guidellm.core.result import TextGenerationResult
 
-__all__ = ["Backend", "BackendTypes", "GenerativeResponse"]
+__all__ = ["Backend", "BackendEngine", "GenerativeResponse"]
 
 
-class BackendTypes(Enum):
+class BackendEngine(str, Enum):
+    """
+    Determines the Engine of the LLM Backend.
+    All the implemented backends in the project have the engine.
+
+    NOTE: the `TEST` engine has to be used only for testing purposes.
+    """
+
     TEST = "test"
     OPENAI_SERVER = "openai_server"
 
@@ -33,18 +41,18 @@ class GenerativeResponse:
 
 class Backend(ABC):
     """
-    An abstract base class for generative AI backends.
+    An abstract base class with template methods for generative AI backends.
     """
 
     _registry = {}
 
     @staticmethod
-    def register_backend(backend_type: BackendTypes):
+    def register(backend_type: BackendEngine):
         """
         A decorator to register a backend class in the backend registry.
 
         :param backend_type: The type of backend to register.
-        :type backend_type: BackendTypes
+        :type backend_type: BackendType
         """
 
         def inner_wrapper(wrapped_class: Type["Backend"]):
@@ -54,21 +62,24 @@ def inner_wrapper(wrapped_class: Type["Backend"]):
         return inner_wrapper
 
     @staticmethod
-    def create_backend(backend_type: Union[str, BackendTypes], **kwargs) -> "Backend":
+    def create(backend_type: Union[str, BackendEngine], **kwargs) -> "Backend":
         """
         Factory method to create a backend based on the backend type.
 
         :param backend_type: The type of backend to create.
-        :type backend_type: BackendTypes
+        :type backend_type: BackendType
         :param kwargs: Additional arguments for backend initialization.
         :type kwargs: dict
         :return: An instance of a subclass of Backend.
         :rtype: Backend
         """
+
         logger.info(f"Creating backend of type {backend_type}")
+
         if backend_type not in Backend._registry:
             logger.error(f"Unsupported backend type: {backend_type}")
             raise ValueError(f"Unsupported backend type: {backend_type}")
+
         return Backend._registry[backend_type](**kwargs)
 
     def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
@@ -121,8 +132,10 @@ def available_models(self) -> List[str]:
         :return: A list of available models.
         :rtype: List[str]
         """
-        raise NotImplementedError()
 
+        pass
+
+    @property
     @abstractmethod
     def default_model(self) -> str:
         """
@@ -131,7 +144,8 @@ def default_model(self) -> str:
         :return: The default model.
         :rtype: str
         """
-        raise NotImplementedError()
+
+        pass
 
     @abstractmethod
     def model_tokenizer(self, model: str) -> Optional[str]:
@@ -143,4 +157,5 @@ def model_tokenizer(self, model: str) -> Optional[str]:
         :return: The tokenizer for the model, or None if it cannot be created.
         :rtype: Optional[str]
         """
-        raise NotImplementedError()
+
+        pass

src/guidellm/backend/openai.py

@@ -1,16 +1,18 @@
-from typing import Any, Iterator, List, Optional
+import functools
+import os
+from typing import Any, Dict, Iterator, List, Optional
 
-import openai
 from loguru import logger
+from openai import OpenAI
 from transformers import AutoTokenizer
 
-from guidellm.backend import Backend, BackendTypes, GenerativeResponse
+from guidellm.backend import Backend, BackendEngine, GenerativeResponse
 from guidellm.core.request import TextGenerationRequest
 
 __all__ = ["OpenAIBackend"]
 
 
-@Backend.register_backend(BackendTypes.OPENAI_SERVER)
+@Backend.register(BackendEngine.OPENAI_SERVER)
 class OpenAIBackend(Backend):
     """
     An OpenAI backend implementation for the generative AI result.
@@ -33,34 +35,35 @@ class OpenAIBackend(Backend):
 
     def __init__(
         self,
-        target: Optional[str] = None,
-        host: Optional[str] = None,
-        port: Optional[int] = None,
-        path: Optional[str] = None,
+        openai_api_key: Optional[str] = None,
+        internal_callback_url: Optional[str] = None,
         model: Optional[str] = None,
-        api_key: Optional[str] = None,
-        **request_args,
+        **request_args: Any,
     ):
-        self.target = target
-        self.model = model
-        self.request_args = request_args
-
-        if not self.target:
-            if not host:
-                raise ValueError("Host is required if target is not provided.")
-
-            port_incl = f":{port}" if port else ""
-            path_incl = path if path else ""
-            self.target = f"http://{host}{port_incl}{path_incl}"
+        """
+        Initialize an OpenAI Client
+        """
 
-        openai.api_base = self.target
-        openai.api_key = api_key
+        self.request_args = request_args
 
-        if not model:
-            self.model = self.default_model()
+        if not (_api_key := (openai_api_key or os.getenv("OPENAI_API_KEY", None))):
+            raise ValueError(
+                "`OPENAI_API_KEY` environment variable or --openai-api-key CLI parameter "
+                "must be specify for the OpenAI backend"
+            )
+
+        if not (
+            _base_url := (internal_callback_url or os.getenv("OPENAI_BASE_URL", None))
+        ):
+            raise ValueError(
+                "`OPENAI_BASE_URL` environment variable or --openai-base-url CLI parameter "
+                "must be specify for the OpenAI backend"
+            )
+        self.openai_client = OpenAI(api_key=_api_key, base_url=_base_url)
+        self.model = model or self.default_model
 
         logger.info(
-            f"Initialized OpenAIBackend with target: {self.target} "
+            f"Initialized OpenAIBackend with callback url: {internal_callback_url} "
             f"and model: {self.model}"
         )
 
@@ -75,11 +78,10 @@ def make_request(
         :return: An iterator over the generative responses.
         :rtype: Iterator[GenerativeResponse]
         """
+
         logger.debug(f"Making request to OpenAI backend with prompt: {request.prompt}")
         num_gen_tokens = request.params.get("generated_tokens", None)
-        request_args = {
-            "n": 1,
-        }
+        request_args: Dict = {"n": 1}
 
         if num_gen_tokens:
             request_args["max_tokens"] = num_gen_tokens
@@ -88,8 +90,8 @@ def make_request(
         if self.request_args:
             request_args.update(self.request_args)
 
-        response = openai.Completion.create(
-            engine=self.model,
+        response = self.openai_client.completions.create(
+            model=self.model,
             prompt=request.prompt,
             stream=True,
             **request_args,
@@ -129,21 +131,28 @@ def available_models(self) -> List[str]:
         :return: A list of available models.
         :rtype: List[str]
         """
-        models = [model["id"] for model in openai.Engine.list()["data"]]
+
+        models: list[str] = [
+            model.id for model in self.openai_client.models.list().data
+        ]
         logger.info(f"Available models: {models}")
+
         return models
 
+    @property
+    @functools.lru_cache(maxsize=1)
     def default_model(self) -> str:
         """
         Get the default model for the backend.
 
         :return: The default model.
         :rtype: str
         """
-        models = self.available_models()
-        if models:
+
+        if models := self.available_models():
             logger.info(f"Default model: {models[0]}")
             return models[0]
+
         logger.error("No models available.")
         raise ValueError("No models available.")