Rework user scripts, update DB schema to support embeddings

This commit is contained in:
2023-08-19 15:46:19 -04:00
parent 27188b6fa0
commit 857f3315c2
62 changed files with 1842 additions and 1250 deletions

View File

@@ -7,7 +7,7 @@
to the <a href="https://github.com/simon987/sist2/issues/new/choose" target="_blank">issue tracker on
Github</a>. Thank you!
</b-alert>
<router-view/>
<router-view v-if="$store.state.sist2AdminInfo"/>
</b-container>
</div>
</template>
@@ -71,10 +71,12 @@ html, body {
.info-icon {
width: 1rem;
min-width: 1rem;
margin-right: 0.2rem;
cursor: pointer;
line-height: 1rem;
height: 1rem;
min-height: 1rem;
background-image: url(data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHhtbG5zOnhsaW5rPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5L3hsaW5rIiB4PSIwcHgiIHk9IjBweCIKICAgICB2aWV3Qm94PSIwIDAgNDI2LjY2NyA0MjYuNjY3IiBzdHlsZT0iZW5hYmxlLWJhY2tncm91bmQ6bmV3IDAgMCA0MjYuNjY3IDQyNi42Njc7IiBmaWxsPSIjZmZmIj4KPGc+CiAgICA8Zz4KICAgICAgICA8Zz4KICAgICAgICAgICAgPHJlY3QgeD0iMTkyIiB5PSIxOTIiIHdpZHRoPSI0Mi42NjciIGhlaWdodD0iMTI4Ii8+CiAgICAgICAgICAgIDxwYXRoIGQ9Ik0yMTMuMzMzLDBDOTUuNDY3LDAsMCw5NS40NjcsMCwyMTMuMzMzczk1LjQ2NywyMTMuMzMzLDIxMy4zMzMsMjEzLjMzM1M0MjYuNjY3LDMzMS4yLDQyNi42NjcsMjEzLjMzMwogICAgICAgICAgICAgICAgUzMzMS4yLDAsMjEzLjMzMywweiBNMjEzLjMzMywzODRjLTk0LjA4LDAtMTcwLjY2Ny03Ni41ODctMTcwLjY2Ny0xNzAuNjY3UzExOS4yNTMsNDIuNjY3LDIxMy4zMzMsNDIuNjY3CiAgICAgICAgICAgICAgICBTMzg0LDExOS4yNTMsMzg0LDIxMy4zMzNTMzA3LjQxMywzODQsMjEzLjMzMywzODR6Ii8+CiAgICAgICAgICAgIDxyZWN0IHg9IjE5MiIgeT0iMTA2LjY2NyIgd2lkdGg9IjQyLjY2NyIgaGVpZ2h0PSI0Mi42NjciLz4KICAgICAgICA8L2c+CiAgICA8L2c+CjwvZz4KPC9zdmc+Cg==);
filter: brightness(45%);
display: block;

View File

@@ -139,6 +139,38 @@ class Sist2AdminApi {
deleteTaskLogs(taskId) {
return axios.post(`${this.baseUrl}/api/task/${taskId}/delete_logs`);
}
getUserScripts() {
return axios.get(`${this.baseUrl}/api/user_script`);
}
getUserScript(name) {
return axios.get(`${this.baseUrl}/api/user_script/${name}`);
}
createUserScript(name, template) {
return axios.post(`${this.baseUrl}/api/user_script/${name}`, null, {
params: {
template: template
}
});
}
updateUserScript(name, data) {
return axios.put(`${this.baseUrl}/api/user_script/${name}`, data);
}
deleteUserScript(name) {
return axios.delete(`${this.baseUrl}/api/user_script/${name}`);
}
testUserScript(name, job) {
return axios.get(`${this.baseUrl}/api/user_script/${name}/run`, {
params: {
job: job
}
});
}
}
export default new Sist2AdminApi()

View File

@@ -0,0 +1,34 @@
<template>
<b-progress v-if="loading" striped animated value="100"></b-progress>
<span v-else-if="jobs.length === 0"></span>
<b-form-select v-else :options="jobs" text-field="name" value-field="name"
@change="$emit('change', $event)" :value="$t('selectJob')"></b-form-select>
</template>
<script>
import Sist2AdminApi from "@/Sist2AdminApi";
export default {
name: "JobSelect",
mounted() {
Sist2AdminApi.getJobs().then(resp => {
this._jobs = resp.data;
this.loading = false;
});
},
computed: {
jobs() {
return [
{name: this.$t("selectJob"), disabled: true},
...this._jobs.filter(job => job.index_path)
]
}
},
data() {
return {
loading: true,
_jobs: null
}
}
}
</script>

View File

@@ -0,0 +1,18 @@
<template>
<b-list-group-item action :to="`/userScript/${script.name}`">
<div class="d-flex w-100 justify-content-between">
<h5 class="mb-1">
{{ script.name }}
</h5>
</div>
</b-list-group-item>
</template>
<script>
export default {
name: "UserScriptListItem",
props: ["script"],
}
</script>

View File

@@ -0,0 +1,88 @@
<template>
<b-progress v-if="loading" striped animated value="100"></b-progress>
<b-row v-else>
<b-col cols="6">
<h5>Selected scripts</h5>
<b-list-group>
<b-list-group-item v-for="script in selectedScripts" :key="script"
button
@click="onRemoveScript(script)"
class="d-flex justify-content-between align-items-center">
{{ script }}
<b-button-group>
<b-button variant="light" @click.stop="moveUpScript(script)"></b-button>
<b-button variant="light" @click.stop="moveDownScript(script)"></b-button>
</b-button-group>
</b-list-group-item>
</b-list-group>
</b-col>
<b-col cols="6">
<h5>Available scripts</h5>
<b-list-group>
<b-list-group-item v-for="script in availableScripts" :key="script" button
@click="onSelectScript(script)">
{{ script }}
</b-list-group-item>
</b-list-group>
</b-col>
</b-row>
<!-- <b-checkbox-group v-else :options="scripts" stacked :checked="selectedScripts"-->
<!-- @input="$emit('change', $event)"></b-checkbox-group>-->
</template>
<script>
import Sist2AdminApi from "@/Sist2AdminApi";
export default {
name: "UserScriptPicker",
props: ["selectedScripts"],
data() {
return {
loading: true,
scripts: []
}
},
computed: {
availableScripts() {
return this.scripts.filter(script => !this.selectedScripts.includes(script))
}
},
mounted() {
Sist2AdminApi.getUserScripts().then(resp => {
this.scripts = resp.data.map(script => script.name);
this.loading = false;
});
},
methods: {
onSelectScript(name) {
this.selectedScripts.push(name);
this.$emit("change", this.selectedScripts)
},
onRemoveScript(name) {
this.selectedScripts.splice(this.selectedScripts.indexOf(name), 1);
this.$emit("change", this.selectedScripts);
},
moveUpScript(name) {
const index = this.selectedScripts.indexOf(name);
if (index > 0) {
this.selectedScripts.splice(index, 1);
this.selectedScripts.splice(index - 1, 0, name);
}
this.$emit("change", this.selectedScripts);
},
moveDownScript(name) {
const index = this.selectedScripts.indexOf(name);
if (index < this.selectedScripts.length - 1) {
this.selectedScripts.splice(index, 1);
this.selectedScripts.splice(index + 1, 0, name);
}
this.$emit("change", this.selectedScripts);
}
}
}
</script>
<style scoped>
</style>

View File

@@ -54,8 +54,18 @@ export default {
frontendTab: "Frontend",
backendTab: "Backend",
scripts: "User Scripts",
script: "User Script",
testScript: "Test/debug User Script",
newScriptName: "New script name",
scriptType: "Script type",
scriptCode: "Script code (Python)",
scriptOptions: "User scripts",
gitRepository: "Git repository URL",
extraArgs: "Extra command line arguments",
selectJobs: "Available jobs",
selectJob: "Select a job",
webOptions: {
title: "Web options",
lang: "UI Language",

View File

@@ -6,12 +6,18 @@ import Tasks from "@/views/Tasks";
import Frontend from "@/views/Frontend";
import Tail from "@/views/Tail";
import SearchBackend from "@/views/SearchBackend.vue";
import UserScript from "@/views/UserScript.vue";
Vue.use(VueRouter);
const routes = [
{
path: "/",
path: "/task",
name: "Tasks",
component: Tasks
},
{
path: "/:tab?",
name: "Home",
component: Home
},
@@ -20,11 +26,6 @@ const routes = [
name: "Job",
component: Job
},
{
path: "/task/",
name: "Tasks",
component: Tasks
},
{
path: "/frontend/:name",
name: "Frontend",
@@ -35,6 +36,11 @@ const routes = [
name: "SearchBackend",
component: SearchBackend
},
{
path: "/userScript/:name",
name: "UserScript",
component: UserScript
},
{
path: "/log/:taskId",
name: "Tail",

View File

@@ -1,6 +1,6 @@
<template>
<div>
<b-tabs content-class="mt-3">
<b-tabs content-class="mt-3" v-model="tab" @input="onTabChange($event)">
<b-tab :title="$t('backendTab')">
<b-card>
@@ -25,7 +25,6 @@
<SearchBackendListItem v-for="backend in backends"
:key="backend.name" :backend="backend"></SearchBackendListItem>
</b-list-group>
</b-card>
<br/>
@@ -36,12 +35,12 @@
<b-col>
<b-input id="new-job" v-model="newJobName" :placeholder="$t('newJobName')"></b-input>
<b-popover
:show.sync="showHelp"
target="new-job"
placement="top"
triggers="manual"
variant="primary"
:content="$t('newJobHelp')"
:show.sync="showHelp"
target="new-job"
placement="top"
triggers="manual"
variant="primary"
:content="$t('newJobHelp')"
></b-popover>
</b-col>
<b-col>
@@ -59,6 +58,37 @@
</b-list-group>
</b-card>
</b-tab>
<b-tab :title="$t('scripts')">
<b-progress v-if="scriptsLoading" striped animated value="100"></b-progress>
<b-card v-else>
<b-card-title>{{ $t("scripts") }}</b-card-title>
<label>Select template</label>
<b-form-radio-group stacked :options="scriptTemplates" v-model="scriptTemplate"></b-form-radio-group>
<br>
<b-row>
<b-col>
<b-form-input v-model="newScriptName" :disabled="!scriptTemplate" :placeholder="$t('newScriptName')"></b-form-input>
</b-col>
<b-col>
<b-button variant="primary" @click="createScript()"
:disabled="!scriptNameValid(newScriptName)">
{{ $t("create") }}
</b-button>
</b-col>
</b-row>
<hr/>
<b-list-group>
<UserScriptListItem v-for="script in scripts"
:key="script.name" :script="script"></UserScriptListItem>
</b-list-group>
</b-card>
</b-tab>
<b-tab :title="$t('frontendTab')">
<b-card>
@@ -96,10 +126,11 @@ import {formatBindAddress} from "@/util";
import Sist2AdminApi from "@/Sist2AdminApi";
import FrontendListItem from "@/components/FrontendListItem";
import SearchBackendListItem from "@/components/SearchBackendListItem.vue";
import UserScriptListItem from "@/components/UserScriptListItem.vue";
export default {
name: "Jobs",
components: {SearchBackendListItem, JobListItem, FrontendListItem},
components: {UserScriptListItem, SearchBackendListItem, JobListItem, FrontendListItem},
data() {
return {
jobsLoading: true,
@@ -115,11 +146,24 @@ export default {
backendsLoading: true,
newBackendName: "",
showHelp: false
scripts: [],
scriptTemplates: [],
newScriptName: "",
scriptTemplate: null,
scriptsLoading: true,
showHelp: false,
tab: 0
}
},
mounted() {
this.loading = true;
if (this.$route.params.tab) {
console.log("mounted " + this.$route.params.tab)
window.setTimeout(() => {
this.tab = Math.round(Number(this.$route.params.tab));
}, 1)
}
this.reload();
},
methods: {
@@ -144,11 +188,20 @@ export default {
return /^[a-zA-Z0-9-_,.; ]+$/.test(name);
},
scriptNameValid(name) {
if (this.scripts.some(script => script.name === name)) {
return false;
}
if (name.length > 16) {
return false;
}
return /^[a-zA-Z0-9-_,.; ]+$/.test(name);
},
reload() {
Sist2AdminApi.getJobs().then(resp => {
this.jobs = resp.data;
this.jobsLoading = false;
this.showHelp = this.jobs.length === 0;
});
Sist2AdminApi.getFrontends().then(resp => {
@@ -159,6 +212,11 @@ export default {
this.backends = resp.data;
this.backendsLoading = false;
})
Sist2AdminApi.getUserScripts().then(resp => {
this.scripts = resp.data;
this.scriptTemplates = this.$store.state.sist2AdminInfo.user_script_templates;
this.scriptsLoading = false;
})
},
createJob() {
Sist2AdminApi.createJob(this.newJobName).then(this.reload);
@@ -168,6 +226,14 @@ export default {
},
createBackend() {
Sist2AdminApi.createBackend(this.newBackendName).then(this.reload);
},
createScript() {
Sist2AdminApi.createUserScript(this.newScriptName, this.scriptTemplate).then(this.reload)
},
onTabChange(tab) {
if (this.$route.params.tab != tab) {
this.$router.push({params: {tab: tab}})
}
}
}
}

View File

@@ -30,6 +30,13 @@
<SearchBackendSelect :value="job.index_options.search_backend"
@change="onBackendSelect($event)"></SearchBackendSelect>
</b-card>
<br/>
<h4>{{ $t("scriptOptions") }}</h4>
<b-card>
<UserScriptPicker :selected-scripts="job.user_scripts"
@change="onScriptChange($event)"></UserScriptPicker>
</b-card>
<br/>
@@ -48,10 +55,12 @@ import ScanOptions from "@/components/ScanOptions";
import Sist2AdminApi from "@/Sist2AdminApi";
import JobOptions from "@/components/JobOptions";
import SearchBackendSelect from "@/components/SearchBackendSelect.vue";
import UserScriptPicker from "@/components/UserScriptPicker.vue";
export default {
name: "Job",
components: {
UserScriptPicker,
SearchBackendSelect,
ScanOptions,
JobOptions
@@ -95,6 +104,10 @@ export default {
onBackendSelect(backend) {
this.job.index_options.search_backend = backend;
this.update();
},
onScriptChange(scripts) {
this.job.user_scripts = scripts;
this.update();
}
},
mounted() {

View File

@@ -44,9 +44,6 @@
<label>{{ $t("backendOptions.batchSize") }}</label>
<b-form-input v-model="backend.batch_size" type="number" min="1" @change="update()"></b-form-input>
<label>{{ $t("backendOptions.script") }}</label>
<b-form-textarea v-model="backend.script" rows="6" @change="update()"></b-form-textarea>
</template>
<template v-else>
<label>{{ $t("backendOptions.searchIndex") }}</label>

View File

@@ -92,6 +92,9 @@ export default {
if ("stderr" in message) {
message.level = "ERROR";
message.message = message["stderr"];
} else if ("stdout" in message) {
message.level = "INFO";
message.message = message["stdout"];
} else {
message.level = "ADMIN";
message.message = message["sist2-admin"];

View File

@@ -0,0 +1,117 @@
<template>
<b-progress v-if="loading" striped animated value="100"></b-progress>
<b-card v-else>
<b-card-title>
{{ $route.params.name }}
{{ $t("script") }}
</b-card-title>
<div class="mb-3">
<b-button variant="danger" @click="deleteScript()">{{ $t("delete") }}</b-button>
</div>
<b-card>
<h5>{{ $t("testScript") }}</h5>
<b-row>
<b-col cols="11">
<JobSelect @change="onJobSelect($event)"></JobSelect>
</b-col>
<b-col cols="1">
<b-button :disabled="!selectedTestJob" variant="primary" @click="testScript()">{{ $t("test") }}
</b-button>
</b-col>
</b-row>
</b-card>
<br/>
<label>{{ $t("scriptType") }}</label>
<b-form-select :options="['git', 'simple']" v-model="script.type" @change="update()"></b-form-select>
<template v-if="script.type === 'git'">
<label>{{ $t("gitRepository") }}</label>
<b-form-input v-model="script.git_repository" placeholder="https://github.com/example/example.git"
@change="update()"></b-form-input>
<label>{{ $t("extraArgs") }}</label>
<b-form-input v-model="script.extra_args" @change="update()" class="text-monospace"></b-form-input>
</template>
<template v-if="script.type === 'simple'">
<label>{{ $t("scriptCode") }}</label>
<p>Find sist2-python documentation <a href="https://sist2-python.readthedocs.io/" target="_blank">here</a></p>
<b-textarea rows="15" class="text-monospace" v-model="script.script" @change="update()" spellcheck="false"></b-textarea>
</template>
<template v-if="script.type === 'local'">
<!-- TODO-->
</template>
</b-card>
</template>
<script>
import Sist2AdminApi from "@/Sist2AdminApi";
import JobOptions from "@/components/JobOptions.vue";
import JobCheckboxGroup from "@/components/JobCheckboxGroup.vue";
import JobSelect from "@/components/JobSelect.vue";
export default {
name: "UserScript",
components: {JobSelect, JobCheckboxGroup, JobOptions},
data() {
return {
loading: true,
script: null,
selectedTestJob: null
}
},
methods: {
update() {
Sist2AdminApi.updateUserScript(this.name, this.script);
},
onJobSelect(job) {
this.selectedTestJob = job;
},
deleteScript() {
Sist2AdminApi.deleteUserScript(this.name)
.then(() => {
this.$router.push("/");
})
.catch(err => {
this.$bvToast.toast("Cannot delete user script " +
"because it is referenced by a job", {
title: "Error",
variant: "danger",
toaster: "b-toaster-bottom-right"
});
})
},
testScript() {
Sist2AdminApi.testUserScript(this.name, this.selectedTestJob)
.then(() => {
this.$bvToast.toast(this.$t("runJobConfirmation"), {
title: this.$t("runJobConfirmationTitle"),
variant: "success",
toaster: "b-toaster-bottom-right"
});
})
}
},
mounted() {
Sist2AdminApi.getUserScript(this.name).then(resp => {
this.script = resp.data;
this.loading = false;
});
},
computed: {
name() {
return this.$route.params.name;
},
},
}
</script>

View File

@@ -2,4 +2,6 @@ fastapi
git+https://github.com/simon987/hexlib.git
uvicorn
websockets
pycron
pycron
GitPython
git+https://github.com/simon987/sist2-python.git

View File

@@ -18,12 +18,13 @@ from websockets.exceptions import ConnectionClosed
import cron
from config import LOG_FOLDER, logger, WEBSERVER_PORT, DATA_FOLDER, SIST2_BINARY
from jobs import Sist2Job, Sist2ScanTask, TaskQueue, Sist2IndexTask, JobStatus
from jobs import Sist2Job, Sist2ScanTask, TaskQueue, Sist2IndexTask, JobStatus, Sist2UserScriptTask
from notifications import Subscribe, Notifications
from sist2 import Sist2, Sist2SearchBackend
from state import migrate_v1_to_v2, RUNNING_FRONTENDS, TESSERACT_LANGS, DB_SCHEMA_VERSION, migrate_v3_to_v4, \
get_log_files_to_remove, delete_log_file, create_default_search_backends
from web import Sist2Frontend
from script import UserScript, SCRIPT_TEMPLATES
sist2 = Sist2(SIST2_BINARY, DATA_FOLDER)
db = PersistentState(dbfile=os.path.join(DATA_FOLDER, "state.db"))
@@ -52,7 +53,8 @@ async def home():
async def api():
return {
"tesseract_langs": TESSERACT_LANGS,
"logs_folder": LOG_FOLDER
"logs_folder": LOG_FOLDER,
"user_script_templates": list(SCRIPT_TEMPLATES.keys())
}
@@ -113,8 +115,6 @@ async def update_job(name: str, new_job: Sist2Job):
async def update_frontend(name: str, frontend: Sist2Frontend):
db["frontends"][name] = frontend
# TODO: Check etag
return "ok"
@@ -150,9 +150,21 @@ def _run_job(job: Sist2Job):
db["jobs"][job.name] = job
scan_task = Sist2ScanTask(job, f"Scan [{job.name}]")
index_task = Sist2IndexTask(job, f"Index [{job.name}]", depends_on=scan_task)
index_depends_on = scan_task
script_tasks = []
for script_name in job.user_scripts:
script = db["user_scripts"][script_name]
task = Sist2UserScriptTask(script, job, f"Script <{script_name}> [{job.name}]", depends_on=scan_task)
script_tasks.append(task)
index_depends_on = task
index_task = Sist2IndexTask(job, f"Index [{job.name}]", depends_on=index_depends_on)
task_queue.submit(scan_task)
for task in script_tasks:
task_queue.submit(task)
task_queue.submit(index_task)
@@ -167,6 +179,22 @@ async def run_job(name: str):
return "ok"
@app.get("/api/user_script/{name:str}/run")
def run_user_script(name: str, job: str):
script = db["user_scripts"][name]
if not script:
raise HTTPException(status_code=404)
job = db["jobs"][job]
if not job:
raise HTTPException(status_code=404)
script_task = Sist2UserScriptTask(script, job, f"Script <{name}> [{job.name}]")
task_queue.submit(script_task)
return "ok"
@app.get("/api/job/{name:str}/logs_to_delete")
async def task_history(n: int, name: str):
return get_log_files_to_remove(db, name, n)
@@ -239,7 +267,7 @@ def check_es_version(es_url: str, insecure: bool):
es_url = f"{url.scheme}://{url.hostname}:{url.port}"
else:
auth = None
r = requests.get(es_url, verify=insecure, auth=auth)
r = requests.get(es_url, verify=not insecure, auth=auth)
except SSLError:
return {
"ok": False,
@@ -375,6 +403,59 @@ def create_search_backend(name: str):
return backend
@app.delete("/api/user_script/{name:str}")
def delete_user_script(name: str):
if db["user_scripts"][name] is None:
return HTTPException(status_code=404)
if any(name in job.user_scripts for job in db["jobs"]):
raise HTTPException(status_code=400, detail="in use (job)")
script: UserScript = db["user_scripts"][name]
script.delete_dir()
del db["user_scripts"][name]
return "ok"
@app.post("/api/user_script/{name:str}")
def create_user_script(name: str, template: str):
if db["user_scripts"][name] is not None:
return HTTPException(status_code=400, detail="already exists")
script = SCRIPT_TEMPLATES[template](name)
db["user_scripts"][name] = script
return script
@app.get("/api/user_script")
async def get_user_scripts():
return list(db["user_scripts"])
@app.get("/api/user_script/{name:str}")
async def get_user_script(name: str):
backend = db["user_scripts"][name]
if not backend:
raise HTTPException(status_code=404)
return backend
@app.put("/api/user_script/{name:str}")
async def update_user_script(name: str, script: UserScript):
previous_version: UserScript = db["user_scripts"][name]
if previous_version and previous_version.git_repository != script.git_repository:
script.force_clone = True
db["user_scripts"][name] = script
return "ok"
def tail(filepath: str, n: int):
with open(filepath) as file:
@@ -479,7 +560,8 @@ if __name__ == '__main__':
migrate_v3_to_v4(db)
if db["sist2_admin"]["info"]["version"] != DB_SCHEMA_VERSION:
raise Exception(f"Incompatible database version for {db.dbfile}")
raise Exception(f"Incompatible database {db.dbfile}. "
f"Automatic migration is not available, please delete the database file to continue.")
start_frontends()
cron.initialize(db, _run_job)

View File

@@ -9,9 +9,11 @@ MAX_LOG_SIZE = 1 * 1024 * 1024
SIST2_BINARY = os.environ.get("SIST2_BINARY", "/root/sist2")
DATA_FOLDER = os.environ.get("DATA_FOLDER", "/sist2-admin/")
LOG_FOLDER = os.path.join(DATA_FOLDER, "logs")
SCRIPT_FOLDER = os.path.join(DATA_FOLDER, "scripts")
WEBSERVER_PORT = 8080
os.makedirs(LOG_FOLDER, exist_ok=True)
os.makedirs(SCRIPT_FOLDER, exist_ok=True)
os.makedirs(DATA_FOLDER, exist_ok=True)
logger = logging.Logger("sist2-admin")

View File

@@ -1,13 +1,18 @@
import json
import logging
import os.path
import shlex
import signal
import uuid
from datetime import datetime
from enum import Enum
from io import TextIOWrapper
from logging import FileHandler
from subprocess import Popen
import subprocess
from threading import Lock, Thread
from time import sleep
from typing import List
from uuid import uuid4, UUID
from hexlib.db import PersistentState
@@ -18,6 +23,7 @@ from notifications import Notifications
from sist2 import ScanOptions, IndexOptions, Sist2
from state import RUNNING_FRONTENDS, get_log_files_to_remove, delete_log_file
from web import Sist2Frontend
from script import UserScript
class JobStatus(Enum):
@@ -32,6 +38,8 @@ class Sist2Job(BaseModel):
scan_options: ScanOptions
index_options: IndexOptions
user_scripts: List[str] = []
cron_expression: str
schedule_enabled: bool = False
@@ -182,7 +190,7 @@ class Sist2IndexTask(Sist2Task):
duration = self.ended - self.started
ok = return_code == 0
ok = return_code in (0, 1)
if ok:
self.restart_running_frontends(db, sist2)
@@ -231,6 +239,65 @@ class Sist2IndexTask(Sist2Task):
self._logger.info(json.dumps({"sist2-admin": f"Restart frontend {pid=} {frontend_name=}"}))
class Sist2UserScriptTask(Sist2Task):
def __init__(self, user_script: UserScript, job: Sist2Job, display_name: str, depends_on: Sist2Task = None):
super().__init__(job, display_name, depends_on=depends_on.id if depends_on else None)
self.user_script = user_script
def run(self, sist2: Sist2, db: PersistentState):
super().run(sist2, db)
try:
self.user_script.setup(self.log_callback)
except Exception as e:
logger.error(f"Setup for {self.user_script.name} failed: ")
logger.exception(e)
self.log_callback({"sist2-admin": f"Setup for {self.user_script.name} failed: {e}"})
return -1
executable = self.user_script.get_executable()
index_path = os.path.join(DATA_FOLDER, self.job.index_path)
extra_args = self.user_script.extra_args
args = [
executable,
index_path,
*shlex.split(extra_args)
]
self.log_callback({"sist2-admin": f"Starting user script with {executable=}, {index_path=}, {extra_args=}"})
proc = Popen(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=self.user_script.script_dir())
self.pid = proc.pid
t_stderr = Thread(target=self._consume_logs, args=(self.log_callback, proc, "stderr", False))
t_stderr.start()
self._consume_logs(self.log_callback, proc, "stdout", True)
self.ended = datetime.utcnow()
return 0
@staticmethod
def _consume_logs(logs_cb, proc, stream, wait):
pipe_wrapper = TextIOWrapper(getattr(proc, stream), encoding="utf8", errors="ignore")
try:
for line in pipe_wrapper:
if line.strip() == "":
continue
if line.startswith("$PROGRESS"):
progress = json.loads(line[len("$PROGRESS "):])
logs_cb({"progress": progress})
continue
logs_cb({stream: line})
finally:
if wait:
proc.wait()
pipe_wrapper.close()
class TaskQueue:
def __init__(self, sist2: Sist2, db: PersistentState, notifications: Notifications):
self._lock = Lock()

View File

@@ -0,0 +1,126 @@
import os
import shutil
import stat
import subprocess
from enum import Enum
from git import Repo
from pydantic import BaseModel
from config import SCRIPT_FOLDER
class ScriptType(Enum):
LOCAL = "local"
SIMPLE = "simple"
GIT = "git"
def set_executable(file):
os.chmod(file, os.stat(file).st_mode | stat.S_IEXEC)
def _initialize_git_repository(url, path, log_cb, force_clone):
log_cb({"sist2-admin": f"Cloning {url}"})
if force_clone or not os.path.exists(os.path.join(path, ".git")):
if force_clone:
shutil.rmtree(path, ignore_errors=True)
Repo.clone_from(url, path)
else:
repo = Repo(path)
repo.remote("origin").pull()
setup_script = os.path.join(path, "setup.sh")
if setup_script:
log_cb({"sist2-admin": f"Executing setup script {setup_script}"})
set_executable(setup_script)
result = subprocess.run([setup_script], cwd=path, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
for line in result.stdout.split(b"\n"):
if line:
log_cb({"stdout": line.decode()})
log_cb({"stdout": f"Executed setup script {setup_script}, return code = {result.returncode}"})
if result.returncode != 0:
raise Exception("Error when running setup script!")
log_cb({"sist2-admin": f"Initialized git repository in {path}"})
class UserScript(BaseModel):
name: str
type: ScriptType
git_repository: str = None
force_clone: bool = False
script: str = None
extra_args: str = ""
def script_dir(self):
return os.path.join(SCRIPT_FOLDER, self.name)
def setup(self, log_cb):
os.makedirs(self.script_dir(), exist_ok=True)
if self.type == ScriptType.GIT:
_initialize_git_repository(self.git_repository, self.script_dir(), log_cb, self.force_clone)
self.force_clone = False
elif self.type == ScriptType.SIMPLE:
self._setup_simple()
set_executable(self.get_executable())
def _setup_simple(self):
with open(self.get_executable(), "w") as f:
f.write(
"#!/bin/bash\n"
"python run.py \"$@\""
)
with open(os.path.join(self.script_dir(), "run.py"), "w") as f:
f.write(self.script)
def get_executable(self):
return os.path.join(self.script_dir(), "run.sh")
def delete_dir(self):
shutil.rmtree(self.script_dir(), ignore_errors=True)
SCRIPT_TEMPLATES = {
"CLIP - Generate embeddings to predict the most relevant image based on the text prompt": lambda name: UserScript(
name=name,
type=ScriptType.GIT,
git_repository="https://github.com/simon987/sist2-script-clip",
extra_args="--num-tags=1 --tags-file=general.txt --color=#dcd7ff"
),
"Whisper - Speech to text with OpenAI Whisper": lambda name: UserScript(
name=name,
type=ScriptType.GIT,
git_repository="https://github.com/simon987/sist2-script-whisper",
extra_args="--model=base --num-threads=4 --color=#51da4c --tag"
),
"Hamburger - Simple script example": lambda name: UserScript(
name=name,
type=ScriptType.SIMPLE,
script=
'from sist2 import Sist2Index\n'
'import sys\n'
'\n'
'index = Sist2Index(sys.argv[1])\n'
'for doc in index.document_iter():\n'
' doc.json_data["tag"] = ["hamburger.#00FF00"]\n'
' index.update_document(doc)\n'
'\n'
'index.sync_tag_table()\n'
'index.commit()\n'
'\n'
'print("Done!")\n'
),
"(Blank)": lambda name: UserScript(
name=name,
type=ScriptType.SIMPLE,
script=""
)
}

View File

@@ -41,8 +41,6 @@ class Sist2SearchBackend(BaseModel):
es_insecure_ssl: bool = False
es_index: str = "sist2"
threads: int = 1
script: str = ""
script_file: str = None
batch_size: int = 70
@staticmethod
@@ -74,8 +72,6 @@ class IndexOptions(BaseModel):
f"--es-index={search_backend.es_index}",
f"--batch-size={search_backend.batch_size}"]
if search_backend.script_file:
args.append(f"--script-file={search_backend.script_file}")
if search_backend.es_insecure_ssl:
args.append(f"--es-insecure-ssl")
if self.incremental_index:
@@ -249,13 +245,6 @@ class Sist2:
def index(self, options: IndexOptions, search_backend: Sist2SearchBackend, logs_cb):
if search_backend.script and search_backend.backend_type == SearchBackendType("elasticsearch"):
with NamedTemporaryFile("w", prefix="sist2-admin", suffix=".painless", delete=False) as f:
f.write(search_backend.script)
search_backend.script_file = f.name
else:
search_backend.script_file = None
args = [
self.bin_path,
*options.args(search_backend),

View File

@@ -14,7 +14,7 @@ RUNNING_FRONTENDS: Dict[str, int] = {}
TESSERACT_LANGS = get_tesseract_langs()
DB_SCHEMA_VERSION = "4"
DB_SCHEMA_VERSION = "5"
from pydantic import BaseModel