mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-22 23:31:12 +02:00
feat(capabilities): add pre-flight meter-gate for billed verbs
This commit is contained in:
parent
f257e385db
commit
8fbfa5a6e1
2 changed files with 37 additions and 1 deletions
|
|
@ -6,7 +6,12 @@ from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from sqlalchemy import select
|
from sqlalchemy import select
|
||||||
|
|
||||||
from app.capabilities.types import BillableOutput, BillingUnit, CapabilityContext
|
from app.capabilities.types import (
|
||||||
|
BillableInput,
|
||||||
|
BillableOutput,
|
||||||
|
BillingUnit,
|
||||||
|
CapabilityContext,
|
||||||
|
)
|
||||||
from app.services.token_tracking_service import record_token_usage
|
from app.services.token_tracking_service import record_token_usage
|
||||||
from app.services.web_crawl_credit_service import WebCrawlCreditService
|
from app.services.web_crawl_credit_service import WebCrawlCreditService
|
||||||
|
|
||||||
|
|
@ -16,6 +21,30 @@ if TYPE_CHECKING:
|
||||||
from sqlalchemy.ext.asyncio import AsyncSession
|
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:
|
||||||
|
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
|
||||||
|
await service.check_credits(owner_user_id, estimated_successes)
|
||||||
|
|
||||||
|
|
||||||
async def charge_capability(
|
async def charge_capability(
|
||||||
output: BillableOutput, unit: BillingUnit | None, ctx: CapabilityContext
|
output: BillableOutput, unit: BillingUnit | None, ctx: CapabilityContext
|
||||||
) -> None:
|
) -> None:
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,13 @@ class BillingUnit(StrEnum):
|
||||||
WEB_CRAWL = "web_crawl"
|
WEB_CRAWL = "web_crawl"
|
||||||
|
|
||||||
|
|
||||||
|
class BillableInput(Protocol):
|
||||||
|
"""A billed verb's input that reports its worst-case unit count for pre-flight."""
|
||||||
|
|
||||||
|
@property
|
||||||
|
def estimated_units(self) -> int: ...
|
||||||
|
|
||||||
|
|
||||||
class BillableOutput(Protocol):
|
class BillableOutput(Protocol):
|
||||||
"""A capability output that reports its own billable count."""
|
"""A capability output that reports its own billable count."""
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue