Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,26 +1,21 @@
 import streamlit as st
-import os
-import zipfile
-from langchain_chroma import Chroma  # ✅ Updated import
+from langchain_chroma import Chroma
 from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
 from langchain.prompts import ChatPromptTemplate
 from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnableLambda
-import tempfile
 
 # === Page Setup ===
 st.set_page_config(page_title="Financial QA - ITC Ltd.", layout="wide")
 st.title("📊 Financial Q&A Chatbot (ITC Ltd.)")
 
-# === Step 1: Extract Chroma DB from zip ===
-def load_chroma_db():
-    with zipfile.ZipFile("chroma_db1.zip", 'r') as zip_ref:
-        temp_dir = tempfile.mkdtemp()
-        zip_ref.extractall(temp_dir)
-    embedding = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
-    return Chroma(persist_directory=temp_dir, embedding_function=embedding)
+# === Step 1: Load Chroma DB ===
+embedding = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
 
-vectorstore = load_chroma_db()
+vectorstore = Chroma(
+    persist_directory="src/chroma_db1",  # ✅ Use unzipped path
+    embedding_function=embedding
+)
 
 # === Step 2: MMR Retriever ===
 retriever = vectorstore.as_retriever(
@@ -48,11 +43,13 @@ prompt = ChatPromptTemplate.from_messages([
 
 # === Step 4: LLM Setup ===
 GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
+
 llm = ChatGoogleGenerativeAI(
     api_key=GOOGLE_API_KEY,
     model="gemini-2.0-flash",
     temperature=1
 )
+
 parser = StrOutputParser()
 
 # === Step 5: Helper Functions ===
@@ -60,7 +57,7 @@ def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
 def retrieve_and_answer(question):
-    docs = retriever.invoke(question)  # ✅ Updated to new `invoke()` method
+    docs = retriever.invoke(question)
     context = format_docs(docs)
     final_input = {"question": question, "context": context}
     result = (prompt | llm | parser).invoke(final_input)