app.py

# chainlit run app.py -w
# You can find this code for Chainlit python streaming here (https://docs.chainlit.io/concepts/streaming/python)
# OpenAI Chat completion
from dotenv import load_dotenv
load_dotenv()

import os
import sys
import getpass
import nest_asyncio
# import pandas as pd
import faiss
import openai

import chainlit as cl  # importing chainlit for our app
# https://docs.chainlit.io/api-reference/step-class#update-a-step
# DEPRICATED: from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools

import llama_index
from llama_index.core import Settings
from llama_index.core import VectorStoreIndex
from llama_index.core import StorageContext
from llama_index.vector_stores.faiss import FaissVectorStore
from llama_index.core import set_global_handler
from llama_index.core.node_parser import MarkdownElementNodeParser
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.postprocessor.flag_embedding_reranker import FlagEmbeddingReranker
from llama_parse import LlamaParse

from openai import AsyncOpenAI  # importing openai for API usage

os.environ["CUDA_VISIBLE_DEVICES"] = ""
# GET KEYS
LLAMA_CLOUD_API_KEY= os.getenv('LLAMA_CLOUD_API_KEY')
OPENAI_API_KEY=os.getenv("OPENAI_API_KEY")

"""
os.environ["LLAMA_CLOUD_API_KEY"] = getpass.getpass("LLamaParse API Key:")
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
# os.environ["WANDB_API_KEY"] = getpass.getpass("WandB API Key: ")
"""

nest_asyncio.apply()

# PARSING the pdf file
parser = LlamaParse(
    result_type="markdown",
    verbose=True,
    language="en",
    num_workers=2,
)

nvidia_docs = parser.load_data(["./nvidia_2tables.pdf"])
# Note: nvidia_docs contains only one file (it could contain more). nvidia_docs[0] is the pdf we loaded.
print(nvidia_docs[0].text[:1000])

# Getting Settings out of llama_index.core which is a major part of their v0.10 update!
Settings.llm = OpenAI(model="gpt-3.5-turbo")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# Using MarkdownElementNodeParser to help make sense of our Markdown objects so we can leverage the potentially structured information in the parsed documents.

node_parser = MarkdownElementNodeParser(llm=OpenAI(model="gpt-3.5-turbo"), num_workers=8)

nodes = node_parser.get_nodes_from_documents(documents=[nvidia_docs[0]])
# Let's see what's in the metadata of the nodes:
for nd in nodes:
  print(nd.metadata)
  for k,v in nd:
    if k=='table_df':
      print(nd)
# Now we extract our `base_nodes` and `objects` to create the `VectorStoreIndex`.
base_nodes, objects = node_parser.get_nodes_and_objects(nodes)

# We could use the VectorStoreIndex from llama_index.core
# Or we can use the llama_index FAISS llama-index-vector-stores-faiss
# Trying the faiss, and setting its vectors' dimension.
faiss_dim = 1536
faiss_index = faiss.IndexFlatL2(faiss_dim) # default param overwrite=False, so it will append new vector.
# Parameter overwrite=True suppresses appending a vector. 

# Creating the FaissVectorStore and its recursicve_index_faiss
llama_faiss_vector_store = FaissVectorStore(faiss_index=faiss_index)
storage_context = StorageContext.from_defaults(vector_store=llama_faiss_vector_store)
recursive_index_faiss = VectorStoreIndex(nodes=base_nodes+objects, storage_context=storage_context)

# Now we can build our Recursive Query Engine with reranking!

# We'll need to do a couple steps:
# 1. Initalize our reranker using `FlagEmbeddingReranker` powered by the `BAAI/bge-reranker-large`.
# 2. Set up our recursive query engine!

reranker = FlagEmbeddingReranker(
    top_n=5,
    model="BAAI/bge-reranker-large",
)

recursive_query_engine = recursive_index_faiss.as_query_engine(
    similarity_top_k=15,
    node_postprocessors=[reranker],
    verbose=True
)

"""
# Create pandas dataframe to store query+generated response+added truth
columns=["Query", "Response", "Truth"]
gen_df = pd.DataFrame(columns=columns,dtype='str')
"""

# ChatOpenAI Templates
system_template = """Use the following pieces of context to answer the user's question.
If you don't know the answer, say that you don't know, do not try to make up an answer. 
ALWAYS return a "SOURCES" part in your answer.
The "SOURCES" part should be a reference to the source inside the document from which you got your answer.
You are a helpful assistant who always speaks in a pleasant tone! """

user_template = """ Think through your response step by step."""

#user_query = "Who are the E-VP, Operations - and how old are they?"

#response = recursive_query_engine.query(system_template + user_query + user_template)

#str_resp ="{}".format(response)


def retriever_resp(prompt):
    import time
    response = "this is my response"
    time.sleep(5)
    return response

@cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user
async def main(message: cl.Message):
    settings = cl.user_session.get("settings")

    user_query = message.content
#    prompt = system_template+user_query+user_template
    response = recursive_query_engine.query(system_template + user_query + user_template)
#    response = retriever_resp(prompt)
#    print("AAA",user_query)
    str_resp ="{}".format(response)
    msg = cl.Message(content= str_resp)
    await msg.send()