ai: Refactor the code for JARVI multi platform.

jarvis.py

@@ -26,23 +26,27 @@ os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract
 INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
 INTERNAL_TRAINING_DATA = os.getenv("INTERNAL_TRAINING_DATA")
 
-LINUX_SERVER_HOSTS = [host for host in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if host]
+SYSTEM_PROMPT_MAPPING = json.loads(os.getenv("SYSTEM_PROMPT_MAPPING", "{}"))
+SYSTEM_PROMPT_DEFAULT = os.getenv("DEFAULT_SYSTEM")
+
+LINUX_SERVER_HOSTS = [h for h in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if h]
 LINUX_SERVER_HOSTS_MARKED = set()
 LINUX_SERVER_HOSTS_ATTEMPTS = {}
 
-LINUX_SERVER_PROVIDER_KEYS = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]
+LINUX_SERVER_PROVIDER_KEYS = [k for k in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if k]
 LINUX_SERVER_PROVIDER_KEYS_MARKED = set()
 LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}
 
-LINUX_SERVER_ERRORS = set(map(int, os.getenv("LINUX_SERVER_ERROR").split(",")))
+LINUX_SERVER_ERRORS = set(map(int, os.getenv("LINUX_SERVER_ERROR", "").split(",")))
 
-AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 7)}
+AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 8)}
 RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}
 
 MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
 MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
-MODEL_CHOICES = list(MODEL_MAPPING.values()) if MODEL_MAPPING else []
+MODEL_CHOICES = list(MODEL_MAPPING.values())
 DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))
+DEFAULT_MODEL_KEY = list(MODEL_MAPPING.keys())[0] if MODEL_MAPPING else None
 
 META_TAGS = os.getenv("META_TAGS")
 
@@ -50,157 +54,152 @@ ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS", "[]"))
 
 ACTIVE_CANDIDATE = None
 
+class SessionWithID(requests.Session):
+    def __init__(self):
+        super().__init__()
+        self.session_id = str(uuid.uuid4())
+
+def create_session():
+    return SessionWithID()
+
 def get_available_items(items, marked):
-    available = [item for item in items if item not in marked]
-    random.shuffle(available)
-    return available
+    a = [i for i in items if i not in marked]
+    random.shuffle(a)
+    return a
 
 def marked_item(item, marked, attempts):
     marked.add(item)
     attempts[item] = attempts.get(item, 0) + 1
     if attempts[item] >= 3:
-        def remove_fail():
+        def remove():
             marked.discard(item)
             attempts.pop(item, None)
-        threading.Timer(300, remove_fail).start()
-
-class SessionWithID(requests.Session):
-    def __init__(self):
-        super().__init__()
-        self.session_id = str(uuid.uuid4())
-
-def create_session():
-    return SessionWithID()
+        threading.Timer(300, remove).start()
 
-def get_model_key(display_name):
-    return next((k for k, v in MODEL_MAPPING.items() if v == display_name), list(MODEL_MAPPING.keys())[0] if MODEL_MAPPING else MODEL_CHOICES[0])
+def get_model_key(display):
+    return next((k for k, v in MODEL_MAPPING.items() if v == display), DEFAULT_MODEL_KEY)
 
-def extract_file_content(file_path):
-    ext = Path(file_path).suffix.lower()
-    content = ""
+def extract_file_content(fp):
+    ext = Path(fp).suffix.lower()
+    c = ""
     try:
         if ext == ".pdf":
-            with pdfplumber.open(file_path) as pdf:
-                for page in pdf.pages:
-                    text = page.extract_text()
-                    if text:
-                        content += text + "\n"
-                    for table in page.extract_tables():
-                        table_str = "\n".join([", ".join(row) for row in table if row])
-                        content += "\n" + table_str + "\n"
+            with pdfplumber.open(fp) as pdf:
+                for p in pdf.pages:
+                    t = p.extract_text() or ""
+                    c += t + "\n"
         elif ext in [".doc", ".docx"]:
-            doc = docx.Document(file_path)
-            for para in doc.paragraphs:
-                content += para.text + "\n"
+            d = docx.Document(fp)
+            for para in d.paragraphs:
+                c += para.text + "\n"
         elif ext in [".xlsx", ".xls"]:
-            df = pd.read_excel(file_path)
-            content += df.to_csv(index=False)
+            df = pd.read_excel(fp)
+            c += df.to_csv(index=False)
         elif ext in [".ppt", ".pptx"]:
-            prs = Presentation(file_path)
-            for slide in prs.slides:
-                for shape in slide.shapes:
-                    if hasattr(shape, "text") and shape.text:
-                        content += shape.text + "\n"
-        elif ext in [".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif", ".webp"]:
-            pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
-            image = Image.open(file_path)
-            content += pytesseract.image_to_string(image) + "\n"
+            prs = Presentation(fp)
+            for s in prs.slides:
+                for sh in s.shapes:
+                    if hasattr(sh, "text") and sh.text:
+                        c += sh.text + "\n"
         else:
-            content = Path(file_path).read_text(encoding="utf-8")
+            c = Path(fp).read_text(encoding="utf-8")
     except Exception as e:
-        content = f"{file_path}: {e}"
-    return content.strip()
+        c = f"{fp}: {e}"
+    return c.strip()
 
-async def fetch_response_async(host, provider_key, selected_model, messages, model_config, session_id):
-    timeouts = [60, 80, 120, 240]
-    for timeout in timeouts:
+async def fetch_response_async(host, key, model, msgs, cfg, sid):
+    for t in [60, 80, 120, 240]:
         try:
-            async with httpx.AsyncClient(timeout=timeout) as client:
-                resp = await client.post(host, json={"model": selected_model, "messages": messages, **model_config, "session_id": session_id}, headers={"Authorization": f"Bearer {provider_key}"})
-                if resp.status_code in LINUX_SERVER_ERRORS:
-                    marked_item(provider_key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
+            async with httpx.AsyncClient(timeout=t) as client:
+                r = await client.post(host, json={"model": model, "messages": msgs, **cfg, "session_id": sid}, headers={"Authorization": f"Bearer {key}"})
+                if r.status_code in LINUX_SERVER_ERRORS:
+                    marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
                     return None
-                resp.raise_for_status()
-                resp_json = resp.json()
-                if isinstance(resp_json, dict) and resp_json.get("choices"):
-                    choice = resp_json["choices"][0]
-                    if choice.get("message") and isinstance(choice["message"].get("content"), str):
-                        return choice["message"]["content"]
+                r.raise_for_status()
+                j = r.json()
+                if isinstance(j, dict) and j.get("choices"):
+                    ch = j["choices"][0]
+                    if ch.get("message") and isinstance(ch["message"].get("content"), str):
+                        return ch["message"]["content"]
                 return None
-        except Exception:
+        except:
             continue
-    marked_item(provider_key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
+    marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
     return None
 
-async def chat_with_model_async(history, user_input, selected_model_display, sess):
-    if not get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED) or not get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED):
+async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt):
+    if not get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED) or not get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_ATTEMPTS):
         return RESPONSES["RESPONSE_3"]
     if not hasattr(sess, "session_id"):
         sess.session_id = str(uuid.uuid4())
-    selected_model = get_model_key(selected_model_display)
-    model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)
-    messages = [{"role": "user", "content": user} for user, _ in history] + [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
-    if INTERNAL_TRAINING_DATA and MODEL_CHOICES and selected_model_display == MODEL_CHOICES[0]:
-        messages.insert(0, {"role": "system", "content": INTERNAL_TRAINING_DATA})
-    messages.append({"role": "user", "content": user_input})
+    model_key = get_model_key(model_display)
+    cfg = MODEL_CONFIG.get(model_key, DEFAULT_CONFIG)
+    msgs = [{"role": "user", "content": u} for u, _ in history] + [{"role": "assistant", "content": a} for _, a in history if a]
+    if model_key == DEFAULT_MODEL_KEY and INTERNAL_TRAINING_DATA:
+        prompt = INTERNAL_TRAINING_DATA
+    else:
+        prompt = custom_prompt or SYSTEM_PROMPT_MAPPING.get(model_key, SYSTEM_PROMPT_DEFAULT)
+    msgs.insert(0, {"role": "system", "content": prompt})
+    msgs.append({"role": "user", "content": user_input})
     global ACTIVE_CANDIDATE
     if ACTIVE_CANDIDATE:
-        result = await fetch_response_async(ACTIVE_CANDIDATE[0], ACTIVE_CANDIDATE[1], selected_model, messages, model_config, sess.session_id)
-        if result:
-            return result
+        res = await fetch_response_async(ACTIVE_CANDIDATE[0], ACTIVE_CANDIDATE[1], model_key, msgs, cfg, sess.session_id)
+        if res:
+            return res
         ACTIVE_CANDIDATE = None
     keys = get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED)
-    hosts = get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED)
-    candidates = [(host, key) for host in hosts for key in keys]
-    random.shuffle(candidates)
-    for host, key in candidates:
-        result = await fetch_response_async(host, key, selected_model, messages, model_config, sess.session_id)
-        if result:
-            ACTIVE_CANDIDATE = (host, key)
-            return result
+    hosts = get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_ATTEMPTS)
+    cands = [(h, k) for h in hosts for k in keys]
+    random.shuffle(cands)
+    for h, k in cands:
+        res = await fetch_response_async(h, k, model_key, msgs, cfg, sess.session_id)
+        if res:
+            ACTIVE_CANDIDATE = (h, k)
+            return res
     return RESPONSES["RESPONSE_2"]
 
-async def respond_async(multi_input, history, selected_model_display, sess):
-    message = {"text": multi_input.get("text", "").strip(), "files": multi_input.get("files", [])}
-    if not message["text"] and not message["files"]:
+async def respond_async(multi, history, model_display, sess, custom_prompt):
+    msg = {"text": multi.get("text", "").strip(), "files": multi.get("files", [])}
+    if not msg["text"] and not msg["files"]:
         yield history, gr.MultimodalTextbox(value=None, interactive=True), sess
         return
-    combined_input = ""
-    for file_item in message["files"]:
-        file_path = file_item["name"] if isinstance(file_item, dict) and "name" in file_item else file_item
-        combined_input += f"{Path(file_path).name}\n\n{extract_file_content(file_path)}\n\n"
-    if message["text"]:
-        combined_input += message["text"]
-    history.append([combined_input, ""])
-    ai_response = await chat_with_model_async(history, combined_input, selected_model_display, sess)
+    inp = ""
+    for f in msg["files"]:
+        p = f["name"] if isinstance(f, dict) and "name" in f else f
+        inp += f"{Path(p).name}\n\n{extract_file_content(p)}\n\n"
+    if msg["text"]:
+        inp += msg["text"]
+    history.append([inp, ""])
+    ai = await chat_with_model_async(history, inp, model_display, sess, custom_prompt)
     history[-1][1] = ""
-    def convert_to_string(data):
-        if isinstance(data, (str, int, float)):
-            return str(data)
-        if isinstance(data, bytes):
-            return data.decode("utf-8", errors="ignore")
-        if isinstance(data, (list, tuple)):
-            return "".join(map(convert_to_string, data))
-        if isinstance(data, dict):
-            return json.dumps(data, ensure_ascii=False)
-        return repr(data)
-    for character in ai_response:
-        history[-1][1] += convert_to_string(character)
+    def to_str(d):
+        if isinstance(d, (str, int, float)): return str(d)
+        if isinstance(d, bytes): return d.decode("utf-8", errors="ignore")
+        if isinstance(d, (list, tuple)): return "".join(map(to_str, d))
+        if isinstance(d, dict): return json.dumps(d, ensure_ascii=False)
+        return repr(d)
+    for c in ai:
+        history[-1][1] += to_str(c)
         await asyncio.sleep(0.0001)
         yield history, gr.MultimodalTextbox(value=None, interactive=True), sess
 
-def change_model(new_model_display):
-    return [], create_session(), new_model_display
+def change_model(new):
+    visible = new != MODEL_CHOICES[0]
+    default = SYSTEM_PROMPT_MAPPING.get(get_model_key(new), SYSTEM_PROMPT_DEFAULT)
+    return [], create_session(), new, default, gr.update(value=default, visible=visible)
 
 with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as jarvis:
     user_history = gr.State([])
     user_session = gr.State(create_session())
     selected_model = gr.State(MODEL_CHOICES[0] if MODEL_CHOICES else "")
+    custom_prompt_state = gr.State("")
     chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
     with gr.Row():
         msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)
     with gr.Accordion(AI_TYPES["AI_TYPE_6"], open=False):
         model_dropdown = gr.Dropdown(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
-    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, user_session, selected_model])
-    msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER)
+        system_prompt = gr.Textbox(label=AI_TYPES["AI_TYPE_7"], lines=2, interactive=True, visible=False)
+    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, user_session, selected_model, custom_prompt_state, system_prompt])
+    system_prompt.change(fn=lambda x: x, inputs=[system_prompt], outputs=[custom_prompt_state])
+    msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session, custom_prompt_state], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER)
 jarvis.launch(max_file_size="1mb")