Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,180 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
 import streamlit as st
+from dotenv import load_dotenv
+from huggingface_hub import InferenceClient
 
-"""
-# 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])),
-    ))
+# -----------------------------------------------------------------------------
+# Environment & constants
+# -----------------------------------------------------------------------------
+load_dotenv()
+HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+model_id = "mistralai/Mistral-7B-Instruct-v0.3"
+
+# -----------------------------------------------------------------------------
+# LLM helper
+# -----------------------------------------------------------------------------
+
+def get_llm_hf_inference(model_id=model_id, max_new_tokens: int = 128, temperature: float = 0.1):
+    """Return an InferenceClient wrapper for Hugging Face inference."""
+    client = InferenceClient(model=model_id, token=HUGGINGFACEHUB_API_TOKEN)
+
+    def run(prompt: str) -> str:
+        try:
+            # For future versions with .conversational method
+            response = client.conversational(
+                inputs=prompt,
+                parameters={
+                    "max_new_tokens": max_new_tokens,
+                    "temperature": temperature,
+                },
+            )
+            return response.generated_text
+        except AttributeError:
+            # Fallback for older huggingface_hub clients
+            response = client.post(
+                json={
+                    "inputs": prompt,
+                    "parameters": {
+                        "max_new_tokens": max_new_tokens,
+                        "temperature": temperature,
+                    },
+                },
+                task="conversational"
+            )
+            return response["generated_text"]
+
+    return run
+
+# -----------------------------------------------------------------------------
+# Streamlit page configuration
+# -----------------------------------------------------------------------------
+st.set_page_config(page_title="KolaChatBot", page_icon="🤗")
+st.title("KolaChatBot")
+st.markdown(
+    f"*KolaChatBot utilise l'API Inference de Hugging Face avec le modèle **{model_id}**.*"
+)
+
+# -----------------------------------------------------------------------------
+# Session ‐state initialisation
+# -----------------------------------------------------------------------------
+if "avatars" not in st.session_state:
+    st.session_state.avatars = {"user": "👤", "assistant": "🤗"}
+
+if "user_text" not in st.session_state:
+    st.session_state.user_text = None
+
+if "max_response_length" not in st.session_state:
+    st.session_state.max_response_length = 256
+
+if "system_message" not in st.session_state:
+    st.session_state.system_message = "You are a friendly AI conversing with a human user."
+
+if "starter_message" not in st.session_state:
+    st.session_state.starter_message = "Hello, there! How can I help you today?"
+
+# -----------------------------------------------------------------------------
+# Sidebar settings
+# -----------------------------------------------------------------------------
+with st.sidebar:
+    st.header("Paramètres du système")
+
+    # AI Settings
+    st.session_state.system_message = st.text_area(
+        "System Message", value=st.session_state.system_message
+    )
+    st.session_state.starter_message = st.text_area(
+        "First AI Message", value=st.session_state.starter_message
+    )
+
+    # Model Settings
+    st.session_state.max_response_length = st.number_input(
+        "Max Response Length", value=st.session_state.max_response_length
+    )
+
+    # Avatar Selection
+    st.markdown("*Sélection des avatars :*")
+    col1, col2 = st.columns(2)
+    with col1:
+        st.session_state.avatars["assistant"] = st.selectbox(
+            "Avatar IA", options=["🤗", "💬", "🤖"], index=0
+        )
+    with col2:
+        st.session_state.avatars["user"] = st.selectbox(
+            "Avatar Utilisateur", options=["👤", "👱‍♂️", "👨🏾", "👩", "👧🏾"], index=0
+        )
+
+    # Reset Chat History
+    reset_history = st.button("Réinitialiser l'historique")
+
+# -----------------------------------------------------------------------------
+# Chat history initialisation / reset
+# -----------------------------------------------------------------------------
+if "chat_history" not in st.session_state or reset_history:
+    st.session_state.chat_history = [
+        {"role": "assistant", "content": st.session_state.starter_message}
+    ]
+
+# -----------------------------------------------------------------------------
+# Core inference helper
+# -----------------------------------------------------------------------------
+
+def build_prompt(system_message: str, chat_history: list[dict], user_text: str) -> str:
+    """Format the conversation as a prompt for the LLM."""
+    prompt = f"### SYSTEM:\n{system_message}\n\n"
+    for msg in chat_history:
+        role_tag = "USER" if msg["role"] == "user" else "ASSISTANT"
+        prompt += f"### {role_tag}:\n{msg['content']}\n\n"
+    prompt += f"### USER:\n{user_text}\n\n### ASSISTANT:\n"
+    return prompt
+
+
+def get_response(system_message: str, chat_history: list[dict], user_text: str, max_new_tokens: int = 256):
+    """Generate a response and update chat history."""
+
+    prompt = build_prompt(system_message, chat_history, user_text)
+
+    llm = get_llm_hf_inference(max_new_tokens=max_new_tokens, temperature=0.1)
+    response_text = llm(prompt)
+
+    # Update history
+    chat_history.append({"role": "user", "content": user_text})
+    chat_history.append({"role": "assistant", "content": response_text})
+
+    return response_text, chat_history
+
+# -----------------------------------------------------------------------------
+# Streamlit chat interface
+# -----------------------------------------------------------------------------
+chat_interface = st.container(border=True)
+with chat_interface:
+    output_container = st.container()
+    st.session_state.user_text = st.chat_input(placeholder="Entrez votre message ici…")
+
+# Display chat messages
+with output_container:
+    for message in st.session_state.chat_history:
+        if message["role"] == "system":
+            continue  # Skip system messages
+        with st.chat_message(
+            message["role"], avatar=st.session_state.avatars[message["role"]]
+        ):
+            st.markdown(message["content"])
+
+    # Handle new user message
+    if st.session_state.user_text:
+        # Show the user message immediately
+        with st.chat_message("user", avatar=st.session_state.avatars["user"]):
+            st.markdown(st.session_state.user_text)
+
+        # Generate and display assistant response
+        with st.chat_message(
+            "assistant", avatar=st.session_state.avatars["assistant"]
+        ):
+            with st.spinner("KolaChatBot réfléchit…"):
+                response_text, st.session_state.chat_history = get_response(
+                    system_message=st.session_state.system_message,
+                    user_text=st.session_state.user_text,
+                    chat_history=st.session_state.chat_history,
+                    max_new_tokens=st.session_state.max_response_length,
+                )
+                st.markdown(response_text)