feat: integrate Azure OpenAI support #221
Conversation
…pdate documentation - Added AzureAIClient for interacting with Azure OpenAI models. - Updated README.md to include Azure OpenAI API key, endpoint, and version setup instructions. - Modified configuration to support Azure as a model provider. - Enhanced chat completion handling to include Azure API calls in websocket and streaming contexts. - Updated generator configuration to define Azure models and their parameters.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @mmlmt2604, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces comprehensive support for Azure OpenAI models across the application. It includes the development of a new client for interacting with Azure's API, updates to the core configuration to register Azure as a model provider, and integration into both the standard chat completion and real-time websocket functionalities. Additionally, the documentation has been thoroughly updated to guide users through the setup and usage of Azure OpenAI.
Highlights
- New Azure OpenAI Client: A dedicated
AzureAIClienthas been introduced to facilitate interactions with Azure OpenAI models, supporting both API key and Azure AD token authentication for chat completions and embeddings. - Configuration Integration: The system's configuration (
api/config.pyandapi/config/generator.json) has been updated to recognize Azure as a new model provider, defining its client class and supported models (e.g.,gpt-4o,gpt-4,gpt-35-turbo) along with their default parameters. - API and WebSocket Support: Azure OpenAI models are now integrated into the
simple_chatandwebsocket_wikiAPIs, enabling chat completions, including streaming responses, and incorporating robust error handling for Azure API calls. - Comprehensive Documentation: The
README.mdhas been extensively updated with instructions for setting up Azure OpenAI environment variables, obtaining credentials, and troubleshooting common API errors, along with updates to diagrams and model provider lists.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request integrates Azure OpenAI, including a new AzureAIClient, updates to configuration, documentation, and integration into the chat completion workflows. The README has been updated with setup instructions. The AzureAIClient implementation is functional but has a few areas for improvement regarding input processing, response parsing, state management, and coding conventions.
| def parse_chat_completion( | ||
| self, | ||
| completion: Union[ChatCompletion, Generator[ChatCompletionChunk, None, None]], | ||
| ) -> "GeneratorOutput": | ||
| """Parse the completion, and put it into the raw_response.""" | ||
| log.debug(f"completion: {completion}, parser: {self.chat_completion_parser}") | ||
| try: | ||
| data = self.chat_completion_parser(completion) | ||
| usage = self.track_completion_usage(completion) | ||
| return GeneratorOutput( | ||
| data=None, error=None, raw_response=data, usage=usage | ||
| ) | ||
| except Exception as e: | ||
| log.error(f"Error parsing the completion: {e}") | ||
| return GeneratorOutput(data=None, error=str(e), raw_response=completion) |
There was a problem hiding this comment.
- The
GeneratorOutputis returned withdata=None, while the parseddatais put intoraw_response. The parsed data should be inGeneratorOutput.data. - If
track_completion_usageis called with a streaming generator, it raisesNotImplementedError. This error is then caught here, andGeneratorOutputis returned with an error message. This means that for streaming responses, theGeneratorOutputwill always indicate an error related to usage tracking and won't contain the actual streamed data in thedatafield.
| if self._input_type == "messages": | ||
| system_start_tag = "<START_OF_SYSTEM_PROMPT>" | ||
| system_end_tag = "<END_OF_SYSTEM_PROMPT>" | ||
| user_start_tag = "<START_OF_USER_PROMPT>" | ||
| user_end_tag = "<END_OF_USER_PROMPT>" | ||
| pattern = f"{system_start_tag}(.*?){system_end_tag}{user_start_tag}(.*?){user_end_tag}" | ||
| # Compile the regular expression | ||
| regex = re.compile(pattern) | ||
| # Match the pattern | ||
| match = regex.search(input) | ||
| system_prompt, input_str = None, None | ||
|
|
||
| if match: | ||
| system_prompt = match.group(1) | ||
| input_str = match.group(2) | ||
|
|
||
| else: | ||
| print("No match found.") | ||
| if system_prompt and input_str: | ||
| messages.append({"role": "system", "content": system_prompt}) | ||
| messages.append({"role": "user", "content": input_str}) | ||
| if len(messages) == 0: | ||
| messages.append({"role": "system", "content": input}) | ||
| final_model_kwargs["messages"] = messages |
There was a problem hiding this comment.
The logic for handling input_type seems to incorrectly treat default text input as a system message. Consider revising this to correctly assign roles based on _input_type and handle regex mismatches more explicitly for the "messages" type.
| @@ -365,6 +386,9 @@ You can also mount a .env file to the container: | |||
| echo "GOOGLE_API_KEY=your_google_api_key" > .env | |||
| echo "OPENAI_API_KEY=your_openai_api_key" >> .env | |||
| echo "OPENROUTER_API_KEY=your_openrouter_api_key" >> .env | |||
| echo "AZURE_OPENAI_API_KEY=your_azure_openai_api_key" >> .env | |||
| echo "AZURE_OPENAI_ENDPOINT=your_azure_openai_endpoint" >> .env | |||
| echo "AZURE_OPENAI_VERSION=your_azure_openai_version" >>.env | |||
There was a problem hiding this comment.
Typo in the echo command. It should be >> .env to correctly append to the .env file.
| echo "AZURE_OPENAI_VERSION=your_azure_openai_version" >>.env | |
| echo "AZURE_OPENAI_VERSION=your_azure_openai_version" >> .env |
|
|
||
| __all__ = ["AzureAIClient"] | ||
|
|
||
| # TODO: this overlaps with openai client largely, might need to refactor to subclass openai client to simplify the code |
There was a problem hiding this comment.
Consider refactoring to reduce code duplication with OpenAIClient. A common base class would improve maintainability.
| log_probs = [] | ||
| for c in completion.choices: | ||
| content = c.logprobs.content | ||
| print(content) |
There was a problem hiding this comment.
Remove this print statement or replace it with log.debug().
log.debug(f"Logprobs content: {content}") # Or remove if not needed| # added api_type azure for azure Ai | ||
| self.api_type = "azure" | ||
| self._api_key = api_key | ||
| self._apiversion = api_version |
There was a problem hiding this comment.
Variable name _apiversion should be _api_version to adhere to PEP 8.
self._api_version = api_version| input_str = match.group(2) | ||
|
|
||
| else: | ||
| print("No match found.") |
There was a problem hiding this comment.
Remove this print statement or replace it with appropriate logging.
log.warning("Input string did not match the expected system/user prompt pattern when input_type is 'messages'.")| if "stream" in api_kwargs and api_kwargs.get("stream", False): | ||
| log.debug("streaming call") | ||
| self.chat_completion_parser = handle_streaming_response |
There was a problem hiding this comment.
Modifying self.chat_completion_parser within the call method introduces a side effect that changes the client's state. It's generally better to pass the appropriate parser to parse_chat_completion or have parse_chat_completion determine the parser based on the type of completion object it receives, rather than relying on instance state modified by call parameters.
| obj = super().from_dict(data) | ||
| # recreate the existing clients | ||
| obj.sync_client = obj.init_sync_client() | ||
| obj.async_client = obj.init_async_client() |
There was a problem hiding this comment.
Initialize obj.async_client to None in from_dict to maintain lazy initialization.
obj.async_client = None # Maintain lazy initialization
Client Implementation and update documentation
Added AzureAIClient for interacting with Azure OpenAI models.
Updated README.md to include Azure OpenAI API key, endpoint, and version setup instructions.
Modified configuration to support Azure as a model provider.
Enhanced chat completion handling to include Azure API calls in websocket and streaming contexts.
Updated generator configuration to define Azure models and their parameters.