# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/qdecomp.llm.ipynb.

# %% auto 0
__all__ = ['QUESTION_DECOMPOSITION_SYSTEM_PROMPT_TEMPLATE', 'parse_llm_generation', 'make_question_decomposer']

# %% ../../nbs/qdecomp.llm.ipynb 3
import re
from openai import OpenAI

# %% ../../nbs/qdecomp.llm.ipynb 4
QUESTION_DECOMPOSITION_SYSTEM_PROMPT_TEMPLATE = """Decompose the given question into n sub-questions such that when sub-questions are answered, the original question can be answered correctly.
The second subquestion must refer to the answer of the first question by `#1` as in the examples below. Do not create open-ended sub-questions like "Who is ..." or "How is ...".

Avoid statements like 'Here are the decomposed sub-question, ...' or 'Sure! Here are..' or anything along those lines. Just provide sub-questions as illustrated in the examples below.

Question: Who founded the journal that published The Review of Communication?
Sub-questions:
1. Who published The Review of Communication?
2. Who founded #1?

Question: When was the institute that owned The Collegian founded?
Sub-questions:
1. Which institute does own The Collegian?
2. When #1 founded?

Question: {question}
Sub-questions:
"""

# %% ../../nbs/qdecomp.llm.ipynb 5
def parse_llm_generation(output: str):
    if output.startswith("Sub-questions:"):
        output = output.split("Sub-questions:", 1)[-1]
    enumeration_pattern = r'^\d+\.\s'
    for line in output.splitlines():
        # check if the line starts with enumeration
        if not re.match(enumeration_pattern, line):
            continue
        yield line.split(".", 1)[-1].strip()

# %% ../../nbs/qdecomp.llm.ipynb 6
def make_question_decomposer(
    model: str = "gpt-3.5-turbo",
    client: OpenAI = None,
    completion_kwargs: dict | None = None,
):
    if client is None:
        client = OpenAI()

    if completion_kwargs is None:
        completion_kwargs = {}

    def func(question: str) -> list[str]:
        messages = [
            {
                "role": "system",
                "content": QUESTION_DECOMPOSITION_SYSTEM_PROMPT_TEMPLATE.format(question=question),
            },
        ]
        chat_completion = client.chat.completions.create(
            model=model,
            messages=messages,
            **completion_kwargs,
        )
        text = chat_completion.choices[0].message.content
        return list(parse_llm_generation(text))

    return func
