from typing import Optional, List, Union
from langchain.prompts.chat import ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder, BaseMessagePromptTemplate
from langchain.schema import SystemMessage, BaseMessage


def create_agent_prompt(
    system_message: Optional[SystemMessage] = SystemMessage(content="You are a helpful AI assistant."),
    extra_prompt_messages: Optional[List[BaseMessagePromptTemplate]] = None,
)-> ChatPromptTemplate:
    
    _prompts = extra_prompt_messages or []
    
    messages: List[Union[BaseMessagePromptTemplate, BaseMessage]]
    
    if system_message:
        messages = [system_message]
    else:
        messages = []
        
    messages.extend(
        [
            *_prompts,
            HumanMessagePromptTemplate.from_template("{question}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    return ChatPromptTemplate(messages=messages)