ComposableConsult/juicese

1import os2import gradio as gr3from langchain_community.chat_models import ChatOpenAI4from langchain_community.vectorstores import FAISS5from langchain_community.document_loaders import WebBaseLoader, PyPDFLoader, TextLoader6from langchain_openai import OpenAIEmbeddings7from langchain.text_splitter import RecursiveCharacterTextSplitter8from langchain.chains import RetrievalQA9from langchain.agents import initialize_agent, Tool, AgentType10 11os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")12os.environ["USER_AGENT"] = "JuiceDocsBot/1.0"13 14FAISS_INDEX_DIR = "juice_faiss_index"15 16JUICE_DOC_LINKS = [17    "https://juice-labs.github.io/juice-docs/docs/juice/intro",18    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/setup/systerm-reqs",19    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/setup/installation",20    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app",21    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/cli-login",22    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/reference",23    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/remote-gpu",24    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/concurrent-workloads",25    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/manual-sessions",26    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/gpu-sharing",27    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/advanced-cli",28    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/advanced-cli/advanced-parameters",29    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/advanced-cli/labels",30    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/advanced-cli/taints-tolerates",31    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/advanced-cli/combining-taints-tolerates",32    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/m2m-tokens",33    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/cache",34    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/cuda-compatibility",35    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/known-issues",36    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/cli-app/troubleshooting",37    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/running-applications/docker",38    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/running-applications/jupyterlab",39    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/running-applications/ollama",40    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/running-applications/examples",41    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/performance-latency/example",42    "https://juice-labs.github.io/juice-docs/docs/juice/pro-users/performance-latency/example-workloads",43    "https://juice-labs.github.io/juice-docs/docs/juice/admin/inviting-users",44    "https://juice-labs.github.io/juice-docs/docs/juice/admin/managing-organizations",45    "https://juice-labs.github.io/juice-docs/docs/juice/admin/managing-pools/creating-pool",46    "https://juice-labs.github.io/juice-docs/docs/juice/admin/managing-pools/removing-pool",47    "https://juice-labs.github.io/juice-docs/docs/juice/admin/managing-pools/pool-perms"48    "https://juice-labs.github.io/juice-docs/docs/juice/admin/viewing-gpus"49    "https://juice-labs.github.io/juice-docs/docs/juice/admin/viewing-sessions",50    "https://juice-labs.github.io/juice-docs/docs/juice/admin/support",51    "https://juice-labs.github.io/juice-docs/docs/juice/admin/account-tab",52    "https://juice-labs.github.io/juice-docs/docs/juice/other-juice-docs/cli-glossary",53    "https://juice-labs.github.io/juice-docs/docs/juice/other-juice-docs/troubleshooting-support",54    "https://juice-labs.github.io/juice-docs/docs/juice/other-juice-docs/terms-conditions",55    "https://juice-labs.github.io/juice-docs/docs/contributing", ]56 57 58def load_or_create_vectorstore():59    if os.path.exists(FAISS_INDEX_DIR):60        return FAISS.load_local(FAISS_INDEX_DIR, OpenAIEmbeddings())61    else:62        all_docs = []63        for link in JUICE_DOC_LINKS:64            loader = WebBaseLoader(link)65            docs = loader.load()66            all_docs.extend(docs)67        splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)68        splits = splitter.split_documents(all_docs)69        vectorstore = FAISS.from_documents(splits, OpenAIEmbeddings())70        vectorstore.save_local(FAISS_INDEX_DIR)71        return vectorstore72 73vectorstore = load_or_create_vectorstore()74 75def load_user_docs(uploaded_files):76    docs = []77    for file in uploaded_files:78        ext = os.path.splitext(file.name)[1].lower()79        if ext == ".pdf":80            loader = PyPDFLoader(file.name)81        else:82            loader = TextLoader(file.name)83        docs.extend(loader.load())84    return docs85 86chat_history = []87 88def chat_with_history(user_input, uploaded_files):89    global chat_history90    extra_docs = load_user_docs(uploaded_files) if uploaded_files else []91    retriever = vectorstore.as_retriever()92    if extra_docs:93        splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)94        extra_splits = splitter.split_documents(extra_docs)95        temp_vectorstore = FAISS.from_documents(extra_splits, OpenAIEmbeddings())96        vectorstore.merge_from(temp_vectorstore)97        retriever = vectorstore.as_retriever()98 99    qa_chain = RetrievalQA.from_chain_type(llm=ChatOpenAI(model="gpt-4", streaming=True), chain_type="stuff", retriever=retriever)100    response = qa_chain.run(user_input)101    chat_history.append((user_input, response))102    return chat_history, ""103 104gr.Interface(105    fn=chat_with_history,106    inputs=["text", gr.File(file_types=[".pdf", ".txt"], file_count="multiple")],107    outputs=["chatbot", "text"],108    title="Juice Pre-Sales Engineer Agent"109).launch()110 111 112 113