app.py

import os
import streamlit as st
from datasets import load_dataset
from huggingface_hub import InferenceClient

# Get the API key from the environment variable
api_key = os.getenv("HF_API_KEY")
client = InferenceClient(api_key=api_key)

# Inject custom CSS to change the background color to yellow
st.markdown(
    """
    <style>
    body {
        background-color: yellow;
    }
    </style>
    """,
    unsafe_allow_html=True
)

st.title("Adrega AI Help")
include_context = st.checkbox('Search in Help')

if 'conversation' not in st.session_state:
    st.session_state.conversation = ""
if 'include_context' not in st.session_state: 
    st.session_state.include_context = False

if include_context:
    dataset = load_dataset("andreska/adregadocs", split="test")

    # Function to read the content from the dataset
    def read_dataset(dataset):
        text = []
        for item in dataset:
            text.append(item['text'])
        return "\n".join(text)

    context = read_dataset(dataset)
else:
    context = ""

def handle_submit():
    user_input = st.session_state.user_input
    if user_input:
        if st.session_state.include_context:
            messages = [
                {"role": "system", "content": f"Context: {context}"},
                {"role": "user", "content": user_input}
            ]
        else:
            messages = [
                {"role": "user", "content": user_input}
            ]

        completion = client.chat.completions.create(
            model="Qwen/Qwen2.5-72B-Instruct",
            #model="Qwen/Qwen2.5-Coder-32B-Instruct",
            #model="HuggingFaceTB/SmolLM2-1.7B-Instruct",
            messages=messages,
            max_tokens=500
        )
        
        answer = completion.choices[0].message['content']
        st.session_state.conversation += f"User: {user_input}\nAdrega AI: {answer}\n\n"
        #st.write(f"Adrega AI: {answer}")
    else:
        st.write("Please enter a question.")

if st.button("Ask"): 
    handle_submit()

st.text_input('Ask me a question', key='user_input', on_change=handle_submit)
st.text_area("Conversation History", value=st.session_state.conversation, height=300, max_chars=None)