apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-07-10 11:12:13 +02:00
Code Issues Projects Releases Packages Wiki Activity

chore: cleaup mps v1 billing (#507)

Browse source 
* chore: cleaup mps v1 billing

* chore: remove legacy file upload path

* chore: implement review comments
This commit is contained in:
Abhishek 2026-07-07 18:38:29 +05:30 • committed by GitHub
parent ac01f7775e
commit fdb7f92fcc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
36 changed files with 268 additions and 1319 deletions
Download patch file Download diff file Expand all files Collapse all files

50
api/services/auth/depends.py
View file

@ -34,7 +34,6 @@ async def require_local_auth() -> None:
POSTHOG_ORGANIZATION_GROUP_TYPE = "organization"
POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY = "uses_mps_billing_v2"
async def get_user(
@ -193,7 +192,6 @@ def _sync_created_organization_to_posthog(
organization,
stack_user: dict | None = None,
created_by_provider_id: str | None = None,
uses_mps_billing_v2: bool | None = None,
) -> None:
"""Create/update the PostHog organization group for a newly-created org."""
try:
@ -209,10 +207,6 @@ def _sync_created_organization_to_posthog(
}
if created_by:
properties["created_by_provider_id"] = created_by
if uses_mps_billing_v2 is not None:
properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
uses_mps_billing_v2
)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
@ -231,50 +225,6 @@ def _sync_created_organization_to_posthog(
logger.exception("Failed to sync created organization to PostHog")
def _sync_posthog_organization_group_properties(
*,
organization,
uses_mps_billing_v2: bool | None = None,
) -> None:
"""Update PostHog organization group properties without creating a person."""
try:
organization_id = int(organization.id)
properties = {
"organization_id": organization_id,
"organization_provider_id": getattr(organization, "provider_id", None),
"auth_provider": "stack",
}
if uses_mps_billing_v2 is not None:
properties[POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY] = (
uses_mps_billing_v2
)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
str(organization_id),
properties,
)
except Exception:
logger.exception("Failed to sync organization group properties to PostHog")
def _sync_posthog_organization_mps_billing_v2_status(
organization_id: int,
*,
uses_mps_billing_v2: bool,
) -> None:
"""Update the PostHog organization group with current MPS billing status."""
try:
organization_id = int(organization_id)
group_identify(
POSTHOG_ORGANIZATION_GROUP_TYPE,
str(organization_id),
{POSTHOG_ORGANIZATION_USES_MPS_BILLING_V2_PROPERTY: uses_mps_billing_v2},
)
except Exception:
logger.exception("Failed to sync organization billing status to PostHog")
def _associate_user_with_posthog_organization(
*,
user: UserModel,

48
api/services/gen_ai/embedding/factory.py
View file

@ -2,9 +2,8 @@
Centralizes the provider branching (Azure BYOK / Dograh-managed / OpenAI-compatible
BYOK) that was previously duplicated across document ingestion, the search route,
and the RAG tool, and resolves the MPS billing v2 protocol the same way the voice
path does: attach it only for orgs already on v2, and never create a billing
account to do so.
and the RAG tool, and resolves the MPS correlation id the same way the voice
path does.
"""
from typing import Optional
@ -24,46 +23,23 @@ DEFAULT_AZURE_API_VERSION = "2024-02-15-preview"
async def resolve_embedding_correlation_id(
*,
organization_id: Optional[int],
service_key: Optional[str],
created_by: Optional[str] = None,
) -> Optional[str]:
"""Resolve an MPS correlation id for a managed embedding call made outside a run.
"""Mint an MPS correlation id for a managed embedding call made outside a run.
Mirrors the voice path's gating:
- OSS deployments use a pasted hosted v2 key (v2 by definition), so mint
directly via the bearer endpoint — matching ``_authorize_oss_managed_v2_correlation``.
- Hosted/SaaS: read the org's billing mode (no side effects) and mint only when
it is already v2. Minting for an already-v2 org is a no-op on the account.
Returns ``None`` when the call should be sent without the protocol; MPS accepts
un-gated embedding calls from v1 orgs. Never creates a v2 billing account.
Matches the voice path's ``_authorize_oss_managed_v2_correlation``: the
correlation is minted via the bearer service-key endpoint, so it works for
hosted orgs and OSS keys alike. Returns ``None`` when minting fails; MPS
accepts un-correlated embedding calls.
"""
if not service_key:
return None
# Imported lazily to avoid import-time cycles between the gen_ai and service
# layers (matches the inline-import convention used elsewhere in the app).
from api.constants import DEPLOYMENT_MODE
from api.services.mps_service_key_client import mps_service_key_client
try:
if DEPLOYMENT_MODE == "oss":
minted = await mps_service_key_client.create_correlation_id(
service_key=service_key
)
return minted.get("correlation_id")
if organization_id is None:
return None
status = await mps_service_key_client.get_billing_account_status(
organization_id, created_by=created_by
)
if not status or status.get("billing_mode") != "v2":
return None
minted = await mps_service_key_client.create_correlation_id(
service_key=service_key
)
@ -71,7 +47,7 @@ async def resolve_embedding_correlation_id(
except Exception as e:
logger.warning(
"Could not resolve MPS correlation id for managed embeddings; "
"sending without v2 protocol: {}",
"sending without it: {}",
e,
)
return None
@ -87,8 +63,6 @@ async def build_embedding_service(
endpoint: Optional[str] = None,
api_version: Optional[str] = None,
correlation_id: Optional[str] = None,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
resolve_correlation: bool = False,
) -> BaseEmbeddingService:
"""Construct the right embedding service for a provider/config.
@ -116,11 +90,7 @@ async def build_embedding_service(
if provider == ServiceProviders.DOGRAH.value:
cid = correlation_id
if cid is None and resolve_correlation:
cid = await resolve_embedding_correlation_id(
organization_id=organization_id,
service_key=api_key,
created_by=created_by,
)
cid = await resolve_embedding_correlation_id(service_key=api_key)
return DograhEmbeddingService(
db_client=db_client,
api_key=api_key,

70
api/services/mps_service_key_client.py
View file

@ -241,19 +241,12 @@ class MPSServiceKeyClient:
)
return False
async def check_service_key_usage(
self,
service_key: str,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
) -> dict:
async def check_service_key_usage(self, service_key: str) -> dict:
"""
Check the usage and quota of a service key.
Args:
service_key: The service key to check usage for
organization_id: Organization ID (for authenticated mode)
created_by: User provider ID (for OSS mode)
Returns:
Dictionary containing:
@ -321,39 +314,6 @@ class MPSServiceKeyClient:
response=response,
)
async def get_usage_by_organization(self, organization_id: int) -> dict:
"""
Get aggregated usage for all service keys belonging to an organization (hosted mode).
Args:
organization_id: The organization's ID
Returns:
Dictionary containing total_credits_used and remaining_credits
"""
async with httpx.AsyncClient(timeout=self.timeout) as client:
response = await client.post(
f"{self.base_url}/api/v1/service-keys/usage/organization",
json={"organization_id": organization_id},
headers=self._get_headers(organization_id=organization_id),
)
if response.status_code == 200:
data = response.json()
return {
"total_credits_used": data.get("total_credits_used", 0.0),
"remaining_credits": data.get("remaining_credits", 0.0),
}
else:
logger.error(
f"Failed to get usage by organization: {response.status_code} - {response.text}"
)
raise httpx.HTTPStatusError(
f"Failed to get usage by organization: {response.text}",
request=response.request,
response=response,
)
async def create_credit_purchase_url(
self,
organization_id: int,
@ -422,34 +382,6 @@ class MPSServiceKeyClient:
response=response,
)
async def get_billing_account_status(
self,
organization_id: int,
created_by: Optional[str] = None,
) -> Optional[dict]:
"""Get an existing MPS v2 billing account without creating one."""
async with httpx.AsyncClient(timeout=self.timeout) as client:
response = await client.get(
f"{self.base_url}/api/v1/billing/accounts/{organization_id}/status",
headers=self._get_headers(
organization_id=organization_id,
created_by=created_by,
),
)
if response.status_code == 200:
return response.json()
logger.error(
"Failed to get MPS billing account status: "
f"{response.status_code} - {response.text}"
)
raise httpx.HTTPStatusError(
f"Failed to get MPS billing account status: {response.text}",
request=response.request,
response=response,
)
async def ensure_billing_account_v2(
self,
organization_id: int,

136
api/services/quota_service.py
View file

@ -25,13 +25,13 @@ from api.services.mps_service_key_client import mps_service_key_client
MINIMUM_DOGRAH_CREDITS_FOR_CALL = 0.10
LEGACY_QUOTA_EXCEEDED_MESSAGE = (
OSS_QUOTA_EXCEEDED_MESSAGE = (
"You have exhausted your trial credits. "
"Please email founders@dograh.com for additional Dograh credits "
"or change providers in Models configurations."
"Please sign up on app.dograh.com to create a "
"new service key and set up in your model configurations."
)
BILLING_V2_QUOTA_EXCEEDED_MESSAGE = (
HOSTED_QUOTA_EXCEEDED_MESSAGE = (
"You have exhausted your Dograh credits. "
"Please purchase more credits from /billing "
"or change providers in Models configurations."
@ -60,19 +60,19 @@ def _safe_float(value: Any, default: float = 0.0) -> float:
return default
def _insufficient_billing_v2_quota_result() -> QuotaCheckResult:
def _insufficient_hosted_quota_result() -> QuotaCheckResult:
return QuotaCheckResult(
has_quota=False,
error_code="insufficient_credits",
error_message=BILLING_V2_QUOTA_EXCEEDED_MESSAGE,
error_message=HOSTED_QUOTA_EXCEEDED_MESSAGE,
)
def _insufficient_legacy_quota_result() -> QuotaCheckResult:
def _insufficient_oss_quota_result() -> QuotaCheckResult:
return QuotaCheckResult(
has_quota=False,
error_code="quota_exceeded",
error_message=LEGACY_QUOTA_EXCEEDED_MESSAGE,
error_message=OSS_QUOTA_EXCEEDED_MESSAGE,
)
@ -157,10 +157,10 @@ async def _authorize_hosted_workflow_run_start(
workflow_id: int | None,
workflow_run_id: int | None,
user_config: Any,
) -> tuple[QuotaCheckResult, bool]:
"""Authorize hosted v2 billing and return whether MPS handled enforcement."""
if DEPLOYMENT_MODE == "oss" or organization_id is None:
return QuotaCheckResult(has_quota=True), False
) -> QuotaCheckResult:
"""Authorize a hosted workflow run against the org's MPS billing account."""
if organization_id is None:
return QuotaCheckResult(has_quota=True)
requires_correlation = bool(
workflow_run_id and uses_managed_model_services_v2(user_config)
@ -169,16 +169,13 @@ async def _authorize_hosted_workflow_run_start(
get_dograh_service_api_key(user_config) if requires_correlation else None
)
if requires_correlation and not service_key:
return (
QuotaCheckResult(
has_quota=False,
error_code="invalid_service_key",
error_message=(
"You have invalid keys in your model configuration. "
"Please validate the service keys."
),
return QuotaCheckResult(
has_quota=False,
error_code="invalid_service_key",
error_message=(
"You have invalid keys in your model configuration. "
"Please validate the service keys."
),
True,
)
try:
@ -205,38 +202,28 @@ async def _authorize_hosted_workflow_run_start(
e,
)
if _is_service_key_org_mismatch_error(e):
return (
QuotaCheckResult(
has_quota=False,
error_code="service_key_org_mismatch",
error_message=SERVICE_TOKEN_ORG_MISMATCH_MESSAGE,
),
True,
)
return (
QuotaCheckResult(
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
),
True,
error_code="service_key_org_mismatch",
error_message=SERVICE_TOKEN_ORG_MISMATCH_MESSAGE,
)
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
)
billing_mode = authorization.get("billing_mode")
if billing_mode != "v2":
return QuotaCheckResult(has_quota=True), False
remaining = _safe_float(authorization.get("remaining_credits"))
if (
not authorization.get("allowed", False)
or remaining < MINIMUM_DOGRAH_CREDITS_FOR_CALL
):
logger.warning(
"Insufficient Dograh billing v2 credits for org {}: {:.2f} credits remaining",
"Insufficient Dograh credits for org {}: {:.2f} credits remaining",
organization_id,
remaining,
)
return _insufficient_billing_v2_quota_result(), True
return _insufficient_hosted_quota_result()
try:
await _store_run_correlation_id(
@ -249,35 +236,27 @@ async def _authorize_hosted_workflow_run_start(
workflow_run_id,
e,
)
return (
QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
),
True,
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
)
logger.info(
"Dograh billing v2 run authorization passed for org {}: {:.2f} credits remaining",
"Dograh run authorization passed for org {}: {:.2f} credits remaining",
organization_id,
remaining,
)
return QuotaCheckResult(has_quota=True), True
return QuotaCheckResult(has_quota=True)
async def _authorize_legacy_dograh_keys(
async def _authorize_oss_dograh_keys(
*,
dograh_api_keys: set[str],
organization_id: int | None,
workflow_owner: UserModel,
) -> QuotaCheckResult:
"""Check per-key MPS credits for OSS deployments before a run starts."""
for api_key in dograh_api_keys:
try:
usage = await mps_service_key_client.check_service_key_usage(
api_key,
organization_id=organization_id,
created_by=workflow_owner.provider_id,
)
usage = await mps_service_key_client.check_service_key_usage(api_key)
remaining = usage.get("remaining_credits", 0.0)
# Require at least $0.10 for a short call
@ -286,7 +265,7 @@ async def _authorize_legacy_dograh_keys(
f"Insufficient Dograh credits for key ...{api_key[-8:]}: "
f"${remaining:.2f} remaining"
)
return _insufficient_legacy_quota_result()
return _insufficient_oss_quota_result()
logger.info(
f"Dograh quota check passed for key ...{api_key[-8:]}: "
@ -363,9 +342,9 @@ async def authorize_workflow_run_start(
) -> QuotaCheckResult:
"""Authorize a workflow run before any billable call/text runtime starts.
The workflow organization is the billing subject for hosted v2. The workflow
owner is used only to resolve the effective model configuration and legacy
service-key metadata.
The workflow organization is the billing subject for hosted deployments.
OSS deployments are billed per service key instead. The workflow owner is
used only to resolve the effective model configuration.
"""
try:
workflow = await db_client.get_workflow_by_id(workflow_id)
@ -405,38 +384,27 @@ async def authorize_workflow_run_start(
)
if DEPLOYMENT_MODE != "oss":
hosted_result, hosted_enforced = await _authorize_hosted_workflow_run_start(
return await _authorize_hosted_workflow_run_start(
workflow_owner=workflow_owner,
organization_id=workflow.organization_id,
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
)
if hosted_enforced or not hosted_result.has_quota:
return hosted_result
dograh_api_keys = _dograh_api_keys(user_config)
if not dograh_api_keys:
return QuotaCheckResult(has_quota=True)
legacy_result = await _authorize_legacy_dograh_keys(
dograh_api_keys=dograh_api_keys,
organization_id=(
None if DEPLOYMENT_MODE == "oss" else workflow.organization_id
),
workflow_owner=workflow_owner,
)
if not legacy_result.has_quota:
return legacy_result
if DEPLOYMENT_MODE == "oss":
return await _authorize_oss_managed_v2_correlation(
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
if dograh_api_keys:
oss_result = await _authorize_oss_dograh_keys(
dograh_api_keys=dograh_api_keys,
)
if not oss_result.has_quota:
return oss_result
return QuotaCheckResult(has_quota=True)
return await _authorize_oss_managed_v2_correlation(
workflow_id=workflow.id,
workflow_run_id=workflow_run_id,
user_config=user_config,
)
except Exception as e:
logger.error(f"Error during quota check: {str(e)}")

6
api/services/workflow/audit.py
View file

@ -1,10 +1,8 @@
"""Rule-based audit of a workflow definition's nodes + edges.
Pure, dependency-free helpers derived from `NodeSpec.graph_constraints`.
Lives in tracked code so the regression tests in
`test_workflow_graph_constraints.py` can pin it; the admin cleanup
script in `api/services/admin_utils/local_exec.py` is the production
consumer.
Lives in tracked code so `test_workflow_graph_constraints.py` can pin the
verdicts that one-off cleanup tooling needs to share with runtime validation.
"""
from collections import Counter

3
api/services/workflow/tools/knowledge_base.py
View file

@ -266,8 +266,7 @@ async def _perform_retrieval(
)
# Search runs inside a workflow run: reuse the run's MPS correlation
# id (present only for v2 orgs; None otherwise → sent without the
# protocol). The Dograh-managed path forwards it via request metadata.
# id. The Dograh-managed path forwards it via request metadata.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,

13
api/services/workflow_run_billing.py
View file

@ -32,13 +32,6 @@ def _duration_seconds_from_usage_info(workflow_run) -> float | None:
return duration_seconds if duration_seconds > 0 else None
async def _organization_uses_mps_billing_v2(organization_id: int) -> bool:
account = await mps_service_key_client.get_billing_account_status(
organization_id=organization_id
)
return bool(account and account.get("billing_mode") == "v2")
def _is_usage_not_ready_error(exc: Exception) -> bool:
response = getattr(exc, "response", None)
if getattr(response, "status_code", None) != 409:
@ -79,12 +72,6 @@ async def report_workflow_run_platform_usage(workflow_run) -> None:
return
try:
if not await _organization_uses_mps_billing_v2(organization_id):
logger.debug(
"Not reporting platform usage since org not using mps billing v2"
)
return
result = await mps_service_key_client.report_platform_usage(
organization_id=organization_id,
correlation_id=correlation_id,