Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,136 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from bs4 import BeautifulSoup
+from openai import OpenAI
+from dotenv import load_dotenv
+import os
 
-"""
-# Welcome to Streamlit!
-
-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).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-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])),
-    ))
+# ==== Load API Key from .env ====
+load_dotenv()
+hf_api_key = os.getenv("HUGGINGFACE_API_KEY")
+
+if not hf_api_key:
+    st.error("❌ API key not found. Please set HUGGINGFACE_API_KEY in your .env file.")
+    st.stop()
+
+# ==== Initialize Client ====
+client = OpenAI(
+    base_url="https://router.huggingface.co/v1",
+    api_key=hf_api_key,
+)
+
+# ==== Streamlit Page Setup ====
+st.set_page_config(page_title="OSS ChatGPT", layout="wide")
+st.title("🤖 ChatGPT")
+
+# ==== Sidebar ====
+st.sidebar.title("🛠️ Settings")
+model_choice = st.sidebar.selectbox("Choose a model", [
+    "openai/gpt-oss-20b:hyperbolic",
+    "openai/gpt-oss-120b:hyperbolic"
+])
+
+
+
+if st.sidebar.button("🧹 Clear Chat"):
+    st.session_state.messages = []
+    st.rerun()
+
+# ==== Session Initialization ====
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# ==== LaTeX Rendering Helper ====
+def render_markdown_with_latex(text: str):
+    mathjax_script = """
+    <script type="text/javascript">
+      MathJax = {
+        tex: {
+          inlineMath: [['$', '$']],
+          displayMath: [['$$', '$$']]
+        },
+        svg: {
+          fontCache: 'global'
+        }
+      };
+    </script>
+    <script async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-svg.js"></script>
+    """
+    st.markdown(mathjax_script, unsafe_allow_html=True)
+    st.markdown(text, unsafe_allow_html=True)
+# ... (keep all your existing imports and setup code until the message display section)
+# ... (keep all your existing imports and setup code until the message display section)
+
+# ==== Display Chat History ====
+for i, msg in enumerate(st.session_state.messages):
+    with st.chat_message(msg["role"]):
+        if msg["role"] == "assistant":
+            # Check if content contains thinking pattern
+            if msg["content"].startswith("<think>") and "</think>" in msg["content"]:
+                # Extract thinking and response parts
+                thinking_part = msg["content"].split("<think>")[1].split("</think>")[0].strip()
+                response_part = msg["content"].split("</think>")[1].strip()
+                
+                # Display the main response
+                render_markdown_with_latex(response_part)
+                
+                # Create expander for thinking
+                with st.expander("💭 Show Thinking"):
+                    st.markdown(f"<span style='color: #666; font-style: italic'>{thinking_part}</span>", 
+                               unsafe_allow_html=True)
+            else:
+                render_markdown_with_latex(msg["content"])
+        else:
+            # Display user messages
+            render_markdown_with_latex(msg["content"])
+
+# ==== Chat Input ====
+st.write("Akshay")
+if prompt := st.chat_input("Type your message...",key="unique_chat_input_key"):
+    print(f"=> {prompt}")
+    with st.chat_message("user"):
+        st.markdown(prompt)
+    
+    # Add user message to history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+
+    try:
+        with st.chat_message("assistant"):
+            response_placeholder = st.empty()
+
+            with st.spinner("Thinking..."):
+                completion = client.chat.completions.create(
+                    model=model_choice,
+                    messages=[
+                        {"role": m["role"], "content": m["content"]}
+                        for m in st.session_state.messages
+                    ]
+                )
+                raw = completion.choices[0].message.content
+
+                # Parse thinking and response
+                if "<think>" in raw and "</think>" in raw:
+                    thinking_part = raw.split("<think>")[1].split("</think>")[0].strip()
+                    response_part = raw.split("</think>")[1].strip()
+
+                    # Display main response
+                    clean_response = BeautifulSoup(response_part, "html.parser").get_text()
+                    with st.expander("💭 Show Thinking"):
+                        st.markdown(f"<span style='color: #666; font-style: italic'>{thinking_part}</span>", 
+                               unsafe_allow_html=True)
+                    response_placeholder.markdown(clean_response)
+                    # Store both parts in history
+                    full_content = f"<think>{thinking_part}</think>{clean_response}"
+                    
+                else:
+                    clean_response = BeautifulSoup(raw, "html.parser").get_text()
+                    response_placeholder.markdown(clean_response)
+                    full_content = clean_response
+
+            # Add assistant response to history
+            st.session_state.messages.append({"role": "assistant", "content": full_content})
+
+    except Exception as e:
+        error_msg = f"❌ Error: {str(e)}"
+        with st.chat_message("assistant"):
+            st.error(error_msg)
+        st.session_state.messages.append({"role": "assistant", "content": error_msg})