From 69d5c25e9d822aadbd1fe8a23ea4449acdecf2d8 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Tue, 21 May 2024 20:10:25 -0700
Subject: [PATCH] feat: azure openai pipeline

---
 pipelines/examples/azure_openai_pipeline.py | 55 +++++++++++++++++++++
 pipelines/examples/openai_pipeline.py       |  4 --
 2 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 pipelines/examples/azure_openai_pipeline.py

diff --git a/pipelines/examples/azure_openai_pipeline.py b/pipelines/examples/azure_openai_pipeline.py
new file mode 100644
index 0000000..2e1dcb6
--- /dev/null
+++ b/pipelines/examples/azure_openai_pipeline.py
@@ -0,0 +1,55 @@
+from typing import List, Union, Generator, Iterator
+from schemas import OpenAIChatMessage
+import requests
+
+
+class Pipeline:
+    def __init__(self):
+        # Optionally, you can set the id and name of the pipeline.
+        self.id = "azure_openai_pipeline"
+        self.name = "Azure OpenAI Pipeline"
+        pass
+
+    async def on_startup(self):
+        # This function is called when the server is started.
+        print(f"on_startup:{__name__}")
+        pass
+
+    async def on_shutdown(self):
+        # This function is called when the server is stopped.
+        print(f"on_shutdown:{__name__}")
+        pass
+
+    def get_response(
+        self, user_message: str, messages: List[OpenAIChatMessage], body: dict
+    ) -> Union[str, Generator, Iterator]:
+        # This is where you can add your custom pipelines like RAG.'
+        print(f"get_response:{__name__}")
+
+        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"
+
+        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"
+
+        try:
+            r = requests.post(
+                url=url,
+                json={**body, "model": "gpt-3.5-turbo"},
+                headers=headers,
+                stream=True,
+            )
+
+            r.raise_for_status()
+
+            if data["stream"]:
+                return r.iter_lines()
+            else:
+                return r.json()
+        except Exception as e:
+            return f"Error: {e}"
diff --git a/pipelines/examples/openai_pipeline.py b/pipelines/examples/openai_pipeline.py
index 79c650b..54f99c3 100644
--- a/pipelines/examples/openai_pipeline.py
+++ b/pipelines/examples/openai_pipeline.py
@@ -34,10 +34,6 @@ class Pipeline:
         headers["Authorization"] = f"Bearer {OPENAI_API_KEY}"
         headers["Content-Type"] = "application/json"
 
-        data = {**body, "model": "gpt-3.5-turbo"}
-
-        print(data)
-
         try:
             r = requests.post(
                 url="https://api.openai.com/v1/chat/completions",