I upgraded to gradio 5.0 and I now get some errors when getting a second response because of the history of the messages I think. here is the documentatin on the gradio website:
ChatInterface

gradio.ChatInterface(fn, type="messages", ···)

Description
ChatInterface is Gradio's high-level abstraction for creating chatbot UIs, and allows you to create a web-based demo around a chatbot model in a few lines of code. Only one parameter is required: fn, which takes a function that governs the response of the chatbot based on the user input and chat history. Additional parameters can be used to control the appearance and behavior of the demo.
Example Usage

Basic Example: A chatbot that echoes back the users’s message

import gradio as gr

def echo(message, history):
    return message

demo = gr.ChatInterface(fn=echo, type="messages", examples=["hello", "hola", "merhaba"], title="Echo Bot")
demo.launch()

Custom Chatbot: A gr.ChatInterface with a custom gr.Chatbot that includes a placeholder as well as upvote/downvote buttons. The upvote/downvote buttons are automatically added when a .like() event is attached to a gr.Chatbot. In order to attach event listeners to your custom chatbot, wrap the gr.Chatbot as well as the gr.ChatInterface inside of a gr.Blocks like this:

import gradio as gr

def yes(message, history):
    return "yes"

def vote(data: gr.LikeData):
    if data.liked:
        print("You upvoted this response: " + data.value["value"])
    else:
        print("You downvoted this response: " + data.value["value"])

with gr.Blocks() as demo:
    chatbot = gr.Chatbot(placeholder="<strong>Your Personal Yes-Man</strong><br>Ask Me Anything")
    chatbot.like(vote, None, None)
    gr.ChatInterface(fn=yes, type="messages", chatbot=chatbot)
    
demo.launch()

Initialization
Parameters

fn: Callable

multimodal: bool

type: Literal['messages', 'tuples']

chatbot: Chatbot | None

textbox: Textbox | MultimodalTextbox | None

additional_inputs: str | Component | list[str | Component] | None

additional_inputs_accordion: str | Accordion | None

examples: list[str] | list[MultimodalValue] | list[list] | None

example_labels: list[str] | None

example_icons: list[str] | None

cache_examples: bool | None

cache_mode: Literal['eager', 'lazy'] | None

title: str | None

description: str | None

theme: Theme | str | None

css: str | None

css_paths: str | Path | list[str | Path] | None

js: str | None

head: str | None

head_paths: str | Path | list[str | Path] | None

analytics_enabled: bool | None

autofocus: bool

autoscroll: bool

concurrency_limit: int | None | Literal['default']

fill_height: bool

delete_cache: tuple[int, int] | None

show_progress: Literal['full', 'minimal', 'hidden']

fill_width: bool

submit_btn: str | bool | None

stop_btn: str | bool | None





and here the code you need to fix:
import os
import json
import shutil
import gradio as gr
from itertools import chain

# Import necessary modules
import arxiv
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain_unstructured import UnstructuredLoader
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.retrievers import ArxivRetriever
from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import DirectoryLoader, UnstructuredURLLoader
from langchain.tools import DuckDuckGoSearchRun
from langchain_ollama import ChatOllama

# For LangSmith tracing
from langsmith import traceable
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_API_KEY"] = "<TRUFFLEHOG_REDACTED_LANGSMITH>"  # Using your actual API key
os.environ["LANGCHAIN_PROJECT"] = "l"  # Using your project name

# Initialize embeddings
embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

# Knowledge Base Vector Store
persist_directory = "C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/chroma_db"
vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)

# arXiv Vector Store (Separate)
arxiv_persist_directory = "C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/arxiv_chroma_db"
arxiv_vectorstore = Chroma(persist_directory=arxiv_persist_directory, embedding_function=embeddings)

# Ensure the 'uploads' directories exist
if not os.path.exists("uploads"):
    os.makedirs("uploads")
if not os.path.exists("arxiv_uploads"):
    os.makedirs("arxiv_uploads")

# Function to load or initialize knowledge base JSON
def load_knowledge_base_json():
    if os.path.exists("C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/knowledge_base.json"):
        with open("C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/knowledge_base.json", "r") as f:
            knowledge_base_data = json.load(f)
    else:
        knowledge_base_data = {"files": [], "urls": [], "arxiv_papers": [], "arxiv_files": []}
    return knowledge_base_data

knowledge_base_data = load_knowledge_base_json()

# Function to save knowledge base JSON
def save_knowledge_base_json():
    with open("C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/knowledge_base.json", "w") as f:
        json.dump(knowledge_base_data, f)

# Function to get processed files
def get_processed_files():
    files = os.listdir("uploads")
    return "### Processed Files:\n" + "\n".join(files) if files else "No files processed yet."

# Function to display processed files and URLs
def display_processed_items():
    global knowledge_base_data
    files_text = get_processed_files()
    urls = knowledge_base_data.get("urls", [])
    display_text = files_text + "\n\n### Processed URLs:\n"
    if urls:
        display_text += "\n".join(urls)
    else:
        display_text += "No URLs processed yet."
    return display_text

# Function to update the knowledge base
@traceable
def update_knowledge_base(uploaded_files, website_urls):
    global knowledge_base_data
    if not uploaded_files and not website_urls:
        return "No new documents or URLs to process.", display_processed_items()

    processed_files = []
    skipped_files = []
    docs = []
    new_urls = []

    # Process uploaded files
    if uploaded_files:
        for file in uploaded_files:
            file_path = file.name
            destination_path = os.path.join("uploads", os.path.basename(file_path))
            file_name = os.path.basename(file_path)

            if os.path.exists(destination_path) or file_name in knowledge_base_data["files"]:
                skipped_files.append(file_name)
            else:
                try:
                    shutil.move(file_path, destination_path)
                    processed_files.append(file_name)
                except Exception as e:
                    print(f"Error moving file {file_name}: {str(e)}")

    # Load documents from uploaded files
    if processed_files:
        loader = DirectoryLoader("uploads", glob="**/*")
        documents = loader.load()
        docs.extend(documents)

    # Load documents from URLs
    if website_urls:
        urls = website_urls.split(',')
        for url in urls:
            url = url.strip()
            if url and url not in knowledge_base_data["urls"]:
                loader = UnstructuredURLLoader(urls=[url])
                try:
                    loaded_docs = loader.load()
                    docs.extend(loaded_docs)
                    new_urls.append(url)
                except Exception as e:
                    print(f"Error loading URL {url}: {str(e)}")
            else:
                skipped_files.append(url)

    # Split documents and add to vector store
    if docs:
        text_splitter = CharacterTextSplitter(chunk_size=2500, chunk_overlap=200, length_function=len)
        split_docs = text_splitter.split_documents(docs)
        vectorstore.add_documents(split_docs)
        vectorstore.persist()
        # Update knowledge base data and save to JSON
        knowledge_base_data["files"].extend(processed_files)
        knowledge_base_data["urls"].extend(new_urls)
        save_knowledge_base_json()
        status_message = f"Processed {len(processed_files)} new files and {len(new_urls)} new URLs."
        if skipped_files:
            status_message += f"\nSkipped existing files/URLs: {', '.join(skipped_files)}"
    else:
        status_message = "No new documents or URLs were added."
    return status_message, display_processed_items()

# Function to clear database and uploaded files
def clear_database():
    global vectorstore, knowledge_base_data

    try:
        vectorstore._client._client.close()
    except AttributeError:
        pass

    shutil.rmtree("uploads", ignore_errors=True)
    os.makedirs("uploads", exist_ok=True)

    shutil.rmtree(persist_directory, ignore_errors=True)

    vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)

    # Clear the knowledge base data
    knowledge_base_data = {"files": [], "urls": [], "arxiv_papers": [], "arxiv_files": []}
    # Remove the JSON file
    if os.path.exists("C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/knowledge_base.json"):
        os.remove("C:/Users/theod/Desktop/Coding/AI/!website/tools/rag/knowledge_base.json")

    return "Knowledge base and uploaded files cleared.", display_processed_items()

# Function to process uploaded files in the chat and include their content
def process_uploaded_files(files):
    content = ""
    if files:
        for file in files:
            if file.name.endswith((".pdf", ".txt", ".docx")):
                loader = UnstructuredLoader(file_path=file.name)
                try:
                    docs = loader.load()
                    content += "\n".join([doc.page_content for doc in docs])
                except Exception as e:
                    print(f"Error processing file {file.name}: {str(e)}")
    return content

# Function to count tokens
def count_tokens(text):
    from transformers import AutoTokenizer, AutoModelForCausalLM

    # Load a tokenizer, you can replace this with a different model (e.g., for llama3)
    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

    tokens = tokenizer.tokenize(text)
    token_ids = tokenizer.encode(text)

    return len(tokens)

# Initialize Llama model
llm = ChatOllama(
    model="llama3.1",
    temperature=0,
    base_url="http://localhost:11434",
    streaming=True
)

# Maximum token limit
MAX_TOKENS = 128000  # Example maximum token limit

# Chat with Llama (normal chat) function
@traceable
def chat_with_llama(message, history):
    if isinstance(message, dict):
        text = message.get('text', '')
        files = message.get('files', [])
        
        # Process uploaded files
        file_contents = []
        for file in files:
            file_path = file['path']
            if file_path.lower().endswith(('.pdf', '.txt', '.docx')):
                try:
                    loader = UnstructuredLoader(file_path=file_path)
                    docs = loader.load()
                    file_contents.append("\n".join([doc.page_content for doc in docs]))
                except Exception as e:
                    print(f"Error processing file {file['orig_name']}: {str(e)}")
        
        # Combine text and file contents
        if file_contents:
            text += "\n\nContent from uploaded files:\n" + "\n\n".join(file_contents)
    else:
        text = message

    # Convert history to list of messages
    messages = [SystemMessage(content="You are a helpful assistant.")]
    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append(HumanMessage(content=user_msg))
            if assistant_msg:
                messages.append(AIMessage(content=assistant_msg))
    
    messages.append(HumanMessage(content=text))
    
    # Check if the total tokens exceed the maximum allowed tokens
    total_tokens = sum(count_tokens(msg.content) for msg in messages)
    if total_tokens > MAX_TOKENS:
        # Truncate the history to fit within the token limit
        while total_tokens > MAX_TOKENS and len(messages) > 2:
            # Remove the oldest user-assistant message pair
            messages.pop(1)
            messages.pop(1)
            total_tokens = sum(count_tokens(msg.content) for msg in messages)
        if total_tokens > MAX_TOKENS:
            return "Error: The conversation exceeds the maximum token limit. Please start a new chat."
    
    # Stream the response from the model
    try:
        response_stream = llm.stream(messages)
        partial_response = ""
        for chunk in response_stream:
            partial_response += chunk.content
            yield partial_response
    except Exception as e:
        yield f"Error: {str(e)}"

# Function to generate search queries, now including history
def generate_search_queries(message, history):
    # Combine the conversation history into a formatted string
    history_messages = ""
    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                history_messages += f"User: {user_msg}\n"
            if assistant_msg:
                history_messages += f"Assistant: {assistant_msg}\n"

    query_generation_prompt = PromptTemplate(
        input_variables=["conversation_history", "user_message"],
        template="""
Given the conversation history:
{conversation_history}

And the user's current message:
"{user_message}"

Generate two distinct search queries that would help retrieve the most relevant information to answer the user's question. Be very simple. For a question like "What is the capital of France?" just produce queries like "capital of France" and "France capital".

Provide the search queries as a JSON list, for example:
["search query 1", "search query 2"]

Only respond with the json format of the search queries, nothing else.
"""
    )
    llm_chain = LLMChain(llm=llm, prompt=query_generation_prompt)
    search_queries_json = llm_chain.run(conversation_history=history_messages, user_message=message)
    # Parse the JSON output
    try:
        search_queries = json.loads(search_queries_json)
    except json.JSONDecodeError:
        # If parsing fails, default to using the original message
        search_queries = [message]
    if not isinstance(search_queries, list):
        search_queries = [search_queries]
    return search_queries

# Function to format content
def format_content(search_results):
    content_formatter_prompt = PromptTemplate(
        input_variables=["search_results"],
        template="""
Given the following search results:

{search_results}

Extract the most relevant information that would help answer the user's question. Provide a concise summary, focusing on key facts and data.
"""
    )
    llm_chain = LLMChain(llm=llm, prompt=content_formatter_prompt)
    formatted_content = llm_chain.run(search_results=search_results)
    return formatted_content

# Updated Internet-Enabled LLM Function with Query Generator and Content Formatter
@traceable
def internet_access_llm(message, history, use_wikipedia):
    # Generate search queries using the updated function that includes history
    search_queries = generate_search_queries(message, history)

    print(search_queries)
    
    # Collect search results
    all_search_results = []
    search_tool = DuckDuckGoSearchRun()
    
    for query in search_queries:
        result = search_tool.run(query)
        all_search_results.append(f"Results for query '{query}':\n{result}")
    
    # Optionally search Wikipedia
    if use_wikipedia:
        for query in search_queries:
            wiki_query = f"{query} site:wikipedia.org"
            result = search_tool.run(wiki_query)
            all_search_results.append(f"Wikipedia results for query '{query}':\n{result}")
    
    # Combine all search results
    combined_search_results = "\n\n".join(all_search_results)
    
    # Format the search results
    formatted_content = format_content(combined_search_results)
    
    # Build the conversation history
    messages = [SystemMessage(content="You are an assistant that uses internet search results to provide accurate answers.")]
    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append(HumanMessage(content=user_msg))
                if assistant_msg is not None:
                    messages.append(AIMessage(content=assistant_msg))
    
    # Add the latest message with formatted content
    prompt = f"Based on the following information:\n{formatted_content}\n\nQuestion:\n{message}\n\nAnswer:"
    messages.append(HumanMessage(content=prompt))
    
    # Stream the response from the model
    try:
        response_stream = llm.stream(messages)
        partial_response = ""
        for chunk in response_stream:
            partial_response += chunk.content
            yield partial_response
    except Exception as e:
        yield f"Error: {str(e)}"

# Function to handle knowledge base Q&A (unchanged)
@traceable
def knowledge_base_chat(message, history):
    # Generate search queries, now including history
    search_queries = generate_search_queries(message, history)
    # Perform retrieval for each generated query
    all_results = []
    for query in search_queries:
        results = vectorstore.similarity_search(query, k=4)
        all_results.extend(results)
    if not all_results:
        yield "No relevant information found in the knowledge base."
        return
    # Remove duplicate documents
    unique_results = list({doc.page_content: doc for doc in all_results}.values())
    # Use the LLM to extract important information
    extraction_prompt = PromptTemplate(
        input_variables=["user_message", "document"],
        template="""
Given the user's message:
"{user_message}"

And the following document:

{document}

Extract important information that could help to answer the question. Respond with the relevant parts of the document in original text, omitting irrelevant sections. So respond with the same text from the document, but fix formations and remove unwanded stuff that could confuse the ai. but for the relevate information dont change any language or structure of it.
"""
    )
    extracted_info_list = []
    extraction_chain = LLMChain(llm=llm, prompt=extraction_prompt)
    for doc in unique_results:
        extracted_info = extraction_chain.run(user_message=message, document=doc.page_content)
        extracted_info_list.append(extracted_info)
    # Combine the extracted information
    context = "\n\n".join(extracted_info_list)
    # Build the conversation history
    messages = [SystemMessage(content="You are a helpful assistant that answers questions based on the provided context.")]
    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append(HumanMessage(content=user_msg))
                if assistant_msg is not None:
                    messages.append(AIMessage(content=assistant_msg))
    # Add the latest message with context
    prompt = f"Context:\n{context}\n\nQuestion:\n{message}\n\nAnswer:"
    messages.append(HumanMessage(content=prompt))
    # Stream the response from the model
    try:
        response_stream = llm.stream(messages)
        partial_response = ""
        for chunk in response_stream:
            partial_response += chunk.content
            yield partial_response
    except Exception as e:
        yield f"Error: {str(e)}"

# Function to retrieve arXiv papers and update arXiv knowledge base
@traceable
def arxiv_research(arxiv_query, max_papers, uploaded_files, arxiv_ids_or_links):
    if not arxiv_query and not uploaded_files and not arxiv_ids_or_links:
        return "No query, files, or IDs provided.", display_arxiv_items()

    processed_papers = []
    skipped_papers = []
    docs = []

    # Search arXiv using the query
    if arxiv_query:
        search = arxiv.Search(query=arxiv_query, max_results=int(max_papers))
        for result in search.results():
            paper_id = result.get_short_id()
            paper_title = result.title
            pdf_url = result.pdf_url
            # Check if paper is already processed
            if paper_id in knowledge_base_data.get("arxiv_papers", []):
                skipped_papers.append(paper_title)
                continue
            # Download PDF
            try:
                paper = next(arxiv.Client().results(arxiv.Search(id_list=[paper_id])))
                pdf_path = paper.download_pdf(dirpath="arxiv_uploads")
                # Extract text
                loader = PyPDFLoader(pdf_path)
                paper_doc = loader.load()
                docs.extend(paper_doc)
                processed_papers.append(paper_title)
                # Save paper info
                knowledge_base_data.setdefault("arxiv_papers", []).append(paper_id)
            except Exception as e:
                print(f"Error downloading or processing paper {paper_title}: {str(e)}")

    # Process uploaded files
    if uploaded_files:
        for file in uploaded_files:
            file_path = file.name
            destination_path = os.path.join("arxiv_uploads", os.path.basename(file_path))
            file_name = os.path.basename(file_path)
            if os.path.exists(destination_path) or file_name in knowledge_base_data.get("arxiv_files", []):
                skipped_papers.append(file_name)
            else:
                try:
                    shutil.move(file_path, destination_path)
                    loader = PyPDFLoader(destination_path)
                    paper_doc = loader.load()
                    docs.extend(paper_doc)
                    processed_papers.append(file_name)
                    knowledge_base_data.setdefault("arxiv_files", []).append(file_name)
                except Exception as e:
                    print(f"Error processing file {file_name}: {str(e)}")

    # Process arXiv IDs or URLs
    if arxiv_ids_or_links:
        ids_or_links = arxiv_ids_or_links.split(',')
        for item in ids_or_links:
            item = item.strip()
            if not item:
                continue
            # Extract arXiv ID
            if 'arxiv.org' in item:
                paper_id = item.split('/')[-1]
            else:
                paper_id = item
            if paper_id in knowledge_base_data.get("arxiv_papers", []):
                skipped_papers.append(paper_id)
                continue
            # Download and process paper
            try:
                paper = next(arxiv.Client().results(arxiv.Search(id_list=[paper_id])))
                pdf_path = paper.download_pdf(dirpath="arxiv_uploads")
                loader = PyPDFLoader(pdf_path)
                paper_doc = loader.load()
                docs.extend(paper_doc)
                processed_papers.append(paper.title)
                knowledge_base_data.setdefault("arxiv_papers", []).append(paper_id)
            except Exception as e:
                print(f"Error processing paper {paper_id}: {str(e)}")

    # Split documents and add to arXiv vector store
    if docs:
        text_splitter = CharacterTextSplitter(chunk_size=2500, chunk_overlap=200, length_function=len)
        split_docs = text_splitter.split_documents(docs)
        arxiv_vectorstore.add_documents(split_docs)
        arxiv_vectorstore.persist()
        # Save knowledge base data
        save_knowledge_base_json()
        status_message = f"Processed {len(processed_papers)} new papers."
        if skipped_papers:
            status_message += f"\nSkipped existing papers: {', '.join(skipped_papers)}"
    else:
        status_message = "No new papers were added."
    return status_message, display_arxiv_items()

# Function to display processed arXiv items
def display_arxiv_items():
    arxiv_files = os.listdir("arxiv_uploads")
    arxiv_papers = knowledge_base_data.get("arxiv_papers", [])
    display_text = "### Processed arXiv Papers:\n"
    if arxiv_papers:
        display_text += "\n".join(arxiv_papers)
    else:
        display_text += "No arXiv papers processed yet."
    display_text += "\n\n### Uploaded Files:\n"
    if arxiv_files:
        display_text += "\n".join(arxiv_files)
    else:
        display_text += "No files uploaded yet."
    return display_text

# Function to allow users to download papers
def download_papers():
    files = os.listdir("arxiv_uploads")
    file_paths = [os.path.join("arxiv_uploads", f) for f in files]
    return file_paths

# Function to handle arXiv Knowledge Base Q&A
@traceable
def ask_arxiv_knowledge_base(message, history):
    # Generate search queries
    search_queries = generate_search_queries(message, history)
    # Perform retrieval
    all_results = []
    for query in search_queries:
        results = arxiv_vectorstore.similarity_search(query, k=4)
        all_results.extend(results)
    if not all_results:
        yield "No relevant information found in the arXiv knowledge base."
        return
    # Remove duplicates
    unique_results = list({doc.page_content: doc for doc in all_results}.values())
    # Use LLM to extract information
    extraction_prompt = PromptTemplate(
        input_variables=["user_message", "document"],
        template="""
Given the user's message:
"{user_message}"

And the following document:

{document}

Extract important information that could help to answer the question. Respond with the relevant parts of the document in original text, omitting irrelevant sections.
"""
    )
    extracted_info_list = []
    extraction_chain = LLMChain(llm=llm, prompt=extraction_prompt)
    for doc in unique_results:
        extracted_info = extraction_chain.run(user_message=message, document=doc.page_content)
        extracted_info_list.append(extracted_info)
    # Combine the extracted information
    context = "\n\n".join(extracted_info_list)
    # Also prepare a list of papers used
    papers_used = []
    for doc in unique_results:
        metadata = doc.metadata
        title = metadata.get('Title', 'Unknown title')
        authors = metadata.get('Authors', 'Unknown authors')
        paper_info = f"Title: {title}, Authors: {authors}"
        papers_used.append(paper_info)
    papers_info = "\n".join([f"- {info}" for info in papers_used])
    # Build the conversation history
    messages = [SystemMessage(content="You are a helpful assistant that answers questions based on arXiv research papers.")]
    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append(HumanMessage(content=user_msg))
                if assistant_msg is not None:
                    messages.append(AIMessage(content=assistant_msg))
    # Add the latest message with context
    prompt = f"The following arXiv papers were used:\n{papers_info}\n\nContext:\n{context}\n\nQuestion:\n{message}\n\nAnswer:"
    messages.append(HumanMessage(content=prompt))
    # Stream the response from the model
    try:
        response_stream = llm.stream(messages)
        partial_response = ""
        for chunk in response_stream:
            partial_response += chunk.content
            yield partial_response
    except Exception as e:
        yield f"Error: {str(e)}"

# Initialize the ArxivRetriever
arxiv_retriever = ArxivRetriever(
    load_max_docs=5,
    get_full_documents=True,
)

# Function for the Arxiv Retriever Chat
@traceable
def arxiv_retriever_chat(message, history):
    # Use the ArxivRetriever to get documents
    docs = arxiv_retriever.invoke(message)

    if not docs:
        yield "No relevant papers found."
        return

    # Prepare context from docs
    context = "\n\n".join([doc.page_content for doc in docs])

    # Also prepare a list of papers used
    papers_used = []
    for doc in docs:
        title = doc.metadata.get('Title', 'Unknown title')
        authors = doc.metadata.get('Authors', 'Unknown authors')
        paper_info = f"Title: {title}, Authors: {authors}"
        papers_used.append(paper_info)
    papers_info = "\n".join([f"- {info}" for info in papers_used])

    # Build the conversation history
    messages = [SystemMessage(content="You are a helpful assistant that answers questions based on arXiv papers.")]

    if history:
        for user_msg, assistant_msg in history:
            if user_msg:
                messages.append(HumanMessage(content=user_msg))
                if assistant_msg is not None:
                    messages.append(AIMessage(content=assistant_msg))

    # Add the latest message with context
    prompt = f"The following arXiv papers were used:\n{papers_info}\n\nContext:\n{context}\n\nQuestion:\n{message}\n\nAnswer:"

    messages.append(HumanMessage(content=prompt))

    # Stream the response from the model
    try:
        response_stream = llm.stream(messages)
        partial_response = ""
        for chunk in response_stream:
            partial_response += chunk.content
            yield partial_response
    except Exception as e:
        yield f"Error: {str(e)}"
css = """
.contain { display: flex; flex-direction: column; }
#component-0 { height: 100%; }
#chatbot { flex-grow: 1; overflow-y: auto; }
}
"""

# Gradio interface
with gr.Blocks(css="""
.contain { display: flex; flex-direction: column; }
#component-0 { height: 100%; }
#chatbot { flex-grow: 1; overflow-y: auto; }
""") as rag_interface:
    gr.Markdown("# Multi-Tab Gradio App with Llama and Similarity Search")

    # Chat with Llama Tab
    with gr.Tab("Chat with Llama"):
        gr.Markdown("## Chat Interface")
        chatbot = gr.Chatbot(height=600, type="messages")
        
        chat_interface = gr.ChatInterface(
            fn=chat_with_llama,
            chatbot=chatbot,
            type="messages",
            examples=["Hello", "What's the weather like?", "Tell me a joke"],
            title="Chat with Llama",
            description="Ask me anything!",
            submit_btn="Send",
            stop_btn="Stop",
            autofocus=True,
            # retry_btn=gr.Button("Retry"),
            # undo_btn=gr.Button("Undo"),
            # clear_btn=gr.Button("Clear")
        )

        # chat_interface.retry_btn.click(chat_interface.retry)
        # chat_interface.undo_btn.click(chat_interface.undo)
        # chat_interface.clear_btn.click(chat_interface.clear)

    # Knowledge Base Tab
    with gr.Tab("Knowledge Base"):
        gr.Markdown("## Knowledge Base Management and Chat")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Update Knowledge Base")
                kb_upload_files = gr.File(
                    label="Upload Files",
                    file_types=[".pdf", ".txt", ".docx"],
                    file_count="multiple"
                )
                kb_website_urls = gr.Textbox(
                    label="Website URLs",
                    placeholder="Enter comma-separated URLs"
                )
                kb_update_button = gr.Button("Update Knowledge Base")
                kb_clear_button = gr.Button("Clear Database")
                kb_update_status = gr.Markdown()
                kb_update_button.click(
                    update_knowledge_base,
                    inputs=[kb_upload_files, kb_website_urls],
                    outputs=[kb_update_status, kb_update_status]
                )
                kb_clear_button.click(
                    clear_database,
                    outputs=[kb_update_status, kb_update_status]
                )
                gr.Markdown("### Processed Files and URLs")
                processed_items_display = gr.Markdown(display_processed_items())
            
            with gr.Column(scale=2):
                gr.Markdown("### Chat with Knowledge Base")
                chatbot = gr.Chatbot(height=600, type="messages")
                gr_chat_kb = gr.ChatInterface(
                    fn=knowledge_base_chat,
                    chatbot=chatbot,
                    type="messages",
                    title="Knowledge Base Chat",
                    # submit_btn="Send",
                    # stop_btn="Stop"
                )

    # LLM with Internet Access Tab
    with gr.Tab("LLM with Internet Access"):
        gr.Markdown("## Internet-Enabled LLM with Query Generator and Content Formatter")
        use_wikipedia = gr.Checkbox(label="Use Wikipedia")
        chatbot = gr.Chatbot(height=600, type="messages")
        gr_chat_internet = gr.ChatInterface(
            fn=internet_access_llm,
            chatbot=chatbot,
            type="messages",
            title="LLM with Internet Access",
            additional_inputs=[use_wikipedia],
            # submit_btn="Send",
            # stop_btn="Stop"
        )

    # arXiv Research Papers Tab
    with gr.Tab("arXiv Research Papers"):
        gr.Markdown("## Retrieve arXiv Papers and Chat")
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Retrieve/Add arXiv Papers")
                arxiv_query = gr.Textbox(
                    label="Search Query",
                    placeholder="Enter your arXiv search query"
                )
                max_papers = gr.Slider(
                    label="Number of Papers",
                    minimum=1,
                    maximum=10,
                    step=1,
                    value=5
                )
                arxiv_upload_files = gr.File(
                    label="Upload PDF Files",
                    file_types=[".pdf"],
                    file_count="multiple"
                )
                arxiv_ids_or_links = gr.Textbox(
                    label="arXiv IDs or URLs",
                    placeholder="Enter comma-separated arXiv IDs or URLs"
                )
                arxiv_button = gr.Button("Retrieve/Add Papers")
                arxiv_result = gr.Markdown()
                arxiv_button.click(
                    arxiv_research,
                    inputs=[arxiv_query, max_papers, arxiv_upload_files, arxiv_ids_or_links],
                    outputs=[arxiv_result, arxiv_result]
                )
                gr.Markdown("### Processed arXiv Papers")
                processed_items_display_arxiv = gr.Markdown(display_arxiv_items())
                gr.Markdown("### Download Papers")
                download_button = gr.Button("Download Papers")
                download_files = gr.File()
                download_button.click(
                    download_papers,
                    outputs=[download_files]
                )
            
            with gr.Column(scale=2):
                gr.Markdown("### Chat with arXiv Knowledge Base")
                chatbot = gr.Chatbot(height=600, type="messages")
                gr_chat_arxiv = gr.ChatInterface(
                    fn=ask_arxiv_knowledge_base,
                    chatbot=chatbot,
                    type="messages",
                    title="arXiv Knowledge Base Chat",
                    # submit_btn="Send",
                    # stop_btn="Stop"
                )

    # New Arxiv Retriever Tab
    with gr.Tab("Arxiv Retriever"):
        gr.Markdown("## Ask Questions Based on arXiv Papers")
        gr.Markdown("This tab uses the ArxivRetriever to fetch relevant papers and provide answers.")
        chatbot = gr.Chatbot(height=600, type="messages")
        gr_chat_arxiv_retriever = gr.ChatInterface(
            fn=arxiv_retriever_chat,
            chatbot=chatbot,
            type="messages",
            title="Arxiv Retriever Chat",
            # submit_btn="Send",
            # stop_btn="Stop"
        )

if __name__ == "__main__":
    rag_interface.launch()