app.py

import streamlit as st
from groq import Groq
import os 
from dotenv import load_dotenv
import json

load_dotenv()

st.set_page_config(layout = 'wide')

client = Groq(api_key = os.environ['GROQ_API_KEY'])

with open('prompts/logic.md') as logic_prompt_file : 
    logic_prompt = logic_prompt_file.read()

with open('prompts/artist.md') as artis_prompt_file : 
    artist_prompt = artis_prompt_file.read()

if 'chat_histories' not in st.session_state: st.session_state.chat_histories = {}
if 'current_chat_id' not in st.session_state: st.session_state.current_chat_id = None
if 'chat_counter' not in st.session_state: st.session_state.chat_counter = 0

def start_new_chat() : 

    st.session_state.chat_counter += 1
    new_chat_id = f'Chat {st.session_state.chat_counter}'

    st.session_state.chat_histories[new_chat_id] = []
    st.session_state.current_chat_id = new_chat_id
    st.rerun()

def clear_chat_history() : 

    st.session_state.chat_histories = {}
    st.session_state.chat_counter = 0
    st.session_state.current_chat_id = None

    st.rerun()

with st.sidebar : 

    st.title('Chat Sessions')
    st.button('New Chat' , on_click = start_new_chat , use_container_width = True)

    st.markdown('---')

    chat_ids = list(st.session_state.chat_histories.keys())

    for chat_id in reversed(chat_ids) : 

        if st.button(chat_id , use_container_width = True) : 

            st.session_state.current_chat_id = chat_id
            st.rerun()

    st.markdown('---')

    with st.expander('Danger Zone') : 
        st.button('Clear All' , on_click = clear_chat_history , use_container_width = True,  type = 'primary')

if st.session_state.current_chat_id is None : 

    st.info("Start a new conversation to begin.")
    st.stop()

current_messages = st.session_state.chat_histories[st.session_state.current_chat_id]

for message in current_messages : 

    with st.chat_message(message['role']) : 
        st.markdown(message['content'])

if prompt := st.chat_input('What would you like to ask?'):

    st.session_state.chat_histories[st.session_state.current_chat_id].append({'role' : 'user' , 'content' : prompt})

    with st.chat_message('user') : 
        st.markdown(prompt)

    with st.chat_message('assistant') : 
        with st.spinner('Analyzing intent...') : 

            logic_completion = client.chat.completions.create(

                messages = [
                    {"role" : "system" , "content" : logic_prompt} , 
                    {"role" : "user" , "content" : f'Previous History : {st.session_state.chat_histories[st.session_state.current_chat_id]}\n\nUser says: {prompt}'}
                ] , 
                model = 'llama-3.3-70b-versatile' , 
                response_format = {"type" : "json_object"}
            )

            logic_output = json.loads(logic_completion.choices[0].message.content)

            st.sidebar.write(f"**Intent:** {logic_output.get('intent')}")

            st.sidebar.json(logic_output.get('reason'))

        with st.spinner('Channeling the artist...') : 
            
            creative_input = f"Previous History : {st.session_state.chat_histories[st.session_state.current_chat_id]}\n\nUser says: {prompt}\n\nLogic Engine Results: {json.dumps(logic_output)}"
            
            creative_completion = client.chat.completions.create(
                messages = [
                    {"role" : "system" , "content" : artist_prompt} , 
                    {"role" : "user" , "content" : creative_input}
                ] , 
                model = 'llama-3.3-70b-versatile' , 
                response_format = {"type" : "json_object"}
            )
            
            artist_response_data = json.loads(creative_completion.choices[0].message.content)
            final_text = artist_response_data.get('response', '')
            
            st.markdown(final_text)
            
            if artist_response_data.get('show_image') == 'true' and artist_response_data.get('image') : 
                st.image(artist_response_data['image'])

    st.session_state.chat_histories[st.session_state.current_chat_id].append({'role': 'assistant', 'content': final_text})