fix: send auth credentials with validate service keys

This commit is contained in:
Abhishek Kumar 2026-03-27 00:07:38 +05:30
parent 123114fb94
commit 83f05ab146
9 changed files with 83 additions and 24 deletions

View file

@ -125,7 +125,11 @@ async def update_user_configurations(
try:
validator = UserConfigurationValidator()
await validator.validate(user_configurations)
await validator.validate(
user_configurations,
organization_id=user.selected_organization_id,
created_by=user.provider_id,
)
except ValueError as e:
raise HTTPException(status_code=422, detail=e.args[0])
@ -163,7 +167,11 @@ async def validate_user_configurations(
):
validator = UserConfigurationValidator()
try:
status = await validator.validate(configurations)
status = await validator.validate(
configurations,
organization_id=user.selected_organization_id,
created_by=user.provider_id,
)
await db_client.update_user_configuration_last_validated_at(user.id)
return status
except ValueError as e:

View file

@ -228,7 +228,7 @@ class SignalingManager:
{
"type": "error",
"payload": {
"error_type": "quota_exceeded",
"error_type": quota_result.error_code,
"message": quota_result.error_message,
},
}

View file

@ -14,6 +14,12 @@ from api.schemas.user_configuration import (
from api.services.configuration.registry import ServiceConfig, ServiceProviders
from api.services.mps_service_key_client import mps_service_key_client
AuthContext = TypedDict(
"AuthContext",
{"organization_id": Optional[int], "created_by": Optional[str]},
total=False,
)
class APIKeyStatus(TypedDict):
model: str
@ -43,7 +49,16 @@ class UserConfigurationValidator:
ServiceProviders.SELF_HOSTED.value: self._check_self_hosted_api_key,
}
async def validate(self, configuration: UserConfiguration) -> APIKeyStatusResponse:
async def validate(
self,
configuration: UserConfiguration,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
) -> APIKeyStatusResponse:
self._auth_context: AuthContext = {
"organization_id": organization_id,
"created_by": created_by,
}
status_list = []
status_list.extend(self._validate_service(configuration.llm, "llm"))
@ -165,7 +180,12 @@ class UserConfigurationValidator:
"You provided a Dograh API key (dgr...) instead of a service key. "
"Please use a service key (mps...)."
)
return mps_service_key_client.validate_service_key(api_key)
auth = getattr(self, "_auth_context", {})
return mps_service_key_client.validate_service_key(
api_key,
organization_id=auth.get("organization_id"),
created_by=auth.get("created_by"),
)
def _check_sarvam_api_key(self, model: str, api_key: str) -> bool:
return True

View file

@ -276,7 +276,7 @@ class MPSServiceKeyClient:
"remaining_credits": data.get("remaining_credits", 0.0),
}
else:
logger.error(
logger.warning(
f"Failed to check service key usage: {response.status_code} - {response.text}"
)
raise httpx.HTTPStatusError(
@ -416,7 +416,12 @@ class MPSServiceKeyClient:
response=response,
)
def validate_service_key(self, service_key: str) -> bool:
def validate_service_key(
self,
service_key: str,
organization_id: Optional[int] = None,
created_by: Optional[str] = None,
) -> bool:
"""
Synchronously validate a Dograh service key by checking usage via MPS.
@ -427,7 +432,7 @@ class MPSServiceKeyClient:
response = client.post(
f"{self.base_url}/api/v1/service-keys/usage",
json={"service_key": service_key},
headers=self._get_headers(),
headers=self._get_headers(organization_id, created_by),
)
return response.status_code == 200
except Exception:

View file

@ -20,6 +20,7 @@ class QuotaCheckResult:
has_quota: bool
error_message: str = ""
error_code: str = ""
async def check_dograh_quota(user: UserModel) -> QuotaCheckResult:
@ -76,6 +77,7 @@ async def check_dograh_quota(user: UserModel) -> QuotaCheckResult:
)
return QuotaCheckResult(
has_quota=False,
error_code="quota_exceeded",
error_message=(
"You have exhausted your trial credits. "
"Please email founders@dograh.com for additional Dograh credits "
@ -89,8 +91,16 @@ async def check_dograh_quota(user: UserModel) -> QuotaCheckResult:
)
except Exception as e:
logger.error(f"Failed to check quota for Dograh key: {str(e)}")
error_str = str(e)
if "404" in error_str or "not found" in error_str.lower():
return QuotaCheckResult(
has_quota=False,
error_code="invalid_service_key",
error_message="You have invalid keys in your model configuration. Please validate the service keys.",
)
return QuotaCheckResult(
has_quota=False,
error_code="quota_check_failed",
error_message="Could not verify Dograh credits. Please try again.",
)

View file

@ -6,9 +6,12 @@ inline WebSocket media streaming.
import json
import random
from typing import Any, Dict, List, Optional, TYPE_CHECKING
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import aiohttp
from fastapi import HTTPException
from loguru import logger
from api.enums import WorkflowRunMode
from api.services.telephony.base import (
CallInitiationResult,
@ -16,8 +19,6 @@ from api.services.telephony.base import (
TelephonyProvider,
)
from api.utils.common import get_backend_endpoints
from fastapi import HTTPException
from loguru import logger
if TYPE_CHECKING:
from fastapi import WebSocket