mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-08 22:22:17 +02:00
132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
"""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 (
|
|
BillableInput,
|
|
BillableOutput,
|
|
BillingUnit,
|
|
CapabilityContext,
|
|
)
|
|
from app.config import config
|
|
from app.services.token_tracking_service import record_token_usage
|
|
from app.services.web_crawl_credit_service import WebCrawlCreditService
|
|
from app.utils.captcha import captcha_enabled
|
|
|
|
if TYPE_CHECKING:
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
|
|
async def gate_capability(
|
|
payload: BillableInput, unit: BillingUnit | None, ctx: CapabilityContext
|
|
) -> None:
|
|
"""Pre-flight: block an over-budget owner before the executor runs (03c).
|
|
|
|
Raises ``InsufficientCreditsError`` when the wallet can't cover the input's
|
|
worst-case ``estimated_units``. ``None`` unit = free = no gate.
|
|
"""
|
|
if unit is None:
|
|
return
|
|
if unit is BillingUnit.WEB_CRAWL:
|
|
await _gate_web_crawl(ctx, payload.estimated_units)
|
|
|
|
|
|
async def _gate_web_crawl(ctx: CapabilityContext, estimated_successes: int) -> None:
|
|
"""Reserve the worst-case cost: crawl successes + worst-case captcha attempts.
|
|
|
|
Captcha budget is only reserved when solving is actually enabled — with
|
|
solving off, attempts can never happen, so reserving would wrongly block a
|
|
run for captcha that will never be attempted. Mirrors the indexer path (3d).
|
|
"""
|
|
service = WebCrawlCreditService(ctx.session)
|
|
crawl_on = service.billing_enabled()
|
|
captcha_on = service.captcha_billing_enabled() and captcha_enabled()
|
|
if not crawl_on and not captcha_on:
|
|
return
|
|
owner_user_id = await _resolve_workspace_owner(ctx.session, ctx.workspace_id)
|
|
if owner_user_id is None:
|
|
return
|
|
|
|
required_micros = 0
|
|
if crawl_on:
|
|
required_micros += service.successes_to_micros(estimated_successes)
|
|
if captcha_on:
|
|
worst_case_attempts = estimated_successes * config.CAPTCHA_MAX_ATTEMPTS_PER_URL
|
|
required_micros += service.captcha_solves_to_micros(worst_case_attempts)
|
|
await service.check_balance(owner_user_id, required_micros)
|
|
|
|
|
|
async def charge_capability(
|
|
output: BillableOutput, unit: BillingUnit | None, ctx: CapabilityContext
|
|
) -> None:
|
|
"""Bill the workspace owner for this result's billable successes (03c). ``None`` = free.
|
|
|
|
For crawl-backed verbs this also bills any captcha *attempts* (Phase 3d) as a
|
|
separate per-attempt unit — the solver charges per attempt even when the crawl
|
|
ultimately failed, so it can't ride the per-success crawl meter.
|
|
"""
|
|
if unit is None:
|
|
return
|
|
if unit is BillingUnit.WEB_CRAWL:
|
|
await _charge_web_crawl(ctx, output.billable_units)
|
|
await _charge_captcha(ctx, getattr(output, "captcha_attempts", 0))
|
|
|
|
|
|
async def _charge_web_crawl(ctx: CapabilityContext, successes: int) -> None:
|
|
if successes <= 0:
|
|
return
|
|
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 _charge_captcha(ctx: CapabilityContext, attempts: int) -> None:
|
|
if attempts <= 0:
|
|
return
|
|
service = WebCrawlCreditService(ctx.session)
|
|
if not service.captcha_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_captcha's commit flushes both.
|
|
await record_token_usage(
|
|
ctx.session,
|
|
usage_type="web_crawl_captcha",
|
|
workspace_id=ctx.workspace_id,
|
|
user_id=owner_user_id,
|
|
cost_micros=service.captcha_solves_to_micros(attempts),
|
|
call_details={"attempts": attempts},
|
|
)
|
|
await service.charge_captcha(owner_user_id, attempts)
|
|
|
|
|
|
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()
|