0a8a8fda54
* feat: add credit verification before agent message endpoints
Add credit verification checks to message endpoints to prevent
execution when organizations have insufficient credits.
- Add InsufficientCreditsError exception type
- Add CreditVerificationService that calls step-orchestrator API
- Add credit checks to /agents/{id}/messages endpoints
- Add credit checks to /conversations/{id}/messages endpoint
🐾 Generated with [Letta Code](https://letta.com)
Co-Authored-By: Letta <noreply@letta.com>
* surface error in ade
* do per step instead
* parallel check
* parallel to step
* small fixes
* stage publish api
* fixes
* revert unnecessary frontend changes
* insufficient credits stop reason
---------
Co-authored-by: Letta <noreply@letta.com>
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
import logging
|
|
import os
|
|
|
|
import httpx
|
|
|
|
from letta.errors import InsufficientCreditsError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class CreditVerificationService:
|
|
"""Service for verifying organization credit balance before agent execution."""
|
|
|
|
def __init__(self):
|
|
self.endpoint = os.getenv("STEP_ORCHESTRATOR_ENDPOINT")
|
|
self.auth_key = os.getenv("STEP_COMPLETE_KEY")
|
|
|
|
async def verify_credits(self, organization_id: str) -> bool:
|
|
"""
|
|
Check if an organization has enough credits to proceed.
|
|
|
|
Returns True if credits are available or if the service is not configured.
|
|
Raises InsufficientCreditsError if no credits remain.
|
|
"""
|
|
if not self.endpoint or not self.auth_key:
|
|
return True
|
|
|
|
try:
|
|
headers = {}
|
|
if self.auth_key:
|
|
headers["Authorization"] = f"Bearer {self.auth_key}"
|
|
|
|
async with httpx.AsyncClient(timeout=5.0) as client:
|
|
response = await client.get(
|
|
f"{self.endpoint}/validate/core-organizations/{organization_id}",
|
|
headers=headers,
|
|
)
|
|
response.raise_for_status()
|
|
|
|
data = response.json()
|
|
if not data.get("hasMoreCredits", True):
|
|
raise InsufficientCreditsError()
|
|
|
|
return True
|
|
|
|
except InsufficientCreditsError:
|
|
raise
|
|
except httpx.TimeoutException:
|
|
logger.warning(f"Timeout verifying credits for organization {organization_id}")
|
|
return True
|
|
except httpx.HTTPStatusError as e:
|
|
logger.warning(f"HTTP error verifying credits for organization {organization_id}: {e.response.status_code}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Unexpected error verifying credits for organization {organization_id}: {e}")
|
|
return True
|