Improved Licensing

src/database/client.py

@@ -1184,18 +1184,17 @@ class DatabaseClient:
         self,
         key: str,
         max_employees: int,
-        max_posts_per_day: int,
-        max_researches_per_day: int,
+        daily_token_limit: Optional[int] = None,
         description: Optional[str] = None
     ) -> LicenseKey:
         """Create new license key."""
         data = {
             "key": key,
             "max_employees": max_employees,
-            "max_posts_per_day": max_posts_per_day,
-            "max_researches_per_day": max_researches_per_day,
             "description": description,
         }
+        if daily_token_limit is not None:
+            data["daily_token_limit"] = daily_token_limit
         result = await asyncio.to_thread(
             lambda: self.client.table("license_keys").insert(data).execute()
         )
@@ -1222,6 +1221,16 @@ class DatabaseClient:
         logger.info(f"Marked license key as used: {key}")
         return LicenseKey(**result.data[0])
 
+    async def get_license_key_by_id(self, key_id: UUID) -> Optional[LicenseKey]:
+        """Get license key by UUID."""
+        result = await asyncio.to_thread(
+            lambda: self.client.table("license_keys")
+            .select("*")
+            .eq("id", str(key_id))
+            .execute()
+        )
+        return LicenseKey(**result.data[0]) if result.data else None
+
     async def update_license_key(self, key_id: UUID, updates: Dict[str, Any]) -> LicenseKey:
         """Update license key limits (admin only)."""
         result = await asyncio.to_thread(
@@ -1270,8 +1279,7 @@ class DatabaseClient:
         data = {
             "company_id": str(company_id),
             "date": date_.isoformat(),
-            "posts_created": 0,
-            "researches_created": 0
+            "tokens_used": 0
         }
         result = await asyncio.to_thread(
             lambda: self.client.table("company_daily_quotas")
@@ -1280,47 +1288,22 @@ class DatabaseClient:
         )
         return CompanyDailyQuota(**result.data[0])
 
-    async def increment_company_posts_quota(self, company_id: UUID) -> None:
-        """Increment daily posts count for company."""
+    async def increment_company_tokens(self, company_id: UUID, tokens: int) -> None:
+        """Increment daily token usage for company."""
         try:
             quota = await self.get_company_daily_quota(company_id)
-            new_count = quota.posts_created + 1
+            new_count = quota.tokens_used + tokens
 
-            result = await asyncio.to_thread(
+            await asyncio.to_thread(
                 lambda: self.client.table("company_daily_quotas")
-                .update({"posts_created": new_count})
+                .update({"tokens_used": new_count})
                 .eq("id", str(quota.id))
                 .execute()
             )
 
-            logger.info(f"Incremented posts quota for company {company_id}: {quota.posts_created} -> {new_count}")
-
-            if not result.data:
-                logger.error(f"Failed to increment posts quota - no data returned")
+            logger.info(f"Incremented token quota for company {company_id}: {quota.tokens_used} -> {new_count}")
         except Exception as e:
-            logger.error(f"Error incrementing posts quota for company {company_id}: {e}")
-            raise
-
-    async def increment_company_researches_quota(self, company_id: UUID) -> None:
-        """Increment daily researches count for company."""
-        try:
-            quota = await self.get_company_daily_quota(company_id)
-            new_count = quota.researches_created + 1
-
-            result = await asyncio.to_thread(
-                lambda: self.client.table("company_daily_quotas")
-                .update({"researches_created": new_count})
-                .eq("id", str(quota.id))
-                .execute()
-            )
-
-            logger.info(f"Incremented researches quota for company {company_id}: {quota.researches_created} -> {new_count}")
-
-            if not result.data:
-                logger.error(f"Failed to increment researches quota - no data returned")
-        except Exception as e:
-            logger.error(f"Error incrementing researches quota for company {company_id}: {e}")
-            raise
+            logger.warning(f"Error incrementing token quota for company {company_id}: {e}")
 
     async def get_company_limits(self, company_id: UUID) -> Optional[LicenseKey]:
         """Get company limits from associated license key.
@@ -1340,39 +1323,21 @@ class DatabaseClient:
 
         return LicenseKey(**result.data[0]) if result.data else None
 
-    async def check_company_post_limit(self, company_id: UUID) -> tuple[bool, str]:
-        """Check if company can create more posts today.
+    async def check_company_token_limit(self, company_id: UUID) -> tuple[bool, str, int, int]:
+        """Check if company has token budget remaining today.
 
-        Returns (can_create: bool, error_message: str)
+        Returns (can_proceed: bool, error_message: str, tokens_used: int, daily_limit: int)
         """
         license_key = await self.get_company_limits(company_id)
-        if not license_key:
-            # No license key, use defaults (unlimited)
-            return True, ""
+        if not license_key or license_key.daily_token_limit is None:
+            return True, "", 0, 0  # no limit = unlimited
 
         quota = await self.get_company_daily_quota(company_id)
 
-        if quota.posts_created >= license_key.max_posts_per_day:
-            return False, f"Tageslimit erreicht ({license_key.max_posts_per_day} Posts/Tag). Versuche es morgen wieder."
+        if quota.tokens_used >= license_key.daily_token_limit:
+            return False, f"Tageslimit erreicht ({license_key.daily_token_limit:,} Tokens/Tag). Morgen wieder verfügbar.", quota.tokens_used, license_key.daily_token_limit
 
-        return True, ""
-
-    async def check_company_research_limit(self, company_id: UUID) -> tuple[bool, str]:
-        """Check if company can create more researches today.
-
-        Returns (can_create: bool, error_message: str)
-        """
-        license_key = await self.get_company_limits(company_id)
-        if not license_key:
-            # No license key, use defaults (unlimited)
-            return True, ""
-
-        quota = await self.get_company_daily_quota(company_id)
-
-        if quota.researches_created >= license_key.max_researches_per_day:
-            return False, f"Tageslimit erreicht ({license_key.max_researches_per_day} Researches/Tag). Versuche es morgen wieder."
-
-        return True, ""
+        return True, "", quota.tokens_used, license_key.daily_token_limit
 
     async def check_company_employee_limit(self, company_id: UUID) -> tuple[bool, str]:
         """Check if company can add more employees.

src/database/models.py

@@ -365,8 +365,7 @@ class LicenseKey(DBModel):
 
     # Limits
     max_employees: int = 5
-    max_posts_per_day: int = 10
-    max_researches_per_day: int = 5
+    daily_token_limit: Optional[int] = None  # None = unlimited
 
     # Usage
     used: bool = False
@@ -382,7 +381,6 @@ class CompanyDailyQuota(DBModel):
     id: Optional[UUID] = None
     company_id: UUID
     date: date
-    posts_created: int = 0
-    researches_created: int = 0
+    tokens_used: int = 0
     created_at: Optional[datetime] = None
     updated_at: Optional[datetime] = None