Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: update undertaker to use rate limiter (#9498)

Browse Source
This commit is contained in:
Shubham Naik
2026-02-16 15:07:06 -08:00
committed by Caren Thomas
parent 619e81ed1e
commit e66981c7e8
2 changed files with 17 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
letta/agents/letta_agent_v2.py
View File

@@ -705,10 +705,15 @@ class LettaAgentV2(BaseAgentV2):
async def _check_credits(self) -> bool:
"""Check if the organization still has credits. Returns True if OK or not configured."""
try:
await self.credit_verification_service.verify_credits(self.actor.organization_id)
await self.credit_verification_service.verify_credits(
self.actor.organization_id, self.agent_state.id
)
return True
except InsufficientCreditsError:
self.logger.warning(f"Insufficient credits for organization {self.actor.organization_id}, stopping agent loop")
self.logger.warning(
f"Insufficient credits for organization {self.actor.organization_id}, "
f"agent {self.agent_state.id}, stopping agent loop"
)
return False
@trace_method

16
letta/services/credit_verification_service.py
View File

@@ -15,9 +15,13 @@ class CreditVerificationService:
self.endpoint = os.getenv("STEP_ORCHESTRATOR_ENDPOINT")
self.auth_key = os.getenv("STEP_COMPLETE_KEY")
async def verify_credits(self, organization_id: str) -> bool:
async def verify_credits(self, organization_id: str, agent_id: str) -> bool:
"""
Check if an organization has enough credits to proceed.
Check if an organization has enough credits to proceed with a specific agent.
Args:
organization_id: The organization's core ID
agent_id: The agent's ID (used to determine model-specific costs)
Returns True if credits are available or if the service is not configured.
Raises InsufficientCreditsError if no credits remain.
@@ -32,7 +36,7 @@ class CreditVerificationService:
async with httpx.AsyncClient(timeout=5.0) as client:
response = await client.get(
f"{self.endpoint}/validate/core-organizations/{organization_id}",
f"{self.endpoint}/validate/core-organizations/{organization_id}/agents/{agent_id}",
headers=headers,
)
response.raise_for_status()
@@ -46,11 +50,11 @@ class CreditVerificationService:
except InsufficientCreditsError:
raise
except httpx.TimeoutException:
logger.warning(f"Timeout verifying credits for organization {organization_id}")
logger.warning(f"Timeout verifying credits for organization {organization_id}, agent {agent_id}")
return True
except httpx.HTTPStatusError as e:
logger.warning(f"HTTP error verifying credits for organization {organization_id}: {e.response.status_code}")
logger.warning(f"HTTP error verifying credits for organization {organization_id}, agent {agent_id}: {e.response.status_code}")
return True
except Exception as e:
logger.error(f"Unexpected error verifying credits for organization {organization_id}: {e}")
logger.error(f"Unexpected error verifying credits for organization {organization_id}, agent {agent_id}: {e}")
return True