chatbot.py

import streamlit as st
import re
from typing import List, Dict

# --------------------------- #
# OpenAI / CloudGPT client init
# --------------------------- #
try:
    from src.api.cloudgpt_aoai import get_openai_client
    openai_client = get_openai_client()
except (ImportError, AttributeError, ModuleNotFoundError) as err:
    st.sidebar.warning(f"CloudGPT client unavailable → falling back ({err})")
    import openai, os

    openai_client = openai.AzureOpenAI(
        api_key=os.getenv("AZURE_OPENAI_API_KEY"),
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        api_version=os.getenv("AZURE_OPENAI_VERSION", "2025-02-01-preview"),
    )

# --------------------------- #
# Streamlit setup             #
# --------------------------- #
st.set_page_config(layout="wide", initial_sidebar_state="collapsed")
st.title("System B")

# --------------------------- #
# Model list                  #
# --------------------------- #
cloudgpt_available_models = [
    "gpt-35-turbo-20220309", "gpt-35-turbo-16k-20230613", "gpt-35-turbo-20230613",
    "gpt-35-turbo-1106", "gpt-4-20230321", "gpt-4-20230613", "gpt-4-32k-20230321",
    "gpt-4-32k-20230613", "gpt-4-1106-preview", "gpt-4-0125-preview", "gpt-4-visual-preview",
    "gpt-4-turbo-20240409", "gpt-4o-20240513", "gpt-4o-20240806", "gpt-4o-20241120",
    "gpt-4o-mini-20240718", "gpt-4o-audio-preview-20241217", "o1-mini-20240912",
    "o1-20241217", "o3-mini-20250131", "deepseek-r1-preview",
]

# ---- Hide sidebar & fixed settings ---- #
# Completely hide the Streamlit sidebar via CSS
st.markdown(
    """<style>[data-testid="stSidebar"]{display:none;}</style>""",
    unsafe_allow_html=True,
)

# Fixed defaults (no user controls needed)
MODEL = "deepseek-r1-preview"
ENABLE_REASONING = True
MAX_HISTORY = 10
CONTEXT_SIZE = 8192

# --------------------------- #
# Helper functions            #
# --------------------------- #

def trim_history(hist: List[Dict[str, str]], max_pairs: int):
    while len(hist) > max_pairs * 2:
        hist.pop(0)
        if hist:
            hist.pop(0)

def build_messages(hist: List[Dict[str, str]], user_prompt: str):
    msgs = [{"role": m["role"], "content": m["content"]} for m in hist[-max(0, MAX_HISTORY * 2):]]
    msgs.append({"role": "user", "content": user_prompt})
    return msgs

# --------------------------- #
# Session state               #
# --------------------------- #
if "history" not in st.session_state:
    st.session_state.history = []

# ---- Render past ---- #
for m in st.session_state.history:
    with st.chat_message(m["role"]):
        st.markdown(m["content"])

# --------------------------- #
# Utility to render coloured boxes
# --------------------------- #

def _styled_box(content: str, color: str, header: str | None = None) -> str:
    header_html = f"<strong>{header}</strong><br>" if header else ""
    return (
        f"<div style='background-color:{color};border-radius:6px;padding:10px;'>"
        f"{header_html}{content}</div>"
    )

# --------------------------- #
# Streaming helpers           #
# --------------------------- #

def stream_with_reasoning(msgs: List[Dict[str, str]]) -> str:
    """Stream deepseek‑r1 live with history context and coloured boxes."""
    think_box = st.empty()
    resp_box = st.empty()
    collected_think = ""
    collected_resp = ""
    in_think = False  # becomes True once <think> appears

    stream = openai_client.chat.completions.create(
        model=MODEL,
        messages=msgs,
        temperature=0.6,
        max_tokens=CONTEXT_SIZE,
        stream=True,
    )

    for chunk in stream:
        if not chunk.choices or not chunk.choices[0].delta:
            continue
        delta = chunk.choices[0].delta.content or ""

        # Switch between think / answer blocks on the fly
        if "<think>" in delta:
            in_think = True
            delta = delta.split("<think>", 1)[1]  # drop anything before tag
        if "</think>" in delta:
            before, after = delta.split("</think>", 1)
            collected_think += before
            think_box.markdown(_styled_box(collected_think, "#FFF9DB", "🧠 System thinking…"), unsafe_allow_html=True)
            in_think = False
            delta = after  # continue with remainder into answer block

        if in_think:
            collected_think += delta
            think_box.markdown(_styled_box(collected_think, "#FFF9DB", "🧠 System thinking…"), unsafe_allow_html=True)
        else:
            collected_resp += delta
            resp_box.markdown(_styled_box(collected_resp, "#E8F4FF"), unsafe_allow_html=True)
    return collected_resp

def stream_standard(msgs):
    collected_resp = ""
    in_think = False
    resp_box = st.empty()

    stream = openai_client.chat.completions.create(
        model=MODEL,
        messages=msgs,
        temperature=0.7,
        max_tokens=CONTEXT_SIZE,
        stream=True,
    )
    for chunk in stream:
        if not chunk.choices or not chunk.choices[0].delta:
            continue
        delta = chunk.choices[0].delta.content or ""

        # Detect start/end of <think>..</think> and suppress tokens inside.
        if "<think>" in delta:
            in_think = True
            # remove part after <think>
            delta = delta.split("<think>")[0]
        if "</think>" in delta:
            # show part after </think>
            after = delta.split("</think>", 1)[1]
            in_think = False
            delta = after
        if not in_think and delta:
            collected_resp += delta
            resp_box.markdown(collected_resp)
    return collected_resp

# --------------------------- #
# Main loop                   #
# --------------------------- #
if user_prompt := st.chat_input("Say something"):
    with st.chat_message("user"):
        st.markdown(user_prompt)
    st.session_state.history.append({"role": "user", "content": user_prompt})
    trim_history(st.session_state.history, MAX_HISTORY)

    with st.chat_message("assistant"):
        if ENABLE_REASONING and MODEL == "deepseek-r1-preview":
            assistant_msg = stream_with_reasoning(build_messages(st.session_state.history[:-1], user_prompt))
        else:
            assistant_msg = stream_standard(build_messages(st.session_state.history[:-1], user_prompt))

    st.session_state.history.append({"role": "assistant", "content": assistant_msg})
    trim_history(st.session_state.history, MAX_HISTORY)