""" Copyright 2024-2026 ReplyPier Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. FAQ generation usage/credits: estimate how many LLM calls a run will make, decide whether those calls are metered (hosted model), and enforce the org's message budget before enqueueing. OSS/self-hosted (no enterprise module) is never metered and never blocked. """ import math from dataclasses import dataclass from typing import List, Optional from uuid import UUID from fastapi import HTTPException from sqlalchemy import text from sqlalchemy.orm import Session from app.core.config import settings from app.core.logger import get_logger from app.knowledge.page_editor import PAGE_ID_EXPR from app.models.knowledge import Knowledge from app.repositories.faq import FAQRepository # Shared hosted-model metering helpers; re-exported here so existing FAQ # callers keep importing them from faq_usage. from app.services.usage_metering import ( # noqa: F401 HOSTED_MODEL_TYPE, is_hosted_model, remaining_message_credits, ) logger = get_logger(__name__) # Rough pages-per-LLM-call heuristic for the confirm-dialog estimate (matches # today's floor batch of ~3 average pages). An estimate, not billing. PAGES_PER_CALL = 3 OVER_BUDGET_MESSAGE = ( "Not enough message credits left in your plan for this generation " "(~{estimated} needed, {remaining} remaining). Upgrade your plan or " "switch to your own AI model key." ) @dataclass class SourceEstimate: id: int name: str source_type: str has_faqs: bool pages: int estimated_calls: int @dataclass class GenerationEstimate: total_sources: int new_sources: int pages: int estimated_calls: int sources: List[SourceEstimate] def generation_is_metered(db: Session, organization_id: UUID) -> bool: """Generation costs credits only on the hosted model — own-key orgs pay their provider directly. FAQ_METER_OWN_KEY flips to always-meter.""" return settings.FAQ_METER_OWN_KEY or is_hosted_model(db, organization_id) def count_source_pages(db: Session, knowledge: Knowledge) -> Optional[int]: """Distinct stored pages for a source, capped like generation itself. None when the vector table is missing/unreadable.""" if not knowledge.schema or not knowledge.table_name: return None query = text( f"SELECT COUNT(*) FROM (SELECT {PAGE_ID_EXPR} " f'FROM {knowledge.schema}."{knowledge.table_name}" ' "WHERE name = :source GROUP BY 1 LIMIT :max_pages) pages" ) try: count = db.execute( query, {"source": knowledge.source, "max_pages": settings.FAQ_MAX_PAGES_PER_SOURCE} ).scalar() return int(count or 0) except Exception as e: logger.warning(f"Page count failed for knowledge {knowledge.id}: {e}") db.rollback() return None def estimate_generation_calls( db: Session, organization_id: UUID, knowledge_ids: Optional[List[int]] = None, count_pages: bool = True, ) -> GenerationEstimate: """Confirm-dialog numbers for a GENERATE_ALL run: which sources it would read (new-only unless explicitly targeted) and roughly how many LLM calls that costs. Unreadable sources count as one call each. Always returns the FULL per-source list (for the generation source picker), each with its own page/call estimate and whether it already has FAQs. The aggregate counts (new_sources/pages/estimated_calls) cover the run's TARGET sources: the explicit `knowledge_ids` selection, or — when none is given — every source that doesn't yet have FAQs. count_pages=False skips the per-source page COUNT over the vector tables (cheap — for surfaces that only need counts, like the Generate button label); per-source estimated_calls is then a 1-per-source lower bound.""" all_sources = ( db.query(Knowledge).filter(Knowledge.organization_id == organization_id).all() ) generated = FAQRepository(db).knowledge_ids_with_faqs(organization_id) existing_ids = {s.id for s in all_sources} target_ids = ( (set(knowledge_ids) & existing_ids) if knowledge_ids else {s.id for s in all_sources if s.id not in generated} ) source_list: List[SourceEstimate] = [] agg_pages = 0 agg_calls = 0 for source in all_sources: source_pages = count_source_pages(db, source) if count_pages else None pages = source_pages or 0 if source_pages is None: # Unreadable now, or page-count skipped (count_pages=False) — assume # one call so the estimate isn't zero. calls = 1 else: # A readable but empty source produces nothing → costs no call. calls = max(1, math.ceil(pages / PAGES_PER_CALL)) if pages else 0 source_list.append( SourceEstimate( id=source.id, name=source.source, source_type=source.source_type.value if hasattr(source.source_type, "value") else str(source.source_type), has_faqs=source.id in generated, pages=pages, estimated_calls=calls, ) ) if source.id in target_ids: agg_pages += pages agg_calls += calls return GenerationEstimate( total_sources=len(all_sources), new_sources=len(target_ids), pages=agg_pages, estimated_calls=agg_calls, sources=source_list, ) def ensure_generation_budget(db: Session, organization_id: UUID, estimated_calls: int) -> None: """402 when a metered run would exceed the org's remaining credits.""" if estimated_calls <= 0 or not generation_is_metered(db, organization_id): return remaining = remaining_message_credits(db, organization_id) if remaining is None: return if estimated_calls > remaining: raise HTTPException( status_code=402, detail=OVER_BUDGET_MESSAGE.format(estimated=estimated_calls, remaining=remaining), )