From 02c29b632e59e84b1e4cf9f6999ae265ea9627fd Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Fri, 5 May 2023 17:52:36 +0530
Subject: [PATCH 1/9] azure support

---
 loopgpt/agent.py          | 10 +++++++++-
 loopgpt/models/openai_.py | 25 +++++++++++++++++--------
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/loopgpt/agent.py b/loopgpt/agent.py
index f07d8f4..4590603 100644
--- a/loopgpt/agent.py
+++ b/loopgpt/agent.py
@@ -33,8 +33,16 @@ def __init__(
         model=None,
         embedding_provider=None,
         temperature=0.8,
+        **kwargs,
     ):
-        if model is None:
+        engine = kwargs.get("engine", None)
+        if engine is not None:
+            import openai
+            if openai.api_type == "azure":
+                model = OpenAIModel(engine)
+            else:
+                raise ValueError("`engine` parameter is only supported for Azure OpenAI API. Please use `model` instead.")
+        elif model is None:
             model = OpenAIModel("gpt-3.5-turbo")
         elif isinstance(model, str):
             model = OpenAIModel(model)
diff --git a/loopgpt/models/openai_.py b/loopgpt/models/openai_.py
index 44c74fa..c4e2994 100644
--- a/loopgpt/models/openai_.py
+++ b/loopgpt/models/openai_.py
@@ -36,14 +36,23 @@ def chat(
         num_retries = 3
         for _ in range(num_retries):
             try:
-                resp = openai.ChatCompletion.create(
-                    model=self.model,
-                    messages=messages,
-                    api_key=api_key,
-                    max_tokens=max_tokens,
-                    temperature=temperature,
-                )["choices"][0]["message"]["content"]
-                return resp
+                if openai.api_type == "azure":
+                    resp = openai.ChatCompletion.create(
+                        engine=self.model,
+                        messages=messages,
+                        api_key=api_key,
+                        max_tokens=max_tokens,
+                        temperature=temperature,
+                    )["choices"][0]["message"]["content"]
+                else:
+                    resp = openai.ChatCompletion.create(
+                        model=self.model,
+                        messages=messages,
+                        api_key=api_key,
+                        max_tokens=max_tokens,
+                        temperature=temperature,
+                    )["choices"][0]["message"]["content"]
+                    return resp
 
             except RateLimitError:
                 logger.warn("Rate limit exceeded. Retrying after 20 seconds.")

From 125ccf31405d20b3e534e9eb781d6625f8c9388d Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Fri, 5 May 2023 18:00:00 +0530
Subject: [PATCH 2/9] add gpt-4-32k

---
 loopgpt/models/openai_.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/loopgpt/models/openai_.py b/loopgpt/models/openai_.py
index c4e2994..23e4b54 100644
--- a/loopgpt/models/openai_.py
+++ b/loopgpt/models/openai_.py
@@ -63,6 +63,7 @@ def count_tokens(self, messages: List[Dict[str, str]]) -> int:
         tokens_per_message, tokens_per_name = {
             "gpt-3.5-turbo": (4, -1),
             "gpt-4": (3, 1),
+            "gpt-4-32k": (3, 1),
         }[self.model]
         enc = tiktoken.encoding_for_model(self.model)
         num_tokens = 0
@@ -78,7 +79,9 @@ def count_tokens(self, messages: List[Dict[str, str]]) -> int:
     def get_token_limit(self) -> int:
         return {
             "gpt-3.5-turbo": 4000,
+            "gpt-35-turbo": 4000,
             "gpt-4": 8000,
+            "gpt-4-32k": 32000,
         }[self.model]
 
     def config(self):

From 1b416d30271877ca95751a670b5747097147a502 Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Fri, 5 May 2023 18:00:41 +0530
Subject: [PATCH 3/9] add gpt-35-turbo

---
 loopgpt/models/openai_.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/loopgpt/models/openai_.py b/loopgpt/models/openai_.py
index 23e4b54..cc9786b 100644
--- a/loopgpt/models/openai_.py
+++ b/loopgpt/models/openai_.py
@@ -62,6 +62,7 @@ def chat(
     def count_tokens(self, messages: List[Dict[str, str]]) -> int:
         tokens_per_message, tokens_per_name = {
             "gpt-3.5-turbo": (4, -1),
+            "gpt-35-turbo": (4, -1),
             "gpt-4": (3, 1),
             "gpt-4-32k": (3, 1),
         }[self.model]

From 05543f6347d967ee9d435a2da4cce6df74ca6521 Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Fri, 5 May 2023 20:05:13 +0530
Subject: [PATCH 4/9] bugfix

---
 loopgpt/models/openai_.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/loopgpt/models/openai_.py b/loopgpt/models/openai_.py
index cc9786b..b694e18 100644
--- a/loopgpt/models/openai_.py
+++ b/loopgpt/models/openai_.py
@@ -52,7 +52,7 @@ def chat(
                         max_tokens=max_tokens,
                         temperature=temperature,
                     )["choices"][0]["message"]["content"]
-                    return resp
+                return resp
 
             except RateLimitError:
                 logger.warn("Rate limit exceeded. Retrying after 20 seconds.")

From 72df463e78b4633d2b9ec1b9cec84cb863033915 Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Sun, 7 May 2023 00:26:58 +0530
Subject: [PATCH 5/9] azure embeddings support, clean up

---
 loopgpt/agent.py                   | 32 +++++++------
 loopgpt/embeddings/__init__.py     |  1 +
 loopgpt/embeddings/azure_openai.py | 28 +++++++++++
 loopgpt/embeddings/openai_.py      | 18 ++++----
 loopgpt/models/__init__.py         |  1 +
 loopgpt/models/azure_openai.py     | 74 ++++++++++++++++++++++++++++++
 loopgpt/models/openai_.py          | 46 +++++--------------
 loopgpt/utils/openai_key.py        | 15 +++++-
 8 files changed, 156 insertions(+), 59 deletions(-)
 create mode 100644 loopgpt/embeddings/azure_openai.py
 create mode 100644 loopgpt/models/azure_openai.py

diff --git a/loopgpt/agent.py b/loopgpt/agent.py
index 4590603..0b1dd1a 100644
--- a/loopgpt/agent.py
+++ b/loopgpt/agent.py
@@ -8,17 +8,18 @@
     AgentStates,
 )
 from loopgpt.memory import from_config as memory_from_config
-from loopgpt.models import OpenAIModel, from_config as model_from_config
+from loopgpt.models import OpenAIModel, AzureOpenAIModel, from_config as model_from_config
 from loopgpt.tools import builtin_tools, from_config as tool_from_config
 from loopgpt.tools.code import ai_function
 from loopgpt.memory.local_memory import LocalMemory
-from loopgpt.embeddings import OpenAIEmbeddingProvider
+from loopgpt.embeddings import OpenAIEmbeddingProvider, AzureOpenAIEmbeddingProvider
 from loopgpt.utils.spinner import spinner
 from loopgpt.loops import cli
 
 
 from typing import *
 
+import openai
 import json
 import time
 import ast
@@ -33,22 +34,23 @@ def __init__(
         model=None,
         embedding_provider=None,
         temperature=0.8,
-        **kwargs,
-    ):
-        engine = kwargs.get("engine", None)
-        if engine is not None:
-            import openai
-            if openai.api_type == "azure":
-                model = OpenAIModel(engine)
-            else:
-                raise ValueError("`engine` parameter is only supported for Azure OpenAI API. Please use `model` instead.")
-        elif model is None:
+    ):  
+        if model is None:
             model = OpenAIModel("gpt-3.5-turbo")
         elif isinstance(model, str):
-            model = OpenAIModel(model)
+            if openai.api_type == "azure":
+                model = AzureOpenAIModel(model)
+            else:
+                model = OpenAIModel(model)
 
-        if embedding_provider is None:
-            embedding_provider = OpenAIEmbeddingProvider()
+        if openai.api_type == "azure":
+            if embedding_provider is None:
+                raise ValueError("You must provide a deployed embedding provider to the `embedding_provider` argument when using the OpenAI Azure API")
+            elif isinstance(embedding_provider, str):
+                embedding_provider = AzureOpenAIEmbeddingProvider(embedding_provider)
+        else:
+            if embedding_provider is None:
+                embedding_provider = OpenAIEmbeddingProvider()
 
         self.name = name
         self.description = description
diff --git a/loopgpt/embeddings/__init__.py b/loopgpt/embeddings/__init__.py
index 0367867..1d0d1af 100644
--- a/loopgpt/embeddings/__init__.py
+++ b/loopgpt/embeddings/__init__.py
@@ -1,5 +1,6 @@
 from loopgpt.embeddings.provider import BaseEmbeddingProvider
 from loopgpt.embeddings.openai_ import OpenAIEmbeddingProvider
+from loopgpt.embeddings.azure_openai import AzureOpenAIEmbeddingProvider
 from loopgpt.embeddings.hf import HuggingFaceEmbeddingProvider
 
 user_providers = {}
diff --git a/loopgpt/embeddings/azure_openai.py b/loopgpt/embeddings/azure_openai.py
new file mode 100644
index 0000000..71b6ace
--- /dev/null
+++ b/loopgpt/embeddings/azure_openai.py
@@ -0,0 +1,28 @@
+import numpy as np
+import openai
+from loopgpt.embeddings.openai_ import OpenAIEmbeddingProvider
+from typing import Optional
+
+from loopgpt.utils.openai_key import get_openai_key
+
+class AzureOpenAIEmbeddingProvider(OpenAIEmbeddingProvider):
+    def __init__(self, deployment_id: str, api_key: Optional[str] = None):
+        self.deployment_id = deployment_id
+        self.api_key = api_key
+    
+    def get(self, text: str):
+        api_key = get_openai_key(self.api_key)
+        return np.array(
+            openai.Embedding.create(input=[text], engine=self.deployment_id, api_key=api_key)[
+                "data"
+            ][0]["embedding"],
+            dtype=np.float32,
+        )
+
+    def config(self):
+        cfg = {"deployment_id": self.deployment_id, "api_key": self.api_key}
+        return cfg
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(config["deployment_id"], config["api_key"])
\ No newline at end of file
diff --git a/loopgpt/embeddings/openai_.py b/loopgpt/embeddings/openai_.py
index a655028..3d2ace7 100644
--- a/loopgpt/embeddings/openai_.py
+++ b/loopgpt/embeddings/openai_.py
@@ -1,17 +1,19 @@
+from typing import Optional
 from loopgpt.embeddings.provider import BaseEmbeddingProvider
+from loopgpt.utils.openai_key import get_openai_key
 import numpy as np
+import openai
 
 
 class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
-    def __init__(self, model: str = "text-embedding-ada-002"):
-        super(OpenAIEmbeddingProvider, self).__init__()
+    def __init__(self, model: str = "text-embedding-ada-002", api_key: Optional[str] = None):
         self.model = model
+        self.api_key = api_key
 
     def get(self, text: str):
-        import openai
-
+        api_key = get_openai_key(self.api_key)
         return np.array(
-            openai.Embedding.create(input=[text], model="text-embedding-ada-002")[
+            openai.Embedding.create(input=[text], model="text-embedding-ada-002", api_key=api_key)[
                 "data"
             ][0]["embedding"],
             dtype=np.float32,
@@ -19,11 +21,9 @@ def get(self, text: str):
 
     def config(self):
         cfg = super().config()
-        cfg.update({"model": self.model})
+        cfg.update({"model": self.model, "api_key": self.api_key})
         return cfg
 
     @classmethod
     def from_config(cls, config):
-        obj = cls()
-        obj.model = config["model"]
-        return obj
+        return cls(config["model"], config.get("api_key"))
diff --git a/loopgpt/models/__init__.py b/loopgpt/models/__init__.py
index a55b328..e1632fc 100644
--- a/loopgpt/models/__init__.py
+++ b/loopgpt/models/__init__.py
@@ -1,6 +1,7 @@
 from loopgpt.models.stable_lm import StableLMModel
 from loopgpt.models.llama_cpp import LlamaCppModel
 from loopgpt.models.openai_ import OpenAIModel
+from loopgpt.models.azure_openai import AzureOpenAIModel
 from loopgpt.models.hf import HuggingFaceModel
 from loopgpt.models.base import *
 
diff --git a/loopgpt/models/azure_openai.py b/loopgpt/models/azure_openai.py
new file mode 100644
index 0000000..da26001
--- /dev/null
+++ b/loopgpt/models/azure_openai.py
@@ -0,0 +1,74 @@
+from typing import List, Dict, Optional
+from loopgpt.models.openai_ import OpenAIModel
+from loopgpt.utils.openai_key import get_openai_key
+from loopgpt.logger import logger
+from time import time
+
+from openai.error import RateLimitError
+import requests
+import openai
+
+def get_deployment_details(endpoint, deployment_id, api_version, api_key):
+    api_key = get_openai_key(api_key)
+    response = requests.get(
+        f"{endpoint}/openai/deployments/{deployment_id}?api-version={api_version}", headers={"api-key": api_key}
+    )
+    return response.json()
+
+def get_deployment_model(endpoint, deployment_id, api_version, api_key):
+    details = get_deployment_details(endpoint, deployment_id, api_version, api_key)
+    model = details["model"]
+
+    return {
+        "gpt-35-turbo": "gpt-3.5-turbo",
+        "gpt-4": "gpt-4",
+        "gpt-4-32k": "gpt-4-32k",
+    }[model]
+
+
+class AzureOpenAIModel(OpenAIModel):
+    def __init__(self, deployment_id: str, api_key: Optional[str] = None):
+        # sanity check
+        assert openai.api_type == "azure", "AzureOpenAIModel can only be used with Azure API"
+
+        self.deployment_id = deployment_id
+        self.api_key = api_key
+        self.endpoint = openai.api_base
+        self.api_version = openai.api_version
+        self.model = get_deployment_model(self.endpoint, self.deployment_id, self.api_version, self.api_key)
+
+    def chat(
+        self,
+        messages: List[Dict[str, str]],
+        max_tokens: Optional[int] = None,
+        temperature: float = 0.8,
+    ) -> str:
+
+        api_key = get_openai_key(self.api_key)
+        num_retries = 3
+        for _ in range(num_retries):
+            try:
+                resp = openai.ChatCompletion.create(
+                    engine=self.deployment_id,
+                    messages=messages,
+                    api_key=api_key,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                )["choices"][0]["message"]["content"]
+                return resp
+
+            except RateLimitError:
+                logger.warn("Rate limit exceeded. Retrying after 20 seconds.")
+                time.sleep(20)
+                continue
+    
+    def config(self):
+        cfg = super().config()
+        cfg.update({
+            "deployment_id": self.deployment_id,
+        })
+        return cfg
+    
+    @classmethod
+    def from_config(cls, config):
+        return cls(config["deployment_id"], config.get("api_key"))
diff --git a/loopgpt/models/openai_.py b/loopgpt/models/openai_.py
index b694e18..b73879e 100644
--- a/loopgpt/models/openai_.py
+++ b/loopgpt/models/openai_.py
@@ -1,21 +1,13 @@
 from typing import *
 from loopgpt.logger import logger
 from loopgpt.models.base import BaseModel
+from loopgpt.utils.openai_key import get_openai_key
 
 import tiktoken
 import time
-import os
 
-
-def _getkey(key: Optional[str] = None):
-    key = key or os.getenv("OPENAI_API_KEY")
-    if key is None:
-        raise ValueError(
-            f"OpenAI API Key not found in the current working directory: {os.getcwd()}. "
-            "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. "
-            "See https://github.com/farizrahman4u/loopgpt#setup-your-openai-api-key- for more details"
-        )
-    return key
+from openai.error import RateLimitError
+import openai
 
 
 class OpenAIModel(BaseModel):
@@ -29,29 +21,17 @@ def chat(
         max_tokens: Optional[int] = None,
         temperature: float = 0.8,
     ) -> str:
-        import openai
-        from openai.error import RateLimitError
-
-        api_key = _getkey(self.api_key)
+        api_key = get_openai_key(self.api_key)
         num_retries = 3
         for _ in range(num_retries):
             try:
-                if openai.api_type == "azure":
-                    resp = openai.ChatCompletion.create(
-                        engine=self.model,
-                        messages=messages,
-                        api_key=api_key,
-                        max_tokens=max_tokens,
-                        temperature=temperature,
-                    )["choices"][0]["message"]["content"]
-                else:
-                    resp = openai.ChatCompletion.create(
-                        model=self.model,
-                        messages=messages,
-                        api_key=api_key,
-                        max_tokens=max_tokens,
-                        temperature=temperature,
-                    )["choices"][0]["message"]["content"]
+                resp = openai.ChatCompletion.create(
+                    model=self.model,
+                    messages=messages,
+                    api_key=api_key,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                )["choices"][0]["message"]["content"]
                 return resp
 
             except RateLimitError:
@@ -62,7 +42,6 @@ def chat(
     def count_tokens(self, messages: List[Dict[str, str]]) -> int:
         tokens_per_message, tokens_per_name = {
             "gpt-3.5-turbo": (4, -1),
-            "gpt-35-turbo": (4, -1),
             "gpt-4": (3, 1),
             "gpt-4-32k": (3, 1),
         }[self.model]
@@ -80,7 +59,6 @@ def count_tokens(self, messages: List[Dict[str, str]]) -> int:
     def get_token_limit(self) -> int:
         return {
             "gpt-3.5-turbo": 4000,
-            "gpt-35-turbo": 4000,
             "gpt-4": 8000,
             "gpt-4-32k": 32000,
         }[self.model]
@@ -97,4 +75,4 @@ def config(self):
 
     @classmethod
     def from_config(cls, config):
-        return cls(config["model"], config.get("api_key", None))
+        return cls(config["model"], config.get("api_key"))
diff --git a/loopgpt/utils/openai_key.py b/loopgpt/utils/openai_key.py
index a34c8ba..e674e30 100644
--- a/loopgpt/utils/openai_key.py
+++ b/loopgpt/utils/openai_key.py
@@ -1,6 +1,7 @@
 import os
-import sys
+import openai
 
+from typing import Optional
 from dotenv import load_dotenv
 
 
@@ -19,3 +20,15 @@ def check_openai_key():
             print(
                 "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. LoopGPT cannot work without it."
             )
+
+def get_openai_key(key: Optional[str] = None):
+    # API Key precedence:
+    key = key or openai.api_key or os.getenv("OPENAI_API_KEY")
+    
+    if key is None:
+        raise ValueError(
+            f"OpenAI API Key not found. "
+            "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. "
+            "See https://github.com/farizrahman4u/loopgpt#setup-your-openai-api-key- for more details"
+        )
+    return key

From 3893f4be74d3cf8c2e90c6d60edb38d4c6e0583f Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Sun, 7 May 2023 03:09:22 +0530
Subject: [PATCH 6/9] fix summarizer

---
 loopgpt/summarizer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/loopgpt/summarizer.py b/loopgpt/summarizer.py
index 045e622..93a4bca 100644
--- a/loopgpt/summarizer.py
+++ b/loopgpt/summarizer.py
@@ -11,8 +11,6 @@
 
 class Summarizer:
     def __init__(self, model: Optional[BaseModel] = None):
-        if isinstance(model, str):
-            model = OpenAIModel(model)
         self._model = model
 
     @property
@@ -20,11 +18,13 @@ def model(self):
         if self._model is None:
             if hasattr(self, "agent"):
                 model = self.agent.model
-                if isinstance(model, OpenAIModel):
+                if type(model) == OpenAIModel:
                     if model.model == "gpt-3.5-turbo":
                         self._model = model
                     else:
                         self._model = OpenAIModel("gpt-3.5-turbo")
+                else:
+                    self._model = model
             else:
                 self._model = OpenAIModel("gpt-3.5-turbo")
         return self._model

From eb32d1bf444920293a086244d6d04ef89f471684 Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Tue, 9 May 2023 23:10:59 +0530
Subject: [PATCH 7/9] update

---
 loopgpt/agent.py               | 24 ++++++++++--------------
 loopgpt/tools/agent_manager.py |  8 +++++---
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/loopgpt/agent.py b/loopgpt/agent.py
index 0b1dd1a..32c66c9 100644
--- a/loopgpt/agent.py
+++ b/loopgpt/agent.py
@@ -34,28 +34,24 @@ def __init__(
         model=None,
         embedding_provider=None,
         temperature=0.8,
-    ):  
-        if model is None:
-            model = OpenAIModel("gpt-3.5-turbo")
-        elif isinstance(model, str):
-            if openai.api_type == "azure":
-                model = AzureOpenAIModel(model)
-            else:
-                model = OpenAIModel(model)
-
+    ):
         if openai.api_type == "azure":
+            if model is None:
+                raise ValueError("You must provide an AzureOpenAIModel to the `model` argument when using the OpenAI Azure API")
             if embedding_provider is None:
                 raise ValueError("You must provide a deployed embedding provider to the `embedding_provider` argument when using the OpenAI Azure API")
-            elif isinstance(embedding_provider, str):
-                embedding_provider = AzureOpenAIEmbeddingProvider(embedding_provider)
-        else:
-            if embedding_provider is None:
-                embedding_provider = OpenAIEmbeddingProvider()
+        
+        if model is None:
+            model = OpenAIModel("gpt-3.5-turbo")
+
+        if embedding_provider is None:
+            embedding_provider = OpenAIEmbeddingProvider()
 
         self.name = name
         self.description = description
         self.goals = goals or []
         self.model = model
+        self.embedding_provider = embedding_provider
         self.temperature = temperature
         self.sub_agents = {}
         self.memory = LocalMemory(embedding_provider=embedding_provider)
diff --git a/loopgpt/tools/agent_manager.py b/loopgpt/tools/agent_manager.py
index 65e2057..7ca298f 100644
--- a/loopgpt/tools/agent_manager.py
+++ b/loopgpt/tools/agent_manager.py
@@ -25,13 +25,15 @@ def resp(self):
 
     def run(self, name="", task="", prompt=""):
         from loopgpt.agent import Agent
-
+        
+        model = self.agent.model
+        emb = self.agent.embedding_provider
         agent = Agent(
-            name=name, description=f"An agent for performing this specific task: {task}"
+            name=name, description=f"An agent for performing this specific task: {task}", model=model, embedding_provider=emb
         )
         agent.tools.clear()
         id = uuid4().hex[:8]
-        self.agent.sub_agents[id] = (Agent(), task)
+        self.agent.sub_agents[id] = (agent, task)
         resp = agent.chat(prompt)
         return {"uuid": id, "resp": resp}
 

From ed1b111cadc6bfd3b52a0484b827c394c3de6469 Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Wed, 10 May 2023 04:09:02 +0530
Subject: [PATCH 8/9] doc upds

---
 docs/source/api/embeddings.rst     |  3 +++
 docs/source/api/models.rst         |  3 +++
 loopgpt/agent.py                   | 10 ++++-----
 loopgpt/embeddings/azure_openai.py | 14 +++++++++++++
 loopgpt/models/azure_openai.py     | 33 ++++++++++++++++++++++++++++++
 5 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/docs/source/api/embeddings.rst b/docs/source/api/embeddings.rst
index 1ef5d97..8097810 100644
--- a/docs/source/api/embeddings.rst
+++ b/docs/source/api/embeddings.rst
@@ -10,5 +10,8 @@ Embedding Providers
 .. automodule:: loopgpt.embeddings.openai_
     :members:
 
+.. automodule:: loopgpt.embeddings.azure_openai
+    :members:
+
 .. automodule:: loopgpt.embeddings.hf
     :members:
diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst
index d1c1846..0dba6e1 100644
--- a/docs/source/api/models.rst
+++ b/docs/source/api/models.rst
@@ -10,6 +10,9 @@ Model
 .. automodule:: loopgpt.models.openai_
     :members:
 
+.. automodule:: loopgpt.models.azure_openai
+    :members:
+
 .. automodule:: loopgpt.models.llama_cpp
     :members:
 
diff --git a/loopgpt/agent.py b/loopgpt/agent.py
index ccea218..7eca84c 100644
--- a/loopgpt/agent.py
+++ b/loopgpt/agent.py
@@ -35,13 +35,13 @@ class Agent:
     :param goals: A list of goals for the agent. Defaults to None.
     :type goals: list, optional
     :param model: The model to use for the agent.
-        Strings are accepted only for OpenAI models. Specify a :class:`BaseModel` object for other models.
+        Strings are accepted only for OpenAI models. Specify a :class:`~loopgpt.models.base.BaseModel` object for other models.
         Defaults to "gpt-3.5-turbo".
-    :type model: str, :class:`BaseModel`, optional
+    :type model: str, :class:`~loopgpt.models.base.BaseModel`, optional
     :param embedding_provider: The embedding provider to use for the agent.
-        Defaults to :class:`OpenAIEmbeddingProvider`.
-        Specify a :class:`BaseEmbeddingProvider` object to use other embedding providers.
-    :type embedding_provider: :class:`BaseEmbeddingProvider`, optional
+        Defaults to :class:`~loopgpt.embeddings.OpenAIEmbeddingProvider`.
+        Specify a :class:`~loopgpt.embeddings.provider.BaseEmbeddingProvider` object to use other embedding providers.
+    :type embedding_provider: :class:`~loopgpt.embeddings.provider.BaseEmbeddingProvider`, optional
     :param temperature: The temperature to use for agent's chat completion. Defaults to 0.8.
     :type temperature: float, optional
     """
diff --git a/loopgpt/embeddings/azure_openai.py b/loopgpt/embeddings/azure_openai.py
index 71b6ace..3ffb7b1 100644
--- a/loopgpt/embeddings/azure_openai.py
+++ b/loopgpt/embeddings/azure_openai.py
@@ -6,7 +6,21 @@
 from loopgpt.utils.openai_key import get_openai_key
 
 class AzureOpenAIEmbeddingProvider(OpenAIEmbeddingProvider):
+    """Creates an Azure OpenAI embedding provider from a deployment ID. Can be created only when ``openai.api_type`` is set to ``azure``.
+
+    :param deployment_id: The deployment ID of the embedding provider.
+    :type deployment_id: str
+    :param api_key: The API key to use for the embedding provider.
+        If not specified, it will be found from ``openai.api_key`` or ``.env`` file or the ``OPENAI_API_KEY`` environment variable.
+    :type api_key: str, optional
+
+    .. note::
+        See :class:`AzureOpenAIModel <loopgpt.models.azure_openai.AzureOpenAIModel>` also.
+    """
     def __init__(self, deployment_id: str, api_key: Optional[str] = None):
+        # sanity check
+        assert openai.api_type == "azure", "AzureOpenAIModel can only be used with Azure API"
+
         self.deployment_id = deployment_id
         self.api_key = api_key
     
diff --git a/loopgpt/models/azure_openai.py b/loopgpt/models/azure_openai.py
index da26001..7235be2 100644
--- a/loopgpt/models/azure_openai.py
+++ b/loopgpt/models/azure_openai.py
@@ -27,6 +27,39 @@ def get_deployment_model(endpoint, deployment_id, api_version, api_key):
 
 
 class AzureOpenAIModel(OpenAIModel):
+    """Creates an Azure OpenAI model from a deployment ID. Can be created only when ``openai.api_type`` is set to ``azure``.
+
+    :param deployment_id: The deployment ID of the model.
+    :type deployment_id: str
+    :param api_key: The API key to use for the model.
+        If not specified, it will be found from ``openai.api_key`` or ``.env`` file or the ``OPENAI_API_KEY`` environment variable.
+    :type api_key: str, optional
+    :raises AssertionError: If ``openai.api_type`` is not set to ``azure``.
+
+    .. note::
+        You will also need an embedding provider deployed (e.g., text-embedding-ada-002) for creating an agent.
+
+    Example:
+
+    .. code-block:: python
+        
+        import os
+        import openai
+        import loopgpt
+        from loopgpt.models import AzureOpenAIModel
+        from loopgpt.embeddings import AzureOpenAIEmbeddingProvider
+
+        openai.api_type = "azure"
+        openai.api_base = "https://<your deployment>.openai.azure.com/"
+        openai.api_version = "2023-03-15-preview"
+        openai.api_key = os.getenv("OPENAI_API_KEY")
+
+        model = AzureOpenAIModel("my-gpt4-deployment")
+        embedding_provider = AzureOpenAIEmbeddingProvider("my-embeddings-deployment")
+
+        agent = loopgpt.Agent(model=model, embedding_provider=embedding_provider)
+        agent.chat("Hello, how are you?")
+    """
     def __init__(self, deployment_id: str, api_key: Optional[str] = None):
         # sanity check
         assert openai.api_type == "azure", "AzureOpenAIModel can only be used with Azure API"

From 9082c0912f6c4cdfd1c05d0a871360488d83803b Mon Sep 17 00:00:00 2001
From: FayazRahman <fayazrahman4u@gmail.com>
Date: Wed, 10 May 2023 04:11:21 +0530
Subject: [PATCH 9/9] black

---
 docs/source/conf.py                | 15 +++++++--------
 loopgpt/agent.py                   | 17 +++++++++++++----
 loopgpt/embeddings/azure_openai.py | 16 ++++++++++------
 loopgpt/embeddings/openai_.py      | 10 ++++++----
 loopgpt/embeddings/provider.py     |  4 ++--
 loopgpt/models/azure_openai.py     | 29 +++++++++++++++++++----------
 loopgpt/models/base.py             |  4 ++--
 loopgpt/tools/agent_manager.py     |  7 +++++--
 loopgpt/utils/openai_key.py        |  3 ++-
 9 files changed, 66 insertions(+), 39 deletions(-)

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 385f45f..de06509 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -6,23 +6,22 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
-project = 'LoopGPT'
-copyright = '2023, Fariz Rahman, Fayaz Rahman'
-author = 'Fariz Rahman, Fayaz Rahman'
-release = '0.0.13'
+project = "LoopGPT"
+copyright = "2023, Fariz Rahman, Fayaz Rahman"
+author = "Fariz Rahman, Fayaz Rahman"
+release = "0.0.13"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = ["sphinx.ext.autodoc"]
 
-templates_path = ['_templates']
+templates_path = ["_templates"]
 exclude_patterns = []
 
 
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = 'alabaster'
-html_static_path = ['_static']
+html_theme = "alabaster"
+html_static_path = ["_static"]
diff --git a/loopgpt/agent.py b/loopgpt/agent.py
index 7eca84c..4c73e30 100644
--- a/loopgpt/agent.py
+++ b/loopgpt/agent.py
@@ -8,7 +8,11 @@
     AgentStates,
 )
 from loopgpt.memory import from_config as memory_from_config
-from loopgpt.models import OpenAIModel, AzureOpenAIModel, from_config as model_from_config
+from loopgpt.models import (
+    OpenAIModel,
+    AzureOpenAIModel,
+    from_config as model_from_config,
+)
 from loopgpt.tools import builtin_tools, from_config as tool_from_config
 from loopgpt.tools.code import ai_function
 from loopgpt.memory.local_memory import LocalMemory
@@ -45,6 +49,7 @@ class Agent:
     :param temperature: The temperature to use for agent's chat completion. Defaults to 0.8.
     :type temperature: float, optional
     """
+
     def __init__(
         self,
         name=DEFAULT_AGENT_NAME,
@@ -56,10 +61,14 @@ def __init__(
     ):
         if openai.api_type == "azure":
             if model is None:
-                raise ValueError("You must provide an AzureOpenAIModel to the `model` argument when using the OpenAI Azure API")
+                raise ValueError(
+                    "You must provide an AzureOpenAIModel to the `model` argument when using the OpenAI Azure API"
+                )
             if embedding_provider is None:
-                raise ValueError("You must provide a deployed embedding provider to the `embedding_provider` argument when using the OpenAI Azure API")
-        
+                raise ValueError(
+                    "You must provide a deployed embedding provider to the `embedding_provider` argument when using the OpenAI Azure API"
+                )
+
         if model is None:
             model = OpenAIModel("gpt-3.5-turbo")
 
diff --git a/loopgpt/embeddings/azure_openai.py b/loopgpt/embeddings/azure_openai.py
index 3ffb7b1..5aa7692 100644
--- a/loopgpt/embeddings/azure_openai.py
+++ b/loopgpt/embeddings/azure_openai.py
@@ -5,6 +5,7 @@
 
 from loopgpt.utils.openai_key import get_openai_key
 
+
 class AzureOpenAIEmbeddingProvider(OpenAIEmbeddingProvider):
     """Creates an Azure OpenAI embedding provider from a deployment ID. Can be created only when ``openai.api_type`` is set to ``azure``.
 
@@ -17,19 +18,22 @@ class AzureOpenAIEmbeddingProvider(OpenAIEmbeddingProvider):
     .. note::
         See :class:`AzureOpenAIModel <loopgpt.models.azure_openai.AzureOpenAIModel>` also.
     """
+
     def __init__(self, deployment_id: str, api_key: Optional[str] = None):
         # sanity check
-        assert openai.api_type == "azure", "AzureOpenAIModel can only be used with Azure API"
+        assert (
+            openai.api_type == "azure"
+        ), "AzureOpenAIModel can only be used with Azure API"
 
         self.deployment_id = deployment_id
         self.api_key = api_key
-    
+
     def get(self, text: str):
         api_key = get_openai_key(self.api_key)
         return np.array(
-            openai.Embedding.create(input=[text], engine=self.deployment_id, api_key=api_key)[
-                "data"
-            ][0]["embedding"],
+            openai.Embedding.create(
+                input=[text], engine=self.deployment_id, api_key=api_key
+            )["data"][0]["embedding"],
             dtype=np.float32,
         )
 
@@ -39,4 +43,4 @@ def config(self):
 
     @classmethod
     def from_config(cls, config):
-        return cls(config["deployment_id"], config["api_key"])
\ No newline at end of file
+        return cls(config["deployment_id"], config["api_key"])
diff --git a/loopgpt/embeddings/openai_.py b/loopgpt/embeddings/openai_.py
index 3d2ace7..bc622ec 100644
--- a/loopgpt/embeddings/openai_.py
+++ b/loopgpt/embeddings/openai_.py
@@ -6,16 +6,18 @@
 
 
 class OpenAIEmbeddingProvider(BaseEmbeddingProvider):
-    def __init__(self, model: str = "text-embedding-ada-002", api_key: Optional[str] = None):
+    def __init__(
+        self, model: str = "text-embedding-ada-002", api_key: Optional[str] = None
+    ):
         self.model = model
         self.api_key = api_key
 
     def get(self, text: str):
         api_key = get_openai_key(self.api_key)
         return np.array(
-            openai.Embedding.create(input=[text], model="text-embedding-ada-002", api_key=api_key)[
-                "data"
-            ][0]["embedding"],
+            openai.Embedding.create(
+                input=[text], model="text-embedding-ada-002", api_key=api_key
+            )["data"][0]["embedding"],
             dtype=np.float32,
         )
 
diff --git a/loopgpt/embeddings/provider.py b/loopgpt/embeddings/provider.py
index 541bfb6..d65cb77 100644
--- a/loopgpt/embeddings/provider.py
+++ b/loopgpt/embeddings/provider.py
@@ -2,8 +2,8 @@
 
 
 class BaseEmbeddingProvider:
-    """Base class for all embedding providers.
-    """
+    """Base class for all embedding providers."""
+
     def get(self, text: str) -> np.ndarray:
         raise NotImplementedError()
 
diff --git a/loopgpt/models/azure_openai.py b/loopgpt/models/azure_openai.py
index 7235be2..c569f5f 100644
--- a/loopgpt/models/azure_openai.py
+++ b/loopgpt/models/azure_openai.py
@@ -8,13 +8,16 @@
 import requests
 import openai
 
+
 def get_deployment_details(endpoint, deployment_id, api_version, api_key):
     api_key = get_openai_key(api_key)
     response = requests.get(
-        f"{endpoint}/openai/deployments/{deployment_id}?api-version={api_version}", headers={"api-key": api_key}
+        f"{endpoint}/openai/deployments/{deployment_id}?api-version={api_version}",
+        headers={"api-key": api_key},
     )
     return response.json()
 
+
 def get_deployment_model(endpoint, deployment_id, api_version, api_key):
     details = get_deployment_details(endpoint, deployment_id, api_version, api_key)
     model = details["model"]
@@ -42,7 +45,7 @@ class AzureOpenAIModel(OpenAIModel):
     Example:
 
     .. code-block:: python
-        
+
         import os
         import openai
         import loopgpt
@@ -60,15 +63,20 @@ class AzureOpenAIModel(OpenAIModel):
         agent = loopgpt.Agent(model=model, embedding_provider=embedding_provider)
         agent.chat("Hello, how are you?")
     """
+
     def __init__(self, deployment_id: str, api_key: Optional[str] = None):
         # sanity check
-        assert openai.api_type == "azure", "AzureOpenAIModel can only be used with Azure API"
+        assert (
+            openai.api_type == "azure"
+        ), "AzureOpenAIModel can only be used with Azure API"
 
         self.deployment_id = deployment_id
         self.api_key = api_key
         self.endpoint = openai.api_base
         self.api_version = openai.api_version
-        self.model = get_deployment_model(self.endpoint, self.deployment_id, self.api_version, self.api_key)
+        self.model = get_deployment_model(
+            self.endpoint, self.deployment_id, self.api_version, self.api_key
+        )
 
     def chat(
         self,
@@ -76,7 +84,6 @@ def chat(
         max_tokens: Optional[int] = None,
         temperature: float = 0.8,
     ) -> str:
-
         api_key = get_openai_key(self.api_key)
         num_retries = 3
         for _ in range(num_retries):
@@ -94,14 +101,16 @@ def chat(
                 logger.warn("Rate limit exceeded. Retrying after 20 seconds.")
                 time.sleep(20)
                 continue
-    
+
     def config(self):
         cfg = super().config()
-        cfg.update({
-            "deployment_id": self.deployment_id,
-        })
+        cfg.update(
+            {
+                "deployment_id": self.deployment_id,
+            }
+        )
         return cfg
-    
+
     @classmethod
     def from_config(cls, config):
         return cls(config["deployment_id"], config.get("api_key"))
diff --git a/loopgpt/models/base.py b/loopgpt/models/base.py
index 5ee8be5..e75ca11 100644
--- a/loopgpt/models/base.py
+++ b/loopgpt/models/base.py
@@ -2,8 +2,8 @@
 
 
 class BaseModel:
-    """Base class for all models.
-    """
+    """Base class for all models."""
+
     def chat(
         self,
         messages: List[Dict[str, str]],
diff --git a/loopgpt/tools/agent_manager.py b/loopgpt/tools/agent_manager.py
index 7ca298f..76b5a2e 100644
--- a/loopgpt/tools/agent_manager.py
+++ b/loopgpt/tools/agent_manager.py
@@ -25,11 +25,14 @@ def resp(self):
 
     def run(self, name="", task="", prompt=""):
         from loopgpt.agent import Agent
-        
+
         model = self.agent.model
         emb = self.agent.embedding_provider
         agent = Agent(
-            name=name, description=f"An agent for performing this specific task: {task}", model=model, embedding_provider=emb
+            name=name,
+            description=f"An agent for performing this specific task: {task}",
+            model=model,
+            embedding_provider=emb,
         )
         agent.tools.clear()
         id = uuid4().hex[:8]
diff --git a/loopgpt/utils/openai_key.py b/loopgpt/utils/openai_key.py
index e674e30..811ea04 100644
--- a/loopgpt/utils/openai_key.py
+++ b/loopgpt/utils/openai_key.py
@@ -21,10 +21,11 @@ def check_openai_key():
                 "Please set the `OPENAI_API_KEY` environment variable or add it to `.env`. LoopGPT cannot work without it."
             )
 
+
 def get_openai_key(key: Optional[str] = None):
     # API Key precedence:
     key = key or openai.api_key or os.getenv("OPENAI_API_KEY")
-    
+
     if key is None:
         raise ValueError(
             f"OpenAI API Key not found. "