Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,158 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+import zipfile
 
-"""
-# Welcome to Streamlit!
+# from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import Chroma
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
+from langchain.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain.schema.output_parser import StrOutputParser
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+# --- Streamlit Setup ---
+st.set_page_config(page_title="📊 ITC Financial Analyst AI", layout="wide")
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+# Custom CSS for better UI
+st.markdown("""
+    <style>
+        .main { background-color: #f9f9f9; }
+        .block-container {
+            padding-top: 2rem;
+            padding-bottom: 2rem;
+        }
+        .stChatMessage {
+            background-color: #ffffff;
+            border: 1px solid #e0e0e0;
+            padding: 1rem;
+            border-radius: 12px;
+            margin-bottom: 1rem;
+        }
+        .stButton button {
+            background-color: #FF6347 !important;
+            color: white !important;
+            border-radius: 8px !important;
+            font-weight: 600;
+        }
+        .source-box {
+            background-color: #f0f0f0;
+            border-left: 5px solid #555;
+            padding: 0.5rem;
+            margin-top: 0.5rem;
+            border-radius: 8px;
+            font-size: 0.9rem;
+        }
+    </style>
+""", unsafe_allow_html=True)
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+st.title("📊 ITC Financial Analysis with AI-Powered Insights")
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
+# Chat history buffer
+memory_buffer = {"chat_history": []}
 
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
+# Clear history
+st.sidebar.markdown("## 🛠️ Options")
+if st.sidebar.button("🔁 End Chat"):
+    memory_buffer["chat_history"] = []
+
+# Load Chroma vector DB from ZIP
+with zipfile.ZipFile('chroma_db1.zip', 'r') as zip_ref:
+    zip_ref.extractall('chroma_db')
+
+# Vector embeddings
+embedding = HuggingFaceEmbeddings(model_name='all-MiniLM-L6-v2')
+vectorstore = Chroma(persist_directory='chroma_db', embedding_function=embedding)
+mmr_retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3, "lambda_mult": 1})
+
+# Helper functions
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+def get_docs_and_context(question):
+    docs = mmr_retriever.get_relevant_documents(question)
+    return {"question": question, "docs": docs, "context": format_docs(docs)}
+
+# Prompt setup
+parallel_chain = RunnableLambda(lambda x: {
+    "question": x["input"],
+    **get_docs_and_context(x["input"])
 })
 
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+chat_prompt = ChatPromptTemplate.from_messages([
+    ("system",
+     """
+     You are a domain-specific AI financial analyst focused on company-level performance evaluation.
+
+     Your task is to analyze and respond to user financial queries *strictly based on the provided transcript data*: {context}.
+
+     Rules:
+     1. ONLY extract facts, figures, and insights that are explicitly available in the transcript.
+     2. If data is *missing or partially available*, clearly state: "The required data is not available in the current transcript." Then provide a generic but relevant explanation based on standard financial principles.
+     3. Maintain numerical accuracy and avoid interpretation beyond data boundaries.
+     4. Prioritize answers relevant to *ITC Ltd.*, but keep response format adaptable to other firms and fiscal years.
+     5. Clearly present year-wise or metric-wise insights using bullet points or structured formats if applicable.
+
+     Your goals:
+     - Ensure 100% fidelity to source transcript.
+     - Do not assume or hallucinate missing numbers.
+     - Use clear, reproducible reasoning steps (e.g., show which line items support your conclusion).
+     - Output should be modular enough to scale across other companies and time periods.
+
+     Respond only to this question from the user.
+     """),
+    MessagesPlaceholder(variable_name="chat_history", optional=True),
+    ("human", "{input}")
+])
+
+GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]
+llm = ChatGoogleGenerativeAI(api_key=GOOGLE_API_KEY, model="gemini-2.0-flash-exp", temperature=1)
+parser = StrOutputParser()
+
+def get_history_from_buffer(_):
+    return memory_buffer['chat_history']
+
+runnable_get_history_from_buffer = RunnableLambda(get_history_from_buffer)
+
+main_chain = (
+    parallel_chain |
+    RunnableLambda(lambda x: {
+        "llm_input": {"input": x["question"], "context": x["context"]},
+        "docs": x["docs"]
+    }) |
+    RunnableLambda(lambda x: {
+        "result": (chat_prompt | llm | parser).invoke(x["llm_input"]),
+        "source_documents": x["docs"]
+    })
+)
+
+chain = RunnablePassthrough.assign(chat_history=runnable_get_history_from_buffer) | main_chain
+
+# Chat history UI
+st.markdown("### 💬 Conversation")
+for msg in memory_buffer["chat_history"]:
+    role = "user" if isinstance(msg, HumanMessage) else "assistant"
+    with st.chat_message(role):
+        st.markdown(msg.content)
+
+# Chat input
+user_input = st.chat_input("Ask about ITC’s performance or any financial metric...")
+
+if user_input:
+    with st.chat_message("user"):
+        st.markdown(user_input)
+
+    memory_buffer["chat_history"].append(HumanMessage(content=user_input))
+    output = chain.invoke({"input": user_input})
+    ai_response = output["result"]
+
+    memory_buffer["chat_history"].append(AIMessage(content=ai_response))
+
+    with st.chat_message("assistant"):
+        st.markdown(ai_response)
+
+        # Show source documents
+        if output.get("source_documents"):
+            st.markdown("**Sources:**")
+            for doc in output["source_documents"]:
+                source = doc.metadata.get("source", "Unknown document")
+                st.markdown(f"<div class='source-box'>📄 {source}</div>", unsafe_allow_html=True)