From 24fae660e6913aac6b52745fb2c2fe128ba2eb79 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Mart=C3=ADnez?= <ivanmartit@gmail.com>
Date: Fri, 2 Feb 2024 16:14:22 +0100
Subject: [PATCH] feat: Add stream information to generate SDKs (#1569)

---
 fern/openapi/openapi.json                      | 18 ++++++++++++++++++
 private_gpt/server/chat/chat_router.py         |  7 +++++++
 .../server/completions/completions_router.py   |  7 +++++++
 3 files changed, 32 insertions(+)

diff --git a/fern/openapi/openapi.json b/fern/openapi/openapi.json
index 6c6635b9..af1646f1 100644
--- a/fern/openapi/openapi.json
+++ b/fern/openapi/openapi.json
@@ -44,6 +44,15 @@
               }
             }
           }
+        },
+        "x-fern-streaming": {
+          "stream-condition": "stream",
+          "response": {
+            "$ref": "#/components/schemas/OpenAICompletion"
+          },
+          "response-stream": {
+            "$ref": "#/components/schemas/OpenAICompletion"
+          }
         }
       }
     },
@@ -86,6 +95,15 @@
               }
             }
           }
+        },
+        "x-fern-streaming": {
+          "stream-condition": "stream",
+          "response": {
+            "$ref": "#/components/schemas/OpenAICompletion"
+          },
+          "response-stream": {
+            "$ref": "#/components/schemas/OpenAICompletion"
+          }
         }
       }
     },
diff --git a/private_gpt/server/chat/chat_router.py b/private_gpt/server/chat/chat_router.py
index e493c348..385ba173 100644
--- a/private_gpt/server/chat/chat_router.py
+++ b/private_gpt/server/chat/chat_router.py
@@ -54,6 +54,13 @@ class ChatBody(BaseModel):
     response_model=None,
     responses={200: {"model": OpenAICompletion}},
     tags=["Contextual Completions"],
+    openapi_extra={
+        "x-fern-streaming": {
+            "stream-condition": "stream",
+            "response": {"$ref": "#/components/schemas/OpenAICompletion"},
+            "response-stream": {"$ref": "#/components/schemas/OpenAICompletion"},
+        }
+    },
 )
 def chat_completion(
     request: Request, body: ChatBody
diff --git a/private_gpt/server/completions/completions_router.py b/private_gpt/server/completions/completions_router.py
index da9583a7..646e54b8 100644
--- a/private_gpt/server/completions/completions_router.py
+++ b/private_gpt/server/completions/completions_router.py
@@ -42,6 +42,13 @@ class CompletionsBody(BaseModel):
     summary="Completion",
     responses={200: {"model": OpenAICompletion}},
     tags=["Contextual Completions"],
+    openapi_extra={
+        "x-fern-streaming": {
+            "stream-condition": "stream",
+            "response": {"$ref": "#/components/schemas/OpenAICompletion"},
+            "response-stream": {"$ref": "#/components/schemas/OpenAICompletion"},
+        }
+    },
 )
 def prompt_completion(
     request: Request, body: CompletionsBody