feat: add support for Azure OpenAI

README.md

@@ -16,24 +16,29 @@ With just a few keystrokes in your terminal by using the OpenAI API or 100% loca
 
 Built with [langchain](https://github.com/langchain-ai/langchain), [lama.cpp](https://github.com/ggerganov/llama.cpp) and [treesitter](https://github.com/tree-sitter/tree-sitter).
 
+<kbd>
+
 ![ezgif-4-53d6e634af](https://github.com/fynnfluegge/doc-comments.ai/assets/16321871/8f2756cb-36f9-43c6-94b1-658b89b49786)
 
+</kbd>
+
 </div>
 
 ## ✨ Features
 
-- 📝 Create documentation comment blocks for all methods in a file
+- 📝 &nbsp;Generate documentation comment blocks for all methods in a file
   - e.g. Javadoc, JSDoc, Docstring, Rustdoc etc.
-- ✍️ Create inline documentation comments in method bodies
-- 🌳 Treesitter integration
-- 💻 Local LLM support
+- ✍️ &nbsp; Generate inline documentation comments in method bodies
+- 🌳&nbsp; Treesitter integration
+- 💻&nbsp; Local LLM support
+- 🌐&nbsp; Azure OpenAI support
 
 > [!NOTE]  
 > Documentations will only be added to files without unstaged changes, so nothing is overwritten.
 
 ## 🚀 Usage
 
-Create documentations for any method in the file with GPT-3.5 Turbo model:
+Create documentations for any method in a file specified by `<RELATIVE_FILE_PATH>` with GPT-3.5-Turbo model:
 
 ```
 aicomments <RELATIVE_FILE_PATH>
@@ -51,10 +56,10 @@ Use GPT-4 model (Default is GPT-3.5):
 aicomments <RELATIVE_FILE_PATH> --gpt4
 ```
 
-Guided mode, confirm documentation generation for each method:
+Use Azure OpenAI:
 
 ```
-aicomments <RELATIVE_FILE_PATH> --guided
+aicomments <RELATIVE_FILE_PATH> --azure-deployment <DEPLOYMENT_NAME>
 ```
 
 Use a local LLM on your machine:
@@ -63,9 +68,15 @@ Use a local LLM on your machine:
 aicomments <RELATIVE_FILE_PATH> --local_model <MODEL_PATH>
 ```
 
+Guided mode, confirm documentation generation for each method:
+
+```
+aicomments <RELATIVE_FILE_PATH> --guided
+```
+
 > [!NOTE]  
 > How to download models from huggingface for local usage see [Local LLM usage](README.md#2-local-llm-usage)
- 
+
 > [!IMPORTANT]  
 > The results by using a local LLM will highly be affected by your selected model. To get similar results compared to GPT-3.5/4 you need to select very large models which require a powerful hardware.
 
@@ -88,14 +99,6 @@ aicomments <RELATIVE_FILE_PATH> --local_model <MODEL_PATH>
 
 ## 🔧 Installation
 
-### 1. OpenAI API usage
-
-Create your personal OpenAI API key and add it as `$OPENAI_API_KEY` to your environment with:
-
-```
-export OPENAI_API_KEY=<YOUR_API_KEY>
-```
-
 Install with `pipx`:
 
 ```
@@ -104,16 +107,36 @@ pipx install doc-comments-ai
 
 > It is recommended to use `pipx` for installation, nonetheless it is also possible to use `pip`.
 
-### 2. Local LLM usage
+### 1. OpenAI usage
+
+Create your personal OpenAI API key and add it as `$OPENAI_API_KEY` to your environment with:
+
+```bash
+export OPENAI_API_KEY = <YOUR_API_KEY>
+```
+
+### 2. Azure OpenAI usage
+
+Add the following variables to your environment:
+
+```bash
+export AZURE_API_BASE = "https://<your-endpoint.openai.azure.com/"
+export AZURE_API_KEY = <YOUR_AZURE_OPENAI_API_KEY>
+export AZURE_API_VERSION = "2023-05-15"
+```
+
+### 3. Local LLM usage
 
 By using a local LLM no API key is required. On first usage of `--local_model` you will be asked for confirmation to intall `llama-cpp-python` with its dependencies.
 The installation process will take care of the hardware-accelerated build tailored to your hardware and OS. For further details see:
 [installation-with-hardware-acceleration](https://github.com/abetlen/llama-cpp-python#installation-with-hardware-acceleration)
 
 To download a model from huggingface for local usage the most convenient way is using the `huggingface-cli`:
+
 ```
 huggingface-cli download TheBloke/CodeLlama-13B-Python-GGUF codellama-13b-python.Q5_K_M.gguf
 ```
+
 This will download the `codellama-13b-python.Q5_K_M` model to `~/.cache/huggingface/`.
 After the download has finished the absolute path of the `.gguf` file is printed to the console which can be used as the value for `--local_model`.
 

doc_comments_ai/app.py

@@ -10,14 +10,6 @@
 
 
 def run():
-    """
-    This is the entry point of the application
-    """
-    api_key = os.environ.get("OPENAI_API_KEY")
-
-    if not api_key:
-        sys.exit("OPENAI_API_KEY not found.")
-
     parser = argparse.ArgumentParser()
     parser.add_argument("dir", nargs="?", default=os.getcwd())
     parser.add_argument(
@@ -40,6 +32,11 @@ def run():
         action="store_true",
         help="User will get asked to confirm the doc generation for each method.",
     )
+    parser.add_argument(
+        "--azure-deployment",
+        type=str,
+        help="Azure OpenAI deployment name.",
+    )
 
     if sys.argv.__len__() < 2:
         sys.exit("Please provide a file")
@@ -54,9 +51,15 @@ def run():
     if utils.has_unstaged_changes(file_name):
         sys.exit(f"File {utils.get_bold_text(file_name)} has unstaged changes")
 
+    if args.azure_deployment:
+        utils.is_azure_openai_environment_available()
+        llm_wrapper = llm.LLM(azure_deployment=args.azure_deployment)
+
     if args.gpt4:
+        utils.is_openai_api_key_available()
         llm_wrapper = llm.LLM(model=GptModel.GPT_4)
     else:
+        utils.is_openai_api_key_available()
         llm_wrapper = llm.LLM(local_model=args.local_model)
 
     generated_doc_comments = {}

doc_comments_ai/llm.py

@@ -21,20 +21,27 @@ def __init__(
         self,
         model: GptModel = GptModel.GPT_35,
         local_model: str | None = None,
+        azure_deployment: str | None = None,
     ):
         max_tokens = 2048 if model == GptModel.GPT_35 else 4096
         if local_model is not None:
             self.install_llama_cpp()
 
             self.llm = LlamaCpp(
                 model_path=local_model,
-                temperature=0.9,
+                temperature=0.8,
                 max_tokens=max_tokens,
                 verbose=False,
             )
+        elif azure_deployment is not None:
+            self.llm = ChatLiteLLM(
+                temperature=0.8,
+                max_tokens=max_tokens,
+                model=f"azure/{azure_deployment}",
+            )
         else:
             self.llm = ChatLiteLLM(
-                temperature=0.9, max_tokens=max_tokens, model=model.value
+                temperature=0.8, max_tokens=max_tokens, model=model.value
             )
         self.template = (
             "Add a detailed doc comment to the following {language} method:\n{code}\n"

doc_comments_ai/utils.py

@@ -1,6 +1,7 @@
 import os
 import re
 import subprocess
+import sys
 
 import tiktoken
 
@@ -135,3 +136,39 @@ def count_tokens(text):
     encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
     tokenized = encoding.encode(text)
     return len(tokenized)
+
+
+def is_openai_api_key_available():
+    """
+    Checks if the OpenAI API key is available in the environment variables.
+
+    Returns:
+        bool: True if the OpenAI API key is available, False otherwise.
+    """
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        sys.exit("OPENAI_API_KEY not found.")
+
+
+def is_azure_openai_environment_available():
+    """
+    This method checks if the Azure OpenAI environment is available by verifying the presence
+    of necessary environment variables: 'AZURE_API_BASE', 'AZURE_API_KEY' and 'AZURE_API_VERSION'.
+
+    The method retrieves these environment variables and checks if they are set. If any
+    of these variables is not set, the method will print a message indicating which variable is missing
+    and then terminate the program, instructing the user to set the missing environment variables.
+
+    :raises SystemExit: If any of the required Azure OpenAI environment variables are not set.
+    """
+    azure_api_base = os.environ.get("AZURE_API_BASE")
+    azure_api_key = os.environ.get("AZURE_API_KEY")
+    azure_api_version = os.environ.get("AZURE_API_VERSION")
+    if not azure_api_base or not azure_api_key or not azure_api_version:
+        if not azure_api_base:
+            print("AZURE_API_BASE not found.")
+        if not azure_api_key:
+            print("AZURE_API_KEY not found.")
+        if not azure_api_version:
+            print("AZURE_API_VERSION not found.")
+        sys.exit("Please set the environment variables for Azure OpenAI deployment.")

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "doc-comments-ai"
-version = "0.1.10"
+version = "0.1.11"
 description = ""
 authors = ["fynnfluegge <[email protected]>"]
 readme = "README.md"