main.py

import os
from src.helpers.inference_helper_functions_v2 import InferenceClient
import streamlit as st
from dotenv import load_dotenv

st.set_page_config(
    page_title="Jimmy", page_icon="😒", initial_sidebar_state="collapsed", layout="wide"
)

load_dotenv()


# Password protection
def check_password():
    def password_entered():
        if st.session_state["password"] == os.getenv("APP_PASSWORD"):
            st.session_state["password_correct"] = True
            del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if not st.session_state.get("password_correct", False):
        st.markdown("\n\n")
        st.text_input(
            "Enter the password",
            type="password",
            on_change=password_entered,
            key="password",
        )
        st.divider()
        st.info("Developed by Milan Mrdenovic - IBM Norway 2026")
        if "password_correct" in st.session_state:
            st.error("😕 Password incorrect")
        return False
    return True


if not check_password():
    st.stop()


@st.cache_resource(show_spinner=False)
def get_client() -> InferenceClient:
    """Build the inference client exactly once.

    @st.cache_resource caches the returned object globally (shared across reruns
    and sessions); since this takes no arguments the cache key is constant, so
    InferenceClient is instantiated a single time for the lifetime of the app.
    """
    region = os.getenv("INF_REGION", "us-east")
    project_id = os.getenv("INF_PROJECT_ID", "")
    # INF_ENDPOINT is a template with {INF_REGION}/{INF_PROJECT_ID} placeholders.
    endpoint = os.getenv("INF_ENDPOINT", "").format(
        INF_REGION=region, INF_PROJECT_ID=project_id
    )
    return InferenceClient(
        provider=os.getenv("INF_PROVIDER", "rhai"),
        api_key=os.getenv("INF_API_KEY", ""),
        url=endpoint,
        project=project_id,
        region=region,
        model_id=os.getenv("INF_DEFAULT_MODEL", ""),
        params={
            "max_completion_tokens": int(os.getenv("INF_MAX_TOKENS", "450")),
            "temperature": float(os.getenv("INF_TEMPERATURE", "0.7")),
            "top_p": float(os.getenv("INF_TOP_P", "1.0")),
        },
    )


def initialize_session_state():
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = [
            {"role": "system", "content": os.getenv("PROMPT_TEMPLATE", "")}
        ]


def show_chat_turns() -> bool:
    """Whether to replay prior turns on screen (SHOW_CHAT_TURNS, default False)."""
    return os.getenv("SHOW_CHAT_TURNS", "False").strip().lower() == "true"


def apply_message_font_scale(scale: float = 1.5) -> None:
    """Enlarge the text inside chat-message bubbles by ``scale`` (1.5 = +50%)."""
    st.markdown(
        f"""
        <style>
        [data-testid="stChatMessage"] [data-testid="stMarkdownContainer"] {{
            font-size: {scale}em;
            line-height: 1.4;
        }}
        </style>
        """,
        unsafe_allow_html=True,
    )


def chat_interface():
    st.subheader("Jimmy")

    apply_message_font_scale(float(os.getenv("MESSAGE_FONT_SCALE", "1.35")))

    replay = show_chat_turns()

    # Replay the visible conversation (skip the system prompt) only when enabled.
    if replay:
        for message in st.session_state.chat_history[1:]:
            avatar = "😒" if message["role"] == "assistant" else None
            with st.chat_message(message["role"], avatar=avatar):
                st.markdown(message["content"])

    user_input = st.chat_input("You:", key="user_input")
    if not user_input:
        return

    st.session_state.chat_history.append({"role": "user", "content": user_input})
    if replay:
        with st.chat_message("user"):
            st.markdown(user_input)

    with st.chat_message("assistant", avatar="😒"):
        stream = get_client().stream_chat_inference(
            messages=st.session_state.chat_history
        )
        content = st.write_stream(stream)
        if not content:
            content = "😒 not feeling it right now."
            st.markdown(content)

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


def main():
    initialize_session_state()
    chat_interface()


if __name__ == "__main__":
    main()