apunkt/dograh
1
0
Fork 0
mirror of https://github.com/dograh-hq/dograh.git synced 2026-06-13 08:15:21 +02:00
Code Issues Projects Releases Packages Wiki Activity

feat: deprecate dograh based quota tracking

Browse source
This commit is contained in:
Abhishek Kumar 2026-06-11 18:36:08 +05:30
parent fde84387f2
commit 7d4e2e06a9
9 changed files with 90 additions and 313 deletions
Download patch file Download diff file Expand all files Collapse all files

27
api/alembic/versions/2159d4ac431a_added_quota_tables.py
View file

@ -18,6 +18,9 @@ branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
DEPRECATED_QUOTA_COMMENT = "Deprecated. MPS owns quota and credit ledger state."
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
# 1) Create the `quota_type` enum *before* we add the column that references it.
@ -34,7 +37,12 @@ def upgrade() -> None:
sa.Column("organization_id", sa.Integer(), nullable=False),
sa.Column("period_start", sa.DateTime(), nullable=False),
sa.Column("period_end", sa.DateTime(), nullable=False),
sa.Column("quota_dograh_tokens", sa.Integer(), nullable=False),
sa.Column(
"quota_dograh_tokens",
sa.Integer(),
nullable=False,
comment=DEPRECATED_QUOTA_COMMENT,
),
sa.Column("used_dograh_tokens", sa.Integer(), nullable=False),
sa.Column("created_at", sa.DateTime(timezone=True), nullable=True),
sa.Column("updated_at", sa.DateTime(timezone=True), nullable=True),
@ -63,7 +71,11 @@ def upgrade() -> None:
op.add_column(
"organizations",
sa.Column(
"quota_type", quota_type_enum, nullable=False, server_default="monthly"
"quota_type",
quota_type_enum,
nullable=False,
server_default="monthly",
comment=DEPRECATED_QUOTA_COMMENT,
),
)
op.add_column(
@ -73,6 +85,7 @@ def upgrade() -> None:
sa.Integer(),
nullable=False,
server_default=sa.text("0"),
comment=DEPRECATED_QUOTA_COMMENT,
),
)
op.add_column(
@ -82,10 +95,17 @@ def upgrade() -> None:
sa.Integer(),
nullable=False,
server_default=sa.text("LEAST(EXTRACT(DAY FROM CURRENT_DATE)::int, 28)"),
comment=DEPRECATED_QUOTA_COMMENT,
),
)
op.add_column(
"organizations", sa.Column("quota_start_date", sa.DateTime(), nullable=True)
"organizations",
sa.Column(
"quota_start_date",
sa.DateTime(),
nullable=True,
comment=DEPRECATED_QUOTA_COMMENT,
),
)
op.add_column(
"organizations",
@ -94,6 +114,7 @@ def upgrade() -> None:
sa.Boolean(),
nullable=False,
server_default=sa.text("false"),
comment=DEPRECATED_QUOTA_COMMENT,
),
)
# ### end Alembic commands ###

10
api/alembic/versions/c425d3445750_add_columns_in_usage_table.py
View file

@ -18,6 +18,9 @@ branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
DEPRECATED_QUOTA_COMMENT = "Deprecated. MPS owns quota and credit ledger state."
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.add_column(
@ -26,7 +29,12 @@ def upgrade() -> None:
)
op.add_column(
"organization_usage_cycles",
sa.Column("quota_amount_usd", sa.Float(), nullable=True),
sa.Column(
"quota_amount_usd",
sa.Float(),
nullable=True,
comment=DEPRECATED_QUOTA_COMMENT,
),
)
# ### end Alembic commands ###

2
api/db/db_client.py
View file

@ -53,7 +53,7 @@ class DBClient(
- UserClient: handles user and user configuration operations
- OrganizationClient: handles organization operations
- OrganizationConfigurationClient: handles organization configuration operations
- OrganizationUsageClient: handles organization usage and quota operations
- OrganizationUsageClient: handles organization usage reporting aggregates
- IntegrationClient: handles integration operations
- WorkflowTemplateClient: handles workflow template operations
- CampaignClient: handles campaign operations

53
api/db/models.py
View file

@ -97,22 +97,44 @@ class OrganizationModel(Base):
provider_id = Column(String, unique=True, index=True, nullable=False)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
# Quota fields
# Deprecated: MPS owns quota and credit ledger state.
quota_type = Column(
Enum("monthly", "annual", name="quota_type"),
nullable=False,
default="monthly",
server_default=text("'monthly'::quota_type"),
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
quota_dograh_tokens = Column(
Integer, nullable=False, default=0, server_default=text("0")
Integer,
nullable=False,
default=0,
server_default=text("0"),
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
quota_reset_day = Column(
Integer, nullable=False, default=1, server_default=text("1")
) # 1-28, only for monthly
quota_start_date = Column(DateTime(timezone=True), nullable=True) # Only for annual
Integer,
nullable=False,
default=1,
server_default=text("1"),
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
quota_start_date = Column(
DateTime(timezone=True),
nullable=True,
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
quota_enabled = Column(
Boolean, nullable=False, default=False, server_default=text("false")
Boolean,
nullable=False,
default=False,
server_default=text("false"),
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
price_per_second_usd = Column(Float, nullable=True)
@ -593,8 +615,9 @@ class WorkflowRunTextSessionModel(Base):
class OrganizationUsageCycleModel(Base):
"""
This model is used to track the usage of Dograh tokens for an organization for a given usage
cycle.
This model is used to track reporting aggregates for an organization for a given
usage cycle. Quota fields on this model are deprecated; MPS owns quota and
credit ledger state.
"""
__tablename__ = "organization_usage_cycles"
@ -603,14 +626,24 @@ class OrganizationUsageCycleModel(Base):
organization_id = Column(Integer, ForeignKey("organizations.id"), nullable=False)
period_start = Column(DateTime(timezone=True), nullable=False)
period_end = Column(DateTime(timezone=True), nullable=False)
quota_dograh_tokens = Column(Integer, nullable=False)
quota_dograh_tokens = Column(
Integer,
nullable=False,
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
used_dograh_tokens = Column(Float, nullable=False, default=0)
total_duration_seconds = Column(
Integer, nullable=False, default=0, server_default=text("0")
)
# New USD tracking fields
used_amount_usd = Column(Float, nullable=True, default=0)
quota_amount_usd = Column(Float, nullable=True)
quota_amount_usd = Column(
Float,
nullable=True,
comment="Deprecated. MPS owns quota and credit ledger state.",
info={"deprecated": True},
)
created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
updated_at = Column(
DateTime(timezone=True),

168
api/db/organization_usage_client.py
View file

@ -23,7 +23,7 @@ from api.schemas.ai_model_configuration import EffectiveAIModelConfiguration
class OrganizationUsageClient(BaseDBClient):
"""Client for managing organization usage and quota operations."""
"""Client for managing organization usage reporting aggregates."""
async def get_or_create_current_cycle(
self, organization_id: int, session=None
@ -49,14 +49,7 @@ class OrganizationUsageClient(BaseDBClient):
self, organization_id: int, session, commit: bool
) -> OrganizationUsageCycleModel:
"""Internal implementation for get_or_create_current_cycle."""
# Get organization to determine quota type
org_result = await session.execute(
select(OrganizationModel).where(OrganizationModel.id == organization_id)
)
org = org_result.scalar_one()
# Calculate current period
period_start, period_end = self._calculate_current_period(org)
period_start, period_end = self._calculate_current_period()
# Try to get existing cycle
cycle_result = await session.execute(
@ -78,7 +71,8 @@ class OrganizationUsageClient(BaseDBClient):
organization_id=organization_id,
period_start=period_start,
period_end=period_end,
quota_dograh_tokens=org.quota_dograh_tokens,
# Deprecated non-null column retained for historical schema compatibility.
quota_dograh_tokens=0,
)
# Handle concurrent inserts gracefully
stmt = stmt.on_conflict_do_nothing(
@ -102,54 +96,6 @@ class OrganizationUsageClient(BaseDBClient):
)
return cycle_result.scalar_one()
async def check_and_reserve_quota(
self, organization_id: int, estimated_tokens: int = 0
) -> bool:
"""
Check if organization has sufficient quota and optionally reserve tokens.
Returns True if quota is available, False otherwise.
This method is fully atomic and safe for concurrent access from multiple processes.
"""
async with self.async_session() as session:
# Get organization
org_result = await session.execute(
select(OrganizationModel).where(OrganizationModel.id == organization_id)
)
org = org_result.scalar_one_or_none()
if not org or not org.quota_enabled:
# No quota enforcement if not enabled
return True
# Get or create current cycle within the same session/transaction
cycle = await self._get_or_create_current_cycle_impl(
organization_id, session, commit=False
)
# Atomic check and update with row-level lock
result = await session.execute(
select(OrganizationUsageCycleModel)
.where(
and_(
OrganizationUsageCycleModel.id == cycle.id,
OrganizationUsageCycleModel.used_dograh_tokens
+ estimated_tokens
<= OrganizationUsageCycleModel.quota_dograh_tokens,
)
)
.with_for_update(skip_locked=False)
)
cycle_locked = result.scalar_one_or_none()
if cycle_locked:
# Update the usage atomically
cycle_locked.used_dograh_tokens += estimated_tokens
await session.commit()
return True
return False
async def update_usage_after_run(
self,
organization_id: int,
@ -188,9 +134,8 @@ class OrganizationUsageClient(BaseDBClient):
await session.commit()
async def get_current_usage(self, organization_id: int) -> dict:
"""Get current period usage information."""
"""Get current reporting-period usage information."""
async with self.async_session() as session:
# Get organization
org_result = await session.execute(
select(OrganizationModel).where(OrganizationModel.id == organization_id)
)
@ -201,42 +146,19 @@ class OrganizationUsageClient(BaseDBClient):
organization_id, session, commit=False
)
# Calculate next refresh date
if org.quota_type == "monthly":
next_refresh = cycle.period_end + relativedelta(days=1)
else: # annual
next_refresh = cycle.period_end + relativedelta(days=1)
result = {
"period_start": cycle.period_start.isoformat(),
"period_end": cycle.period_end.isoformat(),
"used_dograh_tokens": cycle.used_dograh_tokens,
"quota_dograh_tokens": cycle.quota_dograh_tokens,
"percentage_used": (
round(
(cycle.used_dograh_tokens / cycle.quota_dograh_tokens) * 100, 2
)
if cycle.quota_dograh_tokens > 0
else 0
),
"next_refresh_date": next_refresh.date().isoformat(),
"quota_enabled": org.quota_enabled,
"total_duration_seconds": cycle.total_duration_seconds,
}
# Add USD fields if organization has pricing
if org.price_per_second_usd is not None:
result["used_amount_usd"] = cycle.used_amount_usd or 0
result["quota_amount_usd"] = cycle.quota_amount_usd
result["currency"] = "USD"
result["price_per_second_usd"] = org.price_per_second_usd
# Calculate percentage based on USD if available
if cycle.quota_amount_usd and cycle.quota_amount_usd > 0:
result["percentage_used"] = round(
((cycle.used_amount_usd or 0) / cycle.quota_amount_usd) * 100, 2
)
return result
async def get_usage_history(
@ -545,83 +467,11 @@ class OrganizationUsageClient(BaseDBClient):
"currency": "USD",
}
async def update_organization_quota(
self,
organization_id: int,
quota_type: str,
quota_dograh_tokens: int,
quota_reset_day: Optional[int] = None,
quota_start_date: Optional[datetime] = None,
) -> OrganizationModel:
"""Update organization quota settings."""
async with self.async_session() as session:
result = await session.execute(
select(OrganizationModel).where(OrganizationModel.id == organization_id)
)
org = result.scalar_one()
org.quota_type = quota_type
org.quota_dograh_tokens = quota_dograh_tokens
org.quota_enabled = True
if quota_type == "monthly" and quota_reset_day:
org.quota_reset_day = quota_reset_day
elif quota_type == "annual" and quota_start_date:
org.quota_start_date = quota_start_date
await session.commit()
await session.refresh(org)
return org
def _calculate_current_period(
self, org: OrganizationModel
) -> tuple[datetime, datetime]:
"""Calculate the current billing period based on organization settings."""
def _calculate_current_period(self) -> tuple[datetime, datetime]:
"""Calculate the current calendar-month reporting period."""
now = datetime.now(timezone.utc)
if org.quota_type == "monthly":
# Find the start of the current billing month
reset_day = org.quota_reset_day
# Handle month boundaries
if now.day >= reset_day:
period_start = now.replace(
day=reset_day, hour=0, minute=0, second=0, microsecond=0
)
else:
# Previous month
period_start = (now - relativedelta(months=1)).replace(
day=reset_day, hour=0, minute=0, second=0, microsecond=0
)
# End is one month later minus 1 second
period_end = (
period_start + relativedelta(months=1) - relativedelta(seconds=1)
)
else: # annual
if not org.quota_start_date:
# Default to calendar year
period_start = now.replace(
month=1, day=1, hour=0, minute=0, second=0, microsecond=0
)
period_end = (
period_start + relativedelta(years=1) - relativedelta(seconds=1)
)
else:
# Find current annual period
start_date = org.quota_start_date.replace(tzinfo=timezone.utc)
years_diff = now.year - start_date.year
# Adjust for whether we've passed the anniversary
if now.month < start_date.month or (
now.month == start_date.month and now.day < start_date.day
):
years_diff -= 1
period_start = start_date + relativedelta(years=years_diff)
period_end = (
period_start + relativedelta(years=1) - relativedelta(seconds=1)
)
period_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
period_end = period_start + relativedelta(months=1) - relativedelta(seconds=1)
return period_start, period_end

8
api/routes/organization_usage.py
View file

@ -25,14 +25,8 @@ class CurrentUsageResponse(BaseModel):
period_start: str
period_end: str
used_dograh_tokens: float
quota_dograh_tokens: int
percentage_used: float
next_refresh_date: str
quota_enabled: bool
total_duration_seconds: int
# New USD fields
used_amount_usd: Optional[float] = None
quota_amount_usd: Optional[float] = None
currency: Optional[str] = None
price_per_second_usd: Optional[float] = None
@ -134,7 +128,7 @@ class DailyUsageBreakdownResponse(BaseModel):
@router.get("/usage/current-period", response_model=CurrentUsageResponse)
async def get_current_period_usage(user: UserModel = Depends(get_user)):
"""Get current billing period usage for the user's organization."""
"""Get current reporting-period usage for the user's organization."""
if not user.selected_organization_id:
raise HTTPException(status_code=400, detail="No organization selected")

113
ui/src/app/reports/page.tsx
View file

@ -1,14 +1,12 @@
'use client';
import { addDays, format, subDays } from 'date-fns';
import { Calendar, ChevronLeft, ChevronRight, CreditCard, Download, Loader2 } from 'lucide-react';
import { Calendar, ChevronLeft, ChevronRight, Download } from 'lucide-react';
import { useEffect, useState } from 'react';
import { client } from '@/client/client.gen';
import {
getDailyReportApiV1OrganizationsReportsDailyGet,
getDailyRunsDetailApiV1OrganizationsReportsDailyRunsGet,
getModelConfigurationV2ApiV1OrganizationsModelConfigurationsV2Get,
getPreferencesApiV1OrganizationsPreferencesGet,
getWorkflowOptionsApiV1OrganizationsReportsWorkflowsGet
} from '@/client/sdk.gen';
@ -19,7 +17,6 @@ import { Card } from '@/components/ui/card';
import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover';
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
import { Skeleton } from '@/components/ui/skeleton';
import { detailFromError } from '@/lib/apiError';
import { useAuth } from '@/lib/auth';
import { DispositionChart } from './components/DispositionChart';
@ -53,10 +50,6 @@ interface DailyReport {
}>;
}
type CreditPurchaseUrlResponse = {
checkout_url: string;
};
export default function ReportsPage() {
const [selectedDate, setSelectedDate] = useState<Date>(new Date());
const [selectedWorkflow, setSelectedWorkflow] = useState<string>('all');
@ -64,9 +57,6 @@ export default function ReportsPage() {
const [report, setReport] = useState<DailyReport | null>(null);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(null);
const [canBuyCredits, setCanBuyCredits] = useState(false);
const [purchaseLoading, setPurchaseLoading] = useState(false);
const [purchaseError, setPurchaseError] = useState<string | null>(null);
const [timezone, setTimezone] = useState('America/New_York');
const auth = useAuth();
@ -104,50 +94,6 @@ export default function ReportsPage() {
fetchPreferences();
}, [auth.isAuthenticated]);
useEffect(() => {
if (auth.loading) return;
if (!auth.isAuthenticated) {
setCanBuyCredits(false);
return;
}
let cancelled = false;
const fetchCreditPurchaseAvailability = async () => {
try {
const response = await getModelConfigurationV2ApiV1OrganizationsModelConfigurationsV2Get();
if (cancelled) return;
if (response.error || !response.data) {
setCanBuyCredits(false);
return;
}
const configuration = response.data.configuration as {
mode?: unknown;
dograh?: unknown;
} | null;
setCanBuyCredits(
response.data.source === 'organization_v2' &&
configuration?.mode === 'dograh' &&
Boolean(configuration.dograh)
);
} catch (err) {
console.error('Failed to check credit purchase availability:', err);
if (!cancelled) {
setCanBuyCredits(false);
}
}
};
fetchCreditPurchaseAvailability();
return () => {
cancelled = true;
};
}, [auth.loading, auth.isAuthenticated]);
// Fetch report data when date or workflow changes
useEffect(() => {
const fetchReport = async () => {
@ -249,41 +195,6 @@ export default function ReportsPage() {
}
};
const handleBuyCredits = async () => {
if (!auth.isAuthenticated || purchaseLoading) return;
setPurchaseLoading(true);
setPurchaseError(null);
try {
const response = await client.post<
{ 200: CreditPurchaseUrlResponse },
{ detail?: unknown }
>({
url: '/api/v1/organizations/usage/mps-credits/purchase-url',
});
if (response.error) {
throw new Error(
detailFromError(response.error, 'Failed to create credit purchase URL')
);
}
const checkoutUrl = response.data?.checkout_url;
if (!checkoutUrl) {
throw new Error('Failed to create credit purchase URL');
}
window.location.href = checkoutUrl;
} catch (err) {
console.error('Failed to create credit purchase URL:', err);
setPurchaseError(
err instanceof Error ? err.message : 'Failed to create credit purchase URL'
);
setPurchaseLoading(false);
}
};
const isToday = format(selectedDate, 'yyyy-MM-dd') === format(new Date(), 'yyyy-MM-dd');
return (
@ -291,27 +202,7 @@ export default function ReportsPage() {
{/* Header */}
<div className="flex flex-col sm:flex-row justify-between items-start sm:items-center gap-4">
<div className="space-y-2">
<div className="flex flex-wrap items-center gap-3">
<h1 className="text-3xl font-bold">Daily Reports</h1>
{canBuyCredits && (
<Button
size="sm"
onClick={handleBuyCredits}
disabled={purchaseLoading}
className="flex items-center gap-2"
>
{purchaseLoading ? (
<Loader2 className="h-4 w-4 animate-spin" />
) : (
<CreditCard className="h-4 w-4" />
)}
Buy Credits
</Button>
)}
</div>
{purchaseError && (
<p className="text-sm text-red-500">{purchaseError}</p>
)}
<h1 className="text-3xl font-bold">Daily Reports</h1>
</div>
{/* Date Navigation & Workflow Selector */}

2
ui/src/client/sdk.gen.ts
View file

@ -1239,7 +1239,7 @@ export const reactivateServiceKeyApiV1UserServiceKeysServiceKeyIdReactivatePut =
/**
* Get Current Period Usage
*
* Get current billing period usage for the user's organization.
* Get current reporting-period usage for the user's organization.
*/
export const getCurrentPeriodUsageApiV1OrganizationsUsageCurrentPeriodGet = <ThrowOnError extends boolean = false>(options?: Options<GetCurrentPeriodUsageApiV1OrganizationsUsageCurrentPeriodGetData, ThrowOnError>) => (options?.client ?? client).get<GetCurrentPeriodUsageApiV1OrganizationsUsageCurrentPeriodGetResponses, GetCurrentPeriodUsageApiV1OrganizationsUsageCurrentPeriodGetErrors, ThrowOnError>({ url: '/api/v1/organizations/usage/current-period', ...options });

20
ui/src/client/types.gen.ts
View file

@ -1642,22 +1642,6 @@ export type CurrentUsageResponse = {
* Used Dograh Tokens
*/
used_dograh_tokens: number;
/**
* Quota Dograh Tokens
*/
quota_dograh_tokens: number;
/**
* Percentage Used
*/
percentage_used: number;
/**
* Next Refresh Date
*/
next_refresh_date: string;
/**
* Quota Enabled
*/
quota_enabled: boolean;
/**
* Total Duration Seconds
*/
@ -1666,10 +1650,6 @@ export type CurrentUsageResponse = {
* Used Amount Usd
*/
used_amount_usd?: number | null;
/**
* Quota Amount Usd
*/
quota_amount_usd?: number | null;
/**
* Currency
*/