From b94d862cae1ee23e30008871cda679468be0a63e Mon Sep 17 00:00:00 2001
From: Simon Larsen <simonlarsen@oneuptime.com>
Date: Fri, 28 Jun 2024 12:14:49 +0100
Subject: [PATCH] chore: Update LLM Dockerfile and build process

---
 CommonServer/API/CodeRepositoryAPI.ts | 15 +++++++++
 LLM/app.py                            | 48 ++++++++++++++++++++++++++-
 LLM/requirements.txt                  |  1 +
 config.example.env                    |  3 ++
 docker-compose.base.yml               |  1 +
 5 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/CommonServer/API/CodeRepositoryAPI.ts b/CommonServer/API/CodeRepositoryAPI.ts
index 83bb193428..c708968d54 100644
--- a/CommonServer/API/CodeRepositoryAPI.ts
+++ b/CommonServer/API/CodeRepositoryAPI.ts
@@ -23,6 +23,21 @@ export default class CodeRepositoryAPI extends BaseAPI<
   public constructor() {
     super(CodeRepository, CodeRepositoryService);
 
+    this.router.get(
+      `${new this.entityType()
+        .getCrudApiPath()
+        ?.toString()}/is-valid/:secretkey`,
+      CodeRepositoryAuthorization.isAuthorizedRepository,
+      async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
+        try {
+          return Response.sendEmptySuccessResponse(req, res);
+        } catch (err) {
+          next(err);
+        }
+      },
+    );
+  
+
     this.router.get(
       `${new this.entityType()
         .getCrudApiPath()
diff --git a/LLM/app.py b/LLM/app.py
index 6a733084e0..9482ae95bb 100644
--- a/LLM/app.py
+++ b/LLM/app.py
@@ -1,17 +1,46 @@
 import uuid
 import transformers
 import asyncio
+import os
 import torch
+import aiohttp
 from fastapi import FastAPI
 from pydantic import BaseModel
 from contextlib import asynccontextmanager
 from apscheduler.schedulers.background import BackgroundScheduler
 
+# ENV VARS
+ONEUPTIME_URL = os.getenv("ONEUPTIME_URL")
+
+if not ONEUPTIME_URL:
+    ONEUPTIME_URL = "https://oneuptime.com"
+
+print(f"ONEUPTIME_URL: {ONEUPTIME_URL}")
+
 # TODO: Store this in redis down the line. 
 items_pending = {}
 items_processed = {}
 errors = {}
 
+async def validateSecretKey(secretKey):
+    try:
+
+        # If no secret key then return false 
+        if not secretKey:
+            return False
+
+        async with aiohttp.ClientSession() as session:
+            url = f"{ONEUPTIME_URL}/api/code-repository/is-valid/{secretKey}"
+            async with session.get(url) as response:
+                if response.status == 200:
+                    return True
+                else:
+                    return False
+                
+    except Exception as e:
+        print(repr(e))
+        return False
+
 async def job(queue):
     print("Processing queue...")
 
@@ -61,10 +90,12 @@ async def lifespan(app:FastAPI):
 # Declare a Pydantic model for the request body
 class Prompt(BaseModel):
    messages: list
+   secretkey: str
 
 # Declare a Pydantic model for the request body
 class PromptResult(BaseModel):
    id: str
+   secretkey: str
 
 app = FastAPI(lifespan=lifespan)
 
@@ -80,6 +111,13 @@ async def create_item(prompt: Prompt):
         if not prompt:
             return {"error": "Prompt is required"}
         
+        # Validate the secret key
+        is_valid = await validateSecretKey(prompt_status.secretkey)
+
+        if not is_valid:
+            print("Invalid secret key")
+            return {"error": "Invalid secret key"}
+        
         # messages are in str format. We need to convert them fron json [] to list
         messages = prompt.messages
 
@@ -103,6 +141,7 @@ async def create_item(prompt: Prompt):
         print(e)
         return {"error": repr(e)}
 
+# Disable this API in production
 @app.get("/queue-status/")
 async def queue_status():
     try: 
@@ -116,7 +155,14 @@ async def prompt_status(prompt_status: PromptResult):
     try:
         # Log prompt status to console
         print(prompt_status)
-    
+
+        # Validate the secret key
+        is_valid = await validateSecretKey(prompt_status.secretkey)
+
+        if not is_valid:
+            print("Invalid secret key")
+            return {"error": "Invalid secret key"}
+        
         # If not prompt status then return bad request error
         if not prompt_status:
             return {"error": "Prompt status is required"}
diff --git a/LLM/requirements.txt b/LLM/requirements.txt
index c6a6e2b311..3a9190a066 100644
--- a/LLM/requirements.txt
+++ b/LLM/requirements.txt
@@ -8,3 +8,4 @@ pydantic===2.4.2
 
 # Rest of the app
 APScheduler===3.10.4
+aiohttp===3.9.5
diff --git a/config.example.env b/config.example.env
index 4a7eeed77b..108057993a 100644
--- a/config.example.env
+++ b/config.example.env
@@ -240,3 +240,6 @@ ALLOWED_ACTIVE_MONITOR_COUNT_IN_FREE_PLAN=10
 
 # This webhook notifies slack when the new user signs up or is created. 
 NOTIFICATION_WEBHOOK_ON_CREATED_USER= 
+
+# This is the OneUptime Server URL LLM container will use to validate requests. It should be the url where you host the server.
+LLM_ONEUPTIME_SERVER_URL=https://localhost
diff --git a/docker-compose.base.yml b/docker-compose.base.yml
index 58056503ab..238260aa19 100644
--- a/docker-compose.base.yml
+++ b/docker-compose.base.yml
@@ -188,6 +188,7 @@ services:
         environment:
             <<: *common-server-variables
             PORT: 8547
+            ONEUPTIME_URL: ${LLM_ONEUPTIME_SERVER_URL}
         volumes:
             - ./LLM/Models:/app/Models
         logging: