add Azure OpenAI integration (#125)

### Changed

* Removed `kwargs` from LLM classes and replaced with `llm_init_params` to provide parameters that should be passed to the LLM constructor.

### Added

* Azure OpenAI integrations for Discovery and Data Modeling: `AzureOpenAIDiscoveryLLM` and `AzureOpenAIDataModelingLLM`
* Added example that used Azure OpenAI models

CHANGELOG.md

@@ -6,8 +6,13 @@
 
 ### Changed
 
+* Removed `kwargs` from LLM classes and replaced with `llm_init_params` to provide parameters that should be passed to the LLM constructor.
+
 ### Added
 
+* Azure OpenAI integrations for Discovery and Data Modeling: `AzureOpenAIDiscoveryLLM` and `AzureOpenAIDataModelingLLM`
+* Added example that used Azure OpenAI models
+
 ## 0.10.0
 
 ### Fixed

neo4j_runway/llm/base.py

@@ -36,7 +36,6 @@ def __init__(
         client: Instructor,
         is_async: bool = False,
         model_params: Optional[dict[str, Any]] = None,
-        **kwargs: Any,
     ) -> None:
         """
         The base DiscoveryLLM class.
@@ -51,8 +50,6 @@ def __init__(
             An LLM client patched with Instructor.
         is_async : bool
             Whether async calls may be made, by default False
-        kwargs : Any
-            Parameters to pass to the model during initialization.
         """
 
         self.model_name = model_name
@@ -108,7 +105,6 @@ def __init__(
         model_name: str,
         client: Instructor,
         model_params: Optional[dict[str, Any]] = None,
-        **kwargs: Any,
     ) -> None:
         """
         The base DataModelingLLM class.
@@ -121,8 +117,6 @@ def __init__(
             Any parameters to pass to the model, by default None
         client : Instructor
             An LLM client patched with Instructor.
-        kwargs : Any
-             Parameters to pass to the model during initialization.
         """
 
         self.model_name = model_name

neo4j_runway/llm/openai/__init__.py

@@ -1,2 +1,11 @@
+from .azure.data_modeling import AzureOpenAIDataModelingLLM
+from .azure.discovery import AzureOpenAIDiscoveryLLM
 from .data_modeling import OpenAIDataModelingLLM
 from .discovery import OpenAIDiscoveryLLM
+
+__all__ = [
+    "OpenAIDataModelingLLM",
+    "OpenAIDiscoveryLLM",
+    "AzureOpenAIDataModelingLLM",
+    "AzureOpenAIDiscoveryLLM",
+]

neo4j_runway/llm/openai/azure/data_modeling.py

@@ -0,0 +1,84 @@
+# mypy: allow-untyped-defs
+"""
+This file contains the LLM module that interfaces with an Azure hosted OpenAI LLM for data modeing via the Instructor library.
+"""
+
+import os
+from typing import Any, Callable, Dict, Optional
+
+try:
+    import openai
+except ImportError:
+    openai = None  # type: ignore[unused-ignore, assignment]
+
+import instructor
+
+from ...base import BaseDataModelingLLM
+
+
+class AzureOpenAIDataModelingLLM(BaseDataModelingLLM):
+    """
+    Interface for interacting with OpenAI LLMs for data modeling services.
+
+    Attributes
+    ----------
+    model_name : str
+        The name of the deployment model.
+    model_params : Optional[dict[str, Any]], optional
+        Any parameters to pass to the model.
+    client : Instructor
+            An LLM client patched with Instructor.
+    """
+
+    def __init__(
+        self,
+        azure_endpoint: str,
+        api_version: str = "2024-07-01-preview",
+        azure_ad_token_provider: Optional[Callable[[], str]] = None,
+        model: str = "gpt-4o",
+        model_params: Optional[dict[str, Any]] = None,
+        azure_open_ai_key: Optional[str] = None,
+        llm_init_params: Dict[str, Any] = dict(),
+    ) -> None:
+        """
+        Interface for interacting with Azure OpenAI LLMs for data modeling services.
+
+        Parameters
+        ----------
+        azure_endpoint : str
+            The Azure endpoint address.
+        api_version : str, optional
+            The API version to use, by default 2024-07-01-preview
+        azure_ad_token_provider : Optional[Callable[[], str]], optional
+            A token to use for Microsoft Entra ID authentication instead of an API key, by default None
+        model : str
+            The name of the deployment model. By default gpt-4o
+        model_params : Optional[dict[str, Any]], optional
+            Any parameters to pass to the model for a request, by default None
+        azure_open_ai_key : Optional[str], optional
+            Your Azure OpenAI API key if it is not declared in an environment variable as 'AZURE_OPENAI_API_KEY'. By default None
+        llm_init_params : Dict[str, Any], optional
+            Parameters to pass to the model during initialization, by default dict()
+        """
+
+        if openai is None:
+            raise ImportError(
+                "Could not import openai python client. "
+                "Please install it with `pip install openai`."
+            )
+
+        client = instructor.from_openai(
+            openai.AzureOpenAI(
+                api_key=(
+                    azure_open_ai_key
+                    if azure_open_ai_key is not None
+                    else os.environ.get("AZURE_OPENAI_API_KEY")
+                ),
+                api_version=api_version,
+                azure_endpoint=azure_endpoint,
+                azure_ad_token_provider=azure_ad_token_provider,
+                **llm_init_params,
+            )
+        )
+
+        super().__init__(model_name=model, model_params=model_params, client=client)

neo4j_runway/llm/openai/azure/discovery.py

@@ -0,0 +1,110 @@
+"""
+This file contains the LLM module that interfaces with an OpenAI LLM for data discovery.
+"""
+
+import os
+from typing import Any, Callable, Dict, Optional
+
+import instructor
+
+try:
+    import openai
+except ImportError:
+    openai = None  # type: ignore[unused-ignore, assignment]
+
+from ...base import BaseDiscoveryLLM
+
+
+class AzureOpenAIDiscoveryLLM(BaseDiscoveryLLM):
+    """
+     Interface for interacting with different LLMs for data discovery.
+
+     Attributes
+     ----------
+    model_name : str
+         The name of the deployment model.
+     model_params : Optional[dict[str, Any]], optional
+         Any parameters to pass to the model.
+     client : Instructor
+             An LLM client patched with Instructor.
+     is_async : bool
+         Whether the client supports asynchronous API calls.
+    """
+
+    def __init__(
+        self,
+        azure_endpoint: str,
+        api_version: str = "2024-07-01-preview",
+        azure_ad_token_provider: Optional[Callable[[], str]] = None,
+        model: str = "gpt-4o",
+        model_params: Optional[Dict[str, Any]] = None,
+        azure_open_ai_key: Optional[str] = None,
+        enable_async: bool = False,
+        llm_init_params: Dict[str, Any] = dict(),
+    ) -> None:
+        """
+        Interface for interacting with OpenAI LLMs for data discovery.
+
+        Parameters
+        ----------
+        azure_endpoint : str
+            The Azure endpoint address.
+        api_version : str, optional
+            The API version to use, by default 2024-07-01-preview
+        azure_ad_token_provider : Optional[Callable[[], str]], optional
+            A token to use for Microsoft Entra ID authentication instead of an API key, by default None
+        model : str
+            The name of the deployment model. By default gpt-4o
+        model_params : Optional[dict[str, Any]], optional
+            Any parameters to pass to the model for a request, by default None
+        azure_open_ai_key: Union[str, None], optional
+            Your OpenAI API key if it is not declared in an environment variable as 'AZURE_OPENAI_API_KEY'. By default None
+        enable_async : bool
+            Whether to allow asynchronous LLM calls. This may be utilized in multi-csv input to improve response speed. By default False
+        llm_init_params : Dict[str, Any], optional
+            Parameters to pass to the model during initialization, by default dict()
+        """
+
+        if openai is None:
+            raise ImportError(
+                "Could not import openai python client. "
+                "Please install it with `pip install openai`."
+            )
+
+        client: Any
+
+        if enable_async:
+            client = instructor.from_openai(
+                openai.AsyncAzureOpenAI(
+                    api_key=(
+                        azure_open_ai_key
+                        if azure_open_ai_key is not None
+                        else os.environ.get("AZURE_OPENAI_API_KEY")
+                    ),
+                    api_version=api_version,
+                    azure_endpoint=azure_endpoint,
+                    azure_ad_token_provider=azure_ad_token_provider,
+                    **llm_init_params,
+                )
+            )
+        else:
+            client = instructor.from_openai(
+                openai.AzureOpenAI(
+                    api_key=(
+                        azure_open_ai_key
+                        if azure_open_ai_key is not None
+                        else os.environ.get("AZURE_OPENAI_API_KEY")
+                    ),
+                    api_version=api_version,
+                    azure_endpoint=azure_endpoint,
+                    azure_ad_token_provider=azure_ad_token_provider,
+                    **llm_init_params,
+                )
+            )
+
+        super().__init__(
+            model_name=model,
+            model_params=model_params,
+            client=client,
+            is_async=enable_async,
+        )

neo4j_runway/llm/openai/data_modeling.py

@@ -4,7 +4,7 @@
 """
 
 import os
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 try:
     import openai
@@ -26,18 +26,16 @@ class OpenAIDataModelingLLM(BaseDataModelingLLM):
         The name of the model.
     model_params : Optional[dict[str, Any]], optional
         Any parameters to pass to the model.
-    open_ai_key : Optional[str], optional
-        Your OpenAI API key if it is not declared in an environment variable.
-    kwargs : Any
-        Parameters to pass to the model during initialization.
+    client : Instructor
+            An LLM client patched with Instructor.
     """
 
     def __init__(
         self,
         model_name: str = "gpt-4o-2024-05-13",
         model_params: Optional[dict[str, Any]] = None,
         open_ai_key: Optional[str] = None,
-        **kwargs: Any,
+        llm_init_params: Dict[str, Any] = dict(),
     ) -> None:
         """
         Interface for interacting with OpenAI LLMs for data modeling services.
@@ -47,11 +45,11 @@ def __init__(
         model_name : str
             The name of the model. By default gpt-4o-2024-05-13
         model_params : Optional[dict[str, Any]], optional
-            Any parameters to pass to the model, by default None
+            Any parameters to pass to the model for a request, by default None
         open_ai_key : Optional[str], optional
             Your OpenAI API key if it is not declared in an environment variable. By default None
-        kwargs : Any
-            Parameters to pass to the model during initialization.
+        llm_init_params : Dict[str, Any], optional
+            Parameters to pass to the model during initialization, by default dict()
         """
 
         if openai is None:
@@ -67,7 +65,7 @@ def __init__(
                     if open_ai_key is not None
                     else os.environ.get("OPENAI_API_KEY")
                 ),
-                **kwargs,
+                **llm_init_params,
             )
         )
 

neo4j_runway/llm/openai/discovery.py

@@ -3,7 +3,7 @@
 """
 
 import os
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 import instructor
 
@@ -25,12 +25,10 @@ class OpenAIDiscoveryLLM(BaseDiscoveryLLM):
         The name of the model.
     model_params : Optional[dict[str, Any]], optional
         Any parameters to pass to the model.
-    open_ai_key: Union[str, None], optional
-        Your OpenAI API key if it is not declared in an environment variable.
     is_async : bool
         Whether the client supports asynchronous API calls.
-    kwargs : Any
-        Parameters to pass to the model during initialization.
+    client : Instructor
+            An LLM client patched with Instructor.
     """
 
     def __init__(
@@ -39,7 +37,7 @@ def __init__(
         model_params: Optional[dict[str, Any]] = None,
         open_ai_key: Optional[str] = None,
         enable_async: bool = False,
-        **kwargs: Any,
+        llm_init_params: Dict[str, Any] = dict(),
     ) -> None:
         """
         Interface for interacting with OpenAI LLMs for data discovery.
@@ -49,13 +47,13 @@ def __init__(
         model_name : str
             The name of the model. By default gpt-4o-2024-05-13
         model_params : Optional[dict[str, Any]], optional
-            Any parameters to pass to the model, by default None
+            Any parameters to pass to the model for a request, by default None
         open_ai_key: Union[str, None], optional
             Your OpenAI API key if it is not declared in an environment variable. By default None
         enable_async : bool
             Whether to allow asynchronous LLM calls. This may be utilized in multi-csv input to improve response speed. By default False
-        kwargs : Any
-            Parameters to pass to the model during initialization.
+        llm_init_params : Dict[str, Any], optional
+            Parameters to pass to the model during initialization, by default dict()
         """
 
         if openai is None:
@@ -74,7 +72,7 @@ def __init__(
                         if open_ai_key is not None
                         else os.environ.get("OPENAI_API_KEY")
                     ),
-                    **kwargs,
+                    **llm_init_params,
                 )
             )
         else:
@@ -85,7 +83,7 @@ def __init__(
                         if open_ai_key is not None
                         else os.environ.get("OPENAI_API_KEY")
                     ),
-                    **kwargs,
+                    **llm_init_params,
                 )
             )