make the arxiv stuff show what papers it used and also make that the arxiv tab has the same generation of queries as the knowladgebase tab:
import os
import json
import shutil
import gradio as gr
from itertools import chain
Import necessary modules

from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
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_community.tools import DuckDuckGoSearchRun
from langchain_community.document_loaders import DirectoryLoader
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>" # Replace with your actual API key
os.environ["LANGCHAIN_PROJECT"] = "l" # Optional: "default" is used if not set
Initialize embeddings and vector store

embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
persist_directory = "./chroma_db"
vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
Ensure the 'uploads' directory exists

if not os.path.exists("uploads"):
os.makedirs("uploads")
Function to load or initialize knowledge base JSON

def load_knowledge_base_json():
if os.path.exists("knowledge_base.json"):
with open("knowledge_base.json", "r") as f:
knowledge_base_data = json.load(f)
else:
knowledge_base_data = {"files": [], "urls": []}
return knowledge_base_data

knowledge_base_data = load_knowledge_base_json()
Function to save knowledge base JSON

def save_knowledge_base_json():
with open("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 = UnstructuredLoader(web_url=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": []}
# Remove the JSON file
if os.path.exists("knowledge_base.json"):
    os.remove("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):
return len(text.split()) # Simple token count; adjust with tokenizer if needed
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, uploaded_files):
# Process uploaded files and include their content
file_content = process_uploaded_files(uploaded_files)
# Convert history to list of messages
messages = [SystemMessage(content="You are a helpful assistant. The user is able to upload files that get processed as text for you. You always answer the user's question by using the context or your knowledge in the same language the user wants you to talk in.")]
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))
# Append file content to user message if any
if file_content:
message += f"\n\nContent files Uploaded:\n{file_content}"
message += "\n\nWith that knowledge answer to the question as best as possible: "
messages.append(HumanMessage(content=message))
# 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)}"
Updated Function to handle knowledge base Q&A using LLM-generated search queries and extracting important information

@traceable
def knowledge_base_chat(message, history):
# Generate 2 search queries using the LLM
query_generation_prompt = PromptTemplate(
input_variables=["user_message"],
template="""
Given the user's message:
"{user_message}"

Generate two distinct search queries that would help retrieve the most relevant documents from the knowledge base to answer the user's question. Be very simple. for a question like "what is the capital of France?" just do a query with "capital" and for for "france captial" so if there is a text about captails and france isnt specified it should still find it.

Provide the search queries as a JSON list, for example:
["search query 1", "search query 2"]
"""
)
llm_chain = LLMChain(llm=llm, prompt=query_generation_prompt)

search_queries_json = llm_chain.run(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]

# 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 based on their content
unique_results = list({doc.page_content: doc for doc in all_results}.values())

# Use the LLM to extract important information from each document
extraction_prompt = PromptTemplate(
    input_variables=["user_message", "document"],
    template="""

Given the user's message:
"{user_message}"

And the following document:

{document}

You only extract important information that could help to answer the question. You respond with all the stuff that could be useful to answer that in the original text. Only leave out parts that are not relevant or could confuse the ai to write a final answer.
So answer with the text of the following document in original text and output cut out part of the stuff that is not relevant.
"""
)

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 to form the context
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
        # Use yield to stream the response
        yield partial_response
except Exception as e:
    yield f"Error: {str(e)}"

Function for the Internet-Enabled LLM

@traceable
def internet_access_llm(message, history, use_wikipedia):
# Generate a query from the user input
query = message
# Search the internet
search_tool = DuckDuckGoSearchRun()
search_results = search_tool.invoke(query)
# Optionally search Wikipedia
if use_wikipedia:
wikipedia_tool = DuckDuckGoSearchRun(backend="wikipedia")
wiki_results = wikipedia_tool.invoke(query)
search_results += "\n\n" + wiki_results
# Build the conversation history
messages = [SystemMessage(content="You get information from the internet and answer the user's question with that information. Provide an accurate and detailed answer to the user's question using the search results.")]
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 search results
prompt = f"Search Results:\n{search_results}\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 knowledge base

@traceable
def arxiv_research(query, max_papers):
global knowledge_base_data
retriever = ArxivRetriever(load_max_docs=int(max_papers), get_full_documents=True)
docs = retriever.get_relevant_documents(query)
# Add documents to the 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
arxiv_titles = [doc.metadata.get('Title', 'Unknown Title') for doc in docs]
knowledge_base_data["files"].extend(arxiv_titles)
save_knowledge_base_json()
status_message = f"Added {len(docs)} papers to the knowledge base."
else:
status_message = "No documents were added from arXiv."
return status_message, display_processed_items()
Function to handle arXiv Knowledge Base Q&A using similarity search

@traceable
def ask_arxiv_knowledge_base(message, history):
# Search in the vector store using similarity search
results = vectorstore.similarity_search(message, k=3)
if not results:
yield "No relevant information found in the arXiv knowledge base."
return
context = "\n\n".join([doc.page_content for doc in results])
# 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"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)}"
Gradio interface

with gr.Blocks() as demo:
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")
    with gr.Row():
        upload_files = gr.File(
            label="Upload PDFs, DOCX or TXT files",
            file_types=[".pdf", ".txt", ".docx"],
            file_count="multiple",
            visible=True
        )
    gr_chat = gr.ChatInterface(
        fn=chat_with_llama,
        title="Chat with Llama",
        additional_inputs=[upload_files]
    )

# 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())
            # Update the display when the knowledge base is updated or cleared
            kb_update_button.click(
                update_knowledge_base,
                inputs=[kb_upload_files, kb_website_urls],
                outputs=[kb_update_status, processed_items_display]
            )
            kb_clear_button.click(
                clear_database,
                outputs=[kb_update_status, processed_items_display]
            )
        with gr.Column(scale=2):
            gr.Markdown("### Chat with Knowledge Base")
            gr_chat_kb = gr.ChatInterface(
                fn=knowledge_base_chat,
                title="Knowledge Base Chat"
            )

# LLM with Internet Access Tab
with gr.Tab("LLM with Internet Access"):
    gr.Markdown("## Internet-Enabled LLM")
    use_wikipedia = gr.Checkbox(label="Use Wikipedia")
    gr_chat_internet = gr.ChatInterface(
        fn=internet_access_llm,
        title="LLM with Internet Access",
        additional_inputs=[use_wikipedia]
    )

# 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 arXiv Papers")
            arxiv_query = gr.Textbox(
                label="Search Query",
                placeholder="Enter your search query for arXiv papers"
            )
            max_papers = gr.Slider(
                label="Number of Papers",
                minimum=1,
                maximum=10,
                step=1,
                value=5
            )
            arxiv_button = gr.Button("Retrieve Papers")
            arxiv_result = gr.Markdown()
            arxiv_button.click(
                arxiv_research,
                inputs=[arxiv_query, max_papers],
                outputs=[arxiv_result, arxiv_result]
            )
            gr.Markdown("### Processed Files and URLs")
            processed_items_display_arxiv = gr.Markdown(display_processed_items())
            # Update the display when new papers are added
            arxiv_button.click(
                arxiv_research,
                inputs=[arxiv_query, max_papers],
                outputs=[arxiv_result, processed_items_display_arxiv]
            )
        with gr.Column(scale=2):
            gr.Markdown("### Chat with arXiv Knowledge Base")
            gr_chat_arxiv = gr.ChatInterface(
                fn=ask_arxiv_knowledge_base,
                title="arXiv Knowledge Base Chat"
            )

# Launch the app
demo.launch()