jarvis.py

#
# SPDX-FileCopyrightText: Hadad <[email protected]>
# SPDX-License-Identifier: CC-BY-NC-SA-4.0
#

import gradio as gr
import requests
import json
import os
import random
import time
import pytesseract
import pdfplumber
import docx
import pandas as pd
import pptx
import fitz
import io
import uuid
import concurrent.futures

from openai import OpenAI

from optillm.cot_reflection import cot_reflection
from optillm.leap import leap
from optillm.plansearch import plansearch
from optillm.reread import re2_approach
from optillm.rto import round_trip_optimization
from optillm.self_consistency import advanced_self_consistency_approach
from optillm.z3_solver import Z3SymPySolverSystem

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

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"
                    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 process_ai_response(ai_text):
    try:
        result = round_trip_optimization(ai_text)
        result = re2_approach(result)
        result = cot_reflection(result)
        result = advanced_self_consistency_approach(result)
        result = plansearch(result)
        result = leap(result)
        solver = Z3SymPySolverSystem()
        result = solver.solve(result)
        return result
    except Exception:
        return ai_text

def fetch_response(host, provider_key, selected_model, messages, model_config, session_id):
    client = OpenAI(base_url=host, api_key=provider_key)
    data = {"model": selected_model, "messages": messages, **model_config}
    response = client.chat.completions.create(extra_body={"optillm_approach": "rto|re2|cot_reflection|self_consistency|plansearch|leap|z3|bon|moa|mcts|mcp|router|privacy|executecode|json", "session_id": session_id}, **data)
    ai_text = response.choices[0].message.content if response.choices and response.choices[0].message and response.choices[0].message.content else RESPONSES["RESPONSE_2"]
    return process_ai_response(ai_text)

def chat_with_model(history, user_input, selected_model_display, sess):
    if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
        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]
    messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
    messages.append({"role": "user", "content": user_input})
    futures = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=len(LINUX_SERVER_HOSTS)) as executor:
        for host, key in zip(LINUX_SERVER_HOSTS, LINUX_SERVER_PROVIDER_KEYS):
            futures.append(executor.submit(fetch_response, host, key, selected_model, messages, model_config, sess.session_id))
        done, not_done = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
        for future in not_done:
            future.cancel()
        result = list(done)[0].result() if done else RESPONSES["RESPONSE_2"]
    return result

def respond(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
        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 = chat_with_model(history, combined_input, selected_model_display, sess)
    history[-1][1] = ""
    for character in ai_response:
        history[-1][1] += str(character)
        time.sleep(0.0009)
        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])
    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])
    with gr.Row():
        msg = gr.MultimodalTextbox(show_label=False, placeholder=RESPONSES["RESPONSE_5"], interactive=True, file_count="single", file_types=ALLOWED_EXTENSIONS)

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

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