Mutual_Fund/nl2sql/contracts.py

"""NL2SQL 服务与其他 Agent 共享的公共契约。"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


@dataclass(frozen=True)
class DataQueryRequest:
    question: str
    user_id: int
    trace_id: str
    session_id: str | None = None
    caller_agent: str | None = None
    data_scope: dict[str, Any] | None = None
    max_rows: int | None = None
    include_sql: bool = False
    page: int = 1
    page_size: int = 100
    sort_by: str | None = None
    sort_order: str = "asc"

    def __post_init__(self) -> None:
        if not isinstance(self.question, str) or not self.question.strip():
            raise ValueError("question must not be empty")
        if self.user_id <= 0:
            raise ValueError("user_id must be positive")
        if not self.trace_id.strip():
            raise ValueError("trace_id must not be empty")
        if self.max_rows is not None and self.max_rows <= 0:
            raise ValueError("max_rows must be positive")
        if self.page < 1 or self.page_size < 1:
            raise ValueError("page and page_size must be positive")
        if self.sort_order not in {"asc", "desc"}:
            raise ValueError("sort_order must be asc or desc")


@dataclass(frozen=True)
class DataQueryResult:
    query_id: str
    trace_id: str = ""
    columns: list[str] = field(default_factory=list)
    rows: list[dict[str, Any]] = field(default_factory=list)
    row_count: int = 0
    truncated: bool = False
    summary: str | None = None
    answer: str | None = None
    markdown: str | None = None
    chart: dict[str, Any] | None = None
    metric_definitions: list[dict[str, Any]] = field(default_factory=list)
    query_plan: dict[str, Any] | None = None
    sql: str | None = None
    elapsed_ms: float = 0.0
    warnings: list[str] = field(default_factory=list)