Backend OpenAI tests

* base unit tests are added
* OpenAI unit tests are added
* `Backend.submit` with OpenaAI backend integration test is added
* `.env.example` is added as a source of a project configuration
* `README.md` now has project configuring and testing small guide

.env.example

@@ -0,0 +1,9 @@
+# OpenAI compatible server address.
+
+# If you are going to run custom OpenAI API compatible service change this configuration.
+# Could be specified by --openai-base-url CLI parameter
+OPENAI_BASE_URL=http://127.0.0.1:8080
+
+# The OpenAI API Key.
+# Could be specified by --openai-api-key CLI parameter
+OPENAI_API_KEY=invalid

README.md

@@ -1 +1,20 @@
-# guidellm
+# guidellm
+
+# Project configuration
+
+The project is configured with environment variables. Check the example in `.env.example`.
+
+```sh
+# Create .env file and update the configuration
+cp .env.example  .env
+
+# Export all variables
+set -o allexport; source .env; set +o allexport
+```
+
+## Environment Variables
+
+| Variable        | Default Value         | Description                                                                                                                                                    |
+| --------------- | --------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| OPENAI_BASE_URL | http://127.0.0.1:8080 | The host where the `openai` library will make requests to. For running integration tests it is required to have the external OpenAI compatible server running. |
+| OPENAI_API_KEY  | invalid               | [OpenAI Platform](https://platform.openai.com/api-keys) to create a new API key. This value is not used for tests.                                             |

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,23 @@
-import uuid
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from enum import Enum
 from typing import Iterator, List, Optional, Type, Union
 
 from loguru import logger
 
-from guidellm.core.request import TextGenerationRequest
-from guidellm.core.result import TextGenerationResult
+from guidellm.core import TextGenerationRequest, 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,43 +38,46 @@ 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):
+    @classmethod
+    def register(cls, 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"]):
-            Backend._registry[backend_type] = wrapped_class
+            cls._registry[backend_type] = wrapped_class
             return wrapped_class
 
         return inner_wrapper
 
-    @staticmethod
-    def create_backend(backend_type: Union[str, BackendTypes], **kwargs) -> "Backend":
+    @classmethod
+    def create(cls, 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:
+
+        if backend_type not in cls._registry:
             logger.error(f"Unsupported backend type: {backend_type}")
             raise ValueError(f"Unsupported backend type: {backend_type}")
-        return Backend._registry[backend_type](**kwargs)
+
+        return cls._registry[backend_type](**kwargs)
 
     def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
         """
@@ -80,23 +88,23 @@ def submit(self, request: TextGenerationRequest) -> TextGenerationResult:
         :return: The populated result result.
         :rtype: TextGenerationResult
         """
+
         logger.info(f"Submitting request with prompt: {request.prompt}")
-        result_id = str(uuid.uuid4())
-        result = TextGenerationResult(result_id)
+
+        result = TextGenerationResult(request=request)
         result.start(request.prompt)
 
-        for response in self.make_request(request):
+        for response in self.make_request(request):  # GenerativeResponse
             if response.type_ == "token_iter" and response.add_token:
                 result.output_token(response.add_token)
             elif response.type_ == "final":
                 result.end(
-                    response.output,
                     response.prompt_token_count,
                     response.output_token_count,
                 )
-                break
 
         logger.info(f"Request completed with output: {result.output}")
+
         return result
 
     @abstractmethod
@@ -111,7 +119,8 @@ def make_request(
         :return: An iterator over the generative responses.
         :rtype: Iterator[GenerativeResponse]
         """
-        raise NotImplementedError()
+
+        pass
 
     @abstractmethod
     def available_models(self) -> List[str]:
@@ -121,8 +130,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 +142,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 +155,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,19 @@
-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, Stream
+from openai.types import Completion
 from transformers import AutoTokenizer
 
-from guidellm.backend import Backend, BackendTypes, GenerativeResponse
-from guidellm.core.request import TextGenerationRequest
+from guidellm.backend import Backend, BackendEngine, GenerativeResponse
+from guidellm.core 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 +36,37 @@ 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,52 +81,46 @@ 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,
-        }
 
-        if num_gen_tokens:
-            request_args["max_tokens"] = num_gen_tokens
-            request_args["stop"] = None
+        # How many completions to generate for each prompt
+        request_args: Dict = {"n": 1}
+
+        if (num_gen_tokens := request.params.get("generated_tokens", None)) is not None:
+            request_args.update(max_tokens=num_gen_tokens, stop=None)
 
         if self.request_args:
             request_args.update(self.request_args)
 
-        response = openai.Completion.create(
-            engine=self.model,
+        response: Stream[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":
-                    logger.debug("Received final response from OpenAI backend")
-                    yield GenerativeResponse(
-                        type_="final",
-                        output=choice["text"],
-                        prompt=request.prompt,
-                        prompt_token_count=(
-                            request.token_count
-                            if request.token_count
-                            else self._token_count(request.prompt)
-                        ),
-                        output_token_count=(
-                            num_gen_tokens
-                            if num_gen_tokens
-                            else self._token_count(choice["text"])
-                        ),
-                    )
-                    break
-                else:
-                    logger.debug("Received token from OpenAI backend")
-                    yield GenerativeResponse(
-                        type_="token_iter", add_token=choice["text"]
-                    )
+            chunk_content: str = getattr(chunk, "content", "")
+
+            if getattr(chunk, "stop", True) is True:
+                logger.debug("Received final response from OpenAI backend")
+
+                yield GenerativeResponse(
+                    type_="final",
+                    prompt=getattr(chunk, "prompt", request.prompt),
+                    prompt_token_count=(
+                        request.prompt_token_count or self._token_count(request.prompt)
+                    ),
+                    output_token_count=(
+                        num_gen_tokens
+                        if num_gen_tokens
+                        else self._token_count(chunk_content)
+                    ),
+                )
+            else:
+                logger.debug("Received token from OpenAI backend")
+                yield GenerativeResponse(type_="token_iter", add_token=chunk_content)
 
     def available_models(self) -> List[str]:
         """
@@ -129,21 +129,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.")