From 9c1b0d3bf7b644f63b6998d409458af2031a5ef3 Mon Sep 17 00:00:00 2001
From: Shubham Naik <shubham.naik10@gmail.com>
Date: Thu, 7 Nov 2024 11:17:09 -0800
Subject: [PATCH] =?UTF-8?q?feat:=20add=20bash=20script=20to=20allow=20peop?=
 =?UTF-8?q?le=20to=20generate=20openapi=20schema=20with=E2=80=A6=20(#2015)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Shubham Naik <shub@memgpt.ai>
---
 letta/server/generate_openapi_schema.sh | 12 +++++
 letta/server/rest_api/app.py            | 60 +++++++++++++------------
 2 files changed, 44 insertions(+), 28 deletions(-)
 create mode 100755 letta/server/generate_openapi_schema.sh

diff --git a/letta/server/generate_openapi_schema.sh b/letta/server/generate_openapi_schema.sh
new file mode 100755
index 0000000000..f8b299dd97
--- /dev/null
+++ b/letta/server/generate_openapi_schema.sh
@@ -0,0 +1,12 @@
+#!/bin/sh
+echo "Generating OpenAPI schema..."
+
+# check if poetry is installed
+if ! command -v poetry &> /dev/null
+then
+    echo "Poetry could not be found. Please install poetry to generate the OpenAPI schema."
+    exit
+fi
+
+# generate OpenAPI schema
+poetry run python -c 'from letta.server.rest_api.app import app, generate_openapi_schema; generate_openapi_schema(app);'
diff --git a/letta/server/rest_api/app.py b/letta/server/rest_api/app.py
index f8300a9719..fb4c4049e3 100644
--- a/letta/server/rest_api/app.py
+++ b/letta/server/rest_api/app.py
@@ -63,6 +63,37 @@
 log = logging.getLogger("uvicorn")
 
 
+def generate_openapi_schema(app: FastAPI):
+    # Update the OpenAPI schema
+    if not app.openapi_schema:
+        app.openapi_schema = app.openapi()
+
+    openai_docs, letta_docs = [app.openapi_schema.copy() for _ in range(2)]
+
+    openai_docs["paths"] = {k: v for k, v in openai_docs["paths"].items() if k.startswith("/openai")}
+    openai_docs["info"]["title"] = "OpenAI Assistants API"
+    letta_docs["paths"] = {k: v for k, v in letta_docs["paths"].items() if not k.startswith("/openai")}
+    letta_docs["info"]["title"] = "Letta API"
+    letta_docs["components"]["schemas"]["LettaResponse"] = {
+        "properties": LettaResponse.model_json_schema(ref_template="#/components/schemas/LettaResponse/properties/{model}")["$defs"]
+    }
+
+    # Split the API docs into Letta API, and OpenAI Assistants compatible API
+    for name, docs in [
+        (
+            "openai",
+            openai_docs,
+        ),
+        (
+            "letta",
+            letta_docs,
+        ),
+    ]:
+        if settings.cors_origins:
+            docs["servers"] = [{"url": host} for host in settings.cors_origins]
+        Path(f"openapi_{name}.json").write_text(json.dumps(docs, indent=2))
+
+
 def create_application() -> "FastAPI":
     """the application start routine"""
     # global server
@@ -122,34 +153,7 @@ def on_startup():
 
         # Tool.load_default_tools(get_db_session())
 
-        # Update the OpenAPI schema
-        if not app.openapi_schema:
-            app.openapi_schema = app.openapi()
-
-        openai_docs, letta_docs = [app.openapi_schema.copy() for _ in range(2)]
-
-        openai_docs["paths"] = {k: v for k, v in openai_docs["paths"].items() if k.startswith("/openai")}
-        openai_docs["info"]["title"] = "OpenAI Assistants API"
-        letta_docs["paths"] = {k: v for k, v in letta_docs["paths"].items() if not k.startswith("/openai")}
-        letta_docs["info"]["title"] = "Letta API"
-        letta_docs["components"]["schemas"]["LettaResponse"] = {
-            "properties": LettaResponse.model_json_schema(ref_template="#/components/schemas/LettaResponse/properties/{model}")["$defs"]
-        }
-
-        # Split the API docs into Letta API, and OpenAI Assistants compatible API
-        for name, docs in [
-            (
-                "openai",
-                openai_docs,
-            ),
-            (
-                "letta",
-                letta_docs,
-            ),
-        ]:
-            if settings.cors_origins:
-                docs["servers"] = [{"url": host} for host in settings.cors_origins]
-            Path(f"openapi_{name}.json").write_text(json.dumps(docs, indent=2))
+        generate_openapi_schema(app)
 
     @app.on_event("shutdown")
     def on_shutdown():