DataMate/runtime/datamate-python/app/module/generation/service/generation_service.py

import asyncio
import json
import uuid

from langchain_core.language_models import BaseChatModel
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.models.data_synthesis import (
    DataSynthInstance,
    DataSynthesisFileInstance,
    DataSynthesisChunkInstance,
    SynthesisData,
)
from app.db.models.dataset_management import DatasetFiles
from app.db.session import logger
from app.module.generation.schema.generation import Config, SyntheConfig
from app.module.generation.service.prompt import (
    QUESTION_GENERATOR_PROMPT,
    ANSWER_GENERATOR_PROMPT,
)
from app.module.shared.common.document_loaders import load_documents
from app.module.shared.common.text_split import DocumentSplitter
from app.module.shared.util.model_chat import extract_json_substring
from app.module.system.service.common_service import chat, get_model_by_id, get_chat_client


class GenerationService:
    def __init__(self, db: AsyncSession):
        self.db = db
        # 全局并发信号量：保证任意时刻最多 10 次模型调用
        self.question_semaphore = asyncio.Semaphore(10)
        self.answer_semaphore = asyncio.Semaphore(100)

    async def process_task(self, task_id: str):
        """处理数据合成任务入口：根据任务ID加载任务并逐个处理源文件。"""
        synth_task: DataSynthInstance | None = await self.db.get(DataSynthInstance, task_id)
        if not synth_task:
            logger.error(f"Synthesis task {task_id} not found, abort processing")
            return

        logger.info(f"Start processing synthe task {task_id}")

        # 从 synth_config 中读取 max_qa_pairs，全局控制 QA 总量上限；<=0 或异常则视为不限制
        try:
            cfg = Config(**(synth_task.synth_config or {}))
            max_qa_pairs = cfg.max_qa_pairs if (cfg and cfg.max_qa_pairs and cfg.max_qa_pairs > 0) else None
        except Exception:
            max_qa_pairs = None

        # 获取任务关联的文件原始ID列表
        file_ids = await self._get_file_ids_for_task(task_id)
        if not file_ids:
            logger.warning(f"No files associated with task {task_id}, abort processing")
            return

        # 逐个文件处理
        for file_id in file_ids:
            try:
                success = await self._process_single_file(synth_task, file_id, max_qa_pairs=max_qa_pairs)
            except Exception as e:
                logger.exception(f"Unexpected error when processing file {file_id} for task {task_id}: {e}")
                # 确保对应文件任务状态标记为失败
                await self._mark_file_failed(str(synth_task.id), file_id, str(e))
                success = False

            if success:
                # 每处理完一个文件，简单增加 processed_files 计数
                synth_task.processed_files = (synth_task.processed_files or 0) + 1
                await self.db.commit()
                await self.db.refresh(synth_task)

        logger.info(f"Finished processing synthesis task {synth_task.id}")

    # ==================== 高层文件处理流程 ====================
    async def _process_single_file(
        self,
        synth_task: DataSynthInstance,
        file_id: str,
        max_qa_pairs: int | None = None,
    ) -> bool:
        """按 chunk 批量流式处理单个源文件。

        流程：
        1. 切片并将所有 chunk 持久化到 DB 后释放内存；
        2. 从 DB 按 chunk_index 升序批量读取 chunk；
        3. 对批次中的每个 chunk：先生成指定数量的问题，再基于这些问题生成答案；
        4. 每成功处理完一个 chunk（即该 chunk 至少生成一条 QA）就更新一次 processed_chunks；
        5. 全部完成后将文件实例标记为 completed。
        """
        # 解析文件路径与配置
        file_path = await self._resolve_file_path(file_id)
        if not file_path:
            logger.warning(f"File path not found for file_id={file_id}, skip")
            await self._mark_file_failed(str(synth_task.id), file_id, "file_path_not_found")
            return False

        logger.info(f"Processing file_id={file_id}, path={file_path}")

        try:
            config = Config(**(synth_task.synth_config or {}))
        except Exception as e:
            logger.error(f"Invalid synth_config for task={synth_task.id}: {e}")
            await self._mark_file_failed(str(synth_task.id), file_id, "invalid_synth_config")
            return False

        # 1. 加载并切片（仅在此处占用内存）
        chunks = self._load_and_split(
            file_path,
            config.text_split_config.chunk_size,
            config.text_split_config.chunk_overlap,
        )
        if not chunks:
            logger.warning(f"No chunks generated for file_id={file_id}")
            await self._mark_file_failed(str(synth_task.id), file_id, "no_chunks_generated")
            return False

        logger.info(f"File {file_id} split into {len(chunks)} chunks by LangChain")

        # 2. 获取文件实例并持久化 chunk 记录
        file_task = await self._get_or_create_file_instance(
            synthesis_task_id=str(synth_task.id),
            source_file_id=file_id,
        )
        if not file_task:
            logger.error(
                f"DataSynthesisFileInstance not found for task={synth_task.id}, file_id={file_id}"
            )
            await self._mark_file_failed(str(synth_task.id), file_id, "file_instance_not_found")
            return False

        await self._persist_chunks(synth_task, file_task, file_id, chunks)
        total_chunks = len(chunks)
        # 释放内存中的切片
        del chunks

        # 3. 读取问答配置
        question_cfg: SyntheConfig | None = config.question_synth_config
        answer_cfg: SyntheConfig | None = config.answer_synth_config
        if not question_cfg or not answer_cfg:
            logger.error(
                f"Question/Answer synth config missing for task={synth_task.id}, file={file_id}"
            )
            await self._mark_file_failed(str(synth_task.id), file_id, "qa_config_missing")
            return False

        logger.info(
            f"Start QA generation for task={synth_task.id}, file={file_id}, total_chunks={total_chunks}"
        )

        # 为本文件构建模型 client
        question_model = await get_model_by_id(self.db, question_cfg.model_id)
        answer_model = await get_model_by_id(self.db, answer_cfg.model_id)
        question_chat = get_chat_client(question_model)
        answer_chat = get_chat_client(answer_model)

        # 分批次从 DB 读取并处理 chunk
        batch_size = 20
        current_index = 1

        while current_index <= total_chunks:
            end_index = min(current_index + batch_size - 1, total_chunks)
            chunk_batch = await self._load_chunk_batch(
                file_task_id=file_task.id,
                start_index=current_index,
                end_index=end_index,
            )
            if not chunk_batch:
                logger.warning(
                    f"Empty chunk batch loaded for file={file_id}, range=[{current_index}, {end_index}]"
                )
                current_index = end_index + 1
                continue

            # 对本批中的每个 chunk 并发处理（内部受 semaphore 限流）
            async def process_one(chunk: DataSynthesisChunkInstance) -> bool:
                return await self._process_single_chunk_qa(
                    file_task=file_task,
                    chunk=chunk,
                    question_cfg=question_cfg,
                    answer_cfg=answer_cfg,
                    question_chat=question_chat,
                    answer_chat=answer_chat,
                    synth_task_id=str(synth_task.id),
                    max_qa_pairs=max_qa_pairs,
                )

            tasks = [process_one(chunk) for chunk in chunk_batch]
            await asyncio.gather(*tasks, return_exceptions=True)

            current_index = end_index + 1

        # 全部完成
        file_task.status = "completed"
        await self.db.commit()
        await self.db.refresh(file_task)
        return True

    async def _process_single_chunk_qa(
        self,
        file_task: DataSynthesisFileInstance,
        chunk: DataSynthesisChunkInstance,
        question_cfg: SyntheConfig,
        answer_cfg: SyntheConfig,
        question_chat: BaseChatModel,
        answer_chat: BaseChatModel,
        synth_task_id: str,
        max_qa_pairs: int | None = None,
    ) -> bool:
        """处理单个 chunk：生成问题列表，然后为每个问题生成答案并落库。

        为了全局控制 QA 总量：在本方法开始处，根据 synth_task_id 查询当前已落盘的
        SynthesisData 条数，如果 >= max_qa_pairs，则不再对当前 chunk 做任何 QA 生成，
        并将当前文件任务标记为 completed，processed_chunks = total_chunks。

        已经进入后续流程的任务（例如其它协程正在生成答案）允许自然执行完。
        """
        # 如果没有全局上限配置，维持原有行为
        if max_qa_pairs is not None and max_qa_pairs > 0:
            from sqlalchemy import func

            # 统计当前整个任务下已生成的 QA 总数
            result = await self.db.execute(
                select(func.count(SynthesisData.id)).where(
                    SynthesisData.synthesis_file_instance_id.in_(
                        select(DataSynthesisFileInstance.id).where(
                            DataSynthesisFileInstance.synthesis_instance_id == synth_task_id
                        )
                    )
                )
            )
            current_qa_count = int(result.scalar() or 0)

            if current_qa_count >= max_qa_pairs:
                logger.info(
                    "max_qa_pairs reached: current=%s, max=%s, task_id=%s, file_task_id=%s, skip new QA generation for this chunk.",
                    current_qa_count,
                    max_qa_pairs,
                    synth_task_id,
                    file_task.id,
                )
                # 将文件任务标记为已完成，并认为所有 chunk 均已处理
                file_task.status = "completed"
                if file_task.total_chunks is not None:
                    file_task.processed_chunks = file_task.total_chunks
                await self.db.commit()
                await self.db.refresh(file_task)
                return False

        # ---- 下面保持原有逻辑不变 ----
        chunk_index = chunk.chunk_index
        chunk_text = chunk.chunk_content or ""
        if not chunk_text.strip():
            logger.warning(
                f"Empty chunk text for file_task={file_task.id}, chunk_index={chunk_index}"
            )
            # 无论成功或失败，均视为该 chunk 已处理完成
            try:
                await self._increment_processed_chunks(file_task.id, 1)
            except Exception as e:
                logger.exception(
                    f"Failed to increment processed_chunks for file_task={file_task.id}, chunk_index={chunk_index}: {e}"
                )
            return False

        success_any = False

        # 1. 生成问题
        try:
            questions = await self._generate_questions_for_one_chunk(
                chunk_text=chunk_text,
                question_cfg=question_cfg,
                question_chat=question_chat,
            )
        except Exception as e:
            logger.error(
                f"Generate questions failed for file_task={file_task.id}, chunk_index={chunk_index}: {e}"
            )
            questions = []

        if not questions:
            logger.info(
                f"No questions generated for file_task={file_task.id}, chunk_index={chunk_index}"
            )
        else:
            # 2. 针对每个问题生成答案并入库
            qa_success = await self._generate_answers_for_one_chunk(
                file_task=file_task,
                chunk=chunk,
                questions=questions,
                answer_cfg=answer_cfg,
                answer_chat=answer_chat,
            )
            success_any = bool(qa_success)

        # 无论本 chunk 处理是否成功，都增加 processed_chunks 计数，避免任务长时间卡住
        try:
            await self._increment_processed_chunks(file_task.id, 1)
        except Exception as e:
            logger.exception(
                f"Failed to increment processed_chunks for file_task={file_task.id}, chunk_index={chunk_index}: {e}"
            )

        return success_any

    async def _generate_questions_for_one_chunk(
        self,
        chunk_text: str,
        question_cfg: SyntheConfig,
        question_chat: BaseChatModel,
    ) -> list[str]:
        """针对单个 chunk 文本，调用 question_chat 生成问题列表。"""
        number = question_cfg.number or 5
        number = number if number is not None else 5
        number = max(int(len(chunk_text) / 1000 * number), 1)
        template = getattr(question_cfg, "prompt_template", QUESTION_GENERATOR_PROMPT)
        template = template if (template is not None and template.strip() != "") else QUESTION_GENERATOR_PROMPT

        prompt = (
            template
            .replace("{text}", chunk_text)
            .replace("{number}", str(number))
            .replace("{textLength}", str(len(chunk_text)))
        )

        async with self.question_semaphore:
            loop = asyncio.get_running_loop()
            raw_answer = await loop.run_in_executor(
                None,
                chat,
                question_chat,
                prompt,
            )

        # 解析为问题列表
        questions = self._parse_questions_from_answer(
            raw_answer,
        )
        return questions

    async def _generate_answers_for_one_chunk(
        self,
        file_task: DataSynthesisFileInstance,
        chunk: DataSynthesisChunkInstance,
        questions: list[str],
        answer_cfg: SyntheConfig,
        answer_chat: BaseChatModel,
    ) -> bool:
        """为一个 chunk 的所有问题生成答案并写入 SynthesisData。

        返回：是否至少成功写入一条 QA。
        """
        if not questions:
            return False

        chunk_text = chunk.chunk_content or ""
        template = getattr(answer_cfg, "prompt_template", ANSWER_GENERATOR_PROMPT)
        template = template if (template is not None and template.strip() != "") else ANSWER_GENERATOR_PROMPT
        extra_vars = getattr(answer_cfg, "extra_prompt_vars", {}) or {}

        success_flags: list[bool] = []

        async def process_single_question(question: str):
            prompt = template.replace("{text}", chunk_text).replace("{question}", question)
            for k, v in extra_vars.items():
                prompt.replace(f"{{{{{k}}}}}", str(v))
            else:
                prompt_local = prompt

            async with self.answer_semaphore:
                loop = asyncio.get_running_loop()
                answer = await loop.run_in_executor(
                    None,
                    chat,
                    answer_chat,
                    prompt_local,
                )

            # 默认结构：与 ANSWER_GENERATOR_PROMPT 一致，并补充 instruction 字段
            base_obj: dict[str, object] = {
                "input": chunk_text,
                "output": answer,
            }

            # 如果模型已经按照 ANSWER_GENERATOR_PROMPT 返回了 JSON，则尝试解析并在其上增加 instruction
            parsed_obj: dict[str, object] | None = None
            if isinstance(answer, str):
                cleaned = extract_json_substring(answer)
                try:
                  parsed = json.loads(cleaned)
                  if isinstance(parsed, dict):
                      parsed_obj = parsed
                except Exception:
                  parsed_obj = None

            if parsed_obj is not None:
                parsed_obj["instruction"] = question
                data_obj = parsed_obj
            else:
                base_obj["instruction"] = question
                data_obj = base_obj

            record = SynthesisData(
                id=str(uuid.uuid4()),
                data=data_obj,
                synthesis_file_instance_id=file_task.id,
                chunk_instance_id=chunk.id,
            )
            self.db.add(record)
            success_flags.append(True)

        tasks = [process_single_question(q) for q in questions]
        await asyncio.gather(*tasks, return_exceptions=True)

        if success_flags:
            await self.db.commit()
            return True
        return False

    @staticmethod
    def _parse_questions_from_answer(
        raw_answer: str,
    ) -> list[str]:
        """从大模型返回中解析问题数组。"""
        if not raw_answer:
            return []

        cleaned = extract_json_substring(raw_answer)
        try:
            data = json.loads(cleaned)
        except Exception as e:
            logger.error(
                f"Failed to parse question list JSON for task: {e}. "
            )
            return []

        if isinstance(data, list):
            return [str(q) for q in data if isinstance(q, str) and q.strip()]
        # 容错：如果是单个字符串
        if isinstance(data, str) and data.strip():
            return [data.strip()]
        return []


    # ==================== 原有辅助方法（文件路径/切片/持久化等） ====================
    async def _resolve_file_path(self, file_id: str) -> str | None:
        """根据文件ID查询 t_dm_dataset_files 并返回 file_path（仅 ACTIVE 文件）。"""
        result = await self.db.execute(
            select(DatasetFiles).where(DatasetFiles.id == file_id)
        )
        file_obj = result.scalar_one_or_none()
        if not file_obj:
            return None
        return file_obj.file_path

    @staticmethod
    def _load_and_split(file_path: str, chunk_size: int, chunk_overlap: int):
        """使用 LangChain 加载文本并进行切片，直接返回 Document 列表。
        Args:
            file_path: 待切片的文件路径
            chunk_size: 切片大小
            chunk_overlap: 切片重叠大小
        """
        try:
            docs = load_documents(file_path)
            split_docs = DocumentSplitter.auto_split(docs, chunk_size, chunk_overlap)
            return split_docs
        except Exception as e:
            logger.error(f"Error loading or splitting file {file_path}: {e}")
            raise

    async def _persist_chunks(
        self,
        synthesis_task: DataSynthInstance,
        file_task: DataSynthesisFileInstance,
        file_id: str,
        chunks,
    ) -> None:
        """将切片结果保存到 t_data_synthesis_chunk_instances，并更新文件级分块计数。"""
        for idx, doc in enumerate(chunks, start=1):
            # 先复制原始 Document.metadata，再在其上追加任务相关字段，避免覆盖原有元数据
            base_metadata = dict(getattr(doc, "metadata", {}) or {})
            base_metadata.update(
                {
                    "task_id": str(synthesis_task.id),
                    "file_id": file_id
                }
            )

            chunk_record = DataSynthesisChunkInstance(
                id=str(uuid.uuid4()),
                synthesis_file_instance_id=file_task.id,
                chunk_index=idx,
                chunk_content=doc.page_content,
                chunk_metadata=base_metadata,
            )
            self.db.add(chunk_record)

        # 更新文件任务的分块数量
        file_task.total_chunks = len(chunks)
        file_task.status = "processing"

        await self.db.commit()
        await self.db.refresh(file_task)

    async def _get_or_create_file_instance(
        self,
        synthesis_task_id: str,
        source_file_id: str,
    ) -> DataSynthesisFileInstance:
        """根据任务ID和原始文件ID，查找或创建对应的 DataSynthesisFileInstance 记录。

        - 如果已存在（同一任务 + 同一 source_file_id），直接返回；
        - 如果不存在，则创建一条新的文件任务记录，file_name 来自文件路径，
          target_file_location 先复用任务的 result_data_location。
        """
        # 尝试查询已有文件任务记录
        result = await self.db.execute(
            select(DataSynthesisFileInstance).where(
                DataSynthesisFileInstance.synthesis_instance_id == synthesis_task_id,
                DataSynthesisFileInstance.source_file_id == source_file_id,
            )
        )
        file_task = result.scalar_one_or_none()
        return file_task

    async def _mark_file_failed(self, synth_task_id: str, file_id: str, reason: str | None = None) -> None:
        """将指定任务下的单个文件任务标记为失败状态，兜底错误处理。

        - 如果找到对应的 DataSynthesisFileInstance，则更新其 status="failed"。
        - 如果未找到，则静默返回，仅记录日志。
        - reason 参数仅用于日志记录，方便排查。
        """
        try:
            result = await self.db.execute(
                select(DataSynthesisFileInstance).where(
                    DataSynthesisFileInstance.synthesis_instance_id == synth_task_id,
                    DataSynthesisFileInstance.source_file_id == file_id,
                )
            )
            file_task = result.scalar_one_or_none()
            if not file_task:
                logger.warning(
                    f"Failed to mark file as failed: no DataSynthesisFileInstance found for task={synth_task_id}, file_id={file_id}, reason={reason}"
                )
                return

            file_task.status = "failed"
            await self.db.commit()
            await self.db.refresh(file_task)
            logger.info(
                f"Marked file task as failed for task={synth_task_id}, file_id={file_id}, reason={reason}"
            )
        except Exception as e:
            # 兜底日志，避免异常向外传播影响其它文件处理
            logger.exception(
                f"Unexpected error when marking file failed for task={synth_task_id}, file_id={file_id}, original_reason={reason}, error={e}"
            )

    async def _get_file_ids_for_task(self, synth_task_id: str):
        """根据任务ID查询关联的文件原始ID列表"""
        result = await self.db.execute(
            select(DataSynthesisFileInstance.source_file_id)
            .where(DataSynthesisFileInstance.synthesis_instance_id == synth_task_id)
        )
        file_ids = result.scalars().all()
        return file_ids

    # ========== 新增：chunk 计数与批量加载、processed_chunks 安全更新辅助方法 ==========
    async def _count_chunks_for_file(self, synth_file_instance_id: str) -> int:
        """统计指定任务与文件下的 chunk 总数。"""
        from sqlalchemy import func

        result = await self.db.execute(
            select(func.count(DataSynthesisChunkInstance.id)).where(
                DataSynthesisChunkInstance.synthesis_file_instance_id == synth_file_instance_id
            )
        )
        return int(result.scalar() or 0)

    async def _load_chunk_batch(
        self,
        file_task_id: str,
        start_index: int,
        end_index: int,
    ) -> list[DataSynthesisChunkInstance]:
        """按索引范围加载指定文件任务下的一批 chunk 记录（含边界）。"""
        result = await self.db.execute(
            select(DataSynthesisChunkInstance)
            .where(
                DataSynthesisChunkInstance.synthesis_file_instance_id == file_task_id,
                DataSynthesisChunkInstance.chunk_index >= start_index,
                DataSynthesisChunkInstance.chunk_index <= end_index,
            )
            .order_by(DataSynthesisChunkInstance.chunk_index.asc())
        )
        return list(result.scalars().all())

    async def _increment_processed_chunks(self, file_task_id: str, delta: int) -> None:
        result = await self.db.execute(
            select(DataSynthesisFileInstance).where(
                DataSynthesisFileInstance.id == file_task_id,
            )
        )
        file_task = result.scalar_one_or_none()
        if not file_task:
            logger.error(f"Failed to increment processed_chunks: file_task {file_task_id} not found")
            return

        # 原始自增
        new_value = (file_task.processed_chunks or 0) + int(delta)

        # 如果存在 total_chunks，上限为 total_chunks，避免超过
        total = file_task.total_chunks
        if isinstance(total, int) and total >= 0:
            new_value = min(new_value, total)

        file_task.processed_chunks = new_value
        await self.db.commit()
        await self.db.refresh(file_task)