DEV Community

Cover image for Prompting ChatGPT with LangChain
Talles L
Talles L

Posted on

Prompting ChatGPT with LangChain

First, set up an virtual environment with needed packages:

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

Then, set your OpenAI key on an environment variable:

export OPENAI_API_KEY="yourkey"
Enter fullscreen mode Exit fullscreen mode

Create the following file:

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()
Enter fullscreen mode Exit fullscreen mode

Running it:

$ 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.
Enter fullscreen mode Exit fullscreen mode

Documentation:

Top comments (0)