Prompting ChatGPT with LangChain

$ python3 -m venv venv
$ venv/bin/pip install langchain-openai # version 0.1.23 at time of writing

export OPENAI_API_KEY="yourkey"

from typing import Any, DefaultDict, Dict

from langchain_core.chat_history import InMemoryChatMessageHistory
from langchain_core.messages.ai import AIMessage
from langchain_core.runnables import chain
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai.chat_models.base import ChatOpenAI

@chain
def ask_user_prompt(_: None) -> str:
    try:
        return input('   User: ')
    except EOFError:
        print('(exit)')
        exit(0)

@chain
def print_chatgpt_response(message: AIMessage) -> None:
    print(f'ChatGPT: {message.content}')

@chain
def line_break(_: Any) -> Any:
    print()
    return _

def create_session_storage() -> DefaultDict[str, InMemoryChatMessageHistory]:
    from collections import defaultdict
    from langchain_core.chat_history import InMemoryChatMessageHistory

    return defaultdict(InMemoryChatMessageHistory)

def create_session_config() -> Dict[str, Dict[str, str]]:
    from random import randint
    session_id =  str(randint(0, 999999))
    config = {'configurable': {'session_id': session_id}}
    return config


def main():
    print('Send Ctrl+D to exit.')
    print()

    session_storage = create_session_storage()
    session_config = create_session_config()

    llm = ChatOpenAI()
    get_session_history = lambda session_id: session_storage[session_id]

    llm_with_memory = RunnableWithMessageHistory(llm, get_session_history)

    chat_chain = (
        ask_user_prompt |
        llm_with_memory |
        print_chatgpt_response |
        line_break
    )

    while True:
        chat_chain.invoke(None, config=session_config)

if __name__ == '__main__':
    main()

$ venv/bin/python chat.py 
Send Ctrl+D to exit.

   User: tell me a joke
ChatGPT: Why did the scarecrow win an award? Because he was outstanding in his field!

   User: explain it
ChatGPT: The joke plays on the double meaning of the phrase "outstanding in his field." In one sense, it means the scarecrow was excellent at his job of scaring away birds in the field. In another sense, it means he was physically standing out in the field, as scarecrows are typically seen standing in fields to protect crops. The humor comes from the pun and the unexpected twist of giving an award to a scarecrow.