not-again-ai 0.14.0__py3-none-any.whl → 0.15.0__py3-none-any.whl

not_again_ai/llm/gh_models/chat_completion.py

@@ -1,81 +0,0 @@
-import contextlib
-import json
-import time
-from typing import Any
-
-from azure.ai.inference import ChatCompletionsClient
-from azure.ai.inference.models import ChatCompletionsToolDefinition, ChatRequestMessage
-
-
-def chat_completion(
-    messages: list[ChatRequestMessage],
-    model: str,
-    client: ChatCompletionsClient,
-    tools: list[ChatCompletionsToolDefinition] | None = None,
-    max_tokens: int | None = None,
-    temperature: float | None = None,
-    json_mode: bool = False,
-    seed: int | None = None,
-) -> dict[str, Any]:
-    """Gets a response from GitHub Models using the Azure AI Inference SDK.
-    See the available models at https://github.com/marketplace/models
-    Full documentation of the SDK is at: https://learn.microsoft.com/en-us/azure/ai-studio/reference/reference-model-inference-chat-completions
-    And samples at: https://github.com/Azure/azure-sdk-for-python/tree/main/sdk/ai/azure-ai-inference/samples
-
-    Returns:
-        dict[str, Any]: A dictionary with the following keys
-            message (str | dict): The content of the generated assistant message.
-                If json_mode is True, this will be a dictionary.
-            tool_names (list[str], optional): The names of the tools called by the model.
-                If the model does not support tools, a ResponseError is raised.
-            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
-            prompt_tokens (int): The number of tokens in the messages sent to the model.
-            completion_tokens (int): The number of tokens used by the model to generate the completion.
-            response_duration (float): The time, in seconds, taken to generate the response by using the model.
-            system_fingerprint (str, optional): If seed is set, a unique identifier for the model used to generate the response.
-    """
-    response_format = {"type": "json_object"} if json_mode else None
-    start_time = time.time()
-    response = client.complete(  # type: ignore
-        messages=messages,
-        model=model,
-        response_format=response_format,  # type: ignore
-        max_tokens=max_tokens,
-        temperature=temperature,
-        tools=tools,
-        seed=seed,
-    )
-    end_time = time.time()
-    response_duration = end_time - start_time
-
-    response_data = {}
-    finish_reason = response.choices[0].finish_reason
-    response_data["finish_reason"] = finish_reason.value  # type: ignore
-
-    message = response.choices[0].message.content
-    if message and json_mode:
-        with contextlib.suppress(json.JSONDecodeError):
-            message = json.loads(message)
-    response_data["message"] = message
-
-    # Check for tool calls because even if the finish_reason is stop, the model may have called a tool
-    tool_calls = response.choices[0].message.tool_calls
-    if tool_calls:
-        tool_names = []
-        tool_args_list = []
-        for tool_call in tool_calls:
-            tool_names.append(tool_call.function.name)
-            tool_args_list.append(json.loads(tool_call.function.arguments))
-        response_data["tool_names"] = tool_names
-        response_data["tool_args_list"] = tool_args_list
-
-    if seed is not None and hasattr(response, "system_fingerprint"):
-        response_data["system_fingerprint"] = response.system_fingerprint
-
-    usage = response.usage
-    if usage is not None:
-        response_data["completion_tokens"] = usage.completion_tokens
-        response_data["prompt_tokens"] = usage.prompt_tokens
-    response_data["response_duration"] = round(response_duration, 4)
-
-    return response_data

not_again_ai/llm/openai_api/chat_completion.py

@@ -1,339 +0,0 @@
-from collections.abc import Generator
-import contextlib
-import json
-import time
-from typing import Any
-
-from openai import AzureOpenAI, OpenAI
-
-
-def chat_completion(
-    messages: list[dict[str, Any]],
-    model: str,
-    client: OpenAI | AzureOpenAI | Any,
-    tools: list[dict[str, Any]] | None = None,
-    tool_choice: str = "auto",
-    max_tokens: int | None = None,
-    temperature: float = 0.7,
-    json_mode: bool = False,
-    json_schema: dict[str, Any] | None = None,
-    seed: int | None = None,
-    logprobs: tuple[bool, int | None] | None = None,
-    n: int = 1,
-    **kwargs: Any,
-) -> dict[str, Any]:
-    """Get an OpenAI chat completion response: https://platform.openai.com/docs/api-reference/chat/create
-
-    NOTE: Depending on the model, certain parameters may not be supported,
-    particularly for older vision-enabled models like gpt-4-1106-vision-preview.
-    Be sure to check the documentation: https://platform.openai.com/docs/models/gpt-4-turbo-and-gpt-4
-
-    Args:
-        messages (list): A list of messages comprising the conversation so far.
-        model (str): ID of the model to use. See the model endpoint compatibility table:
-            https://platform.openai.com/docs/models/model-endpoint-compatibility
-            for details on which models work with the Chat API.
-        client (OpenAI | AzureOpenAI | Any): An instance of the OpenAI or AzureOpenAI client.
-            If anything else is provided, we assume that it follows the OpenAI spec and call it by passing kwargs directly.
-            For example you can provide something like:
-            ```
-            def custom_client(**kwargs):
-                client = openai_client()
-                completion = client.chat.completions.create(**kwargs)
-                return completion.to_dict()
-            ```
-        tools (list[dict[str, Any]], optional):A list of tools the model may call.
-            Use this to provide a list of functions the model may generate JSON inputs for. Defaults to None.
-        tool_choice (str, optional): The tool choice to use. Can be "auto", "required", "none", or a specific function name.
-            Note the function name cannot be any of "auto", "required", or "none". Defaults to "auto".
-        max_tokens (int, optional): The maximum number of tokens to generate in the chat completion.
-            Defaults to None, which automatically limits to the model's maximum context length.
-        temperature (float, optional): What sampling temperature to use, between 0 and 2.
-            Higher values like 0.8 will make the output more random,
-            while lower values like 0.2 will make it more focused and deterministic. Defaults to 0.7.
-        json_mode (bool, optional): When JSON mode is enabled, the model is constrained to only
-            generate strings that parse into valid JSON object and will return a dictionary.
-            See https://platform.openai.com/docs/guides/text-generation/json-mode
-        json_schema (dict, optional): Enables Structured Outputs which ensures the model will
-            always generate responses that adhere to your supplied JSON Schema.
-            See https://platform.openai.com/docs/guides/structured-outputs/structured-outputs
-        seed (int, optional): If specified, OpenAI will make a best effort to sample deterministically,
-            such that repeated requests with the same `seed` and parameters should return the same result.
-            Determinism is not guaranteed, and you should refer to the `system_fingerprint` response
-            parameter to monitor changes in the backend.
-        logprobs (tuple[bool, int], optional): Whether to return log probabilities of the output tokens or not.
-            If `logprobs[0]` is true, returns the log probabilities of each output token returned in the content of message.
-            `logprobs[1]` is an integer between 0 and 5 specifying the number of most likely tokens to return at each token position,
-            each with an associated log probability. `logprobs[0]` must be set to true if this parameter is used.
-        n (int, optional): How many chat completion choices to generate for each input message.
-            Defaults to 1.
-        **kwargs: Additional keyword arguments to pass to the OpenAI client chat completion.
-
-    Returns:
-        dict[str, Any]: A dictionary with the following keys:
-            finish_reason (str): The reason the model stopped generating further tokens.
-                Can be 'stop', 'length', or 'tool_calls'.
-            tool_names (list[str], optional): The names of the tools called by the model.
-            tool_args_list (list[dict], optional): The arguments of the tools called by the model.
-            message (str | dict): The content of the generated assistant message.
-                If json_mode is True, this will be a dictionary.
-            logprobs (list[dict[str, Any] | list[dict[str, Any]]]): If logprobs[1] is between 1 and 5, each element in the list
-                will be a list of dictionaries containing the token, logprob, and bytes for the top `logprobs[1]` logprobs. Otherwise,
-                this will be a list of dictionaries containing the token, logprob, and bytes for each token in the message.
-            choices (list[dict], optional): A list of chat completion choices if n > 1 where each dict contains the above fields.
-            completion_tokens (int): The number of tokens used by the model to generate the completion.
-                NOTE: If n > 1 this is the sum of all completions.
-            prompt_tokens (int): The number of tokens in the messages sent to the model.
-            system_fingerprint (str, optional): If seed is set, a unique identifier for the model used to generate the response.
-            response_duration (float): The time, in seconds, taken to generate the response from the API.
-    """
-
-    if json_mode and json_schema is not None:
-        raise ValueError("json_schema and json_mode cannot be used together.")
-
-    if json_mode:
-        response_format: dict[str, Any] = {"type": "json_object"}
-    elif json_schema is not None:
-        if isinstance(json_schema, dict):
-            response_format = {"type": "json_schema", "json_schema": json_schema}
-    else:
-        response_format = {"type": "text"}
-
-    kwargs.update(
-        {
-            "messages": messages,
-            "model": model,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-            "response_format": response_format,
-            "n": n,
-        }
-    )
-
-    if tools is not None:
-        kwargs["tools"] = tools
-        if tool_choice not in ["none", "auto", "required"]:
-            kwargs["tool_choice"] = {"type": "function", "function": {"name": tool_choice}}
-        else:
-            kwargs["tool_choice"] = tool_choice
-
-    if seed is not None:
-        kwargs["seed"] = seed
-
-    if logprobs is not None:
-        kwargs["logprobs"] = logprobs[0]
-        if logprobs[0] and logprobs[1] is not None:
-            kwargs["top_logprobs"] = logprobs[1]
-
-    start_time = time.time()
-    if isinstance(client, OpenAI | AzureOpenAI):
-        response = client.chat.completions.create(**kwargs)
-        response = response.to_dict()
-    else:
-        response = client(**kwargs)
-    end_time = time.time()
-    response_duration = end_time - start_time
-
-    response_data: dict[str, Any] = {"choices": []}
-    for response_choice in response["choices"]:
-        response_data_curr = {}
-        finish_reason = response_choice["finish_reason"]
-        response_data_curr["finish_reason"] = finish_reason
-
-        # We first check for tool calls because even if the finish_reason is stop, the model may have called a tool
-        tool_calls = response_choice["message"].get("tool_calls", None)
-        if tool_calls:
-            tool_names = []
-            tool_args_list = []
-            for tool_call in tool_calls:
-                tool_names.append(tool_call["function"]["name"])
-                tool_args_list.append(json.loads(tool_call["function"]["arguments"]))
-            response_data_curr["message"] = response_choice["message"]["content"]
-            response_data_curr["tool_names"] = tool_names
-            response_data_curr["tool_args_list"] = tool_args_list
-        elif finish_reason == "stop" or finish_reason == "length":
-            message = response_choice["message"]["content"]
-            if json_mode or json_schema is not None:
-                with contextlib.suppress(json.JSONDecodeError):
-                    message = json.loads(message)
-            response_data_curr["message"] = message
-
-        if response_choice["logprobs"] and response_choice["logprobs"]["content"] is not None:
-            logprobs_list: list[dict[str, Any] | list[dict[str, Any]]] = []
-            for logprob in response_choice["logprobs"]["content"]:
-                if logprob["top_logprobs"]:
-                    curr_logprob_infos = []
-                    for top_logprob in logprob["top_logprobs"]:
-                        curr_logprob_infos.append(
-                            {
-                                "token": top_logprob["token"],
-                                "logprob": top_logprob["logprob"],
-                                "bytes": top_logprob["bytes"],
-                            }
-                        )
-                    logprobs_list.append(curr_logprob_infos)
-                else:
-                    logprobs_list.append(
-                        {
-                            "token": logprob["token"],
-                            "logprob": logprob["logprob"],
-                            "bytes": logprob["bytes"],
-                        }
-                    )
-
-            response_data_curr["logprobs"] = logprobs_list
-        response_data["choices"].append(response_data_curr)
-
-    usage = response["usage"]
-    if usage is not None:
-        response_data["completion_tokens"] = usage["completion_tokens"]
-        response_data["prompt_tokens"] = usage["prompt_tokens"]
-
-    if seed is not None and response["system_fingerprint"] is not None:
-        response_data["system_fingerprint"] = response["system_fingerprint"]
-
-    response_data["response_duration"] = round(response_duration, 4)
-
-    if len(response_data["choices"]) == 1:
-        response_data.update(response_data["choices"][0])
-        del response_data["choices"]
-
-    return response_data
-
-
-def chat_completion_stream(
-    messages: list[dict[str, Any]],
-    model: str,
-    client: OpenAI | AzureOpenAI | Any,
-    tools: list[dict[str, Any]] | None = None,
-    tool_choice: str = "auto",
-    max_tokens: int | None = None,
-    temperature: float = 0.7,
-    seed: int | None = None,
-    **kwargs: Any,
-) -> Generator[dict[str, Any], None, None]:
-    """Stream a chat completion from the OpenAI API.
-
-    Args:
-        messages (list[dict[str, Any]]): The messages to send to the model.
-        model (str): The model to use for the chat completion.
-        client (OpenAI | AzureOpenAI | Any): The client to use to send the request.
-            If anything else is provided, we assume that it follows the OpenAI spec and call it by passing kwargs directly.
-            For example you can provide something like:
-            ```
-            def custom_client(**kwargs) -> Generator[dict[str, Any], None, None]:  # type: ignore
-                client = openai_client()
-                completion = client.chat.completions.create(**kwargs)
-                for chunk in completion:
-                    yield chunk.to_dict()
-            ```
-        tools (list[dict[str, Any]], optional):A list of tools the model may call.
-            Use this to provide a list of functions the model may generate JSON inputs for. Defaults to None.
-        tool_choice (str, optional): The tool choice to use. Can be "auto", "required", "none", or a specific function name.
-            Note the function name cannot be any of "auto", "required", or "none". Defaults to "auto".
-        max_tokens (int | None): The maximum number of tokens to generate.
-        temperature (float): The temperature to use for the chat completion.
-        seed (int, optional): If specified, OpenAI will make a best effort to sample deterministically,
-            such that repeated requests with the same `seed` and parameters should return the same result.
-            Does not currently return `system_fingerprint`.
-
-    Returns:
-        Generator[dict[str, Any], None, None]: A generator of chunks of the chat completion.
-        Each chunk is a dictionary with the following keys:
-            role (str): The role of the chunk. Can be "assistant", "tool", or "usage".
-            content (str): The content of the chunk.
-            tool_name (str | None): The name of the tool called by the model.
-            tool_call_id (str | None): The ID of the tool call.
-            completion_tokens (int | None): The number of tokens used by the model to generate the completion.
-            prompt_tokens (int | None): The number of tokens in the messages sent to the model.
-    """
-
-    class ChatCompletionStreamParser:
-        def __init__(self) -> None:
-            # Remembers if we are currently streaming an assistant message or tool call
-            self.last_type: str = ""
-            self.last_tool_name: str | None = None
-            self.last_tool_call_id: str | None = None
-
-        def process_chunk(self, chunk: dict[str, Any]) -> dict[str, Any] | None:
-            """Convert the current chunk into a more digestible format
-            {
-                "role": Literal["assistant", "tool", "usage"],
-                "content": str,
-                "tool_name": str | None,
-                "tool_call_id": str | None,
-                "completion_tokens": int | None,
-                "prompt_tokens": int | None,
-            }
-            """
-            processed_chunk: dict[str, Any] = {}
-            if chunk["choices"]:
-                choice = chunk["choices"][0]
-                # This checks if its just a regular message currently being streamed
-                if choice["delta"].get("role", "") and choice["delta"].get("tool_calls", None) is None:
-                    if choice["delta"]["role"] != self.last_type:
-                        self.last_type = choice["delta"]["role"]
-                        processed_chunk["role"] = self.last_type
-                        if not choice["delta"]["content"]:
-                            processed_chunk["content"] = ""
-                        else:
-                            processed_chunk["content"] = choice["delta"]["content"]
-                    else:
-                        processed_chunk["role"] = self.last_type
-                elif choice["delta"].get("tool_calls", None):
-                    # tool_calls will always be present if the model is calling a tool
-                    tool_call = choice["delta"]["tool_calls"][0]
-                    if tool_call["function"].get("name"):
-                        self.last_type = "tool"
-                        self.last_tool_name = tool_call["function"]["name"]
-                        self.last_tool_call_id = tool_call["id"]
-                    processed_chunk["role"] = "tool"
-                    processed_chunk["content"] = tool_call["function"]["arguments"]
-                    processed_chunk["tool_name"] = self.last_tool_name
-                    processed_chunk["tool_call_id"] = self.last_tool_call_id
-                elif choice["delta"].get("content", ""):
-                    # This is the case after the first regular assistant message
-                    processed_chunk["role"] = self.last_type
-                    processed_chunk["content"] = choice["delta"]["content"]
-            else:
-                if chunk.get("usage"):
-                    processed_chunk["role"] = "usage"
-                    processed_chunk["completion_tokens"] = chunk["usage"]["completion_tokens"]
-                    processed_chunk["prompt_tokens"] = chunk["usage"]["prompt_tokens"]
-                else:
-                    return None
-            return processed_chunk
-
-    kwargs.update(
-        {
-            "messages": messages,
-            "model": model,
-            "max_tokens": max_tokens,
-            "temperature": temperature,
-            "stream": True,
-            "stream_options": {"include_usage": True},
-        }
-    )
-
-    if tools is not None:
-        kwargs["tools"] = tools
-        if tool_choice not in ["none", "auto", "required"]:
-            kwargs["tool_choice"] = {"type": "function", "function": {"name": tool_choice}}
-        else:
-            kwargs["tool_choice"] = tool_choice
-
-    if seed is not None:
-        kwargs["seed"] = seed
-
-    if isinstance(client, OpenAI | AzureOpenAI):
-        response = client.chat.completions.create(**kwargs)
-    else:
-        response = client(**kwargs)
-
-    parser = ChatCompletionStreamParser()
-    for chunk in response:
-        if isinstance(client, OpenAI | AzureOpenAI):
-            chunk = chunk.to_dict()
-        processed_chunk = parser.process_chunk(chunk)
-        if processed_chunk:
-            yield processed_chunk

not_again_ai/llm/openai_api/context_management.py

@@ -1,70 +0,0 @@
-import copy
-
-from not_again_ai.llm.openai_api.tokens import load_tokenizer, num_tokens_from_messages, truncate_str
-
-
-def _inject_variable(
-    messages_unformatted: list[dict[str, str]], variable_name: str, variable_text: str
-) -> list[dict[str, str]]:
-    """Injects variables into the messages using Python string formatting."""
-    messages_formatted = copy.deepcopy(messages_unformatted)
-    for message in messages_formatted:
-        message["content"] = message["content"].replace("{{" + variable_name + "}}", variable_text)
-    return messages_formatted
-
-
-def priority_truncation(
-    messages_unformatted: list[dict[str, str]],
-    variables: dict[str, str],
-    priority: list[str],
-    token_limit: int,
-    model: str = "gpt-3.5-turbo-0125",
-) -> list[dict[str, str]]:
-    """Formats messages_unformatted and injects variables into the messages in the order of priority, truncating the messages to fit the token limit.
-
-    Algorithm:
-        0. Checks if all variables in the priority list are in the variables dict. If not, adds the missing variables into priority in any order.
-        1. Iterating over priority:
-            a. Count the current number of tokens in messages_formatted and compute how many tokens remain.
-            b. Count the number of times the variable occurs in messages_formatted.
-            c. Truncate the variable to fit the remaining token budget taking into account the number of times it occurs in the messages.
-            d. Inject the variable text into messages_formatted.
-
-    Args:
-        messages_unformatted: A list of dictionaries where each dictionary
-            represents a message. Each message must have 'role' and 'content'
-            keys with string values, where content is a string with any number of occurrences of {{variable_name}}.
-        variables: A dictionary where each key-value pair represents a variable name and its value to inject.
-        priority: A list of variable names in their order of priority.
-        token_limit: The maximum number of tokens allowed in the messages.
-        model: The model to use for tokenization. Defaults to "gpt-3.5-turbo-0125".
-    """
-    tokenizer = load_tokenizer(model)
-
-    # Check if all variables in the priority list are in the variables dict.
-    # If not, add the missing variables into priority in any order.
-    for var in variables:
-        if var not in priority:
-            priority.append(var)
-
-    messages_formatted = copy.deepcopy(messages_unformatted)
-    for var in priority:
-        # Count the current number of tokens in messages_formatted and compute a remaining token budget.
-        tokenizer = load_tokenizer(model)
-        num_tokens = num_tokens_from_messages(messages_formatted, tokenizer=tokenizer, model=model)
-        remaining_tokens = token_limit - num_tokens
-        if remaining_tokens <= 0:
-            break
-
-        # Count the number of times the variable occurs in messages_formatted (including within the same message).
-        num_var_occurrences = 0
-        for message in messages_formatted:
-            num_var_occurrences += message["content"].count("{{" + var + "}}")
-
-        # Truncate the variable to fit the remaining token budget taking into account the number of times it occurs in the messages.
-        truncated_var = truncate_str(variables[var], remaining_tokens // num_var_occurrences, tokenizer=tokenizer)
-
-        # Inject the variable text into messages_formatted.
-        messages_formatted = _inject_variable(messages_formatted, var, truncated_var)
-
-    return messages_formatted

not_again_ai/llm/openai_api/embeddings.py

@@ -1,62 +0,0 @@
-from typing import Any
-
-from openai import OpenAI
-
-
-def embed_text(
-    text: str | list[str],
-    client: OpenAI,
-    model: str = "text-embedding-3-large",
-    dimensions: int | None = None,
-    encoding_format: str = "float",
-    **kwargs: Any,
-) -> list[float] | str | list[list[float]] | list[str]:
-    """Generates an embedding vector for a given text using OpenAI's API.
-
-    Args:
-        text (str | list[str]): The input text to be embedded. Each text should not exceed 8191 tokens, which is the max for V2 and V3 models
-        client (OpenAI): The OpenAI client used to interact with the API.
-        model (str, optional): The ID of the model to use for embedding.
-            Defaults to "text-embedding-3-large".
-            Choose from text-embedding-3-small, text-embedding-3-large, text-embedding-ada-002.
-            See https://platform.openai.com/docs/models/embeddings for more details.
-        dimensions (int | None, optional): The number of dimensions for the output embeddings.
-            This is only supported in "text-embedding-3" and later models. Defaults to None.
-        encoding_format (str, optional): The format for the returned embeddings. Can be either "float" or "base64".
-            Defaults to "float".
-
-    Returns:
-        list[float] | str | list[list[float]] | list[str]: The embedding vector represented as a list of floats or base64 encoded string.
-            If multiple text inputs are provided, a list of embedding vectors is returned.
-            The length and format of the vector depend on the model, encoding_format, and dimensions.
-
-    Raises:
-        ValueError: If 'text-embedding-ada-002' model is used and dimensions are specified,
-            as this model does not support specifying dimensions.
-
-    Example:
-        client = OpenAI()
-        embedding = embed_text("Example text", client, model="text-embedding-ada-002")
-    """
-    if model == "text-embedding-ada-002" and dimensions:
-        # text-embedding-ada-002 does not support dimensions
-        raise ValueError("text-embedding-ada-002 does not support dimensions")
-
-    kwargs = {
-        "model": model,
-        "input": text,
-        "encoding_format": encoding_format,
-    }
-    if dimensions:
-        kwargs["dimensions"] = dimensions
-
-    response = client.embeddings.create(**kwargs)
-
-    responses = []
-    for embedding in response.data:
-        responses.append(embedding.embedding)
-
-    if len(responses) == 1:
-        return responses[0]
-
-    return responses

not_again_ai/llm/openai_api/openai_client.py

@@ -1,78 +0,0 @@
-from typing import Literal
-
-from azure.identity import DefaultAzureCredential, get_bearer_token_provider
-from openai import AzureOpenAI, OpenAI
-
-
-class InvalidOAIAPITypeError(Exception):
-    """Raised when an invalid OAIAPIType string is provided."""
-
-    pass
-
-
-def openai_client(
-    api_type: Literal["openai", "azure_openai"] = "openai",
-    api_key: str | None = None,
-    organization: str | None = None,
-    aoai_api_version: str = "2024-06-01",
-    azure_endpoint: str | None = None,
-    timeout: float | None = None,
-    max_retries: int | None = None,
-) -> OpenAI | AzureOpenAI:
-    """Create an OpenAI or Azure OpenAI client instance based on the specified API type and other provided parameters.
-
-    Azure OpenAI requires RBAC authentication. You must be signed in with the Azure CLI and have correct role assigned.
-    See https://techcommunity.microsoft.com/t5/microsoft-developer-community/using-keyless-authentication-with-azure-openai/ba-p/4111521
-
-    Args:
-        api_type (str, optional): Type of the API to be used. Accepted values are 'openai' or 'azure_openai'.
-            Defaults to 'openai'.
-        api_key (str, optional): The API key to authenticate the client. If not provided,
-            OpenAI automatically uses `OPENAI_API_KEY` from the environment.
-        organization (str, optional): The ID of the organization. If not provided,
-            OpenAI automotically uses `OPENAI_ORG_ID` from the environment.
-        aoai_api_version (str, optional): Only applicable if using Azure OpenAI https://learn.microsoft.com/azure/ai-services/openai/reference#rest-api-versioning
-        azure_endpoint (str, optional): The endpoint to use for Azure OpenAI.
-            If not provided, will be read from the `AZURE_OPENAI_ENDPOINT` environment variable.
-        timeout (float, optional): By default requests time out after 10 minutes.
-        max_retries (int, optional): Certain errors are automatically retried 2 times by default,
-            with a short exponential backoff. Connection errors (for example, due to a network connectivity problem),
-            408 Request Timeout, 409 Conflict, 429 Rate Limit, and >=500 Internal errors are all retried by default.
-
-    Returns:
-        OpenAI: An instance of the OpenAI client.
-
-    Raises:
-        InvalidOAIAPITypeError: If an invalid API type string is provided.
-        NotImplementedError: If the specified API type is recognized but not yet supported (e.g., 'azure_openai').
-
-    Examples:
-        >>> client = openai_client(api_type="openai", api_key="YOUR_API_KEY")
-    """
-    if api_type not in ["openai", "azure_openai"]:
-        raise InvalidOAIAPITypeError(f"Invalid OAIAPIType: {api_type}. Must be 'openai' or 'azure_openai'.")
-
-    if api_type == "openai":
-        args = {
-            "api_key": api_key,
-            "organization": organization,
-            "timeout": timeout,
-            "max_retries": max_retries,
-        }
-        # Remove any None values in order to use the default values.
-        filtered_args = {k: v for k, v in args.items() if v is not None}
-        return OpenAI(**filtered_args)  # type: ignore
-    elif api_type == "azure_openai":
-        azure_credential = DefaultAzureCredential()
-        ad_token_provider = get_bearer_token_provider(azure_credential, "https://cognitiveservices.azure.com/.default")
-        args = {
-            "api_version": aoai_api_version,
-            "azure_endpoint": azure_endpoint,
-            "azure_ad_token_provider": ad_token_provider,  # type: ignore
-            "timeout": timeout,
-            "max_retries": max_retries,
-        }
-        filtered_args = {k: v for k, v in args.items() if v is not None}
-        return AzureOpenAI(**filtered_args)  # type: ignore
-    else:
-        raise NotImplementedError(f"API type '{api_type}' is invalid.")