Generate augmented dialogs on one given dialog

DialogAugmenter is a class used to augment an original dialog by paraphrasing dialog turns while maintaining dialog structure and flow of the conversation.

from langchain_openai import ChatOpenAI

from dialog2graph.datasets.augment_dialogs.augmentation import DialogAugmenter
from dialog2graph.pipelines.model_storage import ModelStorage

First, LLM models should be configured for further use. So, ModelStorage instance is created to which choosen LLMs are added for dialog generation (i.e. dialog augmentation) and formatting LLM’s output. More information on ModelStorage usage may be found in this userguide.

model_storage = ModelStorage()
model_storage.add(
        key="generation-llm",
        config={"model_name": "gpt-4o-mini-2024-07-18", "temperature": 0.7},
        model_type=ChatOpenAI
    )
model_storage.add(
        key="formatting-llm",
        config={"model_name": "gpt-3.5-turbo", "temperature": 0.7},
        model_type=ChatOpenAI
    )

Then, DialogAugmenter instance is created to use invoke method for getting augmented dialogs. DialogAugmenter takes previously created ModelStorage instance and the names given to models should be leveraged for new dialog generation and dialog formatting.

augmenter = DialogAugmenter(
        model_storage=model_storage,
        generation_llm="generation-llm",
        formatting_llm="formatting-llm"
    )
result = augmenter.invoke(
    dialog=dialog, # original dialog as a Dialog object
    topic=topic, # topic of the original dialog as a string
    prompt=augmentation_prompt # prompt for dialog augmentation as a string
    )