app.py

import os
import chainlit as cl
from dotenv import load_dotenv
from langchain_anthropic import ChatAnthropic
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.chat_history import InMemoryChatMessageHistory

# Load environment variables (Hugging Face Spaces injects secrets automatically)
load_dotenv()

# Verify API key is available
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
if not ANTHROPIC_API_KEY:
    print("WARNING: ANTHROPIC_API_KEY not found in environment variables!")

# Load the system prompt from prompt.md


def load_system_prompt():
    """Load the system prompt from prompt.md file."""
    try:
        with open("prompt.md", "r", encoding="utf-8") as f:
            prompt_content = f.read()
            # Verify the prompt loaded correctly (basic sanity check)
            if len(prompt_content) < 100:
                print("WARNING: prompt.md seems too short. Using fallback prompt.")
                return "You are an Adaptive Coach guiding a small team through the Design Thinking process."
            return prompt_content
    except FileNotFoundError:
        print("WARNING: prompt.md not found. Using fallback prompt.")
        return "You are an Adaptive Coach guiding a small team through the Design Thinking process."
    except UnicodeDecodeError as e:
        print(
            f"WARNING: Encoding error reading prompt.md: {e}. Using fallback prompt.")
        return "You are an Adaptive Coach guiding a small team through the Design Thinking process."


# Initialize chain components (outside of handlers to avoid ContextVar issues)
SYSTEM_PROMPT = load_system_prompt()


def get_chain():
    """Create and return the LCEL chain with conversation history support."""
    if not ANTHROPIC_API_KEY:
        return None

    # Initialize the ChatAnthropic model with streaming enabled
    model = ChatAnthropic(
        model="claude-sonnet-4-5-20250929",  # Using the standard Claude 3.5 Sonnet model
        streaming=True,
        temperature=0.7
    )

    # Create the prompt template with system message and conversation history placeholder
    # Always include MessagesPlaceholder - we'll pass empty list if no history
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])

    # Build the LCEL chain: Prompt | Model | Parser
    return prompt | model | StrOutputParser()


@cl.on_chat_start
async def on_chat_start():
    """Initialize the chain when a chat session starts."""
    # Initialize chat history for this session
    chat_history = InMemoryChatMessageHistory()
    cl.user_session.set("chat_history", chat_history)

    # Mark that this is a new session (for welcome message)
    cl.user_session.set("first_message", True)

    # Try to send welcome message here (works locally and sometimes in deployment)
    # If it fails due to ContextVar issues, we'll send it on first message instead
    try:
        if ANTHROPIC_API_KEY:
            welcome_msg = "What are you working on?"
            await cl.Message(content=welcome_msg).send()

            # Add welcome message to chat history
            from langchain_core.messages import AIMessage
            chat_history.add_message(AIMessage(content=welcome_msg))
            cl.user_session.set("chat_history", chat_history)

            # Mark that welcome message was sent
            cl.user_session.set("welcome_sent", True)
    except (LookupError, RuntimeError, Exception) as e:
        # ContextVar or other initialization error - will send welcome on first message instead
        print(f"Could not send welcome message in on_chat_start: {e}")
        cl.user_session.set("welcome_sent", False)


def is_simple_greeting(text: str) -> bool:
    """Check if the message is just a simple greeting."""
    text_lower = text.strip().lower()
    simple_greetings = ["hi", "hello", "hey",
                        "hey there", "hi there", "greetings"]
    return text_lower in simple_greetings or len(text_lower) <= 3


@cl.on_message
async def on_message(message: cl.Message):
    """Handle incoming messages and stream responses."""
    # Get or initialize chat history
    chat_history = cl.user_session.get("chat_history")
    if chat_history is None:
        chat_history = InMemoryChatMessageHistory()
        cl.user_session.set("chat_history", chat_history)

    # Check if this is the first message and send welcome message (if not already sent in on_chat_start)
    is_first = cl.user_session.get("first_message", False)
    welcome_sent = cl.user_session.get("welcome_sent", False)

    if is_first:
        cl.user_session.set("first_message", False)

        # Check API key
        if not ANTHROPIC_API_KEY:
            await cl.Message(
                content="❌ Error: ANTHROPIC_API_KEY not configured. Please set it in Space settings → Variables and secrets."
            ).send()
            return

        # Send welcome message only if it wasn't sent in on_chat_start
        if not welcome_sent:
            welcome_msg = "What are you working on?"
            await cl.Message(content=welcome_msg).send()

            # Add welcome message to chat history and persist back to session
            from langchain_core.messages import AIMessage
            chat_history.add_message(AIMessage(content=welcome_msg))
            # Persist changes
            cl.user_session.set("chat_history", chat_history)

        # If the user's first message is just a greeting, don't process it
        # Otherwise, continue to process their meaningful response
        if is_simple_greeting(message.content):
            return

    # Check API key
    if not ANTHROPIC_API_KEY:
        await cl.Message(
            content="❌ Error: ANTHROPIC_API_KEY not configured. Please set it in Space settings → Variables and secrets."
        ).send()
        return

    # Get the chain (it supports conversation history via MessagesPlaceholder)
    chain = get_chain()
    if chain is None:
        await cl.Message(content="❌ Error: Chain not initialized. Please refresh the page.").send()
        return

    # Create a message object for streaming
    msg = cl.Message(content="")
    await msg.send()

    # Prepare input with chat history (previous messages only, current message goes in {input})
    input_dict = {
        "input": message.content,
        "chat_history": chat_history.messages  # All previous messages
    }

    # Invoke the chain asynchronously with streaming
    try:
        full_response = ""
        async for chunk in chain.astream(input_dict):
            await msg.stream_token(chunk)
            full_response += chunk
        await msg.update()

        # Add both user message and assistant response to chat history AFTER processing
        from langchain_core.messages import HumanMessage, AIMessage
        chat_history.add_message(HumanMessage(content=message.content))
        chat_history.add_message(AIMessage(content=full_response))

        # Persist changes back to session
        cl.user_session.set("chat_history", chat_history)

    except Exception as e:
        # Log the full error for debugging
        error_msg = f"❌ Error: {str(e)}"
        print(f"Error in on_message: {error_msg}")
        import traceback
        traceback.print_exc()
        await cl.Message(
            content=f"{error_msg}\n\nPlease check the Space logs for more details."
        ).send()