From a091705aea26d86cf45f06e7ad28bebdfc12ee8a Mon Sep 17 00:00:00 2001
From: "K.Y" <beidongjiedeguang@gmail.com>
Date: Thu, 16 Nov 2023 18:52:19 +0800
Subject: [PATCH] Consider all openai headers (#93)

---
 Examples/beta.py               | 18 ++++++++++++++++++
 openai_forward/forward/base.py | 18 +++++++++++++-----
 2 files changed, 31 insertions(+), 5 deletions(-)
 create mode 100644 Examples/beta.py

diff --git a/Examples/beta.py b/Examples/beta.py
new file mode 100644
index 0000000..36840a5
--- /dev/null
+++ b/Examples/beta.py
@@ -0,0 +1,18 @@
+from openai import OpenAI
+from sparrow import yaml_load
+
+config = yaml_load("config.yaml", rel_path=True)
+print(f"{config=}")
+
+client = OpenAI(
+    api_key=config['api_key'],
+    base_url=config['api_base'],
+)
+
+my_assistant = client.beta.assistants.create(
+    instructions="You are a personal math tutor. When asked a question, write and run Python code to answer the question.",
+    name="Math Tutor",
+    tools=[{"type": "code_interpreter"}],
+    model="gpt-4",
+)
+print(my_assistant)
diff --git a/openai_forward/forward/base.py b/openai_forward/forward/base.py
index 5aacd32..d2ca1bb 100644
--- a/openai_forward/forward/base.py
+++ b/openai_forward/forward/base.py
@@ -138,11 +138,19 @@ def prepare_client(self, request: Request, return_origin_header=False) -> dict:
         )
         url = f"{self.BASE_URL}{url_path}?{request.url.query}"
 
-        headers = dict(request.headers)
-        auth = headers.get("authorization", "")
-        content_type = headers.get("content-type", "application/json")
-        if not return_origin_header:
-            headers = {"Content-Type": content_type, "Authorization": auth}
+        auth = request.headers.get("Authorization", "")
+
+        if return_origin_header:
+            headers = request.headers
+        else:
+            headers = {
+                "content-type": request.headers.get("content-type", "application/json"),
+                "authorization": auth,
+            }
+            for key, value in request.headers.items():
+                if key.startswith("openai"):
+                    headers[key] = value
+            print(f"{headers=}")
 
         return {
             'auth': auth,