mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-06 22:12:12 +02:00
feat(capabilities): bill workspace owner at executor
This commit is contained in:
parent
3e6007431d
commit
bfe3117302
1 changed files with 58 additions and 0 deletions
58
surfsense_backend/app/capabilities/billing.py
Normal file
58
surfsense_backend/app/capabilities/billing.py
Normal file
|
|
@ -0,0 +1,58 @@
|
||||||
|
"""Charge the workspace owner per billable success at the capability executor (03c)."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
from sqlalchemy import select
|
||||||
|
|
||||||
|
from app.capabilities.types import BillableOutput, BillingUnit, CapabilityContext
|
||||||
|
from app.services.token_tracking_service import record_token_usage
|
||||||
|
from app.services.web_crawl_credit_service import WebCrawlCreditService
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from uuid import UUID
|
||||||
|
|
||||||
|
from sqlalchemy.ext.asyncio import AsyncSession
|
||||||
|
|
||||||
|
|
||||||
|
async def charge_capability(
|
||||||
|
output: BillableOutput, unit: BillingUnit, ctx: CapabilityContext
|
||||||
|
) -> None:
|
||||||
|
"""Bill the workspace owner for this result's billable successes (03c)."""
|
||||||
|
units = output.billable_units
|
||||||
|
if units <= 0:
|
||||||
|
return
|
||||||
|
if unit is BillingUnit.WEB_CRAWL:
|
||||||
|
await _charge_web_crawl(ctx, units)
|
||||||
|
|
||||||
|
|
||||||
|
async def _charge_web_crawl(ctx: CapabilityContext, successes: int) -> None:
|
||||||
|
service = WebCrawlCreditService(ctx.session)
|
||||||
|
if not service.billing_enabled():
|
||||||
|
return
|
||||||
|
owner_user_id = await _resolve_workspace_owner(ctx.session, ctx.workspace_id)
|
||||||
|
if owner_user_id is None:
|
||||||
|
return
|
||||||
|
# Stage the audit row before charge_credits' commit flushes both.
|
||||||
|
await record_token_usage(
|
||||||
|
ctx.session,
|
||||||
|
usage_type="web_crawl",
|
||||||
|
workspace_id=ctx.workspace_id,
|
||||||
|
user_id=owner_user_id,
|
||||||
|
cost_micros=service.successes_to_micros(successes),
|
||||||
|
call_details={"successes": successes},
|
||||||
|
)
|
||||||
|
await service.charge_credits(owner_user_id, successes)
|
||||||
|
|
||||||
|
|
||||||
|
async def _resolve_workspace_owner(
|
||||||
|
session: AsyncSession, workspace_id: int
|
||||||
|
) -> UUID | None:
|
||||||
|
"""The ``user_id`` that owns ``workspace_id`` (the crawl payer, not the caller)."""
|
||||||
|
from app.db import Workspace
|
||||||
|
|
||||||
|
result = await session.execute(
|
||||||
|
select(Workspace.user_id).where(Workspace.id == workspace_id)
|
||||||
|
)
|
||||||
|
return result.scalar_one_or_none()
|
||||||
Loading…
Add table
Add a link
Reference in a new issue