from __future__ import annotations import uuid from datetime import datetime, timedelta from decimal import Decimal from typing import Any from sqlalchemy.orm import Session from app.models.agent_run import AgentRun, AgentToolCall from app.schemas.commercial import ( CommercialEntitlementUpsert, CommercialPlanCreate, CommercialSubscriptionCreate, ) from app.services.commercial_admin import CommercialAdminService def seed_meter( db: Session, tenant_id: str, now: datetime, *, basis: str, tool_type: str = "llm", tool_name: str = "chat.completions", hard_limit: Decimal = Decimal("1000"), internal_cost: dict[str, Any] | None = None, subscription_status: str = "active", preflight_quantity: Decimal | None = None, ): admin = CommercialAdminService(db) plan = admin.create_plan( tenant_id, CommercialPlanCreate( plan_code="runtime", name="运行时计量测试套餐", pricing_model="hybrid", billing_interval="monthly", currency="CNY", base_fee=Decimal("100"), included_seats=10, effective_from=now - timedelta(days=30), ), actor_id="platform-admin", ) admin.activate_plan(tenant_id, plan.id, expected_version=plan.version) subscription = admin.create_subscription( tenant_id, CommercialSubscriptionCreate( subscription_key=f"{tenant_id}-runtime", plan_id=plan.id, status=subscription_status, starts_at=now - timedelta(days=10), current_period_start=now - timedelta(days=1), current_period_end=now + timedelta(days=29), seats=5, ), actor_id="platform-admin", ) runtime_meter: dict[str, Any] = { "enabled": True, "tool_type": tool_type, "tool_name": tool_name, "quantity_basis": basis, } if preflight_quantity is not None: runtime_meter["preflight_quantity"] = str(preflight_quantity) if internal_cost is not None: runtime_meter["internal_cost"] = internal_cost entitlement = admin.upsert_entitlement( tenant_id, CommercialEntitlementUpsert( subscription_id=subscription.id, entitlement_key=f"runtime_{basis}", metric_key=f"runtime_{basis}", entitlement_type="metered", unit="call" if basis == "call" else basis, included_quantity=hard_limit, hard_limit_quantity=hard_limit, reset_interval="monthly", overage_policy="block", status="active", effective_from=now - timedelta(days=10), config_json={"runtime_meter": runtime_meter}, ), ) db.flush() return subscription, entitlement def seed_tool_call( db: Session, occurred_at: datetime, *, route_json: dict[str, Any], request_json: dict[str, Any] | None = None, response_json: dict[str, Any] | None = None, duration_ms: int = 0, tool_type: str = "llm", tool_name: str = "chat.completions", user_id: str = "sensitive-user-id", run: AgentRun | None = None, status: str = "succeeded", ) -> tuple[AgentRun, AgentToolCall]: if run is None: run = AgentRun( run_id=f"run-{uuid.uuid4().hex}", agent="user_agent", source="chat", user_id=user_id, route_json=route_json, permission_level="write", status="succeeded", started_at=occurred_at, ) db.add(run) tool_call = AgentToolCall( id=str(uuid.uuid4()), run_id=run.run_id, tool_type=tool_type, tool_name=tool_name, request_json=request_json or {}, response_json=response_json or {}, status=status, duration_ms=duration_ms, created_at=occurred_at, ) db.add(tool_call) db.flush() return run, tool_call def seed_run( db: Session, started_at: datetime, *, route_json: dict[str, Any], status: str = "running", ) -> AgentRun: run = AgentRun( run_id=f"run-{uuid.uuid4().hex}", agent="user_agent", source="chat", user_id="runtime-user", route_json=route_json, permission_level="write", status=status, started_at=started_at, ) db.add(run) db.commit() return run