unified_chat_app.py

import os
import uuid
import json
from datetime import datetime
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain.schema import SystemMessage, HumanMessage, AIMessage
import pandas as pd
import PyPDF2
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Available models for each provider
OPENAI_MODELS = [
    "gpt-5",
    "gpt-4o",
    "gpt-4o-mini",
    "gpt-4-turbo",
    "gpt-4",
    "gpt-3.5-turbo"
]

ANTHROPIC_MODELS = [
    "claude-haiku-4-5-20251001",
    "claude-sonnet-4-5-20250929",
    "claude-opus-4-1-20250805",
    "claude-sonnet-4-20250514",
    "claude-3-7-sonnet-20250219",
    "claude-3-5-sonnet-20240620",
    "claude-3-opus-20240229",
    "claude-3-sonnet-20240229",
    "claude-3-haiku-20240307"
]

# Global variable to maintain entire chat history
# Structure: list of sessions, each with meta and prompts
global_chat_history = []

# Current session metadata tracker
current_meta = {}

def get_current_time():
    """Return current time in yyyymmdd_hhmmss format"""
    return datetime.now().strftime("%Y%m%d_%H%M%S")

def read_pdf(file_path):
    """Extract text from PDF file"""
    with open(file_path, 'rb') as file:
        pdf_reader = PyPDF2.PdfReader(file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() + "\n"
    return text

def read_csv(file_path):
    """Convert CSV to string format"""
    df = pd.read_csv(file_path)
    return df.to_string()

def read_markdown(file_path):
    """Read markdown file"""
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

def read_text(file_path):
    """Read text file"""
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

def process_uploaded_file(file):
    """Process uploaded file based on its extension"""
    if file is None:
        return None, None

    file_ext = os.path.splitext(file.name)[1].lower()

    try:
        if file_ext == '.pdf':
            content = read_pdf(file.name)
            return content, file.name
        elif file_ext == '.csv':
            content = read_csv(file.name)
            return content, file.name
        elif file_ext == '.md':
            content = read_markdown(file.name)
            return content, file.name
        elif file_ext == '.txt':
            content = read_text(file.name)
            return content, file.name
        else:
            return f"Unsupported file type: {file_ext}", None
    except Exception as e:
        return f"Error processing file: {str(e)}", None

def save_chat_history_to_json(session_id):
    """Save chat history as JSON file"""
    output_directory = "_output"
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    timestamp = get_current_time()
    file_name = f"{timestamp}_{session_id}.json"
    file_path = os.path.join(output_directory, file_name)

    with open(file_path, 'w', encoding='utf-8') as file:
        json.dump(global_chat_history, file, ensure_ascii=False, indent=2)

    print(f"Chat history saved to {file_path}")
    return f"Chat history saved to {file_path}"

def save_chat_history_to_markdown(session_id):
    """Save chat history as Markdown file"""
    output_directory = "_output"
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    timestamp = get_current_time()
    file_name = f"{timestamp}_{session_id}.md"
    file_path = os.path.join(output_directory, file_name)

    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(f"# Chat History\n\n")

        # Iterate through all sessions
        for session_idx, session in enumerate(global_chat_history):
            file.write(f"## Session {session_idx + 1}\n\n")

            # Write metadata for this session
            meta = session.get('meta', {})
            file.write(f"**API Provider**: {meta.get('api_provider', 'N/A')}\n\n")
            file.write(f"**Model**: {meta.get('model', 'N/A')}\n\n")
            file.write(f"**Temperature**: {meta.get('temperature', 'N/A')}\n\n")
            file.write(f"**Max Tokens**: {meta.get('max_tokens', 'N/A')}\n\n")

            # Add uploaded file info if exists
            if meta.get('uploaded_file'):
                file.write(f"**Uploaded File**: {meta.get('uploaded_file')}\n\n")

            # Write prompts and chat history for this session
            for i, prompt in enumerate(session.get("prompts", [])):
                file.write(f"### Prompt {i+1}\n")
                file.write(f"{prompt['prompt']}\n\n")
                for entry in prompt["history"]:
                    file.write(f"#### User Message\n")
                    file.write(f"**Time**: {entry['user_message_time']}\n\n")
                    file.write(f"{entry['user_message']}\n\n")
                    file.write(f"#### Bot Response\n")
                    file.write(f"**Time**: {entry['bot_response_time']}\n\n")
                    file.write(f"{entry['bot_response']}\n\n")
                    file.write(f"---\n\n")

            file.write(f"\n---\n\n")

    print(f"Chat history saved to {file_path}")
    return f"Chat history saved to {file_path}"

def get_llm(api_provider, model_name, temperature, max_tokens):
    """Initialize and return the appropriate LLM based on API provider"""
    if api_provider == "OpenAI":
        return ChatOpenAI(
            api_key=os.getenv("OPENAI_API_KEY"),
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens
        )
    else:  # Anthropic
        return ChatAnthropic(
            anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"),
            model=model_name,
            temperature=temperature,
            max_tokens=max_tokens
        )

def check_meta_changed(api_provider, model_name, temperature, max_tokens):
    """Check if meta parameters have changed"""
    global current_meta

    new_meta = {
        "api_provider": api_provider,
        "model": model_name,
        "temperature": temperature,
        "max_tokens": max_tokens
    }

    # Check if meta has changed
    if current_meta != new_meta:
        current_meta = new_meta.copy()
        return True
    return False

def response(message, history, system_message, api_provider, model_name, temperature, max_tokens, file_data=None):
    """Get response from the selected LLM"""
    global global_chat_history

    # Extract filename if file is uploaded
    uploaded_filename = None
    if isinstance(file_data, tuple) and file_data[0] and file_data[1]:
        uploaded_filename = file_data[1]

    # Check if meta has changed
    meta_changed = check_meta_changed(api_provider, model_name, temperature, max_tokens)

    # If meta changed or no sessions exist, create a new session
    if meta_changed or len(global_chat_history) == 0:
        global_chat_history.append({
            "meta": {
                "api_provider": api_provider,
                "model": model_name,
                "temperature": temperature,
                "max_tokens": max_tokens,
                "uploaded_file": uploaded_filename
            },
            "prompts": []
        })
    else:
        # Update uploaded_file in current session meta if file is uploaded
        if uploaded_filename:
            global_chat_history[-1]["meta"]["uploaded_file"] = uploaded_filename

    # Initialize LLM
    llm = get_llm(api_provider, model_name, temperature, max_tokens)

    # Build message history
    history_langchain_format = [SystemMessage(content=system_message)]
    for human, ai in history:
        history_langchain_format.append(HumanMessage(content=human))
        history_langchain_format.append(AIMessage(content=ai))

    # Prepare message for LLM (with file content) and original message (without file content)
    original_message = message
    message_for_llm = message

    # If there's file content, add it to the message for LLM only
    if isinstance(file_data, tuple) and file_data[0] and file_data[1]:
        file_content, filename = file_data
        message_for_llm = f"Here's the content of the uploaded file:\nUploaded file: {filename}\n\n{file_content}\n\nMy question about this content is: {message}"

    history_langchain_format.append(HumanMessage(content=message_for_llm))

    # Get response from LLM
    llm_response = llm.invoke(history_langchain_format)

    user_message_time = get_current_time()
    bot_response_time = get_current_time()

    # Get current session (last one in the list)
    current_session = global_chat_history[-1]

    # Add to global chat history
    if not current_session["prompts"] or current_session["prompts"][-1]["prompt"] != system_message:
        current_session["prompts"].append({
            "prompt": system_message,
            "history": []
        })

    # Save original message (without file content) to history
    current_session["prompts"][-1]["history"].append({
        "user_message": original_message,
        "user_message_time": user_message_time,
        "bot_response": llm_response.content,
        "bot_response_time": bot_response_time
    })

    return llm_response.content

def generate_session_id():
    """Generate a unique session ID"""
    return str(uuid.uuid4())

def get_model_choices(api_provider):
    """Get available models based on API provider"""
    if api_provider == "OpenAI":
        return gr.update(choices=OPENAI_MODELS, value=OPENAI_MODELS[0])
    else:  # Anthropic
        return gr.update(choices=ANTHROPIC_MODELS, value=ANTHROPIC_MODELS[0])

# Gradio UI setup
with gr.Blocks(title="Unified LLM Chat") as demo:
    gr.Markdown("# Unified LLM Chat Interface")
    gr.Markdown("Chat with OpenAI or Anthropic models with file upload support")

    with gr.Row():
        session_id_input = gr.Textbox(
            value=generate_session_id(),
            label="Session ID",
            placeholder="Enter your session ID",
            scale=2
        )
        api_provider = gr.Dropdown(
            choices=["OpenAI", "Anthropic"],
            value="OpenAI",
            label="API Provider",
            scale=1
        )

    with gr.Row():
        model_name = gr.Dropdown(
            choices=OPENAI_MODELS,
            value=OPENAI_MODELS[0],
            label="Model Name",
            scale=2
        )
        temperature = gr.Slider(
            minimum=0,
            maximum=2,
            value=0.7,
            step=0.1,
            label="Temperature",
            scale=1
        )
        max_tokens = gr.Slider(
            minimum=100,
            maximum=8000,
            value=4000,
            step=100,
            label="Max Tokens",
            scale=1
        )

    chatbot = gr.Chatbot(height=500, autoscroll=True)

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Enter your message here",
            container=False,
            scale=7
        )
        submit_message_button = gr.Button("Submit", scale=1)

    with gr.Row():
        file_upload = gr.File(
            label="Upload File (PDF, MD, CSV, TXT)",
            file_types=[".pdf", ".md", ".csv", ".txt"],
            scale=2
        )
        system_prompt = gr.Textbox(
            "",
            label="System Prompt",
            placeholder="Enter system prompt (optional)",
            scale=5
        )
        update_system_button = gr.Button("Update System Prompt", scale=1)

    with gr.Row():
        save_format = gr.Dropdown(
            choices=["Markdown", "JSON"],
            value="Markdown",
            label="Save Format",
            scale=1
        )
        save_button = gr.Button("Save Chat History", scale=1)
        save_status = gr.Textbox(label="Save Status", scale=2, interactive=False)

    with gr.Row():
        delete_last_button = gr.Button("Delete Last Message ❌")
        clear_chat_button = gr.Button("Clear Chat 💫")

    # States
    state = gr.State(value="")  # System prompt state
    file_data = gr.State(value=None)  # File content state

    # Helper functions
    def submit_message_with_file(message, history, session_id, system_message, api_prov, model, temp, max_tok, current_file_data):
        response_text = response(message, history, system_message, api_prov, model, temp, max_tok, current_file_data)
        history.append((message, response_text))
        return history, "", current_file_data

    def handle_file_upload(file):
        if file is None:
            return None
        content, filename = process_uploaded_file(file)
        return (content, filename) if content and filename else None

    def delete_last_chat(chat_history):
        if len(chat_history) > 0:
            chat_history.pop()
        return chat_history

    def update_system_message(new_prompt):
        return new_prompt

    def save_chat(session_id, save_fmt):
        if save_fmt == "JSON":
            return save_chat_history_to_json(session_id)
        else:
            return save_chat_history_to_markdown(session_id)

    # Event handlers
    msg.submit(
        submit_message_with_file,
        [msg, chatbot, session_id_input, state, api_provider, model_name, temperature, max_tokens, file_data],
        [chatbot, msg, file_data]
    )
    submit_message_button.click(
        submit_message_with_file,
        inputs=[msg, chatbot, session_id_input, state, api_provider, model_name, temperature, max_tokens, file_data],
        outputs=[chatbot, msg, file_data]
    )
    update_system_button.click(update_system_message, inputs=system_prompt, outputs=state)
    save_button.click(save_chat, inputs=[session_id_input, save_format], outputs=save_status)
    file_upload.change(handle_file_upload, inputs=[file_upload], outputs=[file_data])
    delete_last_button.click(fn=delete_last_chat, inputs=chatbot, outputs=chatbot)
    clear_chat_button.click(fn=lambda: [], inputs=None, outputs=chatbot)

    # Update model dropdown when API provider changes
    api_provider.change(get_model_choices, inputs=[api_provider], outputs=[model_name])

if __name__ == "__main__":
    demo.launch(share=True)