From d914bbb5c1931d3d7c844748db570f2b4fc29f97 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sun, 26 Jun 2022 23:31:15 +0200
Subject: [PATCH] Improve API

---
 .gitignore                                    |  1 +
 docs/hello_world.ipynb                        | 26 +++++++++++++++++++
 src/great_ai/__main__.py                      | 21 ++++++++++++---
 .../remote/call_remote_great_ai_async.py      |  3 +++
 4 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 docs/hello_world.ipynb

diff --git a/.gitignore b/.gitignore
index 5603876..679a4eb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@ __pycache__
 .ipynb_checkpoints
 *.egg-info
 build
+tracing_database.json
diff --git a/docs/hello_world.ipynb b/docs/hello_world.ipynb
new file mode 100644
index 0000000..745ac90
--- /dev/null
+++ b/docs/hello_world.ipynb
@@ -0,0 +1,26 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from great_ai import GreatAI\n",
+    "\n",
+    "\n",
+    "@GreatAI.create\n",
+    "def hello_world(name: str) -> str:\n",
+    "    return f\"Hello {name}!\"\n"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/src/great_ai/__main__.py b/src/great_ai/__main__.py
index e24af7f..4d03f3c 100644
--- a/src/great_ai/__main__.py
+++ b/src/great_ai/__main__.py
@@ -2,9 +2,9 @@
 
 import logging
 import re
-import time
 from importlib import import_module, reload
 from pathlib import Path
+from threading import Event
 from typing import Optional
 
 import uvicorn
@@ -65,7 +65,21 @@ def main() -> None:
 
         logger.info(f"Starting uvicorn server with app={app}")
 
-        uvicorn.run(app, **common_config)
+        config = Config(app, **common_config)
+        socket = config.bind_socket()
+        server = GreatAIReload(
+            config, target=uvicorn.Server(config=config).run, sockets=[socket]
+        )
+
+        server.startup()
+        try:
+            Event().wait()
+        finally:
+            server.shutdown()
+            if args.file_name.endswith(".ipynb"):
+                Path(get_script_name_of_notebook(args.file_name)).unlink(
+                    missing_ok=True
+                )
     else:
 
         class EventHandler(PatternMatchingEventHandler):
@@ -106,8 +120,7 @@ def main() -> None:
         observer.start()
 
         try:
-            while True:
-                time.sleep(50)
+            Event().wait()
         finally:
             observer.stop()
             restart_handler.stop_server()
diff --git a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
index 4e99d08..78f059a 100644
--- a/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
+++ b/src/great_ai/great_ai/remote/call_remote_great_ai_async.py
@@ -14,6 +14,9 @@ async def call_remote_great_ai_async(
     if http is None:
         http = HttpClient()
 
+    if base_uri.endswith("/"):
+        base_uri = base_uri[:-1]
+
     url = f"{base_uri}/predict/"
     response = await http.post(
         url=url, data=data, retry_count=retry_count, expected_status=200