Merge pull request #164 from briefercloud/support-azure-ai

Support Azure for OpenAI models

ai/api/llms.py

@@ -1,5 +1,5 @@
 from langchain_aws import BedrockLLM
-from langchain_openai import ChatOpenAI
+from langchain_openai import ChatOpenAI, AzureOpenAI
 from decouple import config
 
 # Add available models here
@@ -10,14 +10,29 @@
     "cohere.command-r-plus-v1:0",
 ]
 
+def str_to_bool(value):
+    if value is None or value == False:
+        return False
+
+    return value.lower() in ['true', '1', 't', 'y', 'yes']
+
 def initialize_llm(model_id=None, openai_api_key=None):
     openai_api_key = openai_api_key or config("OPENAI_API_KEY")
+    use_azure = config("USE_AZURE", default=False, cast=str_to_bool)
 
     if model_id in bedrock_model_ids:
         # Initialize Bedrock using default AWS credentials provider chain
         llm = BedrockLLM(
             model_id=model_id,
         )
+    elif use_azure:
+        # Initialize Azure OpenAI using the environment variables for API key and model name
+        llm = AzureOpenAI(
+            temperature=0,
+            verbose=False,
+            azure_api_key=openai_api_key,
+            model_name=model_id if model_id else config("OPENAI_DEFAULT_MODEL_NAME"),
+        )
     else:
         # Initialize OpenAI using environment variables for API key and model name
         llm = ChatOpenAI(

ai/requirements.txt

@@ -1,21 +1,22 @@
 uvicorn==0.23.2
 fastapi==0.104.1
-langchain==0.2.1
-langchain-community==0.2.1
-langchain-openai==0.1.7
-langchain-aws==0.1.6
-openai==1.27.0
+langchain==0.3.4
+langchain-core==0.3.12
+langchain-community==0.3.3
+langchain-openai==0.2.3
+langchain-aws==0.2.3
+openai==1.52.0
 python-decouple==3.8
 psycopg2-binary==2.9.9
 sqlalchemy-bigquery==1.11.0
 google-cloud-bigquery-storage==2.22.0
 pyarrow
 PyAthena==3.8.2
 oracledb==2.2.0
-boto3==1.34.113
+boto3==1.34.131
 sqlalchemy-redshift==0.8.14
 redshift-connector==2.0.917
 mysqlclient==2.2.4
 trino==0.329.0
 snowflake-connector-python==3.12.2
-snowflake-sqlalchemy==1.6.1
+snowflake-sqlalchemy==1.6.1

chart/Chart.yaml

@@ -2,5 +2,5 @@ apiVersion: v2
 name: briefer
 description: The helm chart for Briefer's open-source version.
 type: application
-version: 0.1.1
-appVersion: '0.1.1'
+version: 0.1.2
+appVersion: '0.1.2'

chart/templates/ai-deployment.yaml

@@ -46,6 +46,8 @@ spec:
             {{- end }}
           {{- end }}
           env:
+            - name: USE_AZURE
+              value: '{{ .Values.ai.env.useAzure | default "false" }}'
             - name: PORT
               value: '{{ .Values.ai.env.port | default "8080" }}'
             - name: OPENAI_DEFAULT_MODEL_NAME