feat: allow for adding new replicate models

UKP-SQuARE · May 14, 2024 · a541740 · a541740
1 parent a284aec
commit a541740
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 6 deletions.
diff --git a/frontend/src/api/index.js b/frontend/src/api/index.js
@@ -322,3 +322,8 @@ export function generateCompletionStreamReplicate(params, token) {
     });
     return response;
 }
+
+
+export function replicateModelHealthCheck(model_id){
+    return axios.get(`${BASE_SQUARE_URL}/replicate/${model_id}/status`);
+}
diff --git a/frontend/src/views/PromptingView.vue b/frontend/src/views/PromptingView.vue
@@ -53,10 +53,23 @@
                       <div class="tab-pane fade" id="pills-replicate" role="tabpanel" aria-labelledby="pills-replicate-tab">
                         <label for="selectedModel" class="form-label mt-2">Model</label>
                         <select v-model="chatConfig.selectedModel" class="form-select" id="selectedModel">
+                          <option value="add-new">Add new model...</option>
                           <option v-for="model in replicateModels" :key="model" :value="model">
                             {{ model }}
                           </option>
                         </select>
+                        <div v-if="chatConfig.selectedModel == 'add-new'" class="mt-3">
+
+                          <div class="row">
+                            <div class="col-9 pe-0">
+                              <input type="text" class="form-control" placeholder="Model Identifier" v-model="newReplicateModel" />
+                            </div>
+                            <div class="col-3 ps-1">
+                              <button type="button" class="btn btn-primary w-100" @click="addNewReplicateModel" >Add</button>
+                            </div>
+                          </div>
+                        </div>
+
                         <div class="row mt-3">
                           <div>
                             <label for="replicate-key" class="form-label">
@@ -560,7 +573,8 @@ import {
   getOpenAIModels,
   getLocalLLMs, 
   getAlternatives, 
-  getReplicateModels
+  getReplicateModels,
+  replicateModelHealthCheck
 } from '@/api';
 import { 
   CustomChatModel, 
@@ -677,12 +691,11 @@ export default {
       sentence: 'The mechanical doll next itself loose.', 
       answer: '0'
     }], 
-
     alternativesWaiting: false,
-
     sensitivityLog: "",
-
     isGenerating: false,
+
+    newReplicateModel: "",
   }),
 
   mounted() {
@@ -759,6 +772,22 @@ export default {
   },
 
   methods: {
+    async addNewReplicateModel(){
+      if (this.newReplicateModel === "") return;
+      if (!this.replicateModels.includes(this.newReplicateModel)){
+        try {
+          await replicateModelHealthCheck(this.newReplicateModel);
+        } catch (e) {
+          this.showErrorToast("Model not found.");
+          return;
+        }
+        console.log("new model added")
+        this.replicateModels.push(this.newReplicateModel);
+      }
+      this.chatConfig.selectedModel = this.newReplicateModel;
+      this.newReplicateModel = "";
+    },
+
     showPopover(event) {
       // this will create a popover once and then it will be reused by Popover automatically
       if (!this.sensitivity_popover) {

diff --git a/replicate/api.http b/replicate/api.http
@@ -82,4 +82,10 @@ Authorization: Bearer {{token}}
     "temperature": 0.7,
     "max_new_tokens": 100,
     "prompt": "What is the theory of relativity?"
-}
+}
+
+
+###
+# @name status
+# @prompt model_id
+GET {{domain}}/{{model_id}}/status HTTP/1.1
diff --git a/replicate/app/app.py b/replicate/app/app.py
@@ -9,6 +9,7 @@
 from replicate.client import Client
 from replicate.exceptions import ReplicateError
 from jinja2 import Template
+import requests
 
 
 router = APIRouter()
@@ -149,6 +150,28 @@ async def generate_completion(request: Request, token: str = Depends(get_token))
     })
 
 
+@router.get("/{identifier}/status")
+@router.get("/{username}/{identifier}/status")
+async def status_check(identifier: str, username: str):
+    if username is not None:
+        model_id = f"{username}/{identifier}"
+    else:
+        model_id = identifier
+
+    url = f"https://replicate.com/{model_id}/status"
+
+    print("Checking status of model: ", url)
+
+    response = requests.get(url)
+
+    if response.status_code == 200:
+        return JSONResponse({
+            "status": "Model is available"
+        })
+    else:
+        raise HTTPException(status_code=404, detail="Model not found")
+
+
 def get_app() -> FastAPI:
     fast_app = FastAPI(
         title="Replicate.ai Client",

diff --git a/replicate/requirements.txt b/replicate/requirements.txt
@@ -2,4 +2,5 @@ replicate==0.25.2
 jinja2==3.1.4
 sentencepiece==0.2.0
 protobuf==5.26.1
-fastapi==0.111.0
+fastapi==0.111.0
+requests==2.31.0