Agentswarm/benchmark/export/__init__.py

"""Benchmark metrics 导出器（阶段2/3 数据采集落库）。

把每次真实 run 终态采集到的 `SwarmMetrics` 快照导出到外部留存，供后续【经验标定】
（O / BASE_COEFFICIENTS / S_gain）累积真实用户数据。

目标可组合，由 `BENCHMARK_EXPORT_TARGET`（逗号分隔）选择：
  - `none`（默认）  → NoopExporter：不导出、无依赖、无凭据（指标仍落 run.metadata Redis）。
  - `cosmos`        → CosmosExporter：写 Cosmos DB for NoSQL（结构化指标，便于查询/排行榜）。
  - `blob`          → BlobExporter：写 Azure Blob 存储账户（全量快照归档）。
例：`BENCHMARK_EXPORT_TARGET=cosmos,blob`。

**凭据一律经环境/Secret 注入，绝不写进代码、日志或提交**（组织安全规则）：
  Cosmos: BENCHMARK_COSMOS_CONNECTION_STRING（必填）
          BENCHMARK_COSMOS_DATABASE（默认 benchmark）/ BENCHMARK_COSMOS_CONTAINER（默认 selfcert）
  Blob:   AZURE_STORAGE_CONNECTION_STRING（或 BENCHMARK_BLOB_ACCOUNT_URL + Managed/Workload Identity）
          BENCHMARK_BLOB_CONTAINER（默认 benchmark-selfcert）
"""
from __future__ import annotations

import json
import logging
import os
from typing import List, Optional

logger = logging.getLogger(__name__)


class MetricsExporter:
    """导出一次 run 的 benchmark 指标快照。子类实现 export()。构造不连网（懒连接）。"""

    name = "abstract"

    def export(self, swarm_id: str, payload: dict) -> None:
        raise NotImplementedError


class NoopExporter(MetricsExporter):
    """默认：不导出（指标已落 run.metadata）。无外部依赖、无凭据。"""

    name = "noop"

    def export(self, swarm_id: str, payload: dict) -> None:
        logger.debug("benchmark export disabled; run %s metrics persisted to run.metadata only", swarm_id)


class CosmosExporter(MetricsExporter):
    """Azure Cosmos DB for NoSQL（SQL API）。一 run 一文档，id=swarm_id，分区键 /swarm_id。"""

    name = "cosmos"

    def __init__(self, connection_string: str, database: str, container: str):
        if not connection_string:
            raise RuntimeError("CosmosExporter 需 BENCHMARK_COSMOS_CONNECTION_STRING")
        self.connection_string = connection_string
        self.database = database
        self.container = container

    @staticmethod
    def build_doc(payload: dict) -> dict:
        """Cosmos 文档：需要字符串 id；分区键字段 swarm_id 已在 payload 内。"""
        return {"id": str(payload.get("swarm_id")), **payload}

    def _container_client(self):
        from azure.cosmos import CosmosClient  # lazy: keep core import-light & key-free
        client = CosmosClient.from_connection_string(self.connection_string)
        db = client.get_database_client(self.database)
        return db.get_container_client(self.container)

    def export(self, swarm_id: str, payload: dict) -> None:
        self._container_client().upsert_item(self.build_doc(payload))
        logger.info("benchmark metrics → cosmos %s/%s (id=%s)", self.database, self.container, swarm_id)


class BlobExporter(MetricsExporter):
    """Azure Blob 存储账户（懒加载 SDK）。归档路径 <scenario>/<swarm_id>.json。"""

    name = "blob"

    def __init__(self, container: str, *, connection_string: Optional[str] = None,
                 account_url: Optional[str] = None):
        if not (connection_string or account_url):
            raise RuntimeError("BlobExporter 需 AZURE_STORAGE_CONNECTION_STRING 或 BENCHMARK_BLOB_ACCOUNT_URL")
        self.container = container
        self.connection_string = connection_string
        self.account_url = account_url

    def _client(self):
        from azure.storage.blob import BlobServiceClient  # lazy
        if self.connection_string:
            return BlobServiceClient.from_connection_string(self.connection_string)
        from azure.identity import DefaultAzureCredential  # Workload/Managed Identity, no secret
        return BlobServiceClient(account_url=self.account_url, credential=DefaultAzureCredential())

    def export(self, swarm_id: str, payload: dict) -> None:
        scenario = payload.get("scenario") or "unknown"
        blob_name = f"{scenario}/{swarm_id}.json"
        client = self._client().get_blob_client(container=self.container, blob=blob_name)
        client.upload_blob(json.dumps(payload, ensure_ascii=False), overwrite=True)
        logger.info("benchmark metrics → blob %s/%s", self.container, blob_name)


def _build_target(target: str) -> Optional[MetricsExporter]:
    """Build one exporter from a target name; None (with a warning) if misconfigured/unknown."""
    if target in ("none", "off", ""):
        return NoopExporter()
    try:
        if target == "cosmos":
            return CosmosExporter(
                connection_string=os.getenv("BENCHMARK_COSMOS_CONNECTION_STRING", ""),
                database=os.getenv("BENCHMARK_COSMOS_DATABASE", "benchmark"),
                container=os.getenv("BENCHMARK_COSMOS_CONTAINER", "selfcert"),
            )
        if target == "blob":
            return BlobExporter(
                container=os.getenv("BENCHMARK_BLOB_CONTAINER", "benchmark-selfcert"),
                connection_string=os.getenv("AZURE_STORAGE_CONNECTION_STRING"),
                account_url=os.getenv("BENCHMARK_BLOB_ACCOUNT_URL"),
            )
    except Exception as exc:
        # Misconfigured target (e.g. missing creds) must not break the others or the run.
        logger.warning("benchmark export target '%s' disabled: %s", target, exc)
        return None
    logger.warning("unknown BENCHMARK_EXPORT_TARGET entry '%s'; ignored", target)
    return None


def get_exporters() -> List[MetricsExporter]:
    """Parse BENCHMARK_EXPORT_TARGET (comma list) → exporters. Default → cosmos,blob."""
    raw = (os.getenv("BENCHMARK_EXPORT_TARGET") or "cosmos,blob").strip().lower()
    targets = [t.strip() for t in raw.split(",") if t.strip()] or ["none"]
    exporters = [e for e in (_build_target(t) for t in targets) if e is not None]
    return exporters or [NoopExporter()]