app.py

import streamlit as st
from together import Together
import os
import re  
import io
import fitz  # PyMuPDF for PDF text extraction

API_KEY = os.getenv("TOGETHER_API_KEY")
if not API_KEY:
    st.error("API key is missing! Set TOGETHER_API_KEY in your environment variables.")
    st.stop()

@st.cache_resource
def get_client():
    return Together(api_key=API_KEY)

if "messages" not in st.session_state:
    st.session_state.messages = []

st.title("DeepSeek R1 Chatbot")

# Function to extract text from a PDF file
def extract_pdf_text(uploaded_file):
    # Convert the uploaded file into a byte stream
    pdf_bytes = uploaded_file.read()
    # Open the byte stream as a PDF document
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")  # Use stream and specify filetype as 'pdf'
    text = ""
    for page in doc:
        text += page.get_text()
    return text

# File uploader for PDFs
uploaded_file = st.file_uploader("Upload a PDF", type=["pdf"])

if uploaded_file:
    pdf_text = extract_pdf_text(uploaded_file)
    st.session_state.messages.append({"role": "user", "content": pdf_text})
    with st.chat_message("user"):
        st.markdown(pdf_text)

def display_response(content):
    latex_pattern = r"\$\$([\s\S]*?)\$\$|\$(.*?)\$"
    matches = list(re.finditer(latex_pattern, content))

    last_end = 0
    formatted_response = ""  

    for match in matches:
        text_part = content[last_end : match.start()]
        formatted_response += text_part  
        st.markdown(text_part)  

        latex_expr = match.group(1) or match.group(2)
        if latex_expr:
            formatted_response += f"{latex_expr}\n"  
            st.latex(latex_expr)  # Renders LaTeX as mathematical notation

        last_end = match.end()

    remaining_text = content[last_end:]
    formatted_response += remaining_text  
    st.markdown(remaining_text)  

    # Streamlit Copy Button (copies pure math notation)
    st.code(formatted_response, language="text")  

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        if message["role"] == "assistant":
            display_response(message["content"])
        else:
            st.markdown(message["content"])

if prompt := st.chat_input("Ask something..."):
    client = get_client()
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        response_placeholder = st.empty()

        try:
            system_prompt = {"role": "system", "content": "Format all mathematical expressions in proper notation."}
            response = client.chat.completions.create(
                model="deepseek-ai/DeepSeek-R1",
                messages=[system_prompt] + st.session_state.messages,  
                max_tokens=4096,
                temperature=0.7,
                top_p=0.9,
                stream=False
            )

            full_response = response.choices[0].message.content  
            response_placeholder.empty()
            display_response(full_response)  

            st.session_state.messages.append({"role": "assistant", "content": full_response})

        except Exception as e:
            st.error(f"Error: {str(e)}")