Update app.py

app.py

@@ -1,106 +1,108 @@
-import os 
-import base64
-from io import BytesIO
-from PIL import Image
-import streamlit as st
-from app_config import SYSTEM_PROMPT,MODEL,MAX_TOKENS,TRANSFORMER_MODEL
-from langchain.memory import ConversationSummaryBufferMemory
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_groq import ChatGroq
-from streamlit_pdf_viewer import pdf_viewer
-from pydantic import BaseModel
-from langchain.chains import LLMChain
-from langchain.prompts import ChatPromptTemplate
-from langchain_community.vectorstores import FAISS
-from sentence_transformers import SentenceTransformer 
-from typing import Any
-
-st.title("Hitachi Support Bot")
-
-class Element(BaseModel):
-    type: str
-    text: Any
-
-# llm = ChatGoogleGenerativeAI(
-#     model=MODEL,
-#     max_tokens=MAX_TOKENS
-# )
-llm = ChatGroq(model=MODEL,api_key='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')
-
-
-prompt = ChatPromptTemplate.from_template(SYSTEM_PROMPT)
-qa_chain = LLMChain(llm=llm,prompt=prompt)
-embeddings = HuggingFaceEmbeddings(model_name=TRANSFORMER_MODEL)
-db = FAISS.load_local("faiss_index",embeddings,allow_dangerous_deserialization=True)
-
-st.markdown(
-    """
-<style>
-    .st-emotion-cache-janbn0 {
-        flex-direction: row-reverse;
-        text-align: right;
-    }
-</style>
-""",
-    unsafe_allow_html=True,
-)
-
-def response_generator(question):
-    relevant_docs = db.similarity_search_with_relevance_scores(question,k=5)
-    context = ""
-    relevant_images = []
-    for d,score in relevant_docs:
-        if score > 0:
-            if d.metadata['type'] == 'text':
-                context += str(d.metadata['original_content'])
-            elif d.metadata['type'] == 'table':
-                context += str(d.metadata['original_content'])
-            elif d.metadata['type'] == 'image':
-                context += d.page_content
-                relevant_images.append(d.metadata['original_content'])
-    result =  qa_chain.run({'context':context,"question":question})
-    return result,relevant_images
-
-with st.sidebar:
-    st.header("Hitachi Support Bot")
-    button = st.toggle("View Doc file.")
-
-if button:
-    pdf_viewer("GPT OUTPUT.pdf")
-else:
-    if "messages" not in st.session_state:
-        st.session_state.messages=[{"role": "system", "content": SYSTEM_PROMPT}]
-        
-    if "llm" not in st.session_state:
-        st.session_state.llm = llm
-    if "rag_memory" not in st.session_state:
-        st.session_state.rag_memory = ConversationSummaryBufferMemory(llm=st.session_state.llm, max_token_limit= 5000)
-
-    container =  st.container(height=700)
-    for message in st.session_state.messages:
-        if message["role"] != "system":
-            if message["role"] == "user":
-                with container.chat_message(message["role"]):
-                    st.write(message["content"])
-            if message["role"] == "assistant":
-                with container.chat_message(message["role"]):
-                    st.write(message["content"])
-                    for i in range(len(message["images"])):
-                        st.image(Image.open(BytesIO(base64.b64decode(message["images"][i].encode('utf-8')))))
-
-    if prompt := st.chat_input("Enter your query here... "):
-        with container.chat_message("user"):
-            st.write(prompt)
-        st.session_state.messages.append({"role":"user" , "content":prompt})
-        with container.chat_message("assistant"):  
-            response,images = response_generator(prompt)
-            st.write(response)
-            for i in range(len(images)):
-                st.markdown("""---""")
-                st.image(Image.open(BytesIO(base64.b64decode(images[i].encode('utf-8')))))
-            st.markdown("""---""")
-
-                
-        st.session_state.rag_memory.save_context({'input': prompt}, {'output': response})
+import os 
+import base64
+from io import BytesIO
+from PIL import Image
+import streamlit as st
+from app_config import SYSTEM_PROMPT,MODEL,MAX_TOKENS,TRANSFORMER_MODEL
+from langchain.memory import ConversationSummaryBufferMemory
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from streamlit_pdf_viewer import pdf_viewer
+from pydantic import BaseModel
+from langchain.chains import LLMChain
+from langchain.prompts import ChatPromptTemplate
+from langchain_community.vectorstores import FAISS
+from sentence_transformers import SentenceTransformer 
+from typing import Any
+
+st.title("Hitachi Support Bot")
+
+class Element(BaseModel):
+    type: str
+    text: Any
+
+# llm = ChatGoogleGenerativeAI(
+#     model=MODEL,
+#     max_tokens=MAX_TOKENS
+# )
+# llm = ChatGroq(model=MODEL,api_key='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')
+llm = ChatGroq(model=MODEL,api_key=os.getenv('API_KEY'))
+
+
+
+prompt = ChatPromptTemplate.from_template(SYSTEM_PROMPT)
+qa_chain = LLMChain(llm=llm,prompt=prompt)
+embeddings = HuggingFaceEmbeddings(model_name=TRANSFORMER_MODEL)
+db = FAISS.load_local("faiss_index",embeddings,allow_dangerous_deserialization=True)
+
+st.markdown(
+    """
+<style>
+    .st-emotion-cache-janbn0 {
+        flex-direction: row-reverse;
+        text-align: right;
+    }
+</style>
+""",
+    unsafe_allow_html=True,
+)
+
+def response_generator(question):
+    relevant_docs = db.similarity_search_with_relevance_scores(question,k=5)
+    context = ""
+    relevant_images = []
+    for d,score in relevant_docs:
+        if score > 0:
+            if d.metadata['type'] == 'text':
+                context += str(d.metadata['original_content'])
+            elif d.metadata['type'] == 'table':
+                context += str(d.metadata['original_content'])
+            elif d.metadata['type'] == 'image':
+                context += d.page_content
+                relevant_images.append(d.metadata['original_content'])
+    result =  qa_chain.run({'context':context,"question":question})
+    return result,relevant_images
+
+with st.sidebar:
+    st.header("Hitachi Support Bot")
+    button = st.toggle("View Doc file.")
+
+if button:
+    pdf_viewer("GPT OUTPUT.pdf")
+else:
+    if "messages" not in st.session_state:
+        st.session_state.messages=[{"role": "system", "content": SYSTEM_PROMPT}]
+        
+    if "llm" not in st.session_state:
+        st.session_state.llm = llm
+    if "rag_memory" not in st.session_state:
+        st.session_state.rag_memory = ConversationSummaryBufferMemory(llm=st.session_state.llm, max_token_limit= 5000)
+
+    container =  st.container(height=700)
+    for message in st.session_state.messages:
+        if message["role"] != "system":
+            if message["role"] == "user":
+                with container.chat_message(message["role"]):
+                    st.write(message["content"])
+            if message["role"] == "assistant":
+                with container.chat_message(message["role"]):
+                    st.write(message["content"])
+                    for i in range(len(message["images"])):
+                        st.image(Image.open(BytesIO(base64.b64decode(message["images"][i].encode('utf-8')))))
+
+    if prompt := st.chat_input("Enter your query here... "):
+        with container.chat_message("user"):
+            st.write(prompt)
+        st.session_state.messages.append({"role":"user" , "content":prompt})
+        with container.chat_message("assistant"):  
+            response,images = response_generator(prompt)
+            st.write(response)
+            for i in range(len(images)):
+                st.markdown("""---""")
+                st.image(Image.open(BytesIO(base64.b64decode(images[i].encode('utf-8')))))
+            st.markdown("""---""")
+
+                
+        st.session_state.rag_memory.save_context({'input': prompt}, {'output': response})
         st.session_state.messages.append({"role":"assistant" , "content":response,'images':images})