Playing around with langchain

sqlgen.py

from langchain.schema import StrOutputParser
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.chat_models import ChatOpenAI

llm = ChatOpenAI(api_key='MY_COOL_API_KEY')


PROMPT_SYSTEM = '''
    You are an experienced SQL developer and database specialist.
    You will receive a schema and a question.
    You will convert questions to SQL queries based on the schema.
    You will only output the SQL query.
'''

PROMPT_USER = '''
    # Table Schema
    ```sql
    {schema}
    ```
                                               
    # Question
    > {question}
'''

def get_schema(_) -> str:
    return '''
        CREATE TABLE users (
            id INTEGER PRIMARY KEY, 
            name TEXT, 
            birthday DATETIME
        )
    '''


prompt = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template(PROMPT_SYSTEM),
    HumanMessagePromptTemplate.from_template(PROMPT_USER),
])

chain = (
    RunnablePassthrough.assign(schema=get_schema)
    | prompt
    | llm
    | StrOutputParser()
)

res = chain.invoke({'question': 'Show me the month with most users born in'})
print(res)
# SELECT strftime('%m', birthday) AS birth_month, COUNT(*) AS num_users
# FROM users
# GROUP BY birth_month
# ORDER BY num_users DESC
# LIMIT 1;