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

4
api/routes/knowledge_base.py
View file

@ -402,7 +402,7 @@ async def search_chunks(
)
# Manual search runs outside any workflow run, so resolve the MPS
# correlation id here (mint only for orgs already on v2; never create one).
# correlation id here.
embedding_service = await build_embedding_service(
db_client=db_client,
provider=embeddings_provider,
@ -411,8 +411,6 @@ async def search_chunks(
base_url=embeddings_base_url,
endpoint=embeddings_endpoint,
api_version=embeddings_api_version,
organization_id=user.selected_organization_id,
created_by=str(user.provider_id),
resolve_correlation=True,
)

16
api/routes/organization.py
View file

@ -42,7 +42,6 @@ from api.schemas.telephony_phone_number import (
ProviderSyncStatus,
)
from api.services.auth.depends import (
_sync_posthog_organization_mps_billing_v2_status,
get_user,
get_user_with_selected_organization,
)
@ -391,22 +390,21 @@ async def migrate_model_configuration_v2(
except ValueError as exc:
raise HTTPException(status_code=422, detail=exc.args[0])
billing_account_status = None
if DEPLOYMENT_MODE != "oss":
try:
billing_account_status = await ensure_hosted_mps_billing_account_v2(
await ensure_hosted_mps_billing_account_v2(
organization_id,
created_by=str(user.provider_id),
)
except Exception as exc:
logger.error(
"Failed to initialize MPS billing v2 account for organization {}: {}",
"Failed to initialize MPS billing account for organization {}: {}",
organization_id,
exc,
)
raise HTTPException(
status_code=502,
detail="Failed to initialize MPS billing v2 account",
detail="Failed to initialize MPS billing account",
)
await upsert_organization_ai_model_configuration_v2(
@ -417,14 +415,6 @@ async def migrate_model_configuration_v2(
organization_id=organization_id,
fallback_user_config=legacy,
)
if DEPLOYMENT_MODE != "oss":
_sync_posthog_organization_mps_billing_v2_status(
organization_id,
uses_mps_billing_v2=bool(
billing_account_status
and billing_account_status.get("billing_mode") == "v2"
),
)
return await _model_configuration_v2_response(
user=user,
configuration=configuration,

95
api/routes/organization_usage.py
View file

@ -1,6 +1,6 @@
import json
from datetime import datetime, timedelta
from typing import Any, Dict, List, Literal, Optional
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import StreamingResponse
@ -29,12 +29,6 @@ class CurrentUsageResponse(BaseModel):
price_per_second_usd: Optional[float] = None
class MPSCreditsResponse(BaseModel):
total_credits_used: float
remaining_credits: float
total_quota: float
class MPSCreditPurchaseUrlResponse(BaseModel):
checkout_url: str
@ -69,7 +63,6 @@ class MPSCreditLedgerEntryResponse(BaseModel):
class MPSBillingCreditsResponse(BaseModel):
billing_version: Literal["legacy", "v2"]
total_credits_used: float = 0.0
remaining_credits: float = 0.0
total_quota: float = 0.0
@ -162,69 +155,13 @@ async def get_current_period_usage(user: UserModel = Depends(get_user)):
raise HTTPException(status_code=500, detail=str(e))
@router.get("/usage/mps-credits", response_model=MPSCreditsResponse)
async def get_mps_credits(user: UserModel = Depends(get_user)):
"""Get aggregated usage and quota from MPS.
OSS users: queries by provider_id (created_by).
Hosted users: queries by organization_id.
"""
try:
if DEPLOYMENT_MODE == "oss":
usage = await mps_service_key_client.get_usage_by_created_by(
str(user.provider_id)
)
else:
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
usage = await mps_service_key_client.get_usage_by_organization(
user.selected_organization_id
)
total_used = usage.get("total_credits_used", 0.0)
total_remaining = usage.get("remaining_credits", 0.0)
return MPSCreditsResponse(
total_credits_used=total_used,
remaining_credits=total_remaining,
total_quota=total_used + total_remaining,
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Failed to fetch MPS credits: {e}")
raise HTTPException(status_code=500, detail=str(e))
async def _get_mps_billing_account_status(
user: UserModel, organization_id: int
) -> Optional[dict]:
return await mps_service_key_client.get_billing_account_status(
organization_id=organization_id,
created_by=str(user.provider_id),
)
def _is_mps_billing_v2(account: Optional[dict]) -> bool:
return bool(account and account.get("billing_mode") == "v2")
async def _legacy_mps_credits_response(user: UserModel) -> MPSBillingCreditsResponse:
if DEPLOYMENT_MODE == "oss":
usage = await mps_service_key_client.get_usage_by_created_by(
str(user.provider_id)
)
else:
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
usage = await mps_service_key_client.get_usage_by_organization(
user.selected_organization_id
)
async def _oss_mps_credits_response(user: UserModel) -> MPSBillingCreditsResponse:
"""Aggregate per-key MPS credits for OSS deployments (no billing account)."""
usage = await mps_service_key_client.get_usage_by_created_by(str(user.provider_id))
total_used = float(usage.get("total_credits_used", 0.0))
total_remaining = float(usage.get("remaining_credits", 0.0))
return MPSBillingCreditsResponse(
billing_version="legacy",
total_credits_used=total_used,
remaining_credits=total_remaining,
total_quota=total_used + total_remaining,
@ -237,16 +174,15 @@ async def get_billing_credits(
limit: int = Query(50, ge=1, le=100),
user: UserModel = Depends(get_user),
):
"""Return legacy MPS credits or paginated v2 billing ledger details for the org."""
"""Return per-key MPS credits (OSS) or the org's paginated billing ledger."""
try:
if DEPLOYMENT_MODE == "oss" or not user.selected_organization_id:
return await _legacy_mps_credits_response(user)
if DEPLOYMENT_MODE == "oss":
return await _oss_mps_credits_response(user)
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")
organization_id = user.selected_organization_id
account_status = await _get_mps_billing_account_status(user, organization_id)
if not _is_mps_billing_v2(account_status):
return await _legacy_mps_credits_response(user)
ledger = await mps_service_key_client.get_credit_ledger(
organization_id=organization_id,
page=page,
@ -287,7 +223,6 @@ async def get_billing_credits(
total_debits = float(ledger["total_debits_credits"])
return MPSBillingCreditsResponse(
billing_version="v2",
total_credits_used=total_debits,
remaining_credits=balance,
total_quota=balance + total_debits,
@ -350,7 +285,7 @@ async def get_billing_credits(
async def create_mps_credit_purchase_url(
user: UserModel = Depends(get_user_with_selected_organization),
):
"""Create a checkout URL for organizations using Dograh-managed MPS v2."""
"""Create a checkout URL for purchasing organization credits."""
if DEPLOYMENT_MODE == "oss":
raise HTTPException(
status_code=404,
@ -359,14 +294,6 @@ async def create_mps_credit_purchase_url(
organization_id = user.selected_organization_id
assert organization_id is not None
account_status = await _get_mps_billing_account_status(user, organization_id)
if not _is_mps_billing_v2(account_status):
raise HTTPException(
status_code=403,
detail=(
"Credit purchases are available only for organizations using billing v2"
),
)
try:
session = await mps_service_key_client.create_credit_purchase_url(