app.py

import os
import streamlit as st
from llama_index import (
    GPTVectorStoreIndex,
    SimpleDirectoryReader,
    ServiceContext,
    StorageContext,
    LLMPredictor,
    load_index_from_storage,
)
from langchain.chat_models import ChatOpenAI

index_name = "./saved_index"
documents_folder = "./documents"


@st.cache_resource
def initialize_index(index_name, documents_folder):
    llm_predictor = LLMPredictor(
        llm=ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
    )
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)
    if os.path.exists(index_name):
        index = load_index_from_storage(
            StorageContext.from_defaults(persist_dir=index_name),
            service_context=service_context,
        )
    else:
        documents = SimpleDirectoryReader(documents_folder).load_data()
        index = GPTVectorStoreIndex.from_documents(
            documents, service_context=service_context
        )
        index.storage_context.persist(persist_dir=index_name)

    return index


@st.cache_data(max_entries=200, persist=True)
def query_index(_index, query_text):
    if _index is None:
        return "Please initialize the index!"
    response = _index.as_query_engine().query(query_text)
    return str(response)


st.title("PQL Chat Demo")
st.header("Team JAFL Miners 🦾")
st.write(
    "Ask for a definition or generate a query in Process Mining language PQL. You can check out the documentation for this language [here](https://docs.celonis.com/en/pql---process-query-language.html). Your query will be answered using the files in the documents folder as context, using embeddings from text-ada-002 and LLM completions from gpt-3.5-turbo."
)

index = None
api_key = st.text_input("Enter your OpenAI API key here:", type="password")
if api_key:
    os.environ["OPENAI_API_KEY"] = api_key
    index = initialize_index(index_name, documents_folder)


if index is None:
    st.warning("Please enter your api key first.")

text = st.text_input("Query text:", value="Insert PQL here to generate Definitions or ask a PQL related question")

sub_prompt_instruction = sub_prompt_instruction = """
Generate a list of definitions and corresponding PQL queries based on the provided list of PQL queries. Ensure that the output maintains the order of the queries.
"""

combined_prompt = f"{sub_prompt_instruction} {text}"


if st.button("Run Query") and text is not None:
    response = query_index(index, combined_prompt)
    st.markdown(response)

    llm_col, embed_col = st.columns(2)
    with llm_col:
        st.markdown(
            f"LLM Tokens Used: {index.service_context.llm_predictor._last_token_usage}"
        )

    with embed_col:
        st.markdown(
            f"Embedding Tokens Used: {index.service_context.embed_model._last_token_usage}"
        )