feat(llm): add uipath openai llm

Author: ionmincu

pyproject.toml

@@ -1,11 +1,13 @@
 [project]
 name = "uipath-llamaindex"
-version = "0.0.23"
+version = "0.0.24"
 description = "UiPath LlamaIndex SDK"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.10"
 dependencies = [
     "llama-index>=0.12.38",
+    "llama-index-embeddings-azure-openai>=0.3.8",
+    "llama-index-llms-azure-openai>=0.3.2",
     "openinference-instrumentation-llama-index>=4.3.0",
     "uipath>=2.0.64",
 ]

src/uipath_llamaindex/embeddings/__init__.py

@@ -0,0 +1,9 @@
+from ._openai_embeddings import (
+    OpenAIEmbeddingModel,
+    UiPathOpenAIEmbedding,
+)
+
+__all__ = [
+    "UiPathOpenAIEmbedding",
+    "OpenAIEmbeddingModel",
+]

src/uipath_llamaindex/embeddings/_openai_embeddings.py

@@ -0,0 +1,48 @@
+import os
+from enum import Enum
+from typing import Any, Union
+
+from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
+
+
+class OpenAIEmbeddingModel(Enum):
+    TEXT_EMBEDDING_3_LARGE = "text-embedding-3-large"
+    TEXT_EMBEDDING_ADA_002 = "text-embedding-ada-002"
+
+
+class UiPathOpenAIEmbedding(AzureOpenAIEmbedding):
+    def __init__(
+        self,
+        model: Union[
+            str, OpenAIEmbeddingModel
+        ] = OpenAIEmbeddingModel.TEXT_EMBEDDING_ADA_002,
+        api_version: str = "2024-10-21",
+        **kwargs: Any,
+    ):
+        default_headers_dict = {
+            "X-UIPATH-STREAMING-ENABLED": "false",
+            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
+            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
+        }
+
+        model_value = model.value if isinstance(model, OpenAIEmbeddingModel) else model
+
+        base_url = os.environ.get(
+            "UIPATH_URL", "EMPTY"
+        ).rstrip("/")
+        
+        if base_url == "EMPTY":
+            raise ValueError(
+                "UIPATH_URL environment variable is not set. Please run uipath auth."
+            )
+
+        defaults = {
+            "model": model_value,
+            "deployment_name": model_value,
+            "azure_endpoint": f"{base_url}/llmgateway_/",
+            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
+            "api_version": api_version,
+            "default_headers": default_headers_dict,
+        }
+        final_kwargs = {**defaults, **kwargs}
+        super().__init__(**final_kwargs)

src/uipath_llamaindex/llms/__init__.py

@@ -0,0 +1,9 @@
+from ._openai_llms import (
+    OpenAIModel,
+    UiPathOpenAI,
+)
+
+__all__ = [
+    "UiPathOpenAI",
+    "OpenAIModel",
+]

src/uipath_llamaindex/llms/_openai_llms.py

@@ -0,0 +1,54 @@
+import os
+from enum import Enum
+from typing import Any, Union
+
+from llama_index.llms.azure_openai import AzureOpenAI
+
+
+class OpenAIModel(Enum):
+    GPT_4_1_2025_04_14 = "gpt-4.1-2025-04-14"
+    GPT_4_1_MINI_2025_04_14 = "gpt-4.1-mini-2025-04-14"
+    GPT_4_1_NANO_2025_04_14 = "gpt-4.1-nano-2025-04-14"
+    GPT_4O_2024_05_13 = "gpt-4o-2024-05-13"
+    GPT_4O_2024_08_06 = "gpt-4o-2024-08-06"
+    GPT_4O_2024_11_20 = "gpt-4o-2024-11-20"
+    GPT_4O_MINI_2024_07_18 = "gpt-4o-mini-2024-07-18"
+    O3_MINI_2025_01_31 = "o3-mini-2025-01-31"
+    TEXT_DAVINCI_003 = "text-davinci-003"
+
+
+# Define your custom AzureOpenAI class with default settings
+class UiPathOpenAI(AzureOpenAI):
+    def __init__(
+        self,
+        model: Union[str, OpenAIModel] = OpenAIModel.GPT_4O_MINI_2024_07_18,
+        api_version: str = "2024-10-21",
+        **kwargs: Any,
+    ):
+        default_headers_dict = {
+            "X-UIPATH-STREAMING-ENABLED": "false",
+            "X-UiPath-LlmGateway-RequestingProduct": "uipath-python-sdk",
+            "X-UiPath-LlmGateway-RequestingFeature": "llama-index-agent",
+        }
+        model_value = model.value if isinstance(model, OpenAIModel) else model
+
+        base_url = os.environ.get(
+            "UIPATH_URL", "EMPTY"
+        ).rstrip("/")
+        
+        if base_url == "EMPTY":
+            raise ValueError(
+                "UIPATH_URL environment variable is not set. Please run uipath auth."
+            )
+
+        defaults = {
+            "model": model_value,
+            "deployment_name": model_value,
+            "azure_endpoint": f"{base_url}/llmgateway_/",
+            "api_key": os.environ.get("UIPATH_ACCESS_TOKEN"),
+            "api_version": api_version,
+            "is_chat_model": True,
+            "default_headers": default_headers_dict,
+        }
+        final_kwargs = {**defaults, **kwargs}
+        super().__init__(**final_kwargs)