jarvis.py

#
# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: Apache-2.0
#

import asyncio
import docx
import gradio as gr
import httpx
import json
import os
import pandas as pd
import pdfplumber
import pytesseract
import random
import requests
import threading
import uuid
import zipfile
import io

from PIL import Image
from pathlib import Path
from pptx import Presentation
from openpyxl import load_workbook

os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract-ocr-eng tesseract-ocr-ind libleptonica-dev libtesseract-dev")

INTERNAL_AI_GET_SERVER = os.getenv("INTERNAL_AI_GET_SERVER")
INTERNAL_TRAINING_DATA = os.getenv("INTERNAL_TRAINING_DATA")

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 = [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(",")))

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())
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")

ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS", "[]"))

ACTIVE_CANDIDATE = None

class SessionWithID(requests.Session):
    def __init__(sess):
        super().__init__()
        sess.session_id = str(uuid.uuid4())

def create_session():
    return SessionWithID()

def ensure_stop_event(sess):
    if not hasattr(sess, "stop_event"):
        sess.stop_event = asyncio.Event()

def get_available_items(items, marked):
    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():
            marked.discard(item)
            attempts.pop(item, None)
        threading.Timer(300, remove).start()

def get_model_key(display):
    return next((k for k, v in MODEL_MAPPING.items() if v == display), DEFAULT_MODEL_KEY)

def extract_pdf_content(fp):
    content = ""
    try:
        with pdfplumber.open(fp) as pdf:
            for page in pdf.pages:
                text = page.extract_text() or ""
                content += text + "\n"
                if page.images:
                    img_obj = page.to_image(resolution=300)
                    for img in page.images:
                        bbox = (img["x0"], img["top"], img["x1"], img["bottom"])
                        cropped = img_obj.original.crop(bbox)
                        ocr_text = pytesseract.image_to_string(cropped)
                        if ocr_text.strip():
                            content += ocr_text + "\n"
                tables = page.extract_tables()
                for table in tables:
                    for row in table:
                        cells = [str(cell) for cell in row if cell is not None]
                        if cells:
                            content += "\t".join(cells) + "\n"
    except Exception as e:
        content += f"{fp}: {e}"
    return content.strip()

def extract_docx_content(fp):
    content = ""
    try:
        doc = docx.Document(fp)
        for para in doc.paragraphs:
            content += para.text + "\n"
        for table in doc.tables:
            for row in table.rows:
                cells = [cell.text for cell in row.cells]
                content += "\t".join(cells) + "\n"
        with zipfile.ZipFile(fp) as z:
            for file in z.namelist():
                if file.startswith("word/media/"):
                    data = z.read(file)
                    try:
                        img = Image.open(io.BytesIO(data))
                        ocr_text = pytesseract.image_to_string(img)
                        if ocr_text.strip():
                            content += ocr_text + "\n"
                    except Exception:
                        pass
    except Exception as e:
        content += f"{fp}: {e}"
    return content.strip()

def extract_excel_content(fp):
    content = ""
    try:
        sheets = pd.read_excel(fp, sheet_name=None)
        for name, df in sheets.items():
            content += f"Sheet: {name}\n"
            content += df.to_csv(index=False) + "\n"
        wb = load_workbook(fp, data_only=True)
        if wb._images:
            for image in wb._images:
                img = image.ref
                if isinstance(img, bytes):
                    try:
                        pil_img = Image.open(io.BytesIO(img))
                        ocr_text = pytesseract.image_to_string(pil_img)
                        if ocr_text.strip():
                            content += ocr_text + "\n"
                    except Exception:
                        pass
    except Exception as e:
        content += f"{fp}: {e}"
    return content.strip()

def extract_pptx_content(fp):
    content = ""
    try:
        prs = Presentation(fp)
        for slide in prs.slides:
            for shape in slide.shapes:
                if hasattr(shape, "text") and shape.text:
                    content += shape.text + "\n"
                if shape.shape_type == 13 and hasattr(shape, "image") and shape.image:
                    try:
                        img = Image.open(io.BytesIO(shape.image.blob))
                        ocr_text = pytesseract.image_to_string(img)
                        if ocr_text.strip():
                            content += ocr_text + "\n"
                    except Exception:
                        pass
            if slide.shapes:
                for shape in slide.shapes:
                    if shape.has_table:
                        table = shape.table
                        for row in table.rows:
                            cells = [cell.text for cell in row.cells]
                            content += "\t".join(cells) + "\n"
    except Exception as e:
        content += f"{fp}: {e}"
    return content.strip()

def extract_file_content(fp):
    ext = Path(fp).suffix.lower()
    if ext == ".pdf":
        return extract_pdf_content(fp)
    elif ext in [".doc", ".docx"]:
        return extract_docx_content(fp)
    elif ext in [".xlsx", ".xls"]:
        return extract_excel_content(fp)
    elif ext in [".ppt", ".pptx"]:
        return extract_pptx_content(fp)
    else:
        try:
            return Path(fp).read_text(encoding="utf-8").strip()
        except Exception as e:
            return f"{fp}: {e}"

async def fetch_response_async(host, key, model, msgs, cfg, sid):
    for t in [1, 2]:
        try:
            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
                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:
            continue
    marked_item(key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
    return None

async def chat_with_model_async(history, user_input, model_display, sess, custom_prompt):
    ensure_stop_event(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_ATTEMPTS):
        return RESPONSES["RESPONSE_3"]
    if not hasattr(sess, "session_id"):
        sess.session_id = str(uuid.uuid4())
        sess.stop_event = asyncio.Event()
    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:
        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_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, history, model_display, sess, custom_prompt):
    ensure_stop_event(sess)
    sess.stop_event.clear()
    msg_input = {"text": multi.get("text", "").strip(), "files": multi.get("files", [])}
    if not msg_input["text"] and not msg_input["files"]:
        yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
        return
    inp = ""
    for f in msg_input["files"]:
        fp = f.get("data", f.get("name", "")) if isinstance(f, dict) else f
        inp += f"{Path(fp).name}\n\n{extract_file_content(fp)}\n\n"
    if msg_input["text"]:
        inp += msg_input["text"]
    history.append([inp, RESPONSES["RESPONSE_8"]])
    yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
    ai = await chat_with_model_async(history, inp, model_display, sess, custom_prompt)
    history[-1][1] = ""
    buffer = []
    last_update = asyncio.get_event_loop().time()
    for char in ai:
        if sess.stop_event.is_set():
            history[-1][1] = RESPONSES["RESPONSE_1"]
            yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess
            sess.stop_event.clear()
            return
        buffer.append(char)
        current_time = asyncio.get_event_loop().time()
        if len(buffer) >= 8 or (current_time - last_update) > 0.04:
            history[-1][1] += "".join(buffer)
            buffer.clear()
            last_update = current_time
            yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
            await asyncio.sleep(0.016)
    if buffer:
        history[-1][1] += "".join(buffer)
        yield history, gr.update(interactive=False, submit_btn=False, stop_btn=True), sess
    yield history, gr.update(value="", interactive=True, submit_btn=True, stop_btn=False), sess

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)

def stop_response(history, sess):
    ensure_stop_event(sess)
    sess.stop_event.set()
    if history:
        history[-1][1] = RESPONSES["RESPONSE_1"]
    new_session = create_session()
    return history, None, new_session

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"])
    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])
        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)
    msg.stop(fn=stop_response, inputs=[user_history, user_session], outputs=[chatbot, msg, user_session])
jarvis.queue(default_concurrency_limit=3).launch(max_file_size="1mb")