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

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

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

LINUX_SERVER_HOSTS = [host for host in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if host]
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_MARKED = set()
LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}

AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 7)}
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 []
DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))

META_TAGS = os.getenv("META_TAGS")

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

ACTIVE_CANDIDATE = None

def get_available_items(items, marked):
    available = [item for item in items if item not in marked]
    random.shuffle(available)
    return available

def marked_item(item, marked, attempts):
    marked.add(item)
    attempts[item] = attempts.get(item, 0) + 1
    if attempts[item] >= 3:
        def remove_fail():
            marked.discard(item)
            attempts.pop(item, None)
        threading.Timer(3600, remove_fail).start()

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

def create_session():
    return SessionWithID()

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 extract_file_content(file_path):
    ext = Path(file_path).suffix.lower()
    content = ""
    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"
        elif ext in [".doc", ".docx"]:
            doc = docx.Document(file_path)
            for para in doc.paragraphs:
                content += para.text + "\n"
        elif ext in [".xlsx", ".xls"]:
            df = pd.read_excel(file_path)
            content += 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"
        else:
            content = Path(file_path).read_text(encoding="utf-8")
    except Exception as e:
        content = f"{file_path}: {e}"
    return content.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:
        try:
            async with httpx.AsyncClient(timeout=timeout) as client:
                data = {"model": selected_model, "messages": messages, **model_config}
                resp = await client.post(host, json={**data, "session_id": session_id}, headers={"Authorization": f"Bearer {provider_key}"})
                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"]
                return RESPONSES["RESPONSE_2"]
        except Exception:
            continue
    marked_item(provider_key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
    return RESPONSES["RESPONSE_2"]

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):
        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})
    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 != RESPONSES["RESPONSE_2"]:
            return result
        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 != RESPONSES["RESPONSE_2"]:
            ACTIVE_CANDIDATE = (host, key)
            return result
    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"]:
        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)
    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)
        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

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 "")
    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], show_progress="full")
    msg.submit(fn=respond_async, inputs=[msg, user_history, selected_model, user_session], outputs=[chatbot, msg, user_session], api_name=INTERNAL_AI_GET_SERVER)
jarvis.launch(max_file_size="1mb")