azureai_client#

AzureOpenAI ModelClient integration.

Functions

get_all_messages_content(completion)

When the n > 1, get all the messages content.

get_first_message_content(completion)

When we only need the content of the first message.

get_probabilities(completion)

Get the probabilities of each token in the completion.

handle_streaming_response(generator)

Handle the streaming response.

parse_stream_response(completion)

Parse the response of the stream API.

Classes

AzureAIClient([api_key, api_version, ...])

A client wrapper for interacting with Azure OpenAI's API.
class AzureAIClient(api_key: str | None = None, api_version: str | None = None, azure_endpoint: str | None = None, credential: DefaultAzureCredential | None = None, chat_completion_parser: Callable[[Completion], Any] = None, input_type: Literal['text', 'messages'] = 'text')[source]#

Bases: ModelClient

A client wrapper for interacting with Azure OpenAI’s API.

This class provides support for both embedding and chat completion API calls. Users can use this class to simplify their interactions with Azure OpenAI models through the Embedder and Generator components.

Initialization:

You can initialize the AzureAIClient with either an API key or Azure Active Directory (AAD) token authentication. It is recommended to set environment variables for sensitive data like API keys.

Parameters:

  • api_key (Optional[str]) – Azure OpenAI API key. Default is None.

  • api_version (Optional[str]) – API version to use. Default is None.

  • azure_endpoint (Optional[str]) – Azure OpenAI endpoint URL. Default is None.

  • credential (Optional[DefaultAzureCredential]) – Azure AD credential for token-based authentication. Default is None.

  • chat_completion_parser (Callable[[Completion], Any]) – Function to parse chat completions. Default is get_first_message_content.

  • input_type (Literal["text", "messages"]) – Format for input, either “text” or “messages”. Default is “text”.

Setup Instructions:

  • Using API Key: Set up the following environment variables: `bash export AZURE_OPENAI_API_KEY="your_api_key" export AZURE_OPENAI_ENDPOINT="your_endpoint" export AZURE_OPENAI_VERSION="your_version" `

  • Using Azure AD Token: Ensure you have configured Azure AD credentials. The DefaultAzureCredential will automatically use your configured credentials.

Example Usage:

from azure.identity import DefaultAzureCredential
from your_module import AzureAIClient  # Adjust import based on your module name

# Initialize with API key
client = AzureAIClient(
    api_key="your_api_key",
    api_version="2023-05-15",
    azure_endpoint="https://your-endpoint.openai.azure.com/"
)

# Or initialize with Azure AD token
client = AzureAIClient(
    api_version="2023-05-15",
    azure_endpoint="https://your-endpoint.openai.azure.com/",
    credential=DefaultAzureCredential()
)

# Example call to the chat completion API
api_kwargs = {
    "model": "gpt-3.5-turbo",
    "messages": [{"role": "user", "content": "What is the meaning of life?"}],
    "stream": True
}
response = client.call(api_kwargs=api_kwargs, model_type=ModelType.LLM)

for chunk in response:
    print(chunk)

Notes: - Ensure that the API key or credentials are correctly set up and accessible to avoid authentication errors. - Use chat_completion_parser to define how to extract and handle the chat completion responses. - The input_type parameter determines how input is formatted for the API call.

References: - [Azure OpenAI API Documentation](https://learn.microsoft.com/en-us/azure/ai-services/openai/overview) - [OpenAI API Documentation](https://platform.openai.com/docs/guides/text-generation)

init_sync_client()[source]#
init_async_client()[source]#
parse_chat_completion(completion: ChatCompletion | Generator[ChatCompletionChunk, None, None]) → GeneratorOutput[source]#

Parse the completion, and put it into the raw_response.

track_completion_usage(completion: ChatCompletion | Generator[ChatCompletionChunk, None, None]) → CompletionUsage[source]#

Track the chat completion usage. Use OpenAI standard API for tracking.

parse_embedding_response(response: CreateEmbeddingResponse) → EmbedderOutput[source]#

Parse the embedding response to a structure LightRAG components can understand.

Should be called in Embedder.

convert_inputs_to_api_kwargs(input: Any | None = None, model_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED) → Dict[source]#

Specify the API input type and output api_kwargs that will be used in _call and _acall methods. Convert the Component’s standard input, and system_input(chat model) and model_kwargs into API-specific format

call(api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED)[source]#

kwargs is the combined input and model_kwargs. Support streaming call.

async acall(api_kwargs: Dict = {}, model_type: ModelType = ModelType.UNDEFINED)[source]#

kwargs is the combined input and model_kwargs

classmethod from_dict(data: Dict[str, Any]) → T[source]#

Create an instance from previously serialized data using to_dict() method.

to_dict() → Dict[str, Any][source]#

Convert the component to a dictionary.

On this page