app.py

from langchain_core.runnables import RunnablePassthrough
# from langchain_text_splitters import RecursiveCharacterTextSplitter
import streamlit as st
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_groq import ChatGroq
from dotenv import load_dotenv, find_dotenv
from PIL import Image
from pathlib import Path
import os, uuid


base_path = Path(__file__).parent
img_path = base_path / "images" 

# Load images
assistant_crs = Image.open(img_path/"assistant_crs.png")
user_crs = Image.open(img_path/"user_crs.png")


def llm_model(model="moonshotai/kimi-k2-instruct-0905"):
    _ = load_dotenv(find_dotenv())
    groq_api_key = os.getenv("GROQ_API_KEY")
    if not groq_api_key :
        try: 
            groq_api_key = st.secrets.get("GROQ_API_KEY")
        except Exception:
            groq_api_key= None
    
    llm = ChatGroq(model=model, groq_api_key=groq_api_key)
    return llm

# --- CONFIGURATION CONSTANTS ---
HISTORY_STORE_KEY = "chat_history_store"    # Dict: session_id -> {name, history, is_placeholder_name}
CURRENT_SESSION_ID_KEY = "current_session_id"   # Tracks the ID of the active session
HISTORY_PLACEHOLDER_KEY = "history"     # Should match "input_history_key" parameterin RunnableWithMessageHistory

# --- 1. PERSISTENCE AND SESSION MANAGEMENT LOGIC

def initialize_session_state():
    """Initializes the required session state variables if they don't exist."""
    # 1. Dictionary to hold all sessions (key: session_id, value: InMemoryChatMessageHistory)
    if HISTORY_STORE_KEY not in st.session_state:
        st.session_state[HISTORY_STORE_KEY] = {}
        # Create a default initial session
        create_new_session("Chat 1")


def create_new_session(session_name: str):
    """Creates a new session and sets it as the active session."""
    new_id = str(uuid.uuid4())
    # Store the new history object and map it to a readable name
    st.session_state[HISTORY_STORE_KEY][new_id] = {
        "name" : session_name,
        "history" : InMemoryChatMessageHistory()
        # "placeholder_name" : 
    }
    st.session_state[CURRENT_SESSION_ID_KEY] = new_id
    # Reset the display buffer to show the new, empty chat
    st.session_state.display_messages = []


def get_session_history(session_id: str) -> InMemoryChatMessageHistory:
    """Retrieves the history object for the given session_ID"""
    if session_id in st.session_state[HISTORY_STORE_KEY]:
        return st.session_state[HISTORY_STORE_KEY][session_id]["history"]
    else:
        # Fallback case, should not hif with proper initialization
        return InMemoryChatMessageHistory()


def delete_session(session_id: str):
    """Deletes the sesson ID from the session state"""
    # 1. Remove the history entry
    if session_id in st.session_state[HISTORY_STORE_KEY]:
        del st.session_state[HISTORY_STORE_KEY][session_id]
    # 2. check if the store is empty
    if not st.session_state[HISTORY_STORE_KEY]:
        # create a new session with default name
        create_new_session("Chat 1")
    else:
        # 3. Current Active session is the first one remaining
        st.session_state[CURRENT_SESSION_ID_KEY] = next(iter(st.session_state[HISTORY_STORE_KEY].keys()))
    # 4. Force rerun to update UI and load the new active chat
    st.rerun()


# --- 2. LANGCHAIN Setup ---
@st.cache_resource
def llm_chain():
    """ Initializes LLM and returns a RunnableWithMessageHistory instance.
    The @st.cache_resource decorator ensures this complex object is only created once.
    """
    # LLM  
    llm = llm_model()

    # prompt
    instruction = """ 
    IDENTITY & OWNERSHIP:
    - NAME: ChatAI 
    - OWNER/CREATOR: Co2fi Rodolphe Segbedji
    - ROLE: You are a sophisticated, high-context Conversational Thought Partner. 
    You are not a static search engine; you are a proactive assistant designed to 
    engage in deep, meaningful, and fluid dialogue. 

    ANTI-HALLUCINATION & INTELLECTUAL HONESTY:
    - If a query is outside your training data or context window, state "I don't have 
   enough information to answer that accurately" rather than guessing.
   - Never fabricate facts, URLs, dates, or technical documentation.
   - If the user provides a premise that is factually incorrect, politely correct 
   the underlying assumption before answering.

"""
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", instruction),
            MessagesPlaceholder(variable_name = HISTORY_PLACEHOLDER_KEY),
            ("human", "{input}"),
        ])
    # parser
    parser = StrOutputParser()
    # chain
    chain = prompt | llm | parser

    # chain_with_message_history
    chain_with_memory = RunnableWithMessageHistory(
        runnable = chain,
        get_session_history = get_session_history,
        input_messages_key = "input",
        history_messages_key = HISTORY_PLACEHOLDER_KEY
    )

    return chain_with_memory

    
# --- 3. STREAMLIT UI AND EXECUTION ---

def main():

    # App title
    st.markdown("# 🧠 ChatAI 💡")
   
    # Initialize all required session state variables
    initialize_session_state()

    # Initialize the state-aware chain (chain with memory)
    chain = llm_chain()

    # Get the active session ID 
    current_session_id = st.session_state[CURRENT_SESSION_ID_KEY]
    current_session_data = st.session_state[HISTORY_STORE_KEY][current_session_id]
    current_session_name = current_session_data["name"] 
    current_history = get_session_history(current_session_id)

    # SideBar UI for Session Management
    with st.sidebar:
        st.header("Chat Sessions")
        # Map of ID to Name for the selectbox
        session_options = {
            k: v["name"] for k, v in st.session_state[HISTORY_STORE_KEY].items()
        }
        # Session selector
        selected_id = st.selectbox(
            "Select a Chat",
            options = list(session_options.keys()),
            format_func = lambda id: session_options[id],
            key = "session_select_box"
        )
        # Logic to switch a session if a different one is selected
        if selected_id != current_session_id:
            st.session_state[CURRENT_SESSION_ID_KEY] = selected_id
            # Force a rerun to load the new chat history
            st.rerun()

        # 2. New session Creator
        new_session_name = st.text_input("➕ New Session")
        if new_session_name and (new_session_name not in [v["name"] for v in st.session_state[HISTORY_STORE_KEY].values()]):
            create_new_session(new_session_name)
            st.rerun()
        if st.button("🗑️ Delete Session"):
            if len(st.session_state[HISTORY_STORE_KEY].keys()) > 1 :
                delete_session(current_session_id)
                st.rerun()
            else:
                st.error("Can't delete the only chat remaining. Create a new chat before deleting it. ")

    # --- MAIN CHAT DISPLAY---
    # Display messages from the current session's history object


    # Display current session messages 
    for message in current_history.messages:
        role = "user" if isinstance(message, HumanMessage) else "assistant"
        avatar=user_crs if role == "user" else assistant_crs
        with st.chat_message(role, avatar=avatar):
            st.markdown(message.content)

    # Handle user input
    user_input = st.chat_input(
        # "Type, or attach files, or record audio",
        "Converse with ChatAI",
        accept_file = "multiple",
        file_type = None,    # allow any file type. You can restrict if you want using a list of file type
        accept_audio = True,
    )

    if user_input:
        text = user_input.text or ""
        files = getattr(user_input, "files", [])
        audio = getattr(user_input, "audio", None)

        # show user message immedialtely
        with st.chat_message("user", avatar=user_crs):
            # Handle text input
            if text:
                st.markdown(text)

            # Handle file upload and show info about files and route by type
            for f in files:
                st.write(f"Uploaded: {f.name} ({f.type})")
                # Document-like types: call your RAG pipeline
                if f.type in ["application/pdf", "text/plain"]:
                    st.write("→ Calling RAG function for document")
                    # e.g. rag_answer = call_rag(f)
                elif f.type.startswith("image/"):
                    st.write("-> Calling image handler")
                    st.image(f)
                    # e.g. img_answer = handle_image(f)
                else:
                    st.write("-> Unsupported file type for special handling")
                
            # Handle audio
            if audio:
                st.write("Recorded audio")
                st.audio(audio)
                # e.g. text = transcribe(audio)
                # st.write(text)

        # Invoke the chain-aware to get AI response
        with st.chat_message("assistant", avatar= assistant_crs):
            with st.spinner("Thinking..."):
                # invoke the chain, passing the current session ID in the config
                ai_response = chain.invoke(
                    {"input": user_input.text},
                    config={"configurable": {"session_id": current_session_id}}
                )
                st.markdown(ai_response)


if __name__ == "__main__" :
    main()