import logging
import sys

import langchain
from extract_100knocks_qa import extract_questions
from langchain.chat_models import ChatOpenAI
from llama_index import (GPTSQLStructStoreIndex, LLMPredictor, ServiceContext,
                         SQLDatabase)
from ruamel.yaml import YAML
from sqlalchemy import create_engine

verbose = True

pgconfig = {
    'host': 'localhost',
    'port': 5432,
    'database': 'dsdojo_db',
    'user': 'padawan',
    'password': 'padawan12345',
}


def main():
    if verbose:
        logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
        logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
        langchain.verbose = True

    # データベースに接続
    database_url = 'postgresql://{user}:{password}@{host}:{port}/{database}'.format(
        **pgconfig)
    engine = create_engine(database_url)

    # LlamaIndexはデフォルトでtext-davinci-003を使うので、gpt-3.5-turboを使うよう設定
    llm = ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0)
    predictor = LLMPredictor(llm)
    service_context = ServiceContext.from_defaults(llm_predictor=predictor)

    # LlamaIndexのtext-to-SQLの準備
    sql_database = SQLDatabase(engine)
    index = GPTSQLStructStoreIndex(
        [],
        service_context=service_context,
        sql_database=sql_database,
    )

    # 問題の一覧を抽出
    questions = extract_questions()

    yaml = YAML()
    yaml.default_style = '|'
    with open('results/result.yaml', 'w', encoding='utf-8') as f:

        # text-to-SQLを実行
        for question in questions:
            try:
                response = index.query(question)
                answer = response.extra_info['sql_query']

                qa = {
                    'question': question,
                    'answer': answer,
                }
            except Exception as e:
                qa = {
                    'question': question,
                    'error': str(e),
                }
            yaml.dump([qa], f)


if __name__ == "__main__":
    main()

langchain.chat_models.ChatOpenAI
[((829, 856), 'sqlalchemy.create_engine', 'create_engine', (['database_url'], {}), '(database_url)\n', (842, 856), False, 'from sqlalchemy import create_engine\n'), ((934, 987), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)'}), "(model_name='gpt-3.5-turbo', temperature=0)\n", (944, 987), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1004, 1021), 'llama_index.LLMPredictor', 'LLMPredictor', (['llm'], {}), '(llm)\n', (1016, 1021), False, 'from llama_index import GPTSQLStructStoreIndex, LLMPredictor, ServiceContext, SQLDatabase\n'), ((1044, 1097), 'llama_index.ServiceContext.from_defaults', 'ServiceContext.from_defaults', ([], {'llm_predictor': 'predictor'}), '(llm_predictor=predictor)\n', (1072, 1097), False, 'from llama_index import GPTSQLStructStoreIndex, LLMPredictor, ServiceContext, SQLDatabase\n'), ((1150, 1169), 'llama_index.SQLDatabase', 'SQLDatabase', (['engine'], {}), '(engine)\n', (1161, 1169), False, 'from llama_index import GPTSQLStructStoreIndex, LLMPredictor, ServiceContext, SQLDatabase\n'), ((1182, 1273), 'llama_index.GPTSQLStructStoreIndex', 'GPTSQLStructStoreIndex', (['[]'], {'service_context': 'service_context', 'sql_database': 'sql_database'}), '([], service_context=service_context, sql_database=\n    sql_database)\n', (1204, 1273), False, 'from llama_index import GPTSQLStructStoreIndex, LLMPredictor, ServiceContext, SQLDatabase\n'), ((1332, 1351), 'extract_100knocks_qa.extract_questions', 'extract_questions', ([], {}), '()\n', (1349, 1351), False, 'from extract_100knocks_qa import extract_questions\n'), ((1364, 1370), 'ruamel.yaml.YAML', 'YAML', ([], {}), '()\n', (1368, 1370), False, 'from ruamel.yaml import YAML\n'), ((520, 579), 'logging.basicConfig', 'logging.basicConfig', ([], {'stream': 'sys.stdout', 'level': 'logging.DEBUG'}), '(stream=sys.stdout, level=logging.DEBUG)\n', (539, 579), False, 'import logging\n'), ((619, 659), 'logging.StreamHandler', 'logging.StreamHandler', ([], {'stream': 'sys.stdout'}), '(stream=sys.stdout)\n', (640, 659), False, 'import logging\n'), ((588, 607), 'logging.getLogger', 'logging.getLogger', ([], {}), '()\n', (605, 607), False, 'import logging\n')]