group_fqcd_jr/app/service/offsite_nl2sql_adapter.py

"""场外基金专用 NL2SQL 适配器。"""

from importlib import import_module
from typing import cast

from sqlalchemy import create_engine
from sqlalchemy.engine import Engine

from app.core.config import Settings, get_settings
from app.core.contracts import RequestContext


class OffsiteNl2SqlAdapter:
    """把场外基金上下文转换为 query_dict 的稳定参数并执行只读查询。"""

    script_path = "nl2sql_yc.py"

    def __init__(
        self,
        db_engine: Engine | None = None,
        settings: Settings | None = None,
    ) -> None:
        self._db_engine = db_engine
        self._settings = settings or get_settings()
        self._owns_engine = db_engine is None

    def query(self, question: str, context: RequestContext) -> dict[str, object]:
        try:
            query_dict = getattr(import_module("nl2sql_yc"), "query_dict", None)
            if not callable(query_dict):
                return {"status": "error", "message": "NL2SQL入口query_dict不存在"}
            db_engine = self._get_db_engine()
            result = cast(dict[str, object], query_dict(
                question,
                self._auth_context(context),
                db_engine=db_engine,
                use_llm=False,
                persist_audit=False,
            ))
            if isinstance(result, dict):
                return result
            return {"status": "error", "message": "NL2SQL返回格式不正确"}
        except (OSError, RuntimeError, ValueError) as exc:
            return {"status": "error", "message": f"NL2SQL调用失败：{type(exc).__name__}"}

    def close(self) -> None:
        """释放适配器自行创建的同步连接池。"""
        if self._owns_engine and self._db_engine is not None:
            self._db_engine.dispose()
            self._db_engine = None

    def _get_db_engine(self) -> Engine:
        if self._db_engine is not None:
            return self._db_engine
        dsn = self._settings.mysql_dsn
        if dsn.startswith("mysql+asyncmy://"):
            dsn = dsn.replace("mysql+asyncmy://", "mysql+pymysql://", 1)
        elif dsn.startswith("mysql://"):
            dsn = dsn.replace("mysql://", "mysql+pymysql://", 1)
        if not dsn:
            raise RuntimeError("未配置数据库连接")
        self._db_engine = create_engine(dsn, pool_pre_ping=True)
        return self._db_engine

    @staticmethod
    def _auth_context(context: RequestContext) -> dict[str, object]:
        return {
            "user_id": int(context.user_id) if context.user_id.isdigit() else None,
            "roles": list(context.roles),
            "allowed_domains": ["market_nav", "trading_account", "product_fee"],
            "customer_scope": context.data_scope if context.data_scope in {
                "self", "own_customers", "all",
            } else "all",
            "max_rows": 100,
            "max_query_seconds": 10,
        }