lpm_kernel/L2/data.py

"""Data processing module for L2 model training.

This module provides the L2DataProcessor class which handles data preparation,
processing, and organization for training L2 models, including conversion between
different formats and extraction of information from notes.
"""

import graphrag
import json
import os
import pandas as pd
import random
import subprocess
import traceback
import yaml
from collections import defaultdict
from datasets import DatasetDict, Dataset
from datetime import datetime
from tqdm import tqdm
from typing import Any, Dict, List

from lpm_kernel.L1.bio import (
    MemoryType,
    Note,
    OBJECT_NOTE_TYPE,
    SUBJECT_NOTE_TYPE
)
from lpm_kernel.L2.data_pipeline.data_prep.context_data.context_generator import ContextGenerator
from lpm_kernel.L2.data_pipeline.data_prep.diversity.diversity_data_generator import DiversityDataGenerator
from lpm_kernel.L2.data_pipeline.data_prep.preference.preference_QA_generate import PreferenceQAGenerator
from lpm_kernel.L2.data_pipeline.data_prep.selfqa.selfqa_generator import SelfQA
from lpm_kernel.L2.note_templates import OBJECTIVE_TEMPLATES, SUBJECTIVE_TEMPLATES
from lpm_kernel.L2.utils import format_timestr
from lpm_kernel.api.services.user_llm_config_service import UserLLMConfigService

from lpm_kernel.configs.logging import get_train_process_logger
logger = get_train_process_logger()

class L2DataProcessor:
    """Data processor for L2 model training.
    
    This class handles the processing and organization of data for training L2 models,
    including conversion between different formats, extraction of information from notes,
    and preparation of training data.
    
    Attributes:
        data_path: Base path for data processing.
        preferred_lang: Preferred language for data processing.
    """

    def __init__(
            self,
            data_path: str = "resources/L2/data_pipeline/raw_data",
            preferred_lang: str = "English",
    ):
        """Initialize the L2DataProcessor.
        
        Args:
            data_path: Base path for data processing. Defaults to "resources/L2/data_pipeline/raw_data".
            preferred_lang: Preferred language for data processing. Defaults to "English".
        """
        self.data_path = data_path
        self.preferred_lang = preferred_lang

    def __call__(self, note_list: List[Note], basic_info: Dict) -> Any:
        """Process a list of notes with basic user information.
        
        This method coordinates the overall data processing workflow, including
        splitting notes by type, refining data, converting to text, and extracting
        entities and relationships.
        
        Args:
            note_list: List of Note objects to process.
            basic_info: Dictionary containing basic user information.
            
        Returns:
            Processing results if any, otherwise None.
        """
        user_info, subjective_memory_notes, objective_memory_notes = self.split_notes_by_type(
            note_list, basic_info
        )

        subjective_notes_remade = self.refine_notes_data_subjective(
            subjective_memory_notes, user_info, self.data_path + "/L1/processed_data/subjective/note_remade.json"
        )

        objective_notes_remade = self.refine_notes_data_objective(
            objective_memory_notes, user_info, self.data_path + "/L1/processed_data/objective/note_remade.json"
        )

        self.json_to_txt_each(
            subjective_notes_remade,
            self.data_path + "/L1/processed_data/subjective",
            file_type="note",
        )

        self.json_to_txt_each(
            objective_notes_remade,
            self.data_path + "/L1/processed_data/objective",
            file_type="note",
        )

        logger.info("Data refinement completed, preparing to extract entities and relations")

        lang = user_info.get("lang", "English")

        if len(subjective_notes_remade) > 0:
            self.graphrag_indexing(
                subjective_notes_remade,
                self.data_path + "/L1/processed_data/subjective",
                self.data_path + "/L1/graphrag_indexing_output/subjective",
                lang,
            )

        if len(objective_notes_remade) > 0:
            self.graphrag_indexing(
                objective_notes_remade,
                self.data_path + "/L1/processed_data/objective",
                self.data_path + "/L1/graphrag_indexing_output/objective",
                lang,
            )
        return

    def gen_subjective_data(
            self,
            note_list: List[Note],
            data_output_base_dir: str,
            preference_output_path: str,
            diversity_output_path: str,
            selfqa_output_path: str,
            global_bio: str,
            topics_path: str,
            entitys_path: str,
            graph_path: str,
            user_name: str,
            config_path: str,
            user_intro: str,
            do_context: bool = False,
    ):
        """Generate subjective data for training.
        
        This method processes various types of subjective data including preferences,
        diversity, self-Q&A, and context data, then merges them into a single output file.
        
        Args:
            note_list: List of Note objects.
            data_output_base_dir: Base directory for output data.
            preference_output_path: Path to save preference data.
            diversity_output_path: Path to save diversity data.
            selfqa_output_path: Path to save self-Q&A data.
            global_bio: User's global biography.
            topics_path: Path to topics data.
            entitys_path: Path to entities data.
            graph_path: Path to graph data.
            user_name: Name of the user.
            config_path: Path to configuration file.
            user_intro: User's introduction.
        """
        preference_output_path = os.path.join(
            data_output_base_dir, preference_output_path
        )
        diversity_output_path = os.path.join(
            data_output_base_dir, diversity_output_path
        )
        selfqa_output_path = os.path.join(data_output_base_dir, selfqa_output_path)
        context_output_path = os.path.join(data_output_base_dir, "context_merged.json")

        logger.info("---" * 30 + "\nPreference data generating\n" + "---" * 30)
        self._gen_preference_data(topics_path, preference_output_path, global_bio)
        logger.info("---" * 30 + "\nPreference data generated\n" + "---" * 30)

        logger.info("---" * 30 + "\nDiversity data generating\n" + "---" * 30)
        self._gen_diversity_data(
            entitys_path,
            note_list,
            graph_path,
            diversity_output_path,
            user_name,
            global_bio,
            config_path,
        )
        logger.info("---" * 30 + "\nDiversity data generated\n" + "---" * 30)

        logger.info("---" * 30 + "\nSelfQA data generating\n" + "---" * 30)
        self._gen_selfqa_data(selfqa_output_path, user_name, user_intro, global_bio)
        logger.info("---" * 30 + "\nSelfQA data generated\n" + "---" * 30)

        if do_context:
            logger.info("---" * 30 + "\nContext data generating\n" + "---" * 30)
            self._gen_context_data(
                note_list,
                entitys_path,
                data_output_base_dir,
                user_name,
                user_intro,
                global_bio
            )
            logger.info("---" * 30 + "\nContext data generated\n" + "---" * 30)
            self._merge_context_data(data_output_base_dir, "context_merged.json")

        # Merge the four specified JSON files
        merged_data = []
        if do_context:
            json_files_to_merge = [
                preference_output_path,
                diversity_output_path,
                selfqa_output_path,
                context_output_path
            ]
        else:
            json_files_to_merge = [
                preference_output_path,
                diversity_output_path,
                selfqa_output_path,
            ]

        logger.info("---" * 30 + "\nMerging JSON files\n" + "---" * 30)

        for file_path in json_files_to_merge:
            if file_path and os.path.exists(file_path):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        file_data = json.load(f)
                        if isinstance(file_data, list):
                            merged_data.extend(file_data)
                            logger.info(f"Added {len(file_data)} items from {file_path}")
                        else:
                            merged_data.append(file_data)
                            logger.info(f"Added 1 item from {file_path}")
                except Exception as e:
                    logger.error(f"Error merging file {file_path}: {str(e)}")
            else:
                if file_path == context_output_path and do_context == False:
                    continue
                logger.warning(f"File not found or path is None: {file_path}")

        # Save the merged data
        merged_output_path = os.path.join(data_output_base_dir, "merged.json")
        with open(merged_output_path, 'w', encoding='utf-8') as f:
            json.dump(merged_data, f, ensure_ascii=False, indent=2)

        logger.info(f"Merged data saved to {merged_output_path} with {len(merged_data)} total items")
        logger.info("---" * 30 + "\nJSON files merged\n" + "---" * 30)

    def _merge_context_data(self, data_output_base_dir: str, context_merged_file_name: str):
        """Merge context_enhanced.json and context_final.jsonl files.
        
        Args:
            data_output_base_dir: Base directory containing the files to merge.
            context_merged_file_name: Name of the output merged file.
        """
        logger.info("---" * 30 + "\nMerging context files\n" + "---" * 30)

        result = []

        # Process context_enhanced.json
        try:
            with open(f"{data_output_base_dir}/context_enhanced.json", 'r', encoding='utf-8') as f:
                enhanced_data = json.load(f)

            for item in enhanced_data:
                try:
                    # Try to parse context_enhanced_need field
                    enhanced_need = json.loads(item["context_enhanced_need"])
                    result.append({
                        "user_request": item["initial_need"],
                        "enhanced_request": enhanced_need["enhanced_request"]
                    })
                except (json.JSONDecodeError, KeyError) as e:
                    logger.warning(f"Failed to process enhanced item: {str(e)}")
                    continue
        except Exception as e:
            logger.error(f"Error processing context_enhanced.json: {str(e)}")

        # Process context_final.jsonl
        try:
            with open(f"{data_output_base_dir}/context_final.jsonl", 'r', encoding='utf-8') as f:
                for line in f:
                    line = line.strip()
                    if not line:  # Skip empty lines
                        continue

                    try:
                        final_item = json.loads(line)
                        try:
                            # Try to parse response field
                            response = json.loads(final_item["response"])
                            result.append({
                                "user_request": final_item["initial_need"],
                                "expert_response": final_item["expert_response"],
                                "user_feedback": response.get("feedback", "") if isinstance(response['feedback'],
                                                                                            str) else
                                response['feedback'][0]
                            })
                        except (json.JSONDecodeError, KeyError) as e:
                            logger.warning(f"Failed to process final item response: {str(e)}")
                            continue
                    except json.JSONDecodeError as e:
                        logger.warning(f"Failed to parse line as JSON: {str(e)}")
                        continue
        except Exception as e:
            logger.error(f"Error processing context_final.jsonl: {str(e)}")

        # Save the merged results
        output_path = f"{data_output_base_dir}/{context_merged_file_name}"
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)

        logger.info(f"Merged context files saved to {output_path}")
        logger.info("---" * 30 + "\nContext files merged\n" + "---" * 30)

    def split_notes_by_type(self, note_list: List[Note], basic_info: Dict):
        """Split notes into subjective and objective categories.
        
        Args:
            note_list: List of Note objects to split.
            basic_info: Dictionary containing basic user information.
            
        Returns:
            Tuple of (user_info, subjective_notes, objective_notes).
        """
        user_info = {
            "username": basic_info["username"],
            "aboutMe": basic_info["aboutMe"],
            "statusBio": basic_info["statusBio"],
            "globalBio": basic_info["globalBio"],
            "lang": basic_info.get("lang", "English"),
        }

        subjective_notes = []
        objective_notes = []
        for note in note_list:
            if note.memory_type in SUBJECT_NOTE_TYPE:
                subjective_notes.append(note)
            elif note.memory_type in OBJECT_NOTE_TYPE:
                objective_notes.append(note)
            else:
                logger.warning(f"Note type not supported: {note.memory_type}")
                continue
        logger.info(
            f"Subjective notes: {len(subjective_notes)}, Objective notes: {len(objective_notes)}"
        )
        return user_info, subjective_notes, objective_notes

    def refine_notes_data_subjective(
            self, note_list: List[Note], user_info: Dict, json_file_remade: str
    ):
        """Refine subjective notes data and save to JSON.
        
        Args:
            note_list: List of Note objects to refine.
            user_info: Dictionary containing user information.
            json_file_remade: Path to save the refined JSON data.
            
        Returns:
            List of refined Note objects.
        """
        data_filtered = []

        lang = user_info.get("lang", "English")

        if lang not in SUBJECTIVE_TEMPLATES:
            lang = "English"

        selected_templates = SUBJECTIVE_TEMPLATES[lang]

        for note in note_list:
            if note.memory_type not in OBJECT_NOTE_TYPE:
                note.create_time = format_timestr(note.create_time)

                basic_template = random.choice(selected_templates["basic"]).format(user_name=user_info["username"])

                if note.insight:
                    # for markdown and doc
                    note.processed = basic_template + str(note.insight)
                else:
                    # for short text
                    note.processed = basic_template + str(note.content)

                if note.title:
                    note.processed += selected_templates["title_suffix"].format(title=note.title)

                data_filtered.append(note)

        logger.info(f"Refined subjective notes: {len(data_filtered)}")

        json_data_filted = [o.to_json() for o in data_filtered]
        file_dir = os.path.dirname(json_file_remade)
        if not os.path.exists(file_dir):
            os.makedirs(file_dir)
        with open(json_file_remade, "w", encoding="utf-8") as file:
            json.dump(json_data_filted, file, ensure_ascii=False, indent=4)

        return data_filtered

    def refine_notes_data_objective(self, note_list: List[Note], user_info: Dict, json_file_remade: str):
        """Refine objective notes data and save to JSON.
        
        Args:
            note_list: List of Note objects to refine.
            user_info: Dictionary containing user information.
            json_file_remade: Path to save the refined JSON data.
            
        Returns:
            List of refined Note objects.
        """
        lang = user_info.get("lang", "English")
        if lang not in OBJECTIVE_TEMPLATES:
            lang = "English"

        templates = OBJECTIVE_TEMPLATES[lang]

        new_item_list = []
        for note in note_list:
            if note.memory_type in OBJECT_NOTE_TYPE:
                if note.insight is None or note.insight == "":
                    continue
                new_item = note.copy()
                if note.content is not None and note.content != "":
                    new_item.processed = random.choice(
                        templates["with_content"]
                    ).format(content=note.content, insight=note.insight)
                else:
                    new_item.processed = random.choice(
                        templates["without_content"]
                    ).format(insight=note.insight)
                new_item_list.append(new_item)

        logger.info(f"Refined objective notes: {len(new_item_list)}")

        file_dir = os.path.dirname(json_file_remade)
        if not os.path.exists(file_dir):
            os.makedirs(file_dir)
        with open(json_file_remade, "w", encoding="utf-8") as f:
            json.dump(new_item_list, f, ensure_ascii=False, indent=4)

        return new_item_list

    def json_to_txt_each(
            self, list_processed_notes: List[Note], txt_file_base: str, file_type: str
    ):
        """Convert processed notes from JSON to individual text files.
        
        Args:
            list_processed_notes: List of processed Note objects.
            txt_file_base: Base directory to save text files.
            file_type: Type of note for naming the output files.
        """
        # Ensure the target directory exists
        if not os.path.exists(txt_file_base):
            os.makedirs(txt_file_base)
            logger.warning("Currently running in function json_to_txt_each")
            logger.warning(f"Specified directory does not exist, created: {txt_file_base}")

        # Clear all existing files in the target directory
        for existing_file in os.listdir(txt_file_base):
            file_path = os.path.join(txt_file_base, existing_file)
            if os.path.isfile(file_path):
                try:
                    os.remove(file_path)
                    logger.info(f"Removed existing file: {file_path}")
                except Exception as e:
                    logger.error(f"Error removing file {file_path}: {str(e)}")

        logger.info(f"Cleared all existing files in {txt_file_base}")

        for no, item in enumerate(tqdm(list_processed_notes)):
            # Build the txt file path
            txt_file = os.path.join(txt_file_base, f"{file_type}_{no}.txt")
            try:
                # Ensure the processed field exists in the item
                if item.processed:
                    with open(txt_file, "w", encoding="utf-8") as tf:
                        tf.write(item.processed)
                else:
                    logger.warning(f"Warning: 'processed' key missing for item {no}")

            except Exception as e:
                logger.error(traceback.format_exc())

    def graphrag_indexing(
            self, note_list: List[Note], graph_input_dir: str, output_dir: str, lang: str
    ):
        """Index notes using GraphRAG.
        
        This method configures and runs GraphRAG indexing on the processed notes,
        creating entity and relation extractions.
        
        Args:
            note_list: List of Note objects to index.
            graph_input_dir: Directory containing input files for indexing.
            output_dir: Directory to save indexing results.
            lang: Language for the prompts.
        """
        GRAPH_CONFIG = os.path.join(
            os.getcwd(), "lpm_kernel/L2/data_pipeline/graphrag_indexing/settings.yaml"
        )

        ENV_CONFIG = os.path.join(
            os.getcwd(), "lpm_kernel/L2/data_pipeline/graphrag_indexing/.env"
        )

        user_llm_config_service = UserLLMConfigService()
        user_llm_config = user_llm_config_service.get_available_llm()

        chat_api_key = user_llm_config.chat_api_key
        chat_base_url = user_llm_config.chat_endpoint
        chat_model_name = user_llm_config.chat_model_name

        embedding_api_key = user_llm_config.embedding_api_key
        embedding_base_url = user_llm_config.embedding_endpoint
        embedding_model_name = user_llm_config.embedding_model_name

        with open(GRAPH_CONFIG, "r", encoding="utf-8") as file:
            settings = yaml.safe_load(file)

        with open(ENV_CONFIG, "w", encoding="utf-8") as file:
            file.write(f"GRAPHRAG_API_KEY={chat_api_key}")

        settings["input"]["base_dir"] = graph_input_dir
        settings["output"]["base_dir"] = output_dir
        settings["reporting"]["base_dir"] = os.path.join(output_dir, "../report")

        settings["models"]["default_chat_model"]["api_base"] = chat_base_url
        settings["models"]["default_chat_model"]["model"] = chat_model_name
        settings["models"]["default_chat_model"]["api_key"] = chat_api_key

        if chat_model_name.startswith("openai"):
            settings["models"]["default_chat_model"]["model"] = chat_model_name.replace("openai/", "")

        if embedding_model_name.startswith("openai"):
            settings["models"]["default_embedding_model"]["model"] = embedding_model_name.replace("openai/", "")
        else:
            settings["models"]["default_embedding_model"]["model"] = embedding_model_name

        settings["models"]["default_embedding_model"]["api_base"] = embedding_base_url
        settings["models"]["default_embedding_model"]["api_key"] = embedding_api_key

        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            logger.warning(f"Specified output directory does not exist, created: {output_dir}.")

        with open(GRAPH_CONFIG, "w", encoding="utf-8") as file:
            yaml.dump(settings, file, default_flow_style=False, allow_unicode=True)

        logger.info(f"Input base_dir has been updated to {graph_input_dir} and saved.")
        logger.info(f"Output base_dir has been updated to {output_dir} and saved.")
        logger.info(
            f"Report base_dir has been updated to {os.path.join(output_dir, 'report')} and saved."
        )

        # Read prompts configuration and modify entity_extraction/summarize_descriptions from "in Chinese" to "in {lang}"
        entity_extraction_path = os.path.join(
            os.getcwd(),
            "lpm_kernel/L2/data_pipeline/graphrag_indexing/prompts/extract_graph.txt",
        )
        with open(entity_extraction_path, "r", encoding="utf-8") as f1:
            entity_extraction = f1.read()
            entity_extraction = entity_extraction.replace("<lang>", lang)
        with open(entity_extraction_path, "w", encoding="utf-8") as f2:
            f2.write(entity_extraction)

        summarize_descriptions_path = os.path.join(
            os.getcwd(),
            "lpm_kernel/L2/data_pipeline/graphrag_indexing/prompts/summarize_descriptions.txt",
        )
        with open(summarize_descriptions_path, "r", encoding="utf-8") as f1:
            summarize_descriptions = f1.read()
            summarize_descriptions = summarize_descriptions.replace("<lang>", lang)
        with open(summarize_descriptions_path, "w", encoding="utf-8") as f2:
            f2.write(summarize_descriptions)

        # Run GraphRAG indexing
        try:
            result = subprocess.run(
                [
                    "bash",
                    os.path.join(
                        os.getcwd(),
                        "lpm_kernel/L2/data_pipeline/data_prep/scripts/graphrag_indexing.sh",
                    ),
                ],
                check=True,
                text=True,
                capture_output=True,
            )
            if result.stderr:
                logger.error(f"subprocess.run graphrag index error: {result.stderr}")
                raise RuntimeError("subprocess.run graphrag index error")
        except Exception as e:
            raise

        """Post-processing"""

        self.creat_mapping(
            output_dir,
            note_list,
            os.path.join(
                os.getcwd(),
                "resources/L2/data_pipeline/raw_data/id_entity_mapping_subjective_v2.json",
            ),
        )

    def creat_mapping(self, graph_dir, note_list, mapped_json_file):
        """Create a mapping between entities and documents.
        
        Args:
            graph_dir: Directory containing GraphRAG output.
            note_list: List of Note objects.
            mapped_json_file: Path to save the mapping file.
        """
        try:
            document = pd.read_parquet(
                os.path.join(graph_dir, "documents.parquet")
            )
            entities = pd.read_parquet(
                os.path.join(graph_dir, "entities.parquet")
            )
        except Exception as e:
            return

        json_data = []

        # show the column names
        logger.info(f"Entity Column names: {entities.columns}")

        logger.info(f"Document Column names: {document.columns}")

        for e_i, e_r in tqdm(entities.iterrows(), total=len(entities)):
            json_item = {}
            json_item["entity_id"] = e_r["id"]
            json_item["entity_name"] = e_r["title"]
            json_item["entity_description"] = e_r["description"]
            json_item["doc_id"] = []
            text_unit_ids = e_r["text_unit_ids"]

            for text_unit_id in text_unit_ids:
                for d_i, d_r in document.iterrows():
                    if text_unit_id in d_r["text_unit_ids"]:
                        if "note" in d_r["title"]:
                            json_item["doc_id"].append(
                                note_list[
                                    int(
                                        d_r["title"]
                                        .replace(".txt", "")
                                        .replace("note_", "")
                                    )
                                ].id
                            )
            json_data.append(json_item)

        with open(os.path.join(mapped_json_file), "w", encoding="utf-8") as file:
            json.dump(json_data, file, ensure_ascii=False, indent=4)

    def _gen_preference_data(self, topics_path, preference_output_path, bio):
        """Generate preference data based on user topics and bio.
        
        Args:
            topics_path: Path to topics data file.
            preference_output_path: Path to save preference data.
            bio: User's bio information.
        """
        processor = PreferenceQAGenerator(
            filename=topics_path, bio=bio, preference_language=self.preferred_lang
        )
        processor.process_clusters(preference_output_path)

    def _gen_diversity_data(
            self,
            entitys_path,
            note_list: List[Note],
            graph_path,
            output_path,
            user_name,
            global_bio,
            config_path,
    ):
        """Generate diversity data based on entities and notes.
        
        Args:
            entitys_path: Path to entities data.
            note_list: List of Note objects.
            graph_path: Path to graph data.
            output_path: Path to save diversity data.
            user_name: Name of the user.
            global_bio: User's global biography.
            config_path: Path to configuration file.
        """
        processor = DiversityDataGenerator(self.preferred_lang)
        processor.generate_data(
            entitys_path, note_list, config_path, graph_path, user_name, global_bio, output_path
        )

    def _gen_selfqa_data(self, output_path, user_name, user_intro, bio):
        """Generate self Q&A data based on user information.
        
        Args:
            output_path: Path to save self Q&A data.
            user_name: Name of the user.
            user_intro: User's introduction.
            bio: User's bio information.
        """
        selfqa = SelfQA(
            user_name=user_name,
            user_input_introduction=user_intro,
            user_global_bio=bio,
            preferred_language=self.preferred_lang,
        )
        q_a_list = selfqa.generate_qa()
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(q_a_list, f, ensure_ascii=False, indent=4)

    def _gen_context_data(
            self,
            note_list: List[Note],
            entitys_path: str,
            data_output_base_dir: str,
            user_name: str,
            user_intro: str,
            global_bio: str
    ):
        """Generate context-based conversation data.
        
        This method generates contextual conversation data in multiple steps:
        1. Generate initial context needs
        2. Enhance context data
        3. Generate expert responses
        4. Generate critic data
        
        Args:
            note_list: List of Note objects.
            entitys_path: Path to entities data.
            data_output_base_dir: Base directory for output data.
            user_name: Name of the user.
            user_intro: User's introduction.
            global_bio: User's global biography.
        """
        context_generator = ContextGenerator(
            preferred_language=self.preferred_lang,
            user_name=user_name,
            user_bio=global_bio
        )

        # 1. Generate initial context needs
        context_generator.generate_context_needs(
            note_list=note_list,
            entity_map_path=os.path.join(entitys_path),
            data_output_base_dir=data_output_base_dir,
            needs_file_name="context_needs.json"
        )
        logger.info("---" * 30 + "\nContext needs generated\n" + "---" * 30)
        logger.info(data_output_base_dir + "/context_needs.json")

        # 2. Generate context enhanced data
        context_generator.generate_context_enhance_data(
            data_output_base_dir=data_output_base_dir,
            needs_file_name="context_needs.json",
            context_enhanced_res_file_name="context_enhanced.json",
            note_list=note_list
        )
        logger.info("---" * 30 + "\nContext enhanced generated\n" + "---" * 30)
        logger.info(data_output_base_dir + "/context_enhanced.json")

        # 3. Generate expert responses
        context_generator.expert_response_generator(
            data_output_base_dir=data_output_base_dir,
            context_enhanced_res_file_name="context_enhanced.json",
            output_file_name="expert_responses.json"
        )
        logger.info("---" * 30 + "\nExpert responses generated\n" + "---" * 30)
        logger.info(data_output_base_dir + "/expert_responses.json")

        # 4. Generate expert responses and critic data
        context_generator.gen_context_critic_data(
            data_output_base_dir=data_output_base_dir,
            expert_response_file_name="expert_responses.json",
            out_file_name="context_final.jsonl"
        )
        logger.info("---" * 30 + "\nContext critic generated\n" + "---" * 30)
        logger.info(data_output_base_dir + "/context_final.jsonl")