Tolan/backend/services/llm_task_service.py

"""
A/B 组 LLM 任务表服务：读取最新任务 prompt / 首句
"""
import logging
import time
from typing import Any, Dict, Optional, TypedDict

from sqlalchemy import desc, func, select

from database import db_manager
from models_db import LlmTask

logger = logging.getLogger(__name__)

_CACHE_TTL_SECONDS = 60.0
_cached_task: Optional[Dict[str, Any]] = None
_cached_at: float = 0.0


class LlmTaskPayload(TypedDict):
    task_name: str
    task_prompt: str
    lumi_first_line: str


class LlmTaskService:
    async def get_latest_task(self) -> Optional[LlmTaskPayload]:
        """取 created_at 最新的一条任务。"""
        try:
            async with db_manager.get_session() as session:
                result = await session.execute(
                    select(LlmTask).order_by(desc(LlmTask.created_at), desc(LlmTask.id)).limit(1)
                )
                row = result.scalar_one_or_none()
            if not row:
                return None
            return {
                "task_name": row.task_name,
                "task_prompt": row.task_prompt,
                "lumi_first_line": row.lumi_first_line,
            }
        except Exception as e:
            logger.warning("[LlmTask] get_latest_task failed: %s", e)
            return None

    async def get_latest_task_cached(self, force_refresh: bool = False) -> Optional[LlmTaskPayload]:
        global _cached_task, _cached_at
        now = time.time()
        if (
            not force_refresh
            and _cached_task is not None
            and (now - _cached_at) < _CACHE_TTL_SECONDS
        ):
            return _cached_task

        task = await self.get_latest_task()
        if task:
            _cached_task = task
            _cached_at = now
        return task

    async def list_tasks(self, page: int = 1, page_size: int = 20) -> Dict[str, Any]:
        page = max(1, page)
        page_size = max(1, min(page_size, 100))
        offset = (page - 1) * page_size
        try:
            async with db_manager.get_session() as session:
                total = int(
                    (await session.execute(select(func.count()).select_from(LlmTask))).scalar_one()
                )
                result = await session.execute(
                    select(LlmTask)
                    .order_by(desc(LlmTask.created_at), desc(LlmTask.id))
                    .offset(offset)
                    .limit(page_size)
                )
                rows = result.scalars().all()
            items = [
                self._row_to_dict(row)
                for row in rows
            ]
            return {"items": items, "total": total, "page": page, "page_size": page_size}
        except Exception as e:
            logger.warning("[LlmTask] list_tasks failed: %s", e)
            raise

    async def create_task(
        self,
        task_name: str,
        task_prompt: str,
        lumi_first_line: str,
    ) -> Dict[str, Any]:
        task_name = task_name.strip()
        task_prompt = task_prompt.strip()
        lumi_first_line = lumi_first_line.strip()
        if not task_name or not task_prompt or not lumi_first_line:
            raise ValueError("task_name, task_prompt, lumi_first_line 均不能为空")
        try:
            async with db_manager.get_session() as session:
                row = LlmTask(
                    task_name=task_name,
                    task_prompt=task_prompt,
                    lumi_first_line=lumi_first_line,
                )
                session.add(row)
                await session.flush()
                await session.refresh(row)
                payload = self._row_to_dict(row)
                await session.commit()
            self.invalidate_cache()
            return payload
        except Exception as e:
            logger.warning("[LlmTask] create_task failed: %s", e)
            raise

    async def get_task_by_id(self, task_id: int) -> Optional[Dict[str, Any]]:
        try:
            async with db_manager.get_session() as session:
                row = await session.get(LlmTask, task_id)
            if not row:
                return None
            return self._row_to_dict(row)
        except Exception as e:
            logger.warning("[LlmTask] get_task_by_id failed: %s", e)
            raise

    async def update_task(
        self,
        task_id: int,
        task_name: str,
        task_prompt: str,
        lumi_first_line: str,
    ) -> Optional[Dict[str, Any]]:
        task_name = task_name.strip()
        task_prompt = task_prompt.strip()
        lumi_first_line = lumi_first_line.strip()
        if not task_name or not task_prompt or not lumi_first_line:
            raise ValueError("task_name, task_prompt, lumi_first_line 均不能为空")
        try:
            async with db_manager.get_session() as session:
                row = await session.get(LlmTask, task_id)
                if not row:
                    return None
                row.task_name = task_name
                row.task_prompt = task_prompt
                row.lumi_first_line = lumi_first_line
                await session.flush()
                await session.refresh(row)
                payload = self._row_to_dict(row)
                await session.commit()
            self.invalidate_cache()
            return payload
        except Exception as e:
            logger.warning("[LlmTask] update_task failed: %s", e)
            raise

    async def delete_task(self, task_id: int) -> bool:
        try:
            async with db_manager.get_session() as session:
                row = await session.get(LlmTask, task_id)
                if not row:
                    return False
                await session.delete(row)
                await session.commit()
            self.invalidate_cache()
            return True
        except Exception as e:
            logger.warning("[LlmTask] delete_task failed: %s", e)
            raise

    @staticmethod
    def _row_to_dict(row: LlmTask) -> Dict[str, Any]:
        return {
            "id": row.id,
            "task_name": row.task_name,
            "task_prompt": row.task_prompt,
            "lumi_first_line": row.lumi_first_line,
            "created_at": row.created_at.isoformat() if row.created_at else None,
            "updated_at": row.updated_at.isoformat() if row.updated_at else None,
        }

    def invalidate_cache(self) -> None:
        global _cached_task, _cached_at
        _cached_task = None
        _cached_at = 0.0


llm_task_service = LlmTaskService()