big dik update: prompts and notify mode

1 files changed, 33 insertions, 10 deletions

diff --git a/supergenerator.py b/supergenerator.py
index bb3efca..e21900b 100644
--- a/supergenerator.py
+++ b/supergenerator.py
@@ -113,15 +113,17 @@ def remove_think_tags(text):
     return result
 
 
-def get_all_users():
+def get_all_users(only_notify=False):
     with shelve.open("users") as db:
-        return {k: db[k] for k in db.keys()}
+        return {k: db[k] for k in db.keys() if not only_notify or db[k].get("notify") == "da"}
 
 
-def get_all_user_ids():
+def get_all_user_ids(only_notify=False):
     ids = []
     with shelve.open("users") as db:
         for user_id in db.keys():
+            if only_notify and db[user_id].get("notify") != "da":
+                continue
             try:
                 ids.append(int(user_id))
             except ValueError:
@@ -129,14 +131,22 @@ def get_all_user_ids():
     return ids
 
 
+
 def set_user_config(
-    user_id, model_name, stream_mode=DEFAULT_STREAM, call_mode=DEFAULT_CALL
+    user_id, 
+    model_name, 
+    stream_mode=DEFAULT_STREAM, 
+    call_mode=DEFAULT_CALL, 
+    prompt_name=DEFAULT_PROMPT,
+    notify_mode=DEFAULT_NOTIFY,
 ):
     with shelve.open("users") as db:
         db[str(user_id)] = {
             "model": model_name,
             "stream": stream_mode,
             "call": call_mode,
+            "prompt": prompt_name,
+            "notify": notify_mode,
         }
 
 
@@ -144,7 +154,9 @@ def ensure_user(
     user_id,
     model_name=DEFAULT_MODEL,
     stream_mode=DEFAULT_STREAM,
-    call_mode=DEFAULT_CALL
+    call_mode=DEFAULT_CALL,
+    prompt_name=DEFAULT_PROMPT,
+    notify_mode=DEFAULT_NOTIFY
 ):
     with shelve.open("users") as db:
         key = str(user_id)
@@ -154,13 +166,15 @@ def ensure_user(
                 "model": model_name,
                 "stream": stream_mode,
                 "call": call_mode,
+                "prompt": prompt_name,
+                "notify": notify_mode,
             }
 
 def get_user_config(user_id):
     with shelve.open("users") as db:
         return db.get(
             str(user_id),
-            {"model": DEFAULT_MODEL, "stream": DEFAULT_STREAM, "call": DEFAULT_CALL},
+            {"model": DEFAULT_MODEL, "stream": DEFAULT_STREAM, "call": DEFAULT_CALL, "prompt": DEFAULT_PROMPT, "notify": DEFAULT_NOTIFY},
         )
 
 
@@ -177,7 +191,16 @@ async def generate(
         else:
             print(f"{RED} -- err model - failed to get prov&model in \"{config.get('model','')}\"{RESET}")
             return "произошла ошибка при парсинге вашей выбранной модели. либо хост бота гандон либо тебе надо сменить модель в настройках!!"
+        
         stream_cf = config["stream"]
+        system_prompt_name = config["prompt"]
+        if system_prompt_name in SYSTEM_PROMPTS:
+            system_prompt = SYSTEM_PROMPTS[system_prompt_name]
+        else:
+            text = "❌произошла ошибка, возможные проблемы:\n" \
+                f"1. у вас выбран удалённый сист промпт (`{system_prompt_name}`)\n" \
+                f"2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
+            return text
 
         if stream_cf == "native" and user_id == chat_id:
             stream = True
@@ -189,14 +212,14 @@ async def generate(
         user_contexts[user_id].append({"role": "user", "content": prompt})
         user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]
 
-        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + user_contexts[
+        messages = [{"role": "system", "content": system_prompt}] + user_contexts[
             user_id
         ]
     else:
         model = DEFAULT_MODEL.split("*", 2)
         stream = False
         stream_cf = "none"
-        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + [
+        messages = [{"role": "system", "content": system_prompt}] + [
             {"role": "user", "content": prompt}
         ]
 
@@ -207,9 +230,9 @@ async def generate(
     try:
         client = providers[model[0]]
     except Exception:
-        text = "произошла ошибка, возможные проблемы:\n" \
+        text = "❌произошла ошибка, возможные проблемы:\n" \
             "1. у вас выбрана удалённая модель или провайдер\n" \
-            "2. конфиг бота настроен неправильно и надо связаться с хостом бота"
+            "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
         return text
         
     text = await openai_generate_text(