app.py

import gradio as gr
from mistralai import Mistral
from langchain_community.tools import TavilySearchResults, JinaSearch
import concurrent.futures
import json
import os
import arxiv 
import fitz  # PyMuPDF
from docx import Document
from PIL import Image
import io
import base64
import mimetypes

# Set environment variables for Tavily API
os.environ["TAVILY_API_KEY"] = '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'

# Mistral client API keys
client_1 = Mistral(api_key='eLES5HrVqduOE1OSWG6C5XyEUeR7qpXQ')
client_2 = Mistral(api_key='VPqG8sCy3JX5zFkpdiZ7bRSnTLKwngFJ')
client_3 = Mistral(api_key='cvyu5Rdk2lS026epqL4VB6BMPUcUMSgt')

# Function to encode images in base64
def encode_image_bytes(image_bytes):
    return base64.b64encode(image_bytes).decode('utf-8')

# Functions to process various file types
def process_file(file_path):
    mime_type, _ = mimetypes.guess_type(file_path)
    if mime_type == 'application/pdf':
        return process_pdf(file_path)
    elif mime_type == 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
        return process_docx(file_path)
    elif mime_type == 'text/plain':
        return process_txt(file_path)
    else:
        print(f"Unsupported file type: {mime_type}")
        return None, []

def process_pdf(file_path):
    text = ""
    images = []
    pdf_document = fitz.open(file_path)
    for page_num in range(len(pdf_document)):
        text += pdf_document[page_num].get_text("text")
        for _, img in enumerate(pdf_document.get_page_images(page_num, full=True)):
            xref = img[0]
            base_image = pdf_document.extract_image(xref)
            image_bytes = base_image["image"]
            image_ext = base_image["ext"]
            base64_image = encode_image_bytes(image_bytes)
            image_data = f"data:image/{image_ext};base64,{base64_image}"
            images.append({"type": "image_url", "image_url": image_data})
    return text, images

def process_docx(file_path):
    doc = Document(file_path)
    text = ""
    images = []
    for paragraph in doc.paragraphs:
        text += paragraph.text + "\n"
    for rel in doc.part.rels.values():
        if "image" in rel.target_ref:
            img_data = rel.target_part.blob
            img = Image.open(io.BytesIO(img_data))
            buffered = io.BytesIO()
            img.save(buffered, format="JPEG")
            image_base64 = encode_image_bytes(buffered.getvalue())
            images.append({"type": "image_url", "image_url": f"data:image/jpeg;base64,{image_base64}"})
    return text, images

def process_txt(file_path):
    with open(file_path, "r", encoding="utf-8") as file:
        text = file.read()
    return text, []

# Search setup function
def setup_search(question):
    try:
        tavily_tool = TavilySearchResults(max_results=20)
        results = tavily_tool.invoke({"query": f"{question}"})
        if isinstance(results, list):
            return results, 'tavily_tool'
    except Exception as e:
        print("Error with TavilySearchResults:", e)
    try:
        jina_tool = JinaSearch()
        results = json.loads(str(jina_tool.invoke({"query": f"{question}"})))
        if isinstance(results, list):
            return results, 'jina_tool'
    except Exception as e:
        print("Error with JinaSearch:", e)
    return [], ''

# Function to extract key topics
def extract_key_topics(content, images=[]):
    prompt = f"""
    Extract the primary themes from the text below. List each theme in as few words as possible, focusing on essential concepts only. Format as a concise, unordered list with no extraneous words.

    ```{content}```

    LIST IN ENGLISH:
    - 
    """
    message_content = [{"type": "text", "text": prompt}] + images
    response = client_1.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": message_content}]
    )
    return response.choices[0].message.content

def search_relevant_articles_arxiv(key_topics, max_articles=100):
    articles_by_topic = {}
    final_topics = []

    def fetch_articles_for_topic(topic):
        topic_articles = []
        try:
            # Fetch articles using arxiv.py based on the topic
            search = arxiv.Search(
                query=topic,
                max_results=max_articles,
                sort_by=arxiv.SortCriterion.Relevance
            )
            for result in search.results():
                article_data = {
                    "title": result.title,
                    "doi": result.doi,
                    "summary": result.summary,
                    "url": result.entry_id,
                    "pdf_url": result.pdf_url
                }
                topic_articles.append(article_data)
            final_topics.append(topic)
        except Exception as e:
            print(f"Error fetching articles for topic '{topic}': {e}")

        return topic, topic_articles

    with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
        # Use threads to fetch articles for each topic
        futures = {executor.submit(fetch_articles_for_topic, topic): topic for topic in key_topics}
        for future in concurrent.futures.as_completed(futures):
            topic, articles = future.result()
            if articles:
                articles_by_topic[topic] = articles

    return articles_by_topic, list(set(final_topics))

# Initialize process for text analysis
def init(content, images=[]):
    key_topics = extract_key_topics(content, images)
    key_topics = [topic.strip("- ") for topic in key_topics.split("\n") if topic]
    articles_by_topic, final_topics = search_relevant_articles_arxiv(key_topics)
    result_json = json.dumps(articles_by_topic, indent=4)
    return final_topics, result_json

# Summarization function
def process_article_for_summary(text, images=[], compression_percentage=30):
    prompt = f"""
    You are a commentator.
    # article:
    {text}

    # Instructions:
    ## Summarize:
    In clear and concise language, summarize the key points and themes presented in the article by cutting it by {compression_percentage} percent in the markdown format.

    """
    message_content = [{"type": "text", "text": prompt}] + images
    response = client_3.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": message_content}]
    )
    return response.choices[0].message.content

# Question answering function
def ask_question_to_mistral(text, question, images=[]):
    prompt = f"Answer the following question without mentioning it or repeating the original text on which the question is asked in style markdown.IN RUSSIAN:\nQuestion: {question}\n\nText:\n{text}"
    message_content = [{"type": "text", "text": prompt}] + images
    search_tool, tool = setup_search(question)
    context = ''
    if search_tool:
        if tool == 'tavily_tool':
            for result in search_tool:
                context += f"{result.get('url', 'N/A')} : {result.get('content', 'No content')} \n"
        elif tool == 'jina_tool':
            for result in search_tool:
                context += f"{result.get('link', 'N/A')} : {result.get('snippet', 'No snippet')} : {result.get('content', 'No content')} \n"
    response = client_2.chat.complete(
        model="pixtral-12b-2409",
        messages=[{"role": "user", "content": f'{message_content}\n\nAdditional Context from Web Search:\n{context}'}]
    )
    return response.choices[0].message.content

# Gradio interface
def gradio_interface(file, task, question, compression_percentage):
    if file:
        text, images = process_file(file.name)
    else:
        text, images = "", []

    topics, articles_json = init(text, images)

    if task == "Summarization":
        summary = process_article_for_summary(text, images, compression_percentage)
        return {"Topics": topics, "Summary": summary, "Articles": articles_json}
    elif task == "Question Answering":
        if question:
            answer = ask_question_to_mistral(text, question, images)
            return {"Topics": topics, "Answer": answer, "Articles": articles_json}
        else:
            return {"Topics": topics, "Answer": "No question provided.", "Articles": articles_json}

with gr.Blocks() as demo:
    gr.Markdown("## Text Analysis: Summarization or Question Answering")
    with gr.Row():
        file_input = gr.File(label="Upload File")
        task_choice = gr.Radio(["Summarization", "Question Answering"], label="Select Task")
        question_input = gr.Textbox(label="Question (for Question Answering)", visible=False)
        compression_input = gr.Slider(label="Compression Percentage (for Summarization)", minimum=10, maximum=90, value=30, visible=False)

    task_choice.change(lambda choice: (gr.update(visible=choice == "Question Answering"), 
                                       gr.update(visible=choice == "Summarization")),
                       inputs=task_choice, outputs=[question_input, compression_input])

    with gr.Row():
        result_output = gr.JSON(label="Results")

    submit_button = gr.Button("Submit")
    submit_button.click(gradio_interface, [file_input, task_choice, question_input, compression_input], result_output)

demo.launch()