Merge branch 'main' into feat/speaches-integration

api/services/campaign/report.py

@@ -1,6 +1,7 @@
 import csv
 import io
-from typing import Any, List
+from datetime import datetime
+from typing import Any, List, Optional
 
 from api.constants import BACKEND_API_ENDPOINT
 from api.db import db_client
@@ -27,12 +28,18 @@ def _collect_extracted_variable_keys(runs: List[Any]) -> list[str]:
     return list(keys)
 
 
-async def generate_campaign_report_csv(campaign_id: int) -> tuple[io.StringIO, str]:
+async def generate_campaign_report_csv(
+    campaign_id: int,
+    start_date: Optional[datetime] = None,
+    end_date: Optional[datetime] = None,
+) -> tuple[io.StringIO, str]:
     """Generate a CSV report for a campaign.
 
     Returns a tuple of (csv_output, filename).
     """
-    runs = await db_client.get_completed_runs_for_report(campaign_id)
+    runs = await db_client.get_completed_runs_for_report(
+        campaign_id, start_date=start_date, end_date=end_date
+    )
 
     # Collect dynamic extracted variable columns
     extracted_var_keys = _collect_extracted_variable_keys(runs)

api/services/configuration/check_validity.py

@@ -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.SPEACHES.value: self._check_speaches_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

api/services/mps_service_key_client.py

@@ -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:

api/services/quota_service.py

@@ -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.",
                 )