From 81173447a7473de71a3587efd83f7b85c9e02e9d Mon Sep 17 00:00:00 2001
From: David Cristofaro <david@dtcristo.com>
Date: Fri, 28 Apr 2023 15:35:52 +1000
Subject: [PATCH] Tweaks

---
 Cargo.lock  |  1 -
 Cargo.toml  |  2 +-
 fly.toml    | 42 +++++++++++++----------------------
 src/main.rs | 63 +++++++++++++++++++++++------------------------------
 4 files changed, 43 insertions(+), 65 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4d627c3..0d195bf 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -423,7 +423,6 @@ dependencies = [
  "axum-extra",
  "eventsource-stream",
  "futures",
- "lazy_static",
  "reqwest",
  "serde",
  "serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index 2389f67..bc885b4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -8,7 +8,7 @@ axum = "0.6"
 axum-extra = "0.7"
 eventsource-stream = "0.2"
 futures = "0.3"
-lazy_static = "1"
+# lazy_static = "1"
 # openai = "1.0.0-alpha.8"
 reqwest = { version = "0.11", features = ["json", "stream"] }
 serde = { version = "1", features = ["derive"] }
diff --git a/fly.toml b/fly.toml
index 64e9f35..cc34af6 100644
--- a/fly.toml
+++ b/fly.toml
@@ -5,35 +5,23 @@
 
 app = "gpt-html-8d2f"
 primary_region = "syd"
-kill_signal = "SIGINT"
-kill_timeout = "5s"
 
-[experimental]
-  auto_rollback = true
-
-[env]
-  PRIMARY_REGION = "syd"
-
-[[services]]
-  protocol = "tcp"
-  internal_port = 9292
-  processes = ["app"]
-
-  [[services.ports]]
-    port = 80
-    handlers = ["http"]
-    force_https = true
-
-  [[services.ports]]
-    port = 443
-    handlers = ["tls", "http"]
-  [services.concurrency]
-    type = "connections"
+[http_service]
+  internal_port = 8080
+  force_https = true
+  auto_stop_machines = true
+  auto_start_machines = true
+  [http_service.concurrency]
     hard_limit = 25
     soft_limit = 20
+    type = "requests"
 
-  [[services.tcp_checks]]
-    interval = "15s"
+[checks]
+  [checks.health]
+    grace_period = "10s"
+    interval = "10s"
+    method = "get"
+    path = "/health"
+    port = 8080
     timeout = "2s"
-    grace_period = "1s"
-    restart_limit = 0
+    type = "http"
diff --git a/src/main.rs b/src/main.rs
index 75f2e42..16595b4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,7 +11,6 @@ use axum_extra::middleware::option_layer;
 use eventsource_stream::Eventsource;
 use futures::future;
 use futures::{StreamExt, TryStreamExt};
-use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 use std::{
     env,
@@ -24,22 +23,11 @@ use tower_http::{services::ServeDir, validate_request::ValidateRequestHeaderLaye
 
 mod error;
 
-lazy_static! {
-    static ref OPENAI_API_KEY: String =
-        env::var("OPENAI_API_KEY").expect("`OPENAI_API_KEY` environment variable should be set");
-    static ref HTTP_BASIC_AUTH_PASSWORD: String = env::var("HTTP_BASIC_AUTH_PASSWORD")
-        .expect("`HTTP_BASIC_AUTH_PASSWORD` environment variable should be set");
-    static ref COMMIT_SHA: String =
-        env::var("COMMIT_SHA").expect("`COMMIT_SHA` environment variable should be set");
-}
-
 #[tokio::main]
 async fn main() {
-    println!("starting server...");
-
-    load_environment();
+    println!("Starting server...");
 
-    let addr = SocketAddr::from(([0, 0, 0, 0], 9292));
+    let addr = SocketAddr::from(([0, 0, 0, 0], 8080));
     Server::bind(&addr)
         .serve(app().into_make_service())
         .with_graceful_shutdown(shutdown_signal())
@@ -97,8 +85,9 @@ async fn shutdown_signal() {
 
 #[derive(Debug, Serialize)]
 struct HealthBody {
-    time: String,
+    time: u64,
     commit_sha: String,
+    basic_auth_enabled: bool,
 }
 
 #[derive(Debug, Serialize)]
@@ -129,35 +118,31 @@ struct Delta {
     content: String,
 }
 
-fn load_environment() {
-    assert!(
-        (*OPENAI_API_KEY).len() > 0,
-        "`OPENAI_API_KEY` should not be empty"
-    );
-    assert!(
-        (*HTTP_BASIC_AUTH_PASSWORD).len() > 0,
-        "`HTTP_BASIC_AUTH_PASSWORD` should not be empty"
-    );
-    assert!((*COMMIT_SHA).len() > 0, "`COMMIT_SHA` should not be empty");
-}
-
 async fn health() -> Result<impl IntoResponse> {
-    println!("health");
+    println!("\n----------");
+    println!("Health");
+    println!("----------");
+
+    env::var("OPENAI_API_KEY").map_err(|_| Error::EnvironmentError)?;
 
     let time = SystemTime::now()
         .duration_since(SystemTime::UNIX_EPOCH)
         .map_err(|_| Error::SystemTimeError)?
-        .as_secs()
-        .to_string();
-    let commit_sha = (*COMMIT_SHA).clone();
-
-    Ok(Json(HealthBody { time, commit_sha }))
+        .as_secs();
+    let commit_sha = env::var("COMMIT_SHA").unwrap_or_else(|_| "unknown".to_string());
+    let basic_auth_enabled = env::var("HTTP_BASIC_AUTH_PASSWORD").is_ok();
+
+    Ok(Json(HealthBody {
+        time,
+        basic_auth_enabled,
+        commit_sha,
+    }))
 }
 
 async fn handler(uri: Uri) -> Result<impl IntoResponse> {
-    println!("");
+    println!("\n----------");
     println!("Fetching: {uri}");
-    println!("---------\n");
+    println!("----------");
 
     let prompt = r#"
 Output a valid HTML document for the webpage that could be located at the URL path provided by the user. Include general navigation anchor tags as well as relative anchor tags to other related pages. Include a minimal amount of inline styles to improve the look of the page. Make the text content quite long with a decent amount of interesting content. Do not use any dummy text on the page.
@@ -185,7 +170,13 @@ Start the reponse with the following exact characters:
     let stream = reqwest::Client::new()
         .post("https://api.openai.com/v1/chat/completions")
         .header("content-type", "application/json")
-        .header("authorization", &format!("Bearer {}", *OPENAI_API_KEY))
+        .header(
+            "authorization",
+            &format!(
+                "Bearer {}",
+                env::var("OPENAI_API_KEY").map_err(|_| Error::EnvironmentError)?
+            ),
+        )
         .body(serde_json::to_string(&body).map_err(|_| Error::SerializationError)?)
         .send()
         .await