one-shot-langchain.py

# imports
import getpass
import os
from langchain import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, 
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate
)

# open ai
os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')
openai = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0.2)

# docs
docs_dir = './books/'
ouput_dir = './output/books/'

# prompt templates
template_sys = "You are a helpful assistant that extracts unique terms from book summaries. Specifically, you must extract all the terms unique to this book along with their definition and return them as a JSON object."
template_base = "The book summary is: {input}"

# get all he md files in the docs dir
docs = [f for f in os.listdir(docs_dir) if f.endswith(".md")]
for doc in docs:

    # load single document from folder
    loader = UnstructuredMarkdownLoader(docs_dir + doc)
    data = loader.load()
    
    # prepare chat prompt
    chat_prompt = ChatPromptTemplate.from_messages([
        SystemMessagePromptTemplate.from_template(template_sys),
        HumanMessagePromptTemplate.from_template(template_base)
    ])
    chat_prompt.format_prompt(input=data).to_messages()

    # prepare request and run chat
    llm_chain = LLMChain(
        prompt=chat_prompt,
        llm=openai,
        verbose=True
    )   
    response = llm_chain.run(data)

    # save respnse to output dir
    with open(ouput_dir + doc.replace('md', 'json'), "w") as f:
        f.write(response)