app.py

import streamlit as st
import os
import PyPDF2
from io import BytesIO
from openai import OpenAI
from huggingface_hub import InferenceClient
from dotenv import load_dotenv

load_dotenv()

# ---------------------
# Utility Functions
# ---------------------


def authenticate():
    """
    A simple authentication mechanism using a password stored in an environment variable (APP_PASSWORD).

    Returns True if the user is authenticated, otherwise stops the Streamlit execution.
    """
    app_password = os.getenv("APP_PASSWORD", None)
    if not app_password:
        st.warning("No password set for the app. Please set the 'APP_PASSWORD' environment variable.")
        return True  # Or return False if you want to block access

    if "authenticated" not in st.session_state:
        st.session_state["authenticated"] = False

    if not st.session_state["authenticated"]:
        st.text_input("Enter your access code:", type="password", key="login_password")
        if st.button("Login"):
            if st.session_state["login_password"] == app_password:
                st.session_state["authenticated"] = True
                st.rerun()
            else:
                st.error("Invalid password. Please try again.")
                st.stop()

    return st.session_state["authenticated"]


def read_pdf(file):
    """
    Reads a PDF file using PyPDF2 and returns the extracted text.
    """
    pdf_reader = PyPDF2.PdfReader(file)
    text = []
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        text.append(page.extract_text())
    return "\n".join(text)

def call_gpt_4o_api(
    messages,
    model,
    temperature,
    max_tokens,
    stream
):
    """
    Calls GPT-4o-compatible API (via OpenAI-like client).
    Expects a list of messages (with "role" and "content" keys),
    including the system message(s) as the first item(s) and
    user/assistant messages subsequently.

    Yields partial (streaming) or complete text.
    """
    client = OpenAI(api_key=os.getenv("OPENAI_API_KEY", ""))

    # remove the second element from messages 
    # (likely the "additional PDF context" system message, or your second system message)
    messages = [messages[0]] + messages[2:]

    if stream:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=True
        )
        partial_text = ""
        for chunk in response:
            delta = chunk.choices[0].delta
            if hasattr(delta, "content") and delta.content:
                partial_text += delta.content
                yield partial_text
    else:
        response = client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            stream=False
        )
        complete_text = response.choices[0].message.content
        yield complete_text

def call_hf_inference(
    messages,
    model_repo,
    temperature=0.7,
    max_tokens=200,
    stream=False
):
    """
    Calls a Hugging Face open-source LLM via the InferenceClient's chat endpoint.
    Expects a list of messages (with "role" and "content"), including
    system and user/assistant roles.

    Yields partial (streaming) or complete text.
    """
    HF_TOKEN = os.getenv("HF_TOKEN", None)
    if not HF_TOKEN:
        raise ValueError("Please set your HF_TOKEN environment variable.")

    client = InferenceClient(api_key=HF_TOKEN)

    # remove the second element from messages
    messages = [messages[0]] + messages[2:]

    response = client.chat.completions.create(
        model=model_repo,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        stream=stream
    )

    if stream:
        partial_text = ""
        for chunk in response:
            delta = chunk.choices[0].delta
            if isinstance(delta, dict):
                chunk_content = delta.get("content", "")
                partial_text += chunk_content
                yield partial_text
    else:
        complete_text = response.choices[0].message["content"]
        yield complete_text

# ---------------------
# Streamlit App
# ---------------------

def main():
    st.set_page_config(page_title="CVI-GPT", layout="centered")
    st.title("CVI-GPT: Conversational Interface")
    
    if not authenticate():
        st.stop()  # or just `return` to end early

    # ---------------------
    # Sidebar: Model & Params
    # ---------------------
    st.sidebar.header("Model & Parameters")
    
    # Model selection
    model_choice = st.sidebar.selectbox(
        "Select Model",
        [
            "meta-llama/Llama-3.3-70B-Instruct",
            "gpt-4o-mini",
            "gpt-4o",
            "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
        ]
    )

    # Temperature & max_tokens
    temperature = st.sidebar.slider("Temperature", 0.0, 1.5, 0.7, 0.1)
    max_tokens = st.sidebar.slider("Max Tokens", 50, 2000, 500, 50)

    # We store the selected model in session_state to detect changes
    if "selected_model" not in st.session_state:
        st.session_state.selected_model = model_choice

    # If the user changes the model, reset the conversation
    if model_choice != st.session_state.selected_model:
        st.session_state.selected_model = model_choice
        st.session_state["messages"] = [
            {"role": "assistant", "content": f"Model changed to `{model_choice}`. How can I help you?"}
        ]

    # System / Instruction Message
    base_instructions = st.sidebar.text_area(
        "System / Instruction Message",
        value=(
            "You are a Helpful Assistant. Respond in a concise, helpful, and markdown-friendly format.\n\n"
            "Formatting Instructions:\n"
            "- Responses should be in markdown.\n"
            "- Use headings, bullet points, bold, italics, etc. for clarity.\n"
            "- Use triple backticks for code blocks.\n"
            "- Provide references or disclaimers when needed."
        ),
        height=200
    )

    # Clear Chat Button
    if st.sidebar.button("Clear Chat"):
        st.session_state["messages"] = [
            {"role": "assistant", "content": "Chat cleared. How can I help you now?"}
        ]

    # ---------------------
    # PDF Upload
    # ---------------------
    st.sidebar.header("Optional: PDF Upload")
    uploaded_file = st.sidebar.file_uploader("Upload a PDF", type=["pdf"])
    pdf_text = ""
    if uploaded_file is not None:
        pdf_text = read_pdf(uploaded_file)
        # We do NOT print the PDF content. Just let user know it's loaded.
        st.sidebar.write("PDF content loaded (not displayed).")
        
    st.sidebar.divider()
    with st.sidebar:
        st.subheader("👨‍💻 Author: *Adrish Maity*", anchor=False)

    # ---------------------
    # Initialize conversation if not present
    # ---------------------
    if "messages" not in st.session_state:
        st.session_state["messages"] = [
            {"role": "assistant", "content": "Hello! How can I help you today?"}
        ]

    # ---------------------
    # Display Conversation
    # ---------------------
    for msg in st.session_state["messages"]:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    # ---------------------
    # Chat Input
    # ---------------------
    if user_input := st.chat_input("Type your question..."):
        # Just store the user's typed text
        user_text = user_input

        st.session_state["messages"].append({"role": "user", "content": user_text})

        # Display user's message
        with st.chat_message("user"):
            st.markdown(user_text)

        # Now build the full conversation:
        # 1) A system message (instructions)
        # 2) A second system message with PDF context if present (kept hidden from UI)
        # 3) All prior conversation
        full_conversation = [{"role": "system", "content": base_instructions}]

        if pdf_text:
            full_conversation[0]["content"] += "\n\n" + "Additional PDF context (user provided):\n" + pdf_text

        full_conversation.extend(st.session_state["messages"])

        # Placeholder for assistant's streaming response
        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            streamed_text = ""

            # Decide how to call the model
            if model_choice in ["gpt-4o", "gpt-4o-mini"]:
                stream_response = call_gpt_4o_api(
                    messages=full_conversation,
                    model=model_choice,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    stream=True
                )
                for partial_output in stream_response:
                    streamed_text = partial_output
                    response_placeholder.markdown(streamed_text)
            else:
                hf_stream = call_hf_inference(
                    messages=full_conversation,
                    model_repo=model_choice,
                    temperature=temperature,
                    max_tokens=max_tokens,
                    stream=True
                )
                for partial_output in hf_stream:
                    streamed_text = partial_output
                    response_placeholder.markdown(streamed_text)

            # Once done, store the final assistant message
            st.session_state["messages"].append({"role": "assistant", "content": streamed_text})

if __name__ == "__main__":
    main()