from datetime import datetime, date
from typing import Optional, Dict
from pathlib import Path
from lib.providers.services import service
from lib.providers.commands import command
from lib.providers.hooks import hook
from lib.utils.debug import debug_box
from .models import CreditTransaction, CreditRatioConfig
from .storage import CreditStorage
from .ledger import CreditLedger
from .conversion import CreditUsageHandler, CreditPolicy
# Global handler for usage tracking
_usage_handler = None
[docs]
class CreditsPlugin:
def __init__(self, base_path: str):
self.base_path = base_path
[docs]
def create_components(self):
"""Create fresh instances of all credit system components"""
storage = CreditStorage(self.base_path)
ratio_config = CreditRatioConfig(self.base_path)
ledger = CreditLedger(storage)
credit_policy = CreditPolicy(ledger, ratio_config, self.base_path)
usage_handler = CreditUsageHandler(ledger, ratio_config, self.base_path)
return storage, ledger, ratio_config, credit_policy, usage_handler
[docs]
def get_base_path(context) -> str:
"""Get the base path for credit data storage"""
return str(Path.cwd() / 'data' / 'credits')
[docs]
@hook()
async def startup(app, context=None):
"""Startup tasks"""
# Initialize components
plugin = CreditsPlugin(get_base_path(context))
_, _, _, _, usage_handler = plugin.create_components()
# Store usage handler globally for the handle_usage service
global _usage_handler
_usage_handler = usage_handler
[docs]
@hook()
async def handle_usage(plugin_id: str, cost_type_id: str, quantity: float,
metadata: dict, context=None, model_id: Optional[str] = None):
"""Handle usage tracking for credits system.
This service is called by the usage plugin after tracking usage.
"""
global _usage_handler
if not _usage_handler:
raise RuntimeError("Credits plugin not properly initialized")
print(cost_type_id)
print("quantity:", quantity)
debug_box("Recording credit usage: {} {} {}".format(plugin_id, cost_type_id, quantity))
await _usage_handler.handle_usage(plugin_id, cost_type_id, quantity,
metadata, context, model_id)
[docs]
@service()
async def allocate_credits(username: str, amount: float,
source: str, reference_id: str,
metadata: Optional[Dict] = None,
context=None) -> float:
"""Allocate credits to a user.
Args:
username: User to allocate credits to
amount: Amount of credits to allocate
source: Source of allocation (e.g., 'purchase', 'admin_grant')
reference_id: External reference (e.g., payment ID)
metadata: Additional information about allocation
context: Request context
Returns:
float: New credit balance
Example:
new_balance = await allocate_credits(
'user123',
1000.0,
'purchase',
'payment_123',
{'payment_method': 'stripe'}
)
"""
plugin = CreditsPlugin(get_base_path(context))
_, ledger, _, _, _ = plugin.create_components()
return await ledger.record_allocation(
username, amount, source, reference_id, metadata
)
[docs]
@service()
async def get_credit_balance(username: str, context=None) -> float:
"""Get current credit balance for a user.
Example:
balance = await get_credit_balance('user123')
"""
plugin = CreditsPlugin(get_base_path(context))
_, ledger, _, _, _ = plugin.create_components()
return await ledger.get_balance(username)
[docs]
@service()
async def check_credits_available(username: str, required_amount: float,
context=None) -> Dict:
"""Check if user has sufficient credits.
Returns:
Dict with 'has_sufficient' and 'current_balance' keys
"""
plugin = CreditsPlugin(get_base_path(context))
_, ledger, _, _, _ = plugin.create_components()
has_sufficient, balance = await ledger.check_credits_available(
username, required_amount
)
return {
'has_sufficient': has_sufficient,
'current_balance': balance
}
[docs]
@service()
async def set_credit_ratio(ratio: float, plugin_id: Optional[str] = None,
cost_type_id: Optional[str] = None,
model_id: Optional[str] = None,
context=None):
"""Set credit ratio for cost conversion.
Args:
ratio: Credits per unit cost
plugin_id: Optional plugin identifier
cost_type_id: Optional cost type identifier
model_id: Optional model identifier
Example:
# Set global default ratio
await set_credit_ratio(100.0)
# Set model-specific ratio
await set_credit_ratio(
90.0,
plugin_id='gpt4',
cost_type_id='gpt4.input_tokens',
model_id='gpt-4-1106-preview'
)
"""
plugin = CreditsPlugin(get_base_path(context))
_, _, ratio_config, _, _ = plugin.create_components()
await ratio_config.set_ratio(ratio, plugin_id, cost_type_id, model_id)
[docs]
@service()
async def get_credit_ratios(context=None) -> Dict:
"""Get current credit ratio configuration."""
plugin = CreditsPlugin(get_base_path(context))
_, _, ratio_config, _, _ = plugin.create_components()
return await ratio_config.get_config()
[docs]
@service()
async def get_credit_report(username: str,
start_date: Optional[str] = None,
end_date: Optional[str] = None,
context=None) -> Dict:
"""Get detailed credit report for a user.
Args:
username: User to get report for
start_date: Optional start date (YYYY-MM-DD)
end_date: Optional end date (YYYY-MM-DD)
Returns:
Dict containing:
- Current balance
- Transaction history
- Usage summary
"""
start = date.fromisoformat(start_date) if start_date else None
end = date.fromisoformat(end_date) if end_date else None
plugin = CreditsPlugin(get_base_path(context))
_, ledger, _, _, _ = plugin.create_components()
transactions = await ledger.get_transactions(username, start, end)
summary = await ledger.get_usage_summary(username, start, end)
current_balance = await ledger.get_balance(username)
return {
'username': username,
'current_balance': current_balance,
'summary': summary,
'transactions': [t.to_dict() for t in transactions]
}
[docs]
@service()
async def estimate_credits(plugin_id: str, cost_type_id: str,
estimated_cost: float,
model_id: Optional[str] = None,
context=None) -> Dict:
"""Estimate credits needed for an operation.
Example:
estimate = await estimate_credits(
'gpt4',
'gpt4.input_tokens',
0.01, # $0.01
'gpt-4-1106-preview'
)
"""
plugin = CreditsPlugin(get_base_path(context))
_, _, _, credit_policy, _ = plugin.create_components()
credits = await credit_policy.estimate_credits_needed(
plugin_id, cost_type_id, estimated_cost, model_id
)
return {
'estimated_cost': estimated_cost,
'credits_required': credits,
'ratio_used': credits / estimated_cost
}