Make plugin work with AzureOAI (#221)

* Make this work better with Azure Open AI * Fix typo * Update README.md Co-authored-by: Derek Legenzoff <delegenz@microsoft.com> --------- Co-authored-by: isafulf <51974293+isafulf@users.noreply.github.com> Co-authored-by: Derek Legenzoff <delegenz@microsoft.com>

Author: 3 people

README.md

@@ -72,6 +72,13 @@ Follow these steps to quickly set up and run the ChatGPT Retrieval Plugin:
  export BEARER_TOKEN=<your_bearer_token>
  export OPENAI_API_KEY=<your_openai_api_key>
 
+ # Optional environment variables used when running Azure OpenAI
+ export OPENAI_API_BASE=https://<AzureOpenAIName>.openai.azure.com/ 
+ export OPENAI_API_TYPE=azure
+ export OPENAI_EMBEDDINGMODEL_DEPLOYMENTID=<Name of text-embedding-ada-002 model deployment>
+ export OPENAI_METADATA_EXTRACTIONMODEL_DEPLOYMENTID=<Name of deployment of model for metatdata>
+ export OPENAI_COMPLETIONMODEL_DEPLOYMENTID=<Name of general model deployment used for completion>
+
  # Add the environment variables for your chosen vector DB.
  # Some of these are optional; read the provider's setup docs in /docs/providers for more information.
 
@@ -237,6 +244,17 @@ The API requires the following environment variables to work:
 | `BEARER_TOKEN` | Yes | This is a secret token that you need to authenticate your requests to the API. You can generate one using any tool or method you prefer, such as [jwt.io](https://jwt.io/). |
 | `OPENAI_API_KEY` | Yes | This is your OpenAI API key that you need to generate embeddings using the `text-embedding-ada-002` model. You can get an API key by creating an account on [OpenAI](https://openai.com/). |
 
+
+### Using the plugin with Azure OpenAI
+
+The Azure Open AI uses URLs that are specific to your resource and references models not by model name but by the deployment id. As a result, you need to set additional environment variables for this case.
+
+In addition to the OPENAI_API_BASE (your specific URL) and OPENAI_API_TYPE (azure), you should also set OPENAI_EMBEDDINGMODEL_DEPLOYMENTID which specifies the model to use for getting embeddings on upsert and query. For this, we recommend deploying text-embedding-ada-002 model and using the deployment name here. 
+
+If you wish to use the data preparation scripts, you will also need to set OPENAI_METADATA_EXTRACTIONMODEL_DEPLOYMENTID, used for metadata extraction and 
+OPENAI_COMPLETIONMODEL_DEPLOYMENTID, used for PII handling.
+
+
 ### Choosing a Vector Database
 
 The plugin supports several vector database providers, each with different features, performance, and pricing. Depending on which one you choose, you will need to use a different Dockerfile and set different environment variables. The following sections provide brief introductions to each vector database provider.

poetry.lock

@@ -10,7 +10,7 @@ packages = [{include = "server"}]
 python = "^3.10"
 fastapi = "^0.92.0"
 uvicorn = "^0.20.0"
-openai = "^0.27.2"
+openai = "^0.27.5"
 python-dotenv = "^0.21.1"
 pydantic = "^1.10.5"
 tenacity = "^8.2.1"

pyproject.toml

@@ -2,7 +2,7 @@
 from services.openai import get_chat_completion
 import json
 from typing import Dict
-
+import os
 
 def extract_metadata_from_document(text: str) -> Dict[str, str]:
  sources = Source.__members__.keys()
@@ -24,8 +24,12 @@ def extract_metadata_from_document(text: str) -> Dict[str, str]:
  {"role": "user", "content": text},
  ]
 
+ # NOTE: Azure Open AI requires deployment id
+ # Read environment variable - if not set - not used
  completion = get_chat_completion(
- messages, "gpt-4"
+ messages,
+ "gpt-4",
+ os.environ.get("OPENAI_METADATA_EXTRACTIONMODEL_DEPLOYMENTID")
  ) # TODO: change to your preferred model name
 
  print(f"completion: {completion}")

services/extract_metadata.py

@@ -1,6 +1,6 @@
 from typing import List
 import openai
-
+import os
 
 from tenacity import retry, wait_random_exponential, stop_after_attempt
 
@@ -20,8 +20,15 @@ def get_embeddings(texts: List[str]) -> List[List[float]]:
  Exception: If the OpenAI API call fails.
  """
  # Call the OpenAI API to get the embeddings
- response = openai.Embedding.create(input=texts, model="text-embedding-ada-002")
+ # NOTE: Azure Open AI requires deployment id
+ deployment = os.environ.get("OPENAI_EMBEDDINGMODEL_DEPLOYMENTID")
 
+ response = {}
+ if deployment == None:
+ response = openai.Embedding.create(input=texts, model="text-embedding-ada-002")
+ else:
+ response = openai.Embedding.create(input=texts, deployment_id=deployment)
+ 
  # Extract the embedding data from the response
  data = response["data"] # type: ignore
 
@@ -33,6 +40,7 @@ def get_embeddings(texts: List[str]) -> List[List[float]]:
 def get_chat_completion(
  messages,
  model="gpt-3.5-turbo", # use "gpt-4" for better results
+ deployment_id = None
 ):
  """
  Generate a chat completion using OpenAI's chat completion API.
@@ -48,10 +56,19 @@ def get_chat_completion(
  Exception: If the OpenAI API call fails.
  """
  # call the OpenAI chat completion API with the given messages
- response = openai.ChatCompletion.create(
- model=model,
- messages=messages,
- )
+ # Note: Azure Open AI requires deployment id
+ response = {}
+ if deployment_id == None:
+ response = openai.ChatCompletion.create(
+ model=model,
+ messages=messages,
+ )
+ else:
+ response = openai.ChatCompletion.create(
+ deployment_id = deployment_id,
+ messages=messages,
+ )
+ 
 
  choices = response["choices"] # type: ignore
  completion = choices[0].message.content.strip()

services/openai.py

@@ -1,3 +1,4 @@
+import os
 from services.openai import get_chat_completion
 
 
@@ -22,6 +23,7 @@ def screen_text_for_pii(text: str) -> bool:
 
  completion = get_chat_completion(
  messages,
+ deployment_id=os.environ.get("OPENAI_COMPLETIONMODEL_DEPLOYMENTID")
  )
 
  if completion.startswith("True"):