feat: add support for Azure OpenAI (#87)

* feat: add support for Azure OpenAI

* code review changes

credentials.yaml.sample

@@ -1,4 +1,18 @@
 customer_credentials:
-  test-customer_id:
-      api_key: sample-api-key
-      model_name: gpt-3.5-turbo
+  testCustomerId:
+    credentialsMap:
+        AZURE_OPENAI:
+            customerId: testCustomerId
+            enabled: true
+            metadata:
+                deploymentName: gpt-4o
+                endpoint: https://myinstance.openai.azure.com/
+            secret: test_secret
+            type: AZURE_OPENAI
+        OPENAI:
+            customerId: testCustomerId
+            enabled: false
+            metadata:
+                model: gpt-3
+            secret: test_secret
+            type: OPENAI

skynet/auth/openai.py

@@ -1,3 +1,5 @@
+from enum import Enum
+
 import aiofiles
 import yaml
 
@@ -10,6 +12,11 @@
 credentials = dict()
 
 
+class CredentialsType(Enum):
+    OPENAI = 'OPENAI'
+    AZURE_OPENAI = 'AZURE_OPENAI'
+
+
 async def open_yaml(file_path):
     try:
         async with aiofiles.open(file_path, mode='r') as file:
@@ -40,4 +47,13 @@ async def setup_credentials():
 
 
 def get_credentials(customer_id):
-    return credentials.get(customer_id, {}) or {}
+    customer_credentials = credentials.get(customer_id, {}) or {}
+    multiple_credentials = customer_credentials.get('credentialsMap')
+
+    if multiple_credentials:
+        result = [val for val in multiple_credentials.values() if val['enabled']]
+        return result[0] if result else {}
+
+    # backwards compatibility
+    customer_credentials.setdefault('type', CredentialsType.OPENAI.value)
+    return customer_credentials

skynet/env.py

@@ -30,6 +30,10 @@ def tobool(val: str | None):
 llama_n_gpu_layers = int(os.environ.get('LLAMA_N_GPU_LAYERS', -1 if is_mac else 40))
 llama_n_batch = int(os.environ.get('LLAMA_N_BATCH', 512))
 
+# azure openai api
+# latest ga version https://learn.microsoft.com/en-us/azure/ai-services/openai/api-version-deprecation#latest-ga-api-release
+azure_openai_api_version = os.environ.get('AZURE_OPENAI_API_VERSION', '2024-02-01')
+
 # openai api
 openai_api_server_path = os.environ.get('OPENAI_API_SERVER_PATH', '/app/llama.cpp/server')
 openai_api_server_port = int(os.environ.get('OPENAI_API_SERVER_PORT', 8002))

skynet/modules/ttt/summaries/jobs.py

@@ -2,7 +2,7 @@
 import time
 import uuid
 
-from skynet.auth.openai import get_credentials
+from skynet.auth.openai import CredentialsType, get_credentials
 
 from skynet.env import job_timeout, modules, redis_exp_seconds, summary_minimum_payload_length
 from skynet.logs import get_logger
@@ -16,7 +16,7 @@
 from skynet.modules.ttt.openai_api.app import restart as restart_openai_api
 
 from .persistence import db
-from .processor import process, process_open_ai
+from .processor import process, process_azure, process_open_ai
 from .v1.models import DocumentMetadata, DocumentPayload, Job, JobId, JobStatus, JobType
 
 log = get_logger(__name__)
@@ -125,15 +125,26 @@ async def run_job(job: Job) -> None:
         exit_task = asyncio.create_task(restart_on_timeout(job))
 
         try:
-            customer_id = job.metadata.customer_id
-            options = get_credentials(customer_id) if customer_id else {}
-            api_key = options.get('secret')
-            model_name = options.get('model')
-
-            if api_key:
-                log.info(f"Forwarding inference to OpenAI for customer {customer_id}")
-
-                result = await process_open_ai(job.payload, job.type, api_key, model_name)
+            customer_id = 'testCustomerId'
+            options = get_credentials(customer_id)
+            secret = options.get('secret')
+            api_type = options.get('type')
+
+            if secret:
+                if api_type == CredentialsType.OPENAI.value:
+                    log.info(f"Forwarding inference to OpenAI for customer {customer_id}")
+
+                    # needed for backwards compatibility
+                    model = options.get('model') or options.get('metadata').get('model')
+                    result = await process_open_ai(job.payload, job.type, secret, model)
+
+                elif api_type == CredentialsType.AZURE_OPENAI.value:
+                    log.info(f"Forwarding inference to Azure openai for customer {customer_id}")
+
+                    metadata = options.get('metadata')
+                    result = await process_azure(
+                        job.payload, job.type, secret, metadata.get('endpoint'), metadata.get('deploymentName')
+                    )
             else:
                 if customer_id:
                     log.info(f'Customer {customer_id} has no API key configured, falling back to local processing')

skynet/modules/ttt/summaries/jobs_test.py

@@ -38,6 +38,7 @@ def run_job_fixture(mocker):
     mocker.patch('skynet.modules.ttt.summaries.jobs.update_job')
     mocker.patch('skynet.modules.ttt.summaries.jobs.process')
     mocker.patch('skynet.modules.ttt.summaries.jobs.process_open_ai')
+    mocker.patch('skynet.modules.ttt.summaries.jobs.process_azure')
     mocker.patch('skynet.modules.ttt.summaries.jobs.db.db')
 
     return mocker
@@ -86,20 +87,58 @@ async def test_run_job_with_open_ai(self, run_job_fixture):
 
         from skynet.modules.ttt.summaries.jobs import process_open_ai, run_job
 
-        run_job_fixture.patch('skynet.modules.ttt.summaries.jobs.get_credentials', return_value={'secret': 'api_key'})
+        secret = 'secret'
+        model = 'gpt-3.5-turbo'
 
-        await run_job(
-            Job(
-                payload=DocumentPayload(
-                    text="Andrew: Hello. Beatrix: Honey? It’s me . . . Andrew: Where are you? Beatrix: At the station. I missed my train."
-                ),
-                metadata=DocumentMetadata(customer_id='test'),
-                type=JobType.SUMMARY,
-                id='job_id',
-            )
+        run_job_fixture.patch(
+            'skynet.modules.ttt.summaries.jobs.get_credentials',
+            return_value={'secret': secret, 'type': 'OPENAI', 'metadata': {'model': model}},
         )
 
-        process_open_ai.assert_called_once()
+        job = Job(
+            payload=DocumentPayload(
+                text="Andrew: Hello. Beatrix: Honey? It’s me . . . Andrew: Where are you? Beatrix: At the station. I missed my train."
+            ),
+            metadata=DocumentMetadata(customer_id='test'),
+            type=JobType.SUMMARY,
+            id='job_id',
+        )
+
+        await run_job(job)
+
+        process_open_ai.assert_called_once_with(job.payload, job.type, secret, model)
+
+    @pytest.mark.asyncio
+    async def test_run_job_with_azure_open_ai(self, run_job_fixture):
+        '''Test that a job is sent for inference to azure openai if there is a customer id with a valid api key.'''
+
+        from skynet.modules.ttt.summaries.jobs import process_azure, run_job
+
+        secret = 'secret'
+        deployment_name = 'gpt-3.5-turbo'
+        endpoint = 'https://myopenai.azure.com'
+
+        run_job_fixture.patch(
+            'skynet.modules.ttt.summaries.jobs.get_credentials',
+            return_value={
+                'secret': secret,
+                'type': 'AZURE_OPENAI',
+                'metadata': {'deploymentName': deployment_name, 'endpoint': endpoint},
+            },
+        )
+
+        job = Job(
+            payload=DocumentPayload(
+                text="Andrew: Hello. Beatrix: Honey? It’s me . . . Andrew: Where are you? Beatrix: At the station. I missed my train."
+            ),
+            metadata=DocumentMetadata(customer_id='test'),
+            type=JobType.SUMMARY,
+            id='job_id',
+        )
+
+        await run_job(job)
+
+        process_azure.assert_called_once_with(job.payload, job.type, secret, endpoint, deployment_name)
 
 
 class TestCanRunNextJob:

skynet/modules/ttt/summaries/processor.py

@@ -2,9 +2,9 @@
 from langchain.prompts import ChatPromptTemplate
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
-from langchain_openai import ChatOpenAI
+from langchain_openai import AzureChatOpenAI, ChatOpenAI
 
-from skynet.env import app_uuid, llama_n_ctx, openai_api_base_url
+from skynet.env import app_uuid, azure_openai_api_version, llama_n_ctx, openai_api_base_url
 from skynet.logs import get_logger
 
 from .prompts.action_items import action_items_conversation_prompt, action_items_text_prompt
@@ -83,3 +83,17 @@ async def process_open_ai(payload: DocumentPayload, job_type: JobType, api_key:
     )
 
     return await process(payload, job_type, llm)
+
+
+async def process_azure(
+    payload: DocumentPayload, job_type: JobType, api_key: str, endpoint: str, deployment_name: str
+) -> str:
+    llm = AzureChatOpenAI(
+        api_key=api_key,
+        api_version=azure_openai_api_version,
+        azure_endpoint=endpoint,
+        azure_deployment=deployment_name,
+        temperature=0,
+    )
+
+    return await process(payload, job_type, llm)