diff --git a/README.md b/README.md
index db5ad6b..a538249 100644
--- a/README.md
+++ b/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).
+
+
+
+
## ✨ Features
-- 📝 Create documentation comment blocks for all methods in a file
+- 📝 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
+- ✍️ Generate inline documentation comments in method bodies
+- 🌳 Treesitter integration
+- 💻 Local LLM support
+- 🌐 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 `` with GPT-3.5-Turbo model:
```
aicomments
@@ -51,10 +56,10 @@ Use GPT-4 model (Default is GPT-3.5):
aicomments --gpt4
```
-Guided mode, confirm documentation generation for each method:
+Use Azure OpenAI:
```
-aicomments --guided
+aicomments --azure-deployment
```
Use a local LLM on your machine:
@@ -63,9 +68,15 @@ Use a local LLM on your machine:
aicomments --local_model
```
+Guided mode, confirm documentation generation for each method:
+
+```
+aicomments --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 --local_model
## 🔧 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=
-```
-
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 =
+```
+
+### 2. Azure OpenAI usage
+
+Add the following variables to your environment:
+
+```bash
+export AZURE_API_BASE = "https://
+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`.
diff --git a/doc_comments_ai/app.py b/doc_comments_ai/app.py
index 6c07780..1968559 100644
--- a/doc_comments_ai/app.py
+++ b/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 = {}
diff --git a/doc_comments_ai/llm.py b/doc_comments_ai/llm.py
index ef99a0f..5a6fde8 100644
--- a/doc_comments_ai/llm.py
+++ b/doc_comments_ai/llm.py
@@ -21,6 +21,7 @@ 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:
@@ -28,13 +29,19 @@ def __init__(
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"
diff --git a/doc_comments_ai/utils.py b/doc_comments_ai/utils.py
index f23256e..37a089b 100644
--- a/doc_comments_ai/utils.py
+++ b/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.")
diff --git a/pyproject.toml b/pyproject.toml
index 73b15ea..ae7a586 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "doc-comments-ai"
-version = "0.1.10"
+version = "0.1.11"
description = ""
authors = ["fynnfluegge "]
readme = "README.md"