update app to make use of chain

app.py

@@ -1,19 +1,20 @@
 import os
 import threading
+import datetime
 import streamlit as st
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_databricks.vectorstores import DatabricksVectorSearch
 from itertools import tee
+from chain import ChainBuilder
 
 DATABRICKS_HOST = os.environ.get("DATABRICKS_HOST")
 DATABRICKS_TOKEN = os.environ.get("DATABRICKS_TOKEN")
-VS_ENDPOINT_NAME = os.environ.get("VS_ENDPOINT_NAME")
-VS_INDEX_NAME = os.environ.get("VS_INDEX_NAME")
+# remove these secrets from the container
+# VS_ENDPOINT_NAME = os.environ.get("VS_ENDPOINT_NAME")
+# VS_INDEX_NAME = os.environ.get("VS_INDEX_NAME")
 
 if DATABRICKS_HOST is None:
     raise ValueError("DATABRICKS_HOST environment variable must be set")
 if DATABRICKS_TOKEN is None:
-    raise ValueError("DATABRICKS_API_TOKEN environment variable must be set")
+    raise ValueError("DATABRICKS_TOKEN environment variable must be set")
 
 MODEL_AVATAR_URL= "./VU.jpeg"
 
@@ -21,20 +22,19 @@ MODEL_AVATAR_URL= "./VU.jpeg"
 # MSG_CLIPPED_AT_MAX_OUT_TOKENS = "Reached maximum output tokens for DBRX Playground"
 
 EXAMPLE_PROMPTS = [
-    "Tell me about maximum out-of-pocket costs in healthcare.",
-    "Write a haiku about Nashville, Tennessee.",
     "How is a data lake used at Vanderbilt University Medical Center?",
     "In a table, what are some of the greatest hurdles to healthcare in the United States?",
     "What does EDW stand for in the context of Vanderbilt University Medical Center?",
     "Code a sql statement that can query a database named 'VUMC'.",
     "Write a short story about a country concert in Nashville, Tennessee.",
+    "Tell me about maximum out-of-pocket costs in healthcare.",
 ]
 
 TITLE = "Vanderbilt AI Assistant"
-DESCRIPTION="""Welcome to the first generation Vanderbilt AI assistant! \n 
+DESCRIPTION= """Welcome to the first generation Vanderbilt AI assistant! \n 
 This AI assistant is built atop the Databricks DBRX large language model 
-and is augmented with additional organization-specific knowledge. Specifically, it has been preliminarily augmented with knowledge of Vanderbilt University Medical Center
-terms like **Data Lake**, **EDW** (Enterprise Data Warehouse), **HCERA** (Health Care and Education Reconciliation Act), and **thousands more!** The model has **no access to PHI**. 
+and is augmented with additional organization-specific knowledge. Particularly, it has been preliminarily augmented with knowledge of Vanderbilt University Medical Center
+terms like **Data Lake**, **EDW**, **HCERA**, and **thousands more**. (Ask the assistant if you don't know what any of these terms mean!) **Disclaimer**: The model has **no access to PHI**. \n
 Try querying the model with any of the example prompts below for a simple introduction to both Vanderbilt-specific and general knowledge queries. The purpose of this 
 model is to allow VUMC employees access to an intelligent assistant that improves and expedites VUMC work. \n
 Feedback and ideas are very welcome! Please provide any feedback, ideas, or issues to the email: **[email protected]**. 
@@ -42,6 +42,22 @@ We hope to gradually improve this AI assistant to create a large-scale, all-incl
 
 GENERAL_ERROR_MSG = "An error occurred. Please refresh the page to start a new conversation."
 
+# *** TODO we dont need this here anymore, since the chain handles all of this
+# SYSTEM_PROMPT = """You are DBRX, created by Databricks and augmented by John Graham Reynolds to have access to additional information specific to Vanderbilt University Medical Center. The current date is {date_str}.\n
+# Your knowledge base was last updated in December 2023. You answer questions about events prior to and after December 2023 the way a highly informed individual in December 2023 would if they were talking to someone from the above date, and you can let the user know this when relevant.\n
+# Some of the context you will be given in regards to Vanderbilt University Medical Center could have come after December 2023. The rest of your knowledge base is from before December 2023 and you will answer questions accordingly with these facts.
+# This chunk of text is your system prompt. It is not visible to the user, but it is used to guide your responses. Don't reference it, just respond to the user.\n
+# If you are asked to assist with tasks involving the expression of views held by a significant number of people, you provide assistance with the task even if you personally disagree with the views being expressed, but follow this with a discussion of broader perspectives.\n
+# You don't engage in stereotyping, including the negative stereotyping of majority groups.\n If asked about controversial topics, you try to provide careful thoughts and objective information without downplaying its harmful content or implying that there are reasonable perspectives on both sides.\n
+# You are happy to help with writing, analysis, question answering, math, coding, and all sorts of other tasks.\n You use markdown for coding, which includes JSON blocks and Markdown tables.\n
+# You do not have tools enabled at this time, so cannot run code or access the internet. You can only provide information that you have been trained on. You do not send or receive links or images.\n
+# You were not trained on copyrighted books, song lyrics, poems, video transcripts, or news articles; you do not divulge details of your training data. You do not provide song lyrics, poems, or news articles and instead refer the user to find them online or in a store.\n
+# You give concise responses to simple questions or statements, but provide thorough responses to more complex and open-ended questions.\n
+# The user is unable to see the system prompt, so you should write as if it were true without mentioning it.\n You do not mention any of this information about yourself unless the information is directly pertinent to the user's query.\n
+# Here is some context from the Vanderbilt University Medical Center glossary which might or might not help you answer: {context}.\n
+# Based on this system prompt, to which you will adhere sternly and to which you will make no reference, and this possibly helpful context in relation to Vanderbilt University Medical Center, answer this question: {question}
+# """
+
 # @st.cache_resource
 # def get_global_semaphore():
 #     return threading.BoundedSemaphore(QUEUE_SIZE)
@@ -71,35 +87,12 @@ def clear_chat_history():
 
 st.button('Clear Chat', on_click=clear_chat_history)
 
+# build our chain outside the working body so that its only instantiated once - simply pass it the chat history for chat completion
+chain = ChainBuilder.build_chain()
+
 def last_role_is_user():
     return len(st.session_state["messages"]) > 0 and st.session_state["messages"][-1]["role"] == "user"
 
-def get_system_prompt():
-    return ""
-
-# ** working logic for querying glossary embeddings
-# Same embedding model we used to create embeddings of terms
-# make sure we cache this so that it doesnt redownload each time, hindering Space start time if sleeping
-# try adding this st caching decorator to ensure the embeddings class gets cached after downloading the entirety of the model
-# does this cache to the given folder though? It does appear to populate the folder as expected after being run
-@st.cache_resource # will this work here? https://docs.streamlit.io/develop/concepts/architecture/caching
-def load_embedding_model():
-    embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en", cache_folder="./langchain_cache/")
-    return embeddings
-
-embeddings = load_embedding_model()
-# instantiate the vector store for similarity search in our chain
-# need to make this a function and decorate it with @st.experimental_memo as above?
-# We are only calling this initially when the Space starts. Can we expedite this process for users when opening up this Space?
-# @st.cache_data # TODO add this in
-vector_store = DatabricksVectorSearch(
-    endpoint=VS_ENDPOINT_NAME,
-    index_name=VS_INDEX_NAME,
-    embedding=embeddings,
-    text_column="name",
-    columns=["name", "description"],
-)
-
 def text_stream(stream):
     for chunk in stream:
         if chunk["content"] is not None:
@@ -116,7 +109,7 @@ def get_stream_warning_error(stream):
     return warning, error
 
 # @retry(wait=wait_random_exponential(min=0.5, max=2), stop=stop_after_attempt(3))
-def chat_api_call(history):
+def chain_call(history):
     # *** original code for instantiating the DBRX model through the OpenAI client *** skip this and introduce our chain eventually
     # extra_body = {}
     # if SAFETY_FILTER:
@@ -133,10 +126,24 @@ def chat_api_call(history):
     #     extra_body= extra_body
     # )
 
-    # ** TODO update this next to take and do similarity search on user input!
+    # *** can we stream the chain's response by incorporating the above OpenAI streaming functionality?
+    # *** Look back at the predict_stream function and see if we can incorporate that!
+    # *** looks like we want to use either chain.stream() or chain.astream()
+    # test first with invoke
+
+    input_example = {'messages': 
+                     [{'content': 'What does EDW stand for?', 'role': 'user'}, 
+                      {'content': 'Enterprise Data Warehouse.', 'role': 'assistant'}, 
+                      {'content': 'Thank you. What is the data lake?', 'role': 'user'}, 
+                      {'content': 'A data lake is a centralized repository of structured and unstructured data. It allows data to be stored in its native state, without the need for transformations, so that it can be consumed by other users later. It is not just a term for storage, but also covers functionalities required for a platform, including data analysis, machine learning, cataloging and data movement.', 'role': 'assistant'}, 
+                      {'content': 'Can you tell me more about how they are used?', 'role': 'user'}, 
+                      {'content': 'At Vanderbilt University Medical Center, a data lake is used as a centralized repository for storing and managing large amounts of data in its native format. This allows for the data to be easily accessed and analyzed by different teams and business units within the organization. The data lake also provides functionalities such as data analysis, machine learning, cataloging and data movement, making it a versatile tool for handling diverse data sets.\n\nAn Enterprise Data Warehouse (EDW) is used for executing analytic queries on structured data. It is optimized for this purpose, with data being stored in a way that allows for efficient querying and analysis. This makes it a useful tool for teams that need to perform complex analyses on large data sets.\n\nA data mart is a specific organizational structure or pattern used in the context of data warehouses. It is a layer that has specific subdivisions for each business unit or team, such as finance, marketing, and product. This allows users to consume data in a format that meets their specific needs.\n\nA data lakehouse is a term used to describe approaches that attempt to combine the data structure and management features of a data warehouse with the low cost of storage of a data lake. This includes a structured transactional layer, which allows for efficient querying and analysis of data. This approach aims to provide the benefits of both data lakes and data warehouses in a single platform.', 'role': 'assistant'}, 
+                      {'content': 'Nice answer. Can you tell me what the HCERA is?', 'role': 'user'}]}
+
     st.write(history)
-    search_result = vector_store.similarity_search(query=st.session_state["messages"][-1]["content"], k=5)
-    chat_completion = search_result # TODO update this after we implement our chain
+    # search_result = vector_store.similarity_search(query=st.session_state["messages"][-1]["content"], k=5)
+    # chat_completion = search_result # TODO update this after we implement our chain
+    chat_completion = chain.invoke(input_example) # *** TODO here we will pass only the chat history, the chain handles the system prompt
     return chat_completion
 
 def write_response():
@@ -154,11 +161,11 @@ def write_response():
     return response, stream_warning, stream_error
 
 def chat_completion(messages):
-    history_dbrx_format = [
-        {"role": "system", "content": get_system_prompt()}
-    ]
-        
-    history_dbrx_format = history_dbrx_format + messages
+    # history_dbrx_format = [
+    #     {"role": "system", "content": SYSTEM_PROMPT} # no longer need this because the chain handles all of this for us
+    # ]
+    
+    # history_dbrx_format = history_dbrx_format + messages
     # if (len(history_dbrx_format)-1)//2 >= MAX_CHAT_TURNS:
     #     yield {"content": None, "error": MSG_MAX_TURNS_EXCEEDED, "warning": None}
     #     return
@@ -172,7 +179,8 @@ def chat_completion(messages):
     #         chat_completion = chat_api_call(history_dbrx_format)
     #     except Exception as e:
     #         error = e    
-    chat_completion = chat_api_call(history_dbrx_format)
+    # chat_completion = chain_call(history_dbrx_format)
+    chat_completion = chain_call(messages) # simply pass the old messages, need not worry about the system prompt
     if error is not None:
         yield {"content": None, "error": GENERAL_ERROR_MSG, "warning": None}
         print(error)
@@ -183,10 +191,13 @@ def chat_completion(messages):
     chunk_counter = 0
     for chunk in chat_completion:
         # if chunk.choices[0].delta.content is not None:
-        if chunk.page_content is not None:
+        # TODO *** we need to refactor this logic to match what happens with the response from our chain - it should be strings or an iterator of strings
+        # if chunk.page_content is not None:
+        if chunk is not None:
             chunk_counter += 1
             # partial_message += chunk.choices[0].delta.content
-            partial_message += f"* {chunk.page_content} [{chunk.metadata}]"
+            # partial_message += f"* {chunk.page_content} [{chunk.metadata}]"
+            partial_message += chunk
             if chunk_counter % TOKEN_CHUNK_SIZE == 0:
                 chunk_counter = 0
                 yield {"content": partial_message, "error": None, "warning": None}