From e25c51ceae8291ef5e6e5b164a75b7c5bfca5fdb Mon Sep 17 00:00:00 2001
From: Flavio Brasil <fwbrasil@gmail.com>
Date: Tue, 5 Dec 2023 14:25:36 -0800
Subject: [PATCH] ais: openai org header

---
 kyo-llm/shared/src/main/scala/kyo/llm/completions.scala | 6 +++++-
 kyo-llm/shared/src/main/scala/kyo/llm/configs.scala     | 5 +++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala b/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala
index 7750a1397..508490606 100644
--- a/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala
+++ b/kyo-llm/shared/src/main/scala/kyo/llm/completions.scala
@@ -58,7 +58,11 @@ object completions {
         Configs.get.map { cfg =>
           Requests[Response](
               _.contentType("application/json")
-                .header("Authorization", s"Bearer $key")
+                .headers(
+                    Map(
+                        "Authorization" -> s"Bearer $key"
+                    ) ++ cfg.apiOrg.map("OpenAI-Organization" -> _)
+                )
                 .post(uri"${cfg.apiUrl}/v1/chat/completions")
                 .body(req)
                 .readTimeout(Duration.Inf)
diff --git a/kyo-llm/shared/src/main/scala/kyo/llm/configs.scala b/kyo-llm/shared/src/main/scala/kyo/llm/configs.scala
index c9b407605..e087345ac 100644
--- a/kyo-llm/shared/src/main/scala/kyo/llm/configs.scala
+++ b/kyo-llm/shared/src/main/scala/kyo/llm/configs.scala
@@ -22,6 +22,7 @@ object configs {
   case class Config(
       apiUrl: String,
       apiKey: Option[String],
+      apiOrg: Option[String],
       model: Model,
       temperature: Double,
       maxTokens: Option[Int],
@@ -50,12 +51,16 @@ object configs {
   object Config {
     val default = {
       val apiKeyProp = "OPENAI_API_KEY"
+      val apiOrgProp = "OPENAI_API_ORG"
       val apiKey =
         Option(System.getenv(apiKeyProp))
           .orElse(Option(System.getProperty(apiKeyProp)))
+      val apiOrg =
+        Option(System.getenv(apiOrgProp))
       Config(
           "https://api.openai.com",
           apiKey,
+          apiOrg,
           Model.gpt4_turbo,
           0.2,
           None,