app.py

#
# Copyright (C) Hadad <[email protected]>
# All rights reserved.
#
# This code is made available under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0) License.
# You are free to share and adapt the code for non-commercial purposes, as long as you provide appropriate credit,
# do not use it for commercial purposes, and distribute your contributions under the same license.
#
# Contributions can be made by directly submitting pull requests.
#
# For inquiries or permission requests, please contact [email protected].
#
# License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
#

import gradio as gr
import requests
import json
import os
import threading
import random
import time
import pytesseract
import pdfplumber
import docx
import pandas as pd
import pptx
import fitz
import io
from pathlib import Path
from PIL import Image
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")

LINUX_SERVER_HOSTS = [host for host in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if host]
LINUX_SERVER_PROVIDER_KEYS = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]

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())
DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))

META_TAGS = os.getenv("META_TAGS")

ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS"))

session = requests.Session()

def get_model_key(display_name):
    return next((k for k, v in MODEL_MAPPING.items() if v == display_name), MODEL_CHOICES[0])

def simulate_streaming_response(text):
    for line in text.splitlines():
        yield line + "\n"
        time.sleep(0.05)

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"
                    tables = page.extract_tables()
                    if tables:
                        for table in 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"]:
            try:
                pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
                image = Image.open(file_path)
                text = pytesseract.image_to_string(image)
                content += text + "\n"
            except Exception as e:
                content += f"{e}\n"
        else:
            content = Path(file_path).read_text(encoding="utf-8")
    except Exception as e:
        content = f"{file_path}: {e}"
    return content.strip()

def chat_with_model(history, user_input, selected_model_display):
    if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
        yield RESPONSES["RESPONSE_3"]
        return
    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]
    messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
    messages.append({"role": "user", "content": user_input})
    data = {"model": selected_model, "messages": messages, **model_config}
    random.shuffle(LINUX_SERVER_PROVIDER_KEYS)
    random.shuffle(LINUX_SERVER_HOSTS)
    for api_key in LINUX_SERVER_PROVIDER_KEYS[:2]:
        for host in LINUX_SERVER_HOSTS[:2]:
            try:
                response = session.post(host, json=data, headers={"Authorization": f"Bearer {api_key}"})
                if response.status_code < 400:
                    ai_text = response.json().get("choices", [{}])[0].get("message", {}).get("content", RESPONSES["RESPONSE_2"])
                    yield from simulate_streaming_response(ai_text)
                    return
            except requests.exceptions.RequestException:
                continue
    yield RESPONSES["RESPONSE_3"]

def respond(multi_input, history, selected_model_display):
    message = {"text": multi_input.get("text", "").strip(), "files": multi_input.get("files", [])}
    if not message["text"] and not message["files"]:
        return history, gr.MultimodalTextbox(value=None, interactive=True)
    combined_input = ""
    for file_item in message["files"]:
        if isinstance(file_item, dict) and "name" in file_item:
            file_path = file_item["name"]
        else:
            file_path = file_item
        file_content = extract_file_content(file_path)
        combined_input += f"{Path(file_path).name}\n\n{file_content}\n\n"
    if message["text"]:
        combined_input += message["text"]
    history.append([combined_input, ""])
    ai_response = ""
    for chunk in chat_with_model(history, combined_input, selected_model_display):
        ai_response += chunk
        history[-1][1] = ai_response
    return history, gr.MultimodalTextbox(value=None, interactive=True)

def change_model(new_model_display):
    return [], new_model_display

with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as demo:
    user_history = gr.State([])
    selected_model = gr.State(MODEL_CHOICES[0])
    chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
    model_dropdown = gr.Dropdown(show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0])
    msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], scale=0, interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)

    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, selected_model])
    msg.submit(fn=respond, inputs=[msg, user_history, selected_model], outputs=[chatbot, msg])

demo.launch(show_api=False, max_file_size="1mb")