module LLM Chat API

The EHRGym Chat API module provides an abstraction layer for interacting with various language models (LLMs), including OpenAI, Azure-hosted models, and locally-hosted vLLM models. It offers utility functions to manage and structure chat-based interactions through a standardized interface.

Utility Functions

  • make_system_message(content: str): Creates a system message dict.
  • make_user_message(content: str): Creates a user message dict.
  • make_assistant_message(content: str): Creates an assistant message dict.

Abstract Classes and Interfaces

AbstractChatModel

An abstract base class defining the interface for chat model implementations.

Methods:

  • __call__(self, messages: list[dict]): Abstract method for generating responses based on message history.
  • get_stats(self): Returns statistics from the model's operation (e.g., retry counts).

BaseModelArgs

Abstract dataclass defining configuration parameters for all model types.

Fields:

  • model_name (str): Name or identifier of the model.
  • max_total_tokens (int): Total token limit.
  • max_input_tokens (int): Input token limit.
  • max_new_tokens (int): Token limit for newly generated content.
  • temperature (Float): Controls randomness in generation.
  • vision_support (Boolean): Indicates if the model supports visual input.
  • log_probs (Boolean): Whether to log token probabilities.

Methods:

  • make_model(self): Abstract method to instantiate specific chat models.

Concrete Model Arguments

OpenAIModelArgs

Configures an OpenAI-hosted chat model.

Methods:

  • make_model(self): Instantiates an OpenAI chat model.

AzureModelArgs

Configures an Azure-hosted chat model.

Fields:

  • deployment_name: Azure-specific deployment name.

Methods:

  • make_model(self): Instantiates an Azure chat model.

vLLMModelArgs

Configures an locally-hosted vLLM model.

Fields:

  • port: Port number for the local vLLM server.

Methods:

  • make_model(self): Instantiates a locally-hosted vLLM chat model.

Chat Model Implementations

ChatModel

Base implementation for interacting with chat-based APIs.

Parameters:

  • port: Model configuration parameters, retry management, logging probabilities, and API credentials.

Methods:

  • __call__(self, messages: list[dict], n_samples: int, temperature: Optional[float]): Executes chat interactions with retry logic and error handling.
  • get_stats(self): Retrieves retry statistics.

Specialized Chat Models

  • OpenAIChatModel: Subclass tilored for OpenAI API interactions.
  • AzureChatModel: Subclass tailored for Azure OpenAI service interactions, incorporating Azure-specific configurations such as deployment names and endpoints.
  • vLLMChatModel: Subclass for local vLLM server interactions, configurable via local ports.

Usage Example

model_args = AzureModelArgs(
    model_name="gpt-4",
    temperature=0.5,
    max_new_tokens=512,
    deployment_name="gpt4-azure-deployment",
    log_probs=False
)
model = model_args.make_model()

response, cost = model([
    make_system_message("You are an assistant."),
    make_user_message("Summarize the key points of diabetes management.")
])