From c4b5f2be476c9f4494d9c105e6b78c6e01849f60 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sat, 1 Jun 2024 22:59:27 -0700
Subject: [PATCH] enh: azure openai pipeline

---
 examples/providers/azure_openai_pipeline.py | 26 ++++++++++++++-------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/examples/providers/azure_openai_pipeline.py b/examples/providers/azure_openai_pipeline.py
index 994c3784..088d8a8b 100644
--- a/examples/providers/azure_openai_pipeline.py
+++ b/examples/providers/azure_openai_pipeline.py
@@ -1,9 +1,19 @@
 from typing import List, Union, Generator, Iterator
 from schemas import OpenAIChatMessage
+from pydantic import BaseModel
 import requests
 
 
 class Pipeline:
+    class Valves(BaseModel):
+        # You can add your custom valves here.
+        AZURE_OPENAI_API_KEY: str = "your-azure-openai-api-key-here"
+        AZURE_OPENAI_ENDPOINT: str = "your-azure-openai-endpoint-here"
+        DEPLOYMENT_NAME: str = "your-deployment-name-here"
+        API_VERSION: str = "2023-10-01-preview"
+        MODEL: str = "gpt-3.5-turbo"
+        pass
+
     def __init__(self):
         # Optionally, you can set the id and name of the pipeline.
         # Assign a unique identifier to the pipeline.
@@ -11,6 +21,7 @@ def __init__(self):
         # The identifier must be an alphanumeric string that can include underscores or hyphens. It cannot contain spaces, special characters, slashes, or backslashes.
         self.id = "azure_openai_pipeline"
         self.name = "Azure OpenAI Pipeline"
+        self.valves = self.Valves()
         pass
 
     async def on_startup(self):
@@ -32,25 +43,22 @@ def pipe(
         print(messages)
         print(user_message)
 
-        AZURE_OPENAI_API_KEY = "your-azure-openai-api-key-here"
-        AZURE_OPENAI_ENDPOINT = "your-azure-openai-endpoint-here"
-        DEPLOYMENT_NAME = "your-deployment-name-here"
-        MODEL = "gpt-3.5-turbo"
+        headers = {
+            "api-key": self.valves.AZURE_OPENAI_API_KEY,
+            "Content-Type": "application/json",
+        }
 
-        headers = {"api-key": AZURE_OPENAI_API_KEY, "Content-Type": "application/json"}
-
-        url = f"{AZURE_OPENAI_ENDPOINT}/openai/deployments/{DEPLOYMENT_NAME}/chat/completions?api-version=2023-10-01-preview"
+        url = f"{self.valves.AZURE_OPENAI_ENDPOINT}/openai/deployments/{self.valves.DEPLOYMENT_NAME}/chat/completions?api-version={self.valves.API_VERSION}"
 
         try:
             r = requests.post(
                 url=url,
-                json={**body, "model": MODEL},
+                json={**body, "model": self.valves.MODEL},
                 headers=headers,
                 stream=True,
             )
 
             r.raise_for_status()
-
             if body["stream"]:
                 return r.iter_lines()
             else: