Merge branch 'feat/openai-llm' into develop

README.md

@@ -269,15 +269,25 @@ cognitive_service:
   # Must be of type "AI services multi-service account"
   endpoint: https://xxx.cognitiveservices.azure.com
 
-openai:
-  api_key: xxx
-  endpoint: https://xxx.openai.azure.com
-  gpt_backup_context: 16385
-  gpt_backup_deployment: gpt-35-turbo-1106
-  gpt_backup_model: gpt-35-turbo-1106
-  gpt_context: 128000
-  gpt_deployment: gpt-4-1106-preview
-  gpt_model: gpt-4-1106-preview
+llm:
+  backup:
+    mode: azure_openai
+    azure_openai:
+      api_key: xxx
+      context: 16385
+      deployment: gpt-35-turbo-0125
+      endpoint: https://xxx.openai.azure.com
+      model: gpt-35-turbo
+      streaming: true
+  primary:
+    mode: azure_openai
+    azure_openai:
+      api_key: xxx
+      context: 128000
+      deployment: gpt-4-0125-preview
+      endpoint: https://xxx.openai.azure.com
+      model: gpt-4
+      streaming: true
 
 ai_search:
   access_key: xxx

bicep/app.bicep

@@ -57,14 +57,27 @@ var config = {
   cognitive_service: {
     endpoint: cognitiveCommunication.properties.endpoint
   }
-  openai: {
-    endpoint: cognitiveOpenai.properties.endpoint
-    gpt_backup_context: gptBackupContext
-    gpt_backup_deployment: gptBackup.name
-    gpt_backup_model: gptBackupModel
-    gpt_context: gptContext
-    gpt_deployment: gpt.name
-    gpt_model: gptModel
+  llm: {
+    backup: {
+      mode: 'azure_openai'
+      azure_openai: {
+        context: gptBackupContext
+        deployment: gptBackup.name
+        endpoint: cognitiveOpenai.properties.endpoint
+        model: gptBackupModel
+        streaming: true
+      }
+    }
+    primary: {
+      mode: 'azure_openai'
+      azure_openai: {
+        context: gptContext
+        deployment: gpt.name
+        endpoint: cognitiveOpenai.properties.endpoint
+        model: gptModel
+        streaming: true
+      }
+    }
   }
   ai_search: {
     access_key: search.listAdminKeys().primaryKey

helpers/config_models/llm.py

@@ -0,0 +1,120 @@
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+from contextlib import asynccontextmanager
+from enum import Enum
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from pydantic import field_validator, SecretStr, BaseModel, ValidationInfo
+from typing import Any, AsyncGenerator, Optional, Tuple, Union
+
+
+class ModeEnum(str, Enum):
+    AZURE_OPENAI = "azure_openai"
+    OPENAI = "openai"
+
+
+class AbstractPlatformModel(BaseModel, frozen=True):
+    _client_kwargs: dict[str, Any] = {
+        # Reliability
+        "max_retries": 3,
+        "timeout": 60,
+    }
+    context: int
+    model: str
+    streaming: bool
+
+
+class AzureOpenaiPlatformModel(AbstractPlatformModel, frozen=True):
+    api_key: Optional[SecretStr] = None
+    deployment: str
+    endpoint: str
+
+    @asynccontextmanager
+    async def instance(
+        self,
+    ) -> AsyncGenerator[Tuple[AsyncAzureOpenAI, AbstractPlatformModel], None]:
+        api_key = self.api_key.get_secret_value() if self.api_key else None
+        token_func = (
+            get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+            if not self.api_key
+            else None
+        )
+        client = AsyncAzureOpenAI(
+            **self._client_kwargs,
+            # Azure deployment
+            api_version="2023-12-01-preview",
+            azure_deployment=self.deployment,
+            azure_endpoint=self.endpoint,
+            # Authentication, either RBAC or API key
+            api_key=api_key,
+            azure_ad_token_provider=token_func,
+        )
+        try:
+            yield client, self
+        finally:
+            await client.close()
+
+
+class OpenaiPlatformModel(AbstractPlatformModel, frozen=True):
+    api_key: SecretStr
+    endpoint: str
+
+    @asynccontextmanager
+    async def instance(
+        self,
+    ) -> AsyncGenerator[Tuple[AsyncOpenAI, AbstractPlatformModel], None]:
+        client = AsyncOpenAI(
+            **self._client_kwargs,
+            # OpenAI deployment
+            base_url=self.endpoint,
+            # Authentication, either API key or no auth
+            api_key=self.api_key.get_secret_value(),
+        )
+        try:
+            yield client, self
+        finally:
+            await client.close()
+
+
+class SelectedPlatformModel(BaseModel, frozen=True):
+    azure_openai: Optional[AzureOpenaiPlatformModel] = None
+    mode: ModeEnum
+    openai: Optional[OpenaiPlatformModel] = None
+
+    @field_validator("azure_openai")
+    def validate_azure_openai(
+        cls,
+        azure_openai: Optional[AzureOpenaiPlatformModel],
+        info: ValidationInfo,
+    ) -> Optional[AzureOpenaiPlatformModel]:
+        if not azure_openai and info.data.get("mode", None) == ModeEnum.AZURE_OPENAI:
+            raise ValueError("Azure OpenAI config required")
+        return azure_openai
+
+    @field_validator("openai")
+    def validate_openai(
+        cls,
+        openai: Optional[OpenaiPlatformModel],
+        info: ValidationInfo,
+    ) -> Optional[OpenaiPlatformModel]:
+        if not openai and info.data.get("mode", None) == ModeEnum.OPENAI:
+            raise ValueError("OpenAI config required")
+        return openai
+
+    def selected(self) -> Union[AzureOpenaiPlatformModel, OpenaiPlatformModel]:
+        platform = (
+            self.azure_openai if self.mode == ModeEnum.AZURE_OPENAI else self.openai
+        )
+        assert platform
+        return platform
+
+
+class LlmModel(BaseModel):
+    backup: SelectedPlatformModel
+    primary: SelectedPlatformModel
+
+    def selected(
+        self, is_backup: bool
+    ) -> Union[AzureOpenaiPlatformModel, OpenaiPlatformModel]:
+        platform = self.backup if is_backup else self.primary
+        return platform.selected()

helpers/config_models/root.py

@@ -7,7 +7,7 @@
 from helpers.config_models.content_safety import ContentSafetyModel
 from helpers.config_models.database import DatabaseModel
 from helpers.config_models.monitoring import MonitoringModel
-from helpers.config_models.openai import OpenAiModel
+from helpers.config_models.llm import LlmModel
 from helpers.config_models.prompts import PromptsModel
 from helpers.config_models.resources import ResourcesModel
 from helpers.config_models.workflow import WorkflowModel
@@ -34,8 +34,8 @@ class RootModel(BaseSettings):
     communication_service: CommunicationServiceModel
     content_safety: ContentSafetyModel
     database: DatabaseModel = DatabaseModel()  # Object is fully defined by default
+    llm: LlmModel
     monitoring: MonitoringModel
-    openai: OpenAiModel
     prompts: PromptsModel = PromptsModel()  # Object is fully defined by default
     resources: ResourcesModel
     workflow: WorkflowModel