Add setting for Azure OpenAI

app.py

@@ -4,7 +4,7 @@ import os
 import re
 import gradio as gr
 from langchain.chains import LLMChain
-from langchain.chat_models import PromptLayerChatOpenAI
+from langchain.chat_models import PromptLayerChatOpenAI, AzureChatOpenAI
 from langchain.memory import ConversationBufferMemory
 from langchain.prompts import PromptTemplate
 from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate
@@ -48,10 +48,39 @@ Use this output format:
 - [textual evidence 2] - outside source"""
 # 7. Make a section for each key point. In each section, make a bullet list of textual evidence from task 3 that relate to that key point.
 
-def load_chain(api_key):
-  api_key = os.environ["OPENAI_API_KEY"] if api_key == "" or api_key.isspace() else api_key
+
+class AzurePromptLayerChatOpenAI(AzureChatOpenAI, PromptLayerChatOpenAI):
+  pass
+
+
+
+def load_chain(api_key, api_type):
+  if api_key == "" or api_key.isspace():
+    if api_type == "OpenAI":
+      api_key = os.environ.get("OPENAI_API_KEY", None)
+    elif api_type == "Azure OpenAI":
+      api_key = os.environ.get("AZURE_OPENAI_API_KEY", None)
+    else:
+      raise RuntimeError("Unknown API type? " + api_type)
+
+
   if api_key:
-    llm = PromptLayerChatOpenAI(temperature=0, model_name="gpt-3.5-turbo", openai_api_key=api_key, pl_tags=["grammar"])
+    shared_args = {
+        "temperature": 0,
+        "model_name": "gpt-3.5-turbo",
+        "openai_api_key": api_key,
+        "pl_tags": ["grammar"],
+    }
+    if api_type == "OpenAI":
+      llm = PromptLayerChatOpenAI(**shared_args)
+    elif api_type == "Azure OpenAI":
+      llm = AzurePromptLayerChatOpenAI(
+        openai_api_base = os.environ.get("AZURE_OPENAI_API_BASE", None),
+        openai_api_version = os.environ.get("AZURE_OPENAI_API_VERSION", "2023-03-15-preview"),
+        deployment_name = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", None),
+        **shared_args
+      )
+
     prompt1 = PromptTemplate(
         input_variables=["content"],
         template=GRAMMAR_PROMPT
@@ -85,7 +114,7 @@ def run(content, chain):
   return chain.run(content)
 
 def run_followup(followup_question, input_vars, chain, chat):
-  
+
   history = [HumanMessage(content=chain.prompt.format(content=m.content)) if isinstance(m, HumanMessage) else m
              for m in chain.memory.chat_memory.messages]
   prompt = ChatPromptTemplate.from_messages([
@@ -189,7 +218,7 @@ def select_diff(evt: gr.SelectData, changes):
     # original, edited = text.split("→")
     return f"Why is it better to change [{original}] to [{edited}]?"
 
-block = gr.Blocks(css="""
+demo = gr.Blocks(css="""
 .diff-component {
   white-space: pre-wrap;
 }
@@ -197,13 +226,14 @@ block = gr.Blocks(css="""
   white-space: normal;
 }
 """)
-with block:
+with demo:
   # api_key = gr.Textbox(
   #     placeholder="Paste your OpenAPI API key here (sk-...)",
   #     show_label=False,
   #     lines=1,
   #     type="password"
   # )
+  api_key = gr.State("")
   gr.HTML("""<div style="display: flex; justify-content: center; align-items: center"><a href="https://thinkcol.com/"><img src="https://thinkcol.com/thinkcol-logo.png" alt="ThinkCol" /></a></div>""")
   gr.Markdown("""Paste a paragraph below, and then choose one of the modes to generate feedback.""")
   content = gr.Textbox(
@@ -291,6 +321,13 @@ with block:
   #   followup_answer_custom = gr.Textbox(
   #       label="Answer"
   #   )
+  with gr.Tab("Settings"):
+    api_type = gr.Radio(
+      ["OpenAI", "Azure OpenAI"],
+      value="OpenAI",
+      label="Server",
+      info="You can try changing this if responses are slow."
+    )
 
   changes = gr.State()
   edited = gr.State()
@@ -301,7 +338,9 @@ with block:
 
   chain_custom = gr.State()
 
-  # api_key.change(load_chain, [api_key], [chain, llm, chain_intro, chain_body1])
+
+  # api_key.change(load_chain, [api_key, api_type], [chain, llm, chain_intro, chain_body1])
+  api_type.change(load_chain, [api_key, api_type], [chain, llm, chain_intro, chain_body1])
 
   inputs = [content, chain]
   outputs = [output_before, output_after, changes, edited]
@@ -321,9 +360,9 @@ with block:
   submit_intro.click(run, [content, chain_intro], output_intro)
   submit_body.click(run_body, [content, title, chain_body1, llm], output_body) # body part A only
   # submit_custom.click(run_custom, [content, llm, prompt], [output_custom, chain_custom]) # TODO standardize api--return memory instead of using chain?
-  
+
   # followup_custom.submit(run_followup, [followup_custom, empty_input, chain_custom, llm], followup_answer_custom)
 
-  block.load(load_chain, gr.State(""), [chain, llm, chain_intro, chain_body1])
+  demo.load(load_chain, [api_key, api_type], [chain, llm, chain_intro, chain_body1])
 
-block.launch(debug=True)
+demo.launch(debug=True)