Enterprise RAG Chat Platform

"""
ChatReadRetrieveRead — Multi-step RAG with query rewriting.

WHAT IT DEMONSTRATES:
  The core retrieval approach for an enterprise RAG chat system on Azure.
  The pipeline rewrites the user's question into an optimal search query,
  retrieves documents from Azure AI Search, then generates a grounded answer —
  all using structured tool calls to the OpenAI API.

WHY IT'S INTERESTING:
  Query rewriting is critical for conversational RAG: the user's raw message
  often contains pronouns or implicit references that degrade search quality.
  By first asking the LLM to produce a search query, then feeding search
  results back as context, the system achieves much higher answer accuracy
  than naive "embed question → retrieve → answer" pipelines.

NOVELTY:
  The architectural decision to use OpenAI tool calling (function calling)
  rather than prompt-engineered JSON extraction is significant: tool calls
  guarantee schema-valid output even when the model is uncertain, eliminating
  the fragile regex/JSON-parse layer that breaks in ~5% of naively prompted
  calls. The multi-retrieval strategy (semantic, vector, hybrid) is managed
  at the query-rewrite stage rather than post-retrieval fusion — meaning
  the model selects the optimal strategy per query rather than always running
  all three and merging results, which keeps latency and cost controlled.
  The .prompty file pattern externalises prompt engineering from application
  code, enabling non-engineers to iterate on prompts without deployments.
"""

import logging
from typing import Any, Optional
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletionMessageParam

logger = logging.getLogger(__name__)


class ChatReadRetrieveReadApproach:
    """
    Three-step RAG: rewrite → retrieve → answer.

    Step 1 — Query rewriting:
        GPT-4o rewrites the conversation into a search query using a
        tool call. The tool definition constrains the output format.

    Step 2 — Document retrieval:
        Azure AI Search returns the top-k documents via semantic / vector
        / hybrid retrieval depending on configuration.

    Step 3 — Answer generation:
        The original question + retrieved chunks are sent to GPT-4o to
        produce a grounded, cited answer.
    """

    def __init__(
        self,
        openai_client: AsyncOpenAI,
        search_client,
        chatgpt_model: str,
        prompt_manager,
        auth_helper,
        content_field: str = "content",
        sourcepage_field: str = "sourcepage",
    ):
        self.openai_client = openai_client
        self.search_client = search_client
        self.chatgpt_model = chatgpt_model
        self.prompt_manager = prompt_manager
        self.auth_helper = auth_helper
        self.content_field = content_field
        self.sourcepage_field = sourcepage_field

        # Prompts loaded from .prompty files — swappable without code changes
        self.query_rewrite_prompt = prompt_manager.load_prompt("chat_query_rewrite.prompty")
        self.query_rewrite_tools = prompt_manager.load_tools("chat_query_rewrite_tools.json")
        self.answer_prompt = prompt_manager.load_prompt("chat_answer_question.prompty")

    async def run(
        self,
        messages: list[ChatCompletionMessageParam],
        auth_claims: dict[str, Any],
        overrides: dict[str, Any],
    ) -> dict:
        # ── Step 1: Rewrite conversation into a search query ────────────────
        rewrite_messages = self.query_rewrite_prompt.format(messages=messages)
        rewrite_response = await self.openai_client.chat.completions.create(
            model=self.chatgpt_model,
            messages=rewrite_messages,
            tools=self.query_rewrite_tools,
            tool_choice="auto",
            temperature=0.0,
        )
        search_query = self._extract_tool_arg(rewrite_response, "search_query")
        logger.info("Rewritten query: %s", search_query)

        # ── Step 2: Retrieve relevant documents ─────────────────────────────
        search_results = await self._search(
            query=search_query,
            auth_claims=auth_claims,
            retrieval_mode=overrides.get("retrieval_mode", "hybrid"),
            top=overrides.get("top", 5),
        )
        context = self._format_context(search_results)

        # ── Step 3: Generate a grounded answer ──────────────────────────────
        answer_messages = self.answer_prompt.format(
            messages=messages, context=context
        )
        answer_response = await self.openai_client.chat.completions.create(
            model=self.chatgpt_model,
            messages=answer_messages,
            temperature=overrides.get("temperature", 0.1),
            max_tokens=overrides.get("max_tokens", 1024),
        )
        return {
            "message": answer_response.choices[0].message,
            "search_query": search_query,
            "sources": [r[self.sourcepage_field] for r in search_results],
        }

    def _extract_tool_arg(self, response, arg_name: str) -> str:
        """Pull a named argument from the first tool call in the response."""
        import json
        tool_call = response.choices[0].message.tool_calls[0]
        args = json.loads(tool_call.function.arguments)
        return args.get(arg_name, "")

    async def _search(self, query: str, auth_claims, retrieval_mode: str, top: int):
        """Delegate to Azure AI Search — semantic, vector, or hybrid."""
        filter_expr = self.auth_helper.build_security_filter(auth_claims)
        async with self.search_client:
            results = await self.search_client.search(
                search_text=query if retrieval_mode != "vector" else None,
                filter=filter_expr,
                top=top,
                query_type="semantic" if retrieval_mode == "semantic" else "simple",
            )
            return [r async for r in results]

    def _format_context(self, search_results: list) -> str:
        return "\n\n".join(
            f"[{r[self.sourcepage_field]}]\n{r[self.content_field]}"
            for r in search_results
        )

"""
Credit limit enforcement decorator for async API endpoints.

WHAT IT DEMONSTRATES:
  A reusable async Python decorator that enforces per-user monthly credit
  limits before allowing expensive LLM API calls to proceed. Integrates
  with Azure AD auth claims and a CosmosDB-backed role manager.

WHY IT'S INTERESTING:
  Decorators are an elegant solution for cross-cutting concerns like rate
  limiting and access control. The decorator extracts auth_claims from
  the wrapped function's arguments (either positional or keyword),
  making it composable with any authenticated endpoint without requiring
  changes to the endpoint's signature.

NOVELTY:
  The typical approach to credit enforcement is a middleware layer that
  intercepts all requests, or manual checks at the top of each endpoint
  function. Both create coupling between the enforcement logic and the
  endpoint code. The decorator pattern here is composable: adding credit
  enforcement to a new endpoint is a one-line annotation change. The
  implementation resolves auth_claims from either positional or keyword
  arguments using inspect.signature(), which is the non-obvious part —
  it makes the decorator work regardless of how the caller passes the
  claims, preventing silent failures if endpoint signatures evolve.
"""

import logging
from functools import wraps
from typing import Callable

from quart import current_app, jsonify

logger = logging.getLogger(__name__)


def credit_limit_required(f: Callable) -> Callable:
    """
    Async decorator: block the request with HTTP 429 if the user has
    exceeded their monthly credit limit.

    Usage:
        @blueprint.route("/api/chat", methods=["POST"])
        @authenticated               # sets auth_claims
        @credit_limit_required       # checks credit before proceeding
        async def chat(auth_claims: dict):
            ...
    """

    @wraps(f)
    async def decorated_function(*args, **kwargs):
        # Locate auth_claims — supports both positional and keyword passing
        auth_claims = kwargs.get("auth_claims")
        if not auth_claims:
            for arg in args:
                if isinstance(arg, dict) and "oid" in arg:
                    auth_claims = arg
                    break

        if not auth_claims:
            return jsonify({"error": "Authentication required"}), 401

        user_id = auth_claims.get("oid")
        if not user_id:
            return jsonify({"error": "User ID missing from auth claims"}), 401

        # Skip enforcement when usage tracking is disabled (local dev, CI)
        from config import CONFIG_USAGE_TRACKING_ENABLED
        if not current_app.config.get(CONFIG_USAGE_TRACKING_ENABLED, False):
            return await f(*args, **kwargs)

        try:
            from core.role_manager import get_role_manager
            role_manager = get_role_manager()
            if not role_manager:
                # Fail open — usage tracking unavailable, allow the request
                logger.warning("Role manager unavailable; skipping credit check")
                return await f(*args, **kwargs)

            # Bootstrap credit fields on first-ever request for this user
            await role_manager.ensure_user_has_credit_limit(user_id)

            if await role_manager.is_user_over_credit_limit(user_id):
                usage = await role_manager.get_user_usage_against_limit(user_id)
                logger.warning("User %s blocked — credit limit exceeded", user_id)
                return jsonify({
                    "error": "Credit limit exceeded",
                    "message": (
                        "You have exceeded your monthly credit limit. "
                        "Contact your administrator or wait for the monthly reset."
                    ),
                    "error_code": "CREDIT_LIMIT_EXCEEDED",
                    "credit_status": {
                        "credit_limit":       usage.get("credit_limit", 0.0),
                        "current_usage":      usage.get("current_usage", 0.0),
                        "remaining_credit":   usage.get("remaining_credit", 0.0),
                        "usage_percentage":   usage.get("usage_percentage", 0.0),
                        "limit_exceeded":     True,
                    },
                }), 429  # HTTP 429 Too Many Requests

            # Under the limit — proceed to the actual endpoint
            return await f(*args, **kwargs)

        except Exception as exc:
            # Fail open on unexpected errors — never block a user due to our bug
            logger.error("Credit check failed for user %s: %s", user_id, exc)
            return await f(*args, **kwargs)

    return decorated_function

"""
Pluggable approach registry — swappable retrieval strategies at runtime.

WHAT IT DEMONSTRATES:
  The architectural pattern that lets the platform support multiple retrieval
  strategies (semantic, vector, hybrid, vision-augmented) as interchangeable
  modules, selected per-request via config without touching application logic.
  Approach classes are registered by name at startup and resolved at request
  time, making the retrieval strategy a runtime decision rather than a
  compile-time one.

WHY IT'S INTERESTING:
  Enterprise RAG systems routinely need to serve different retrieval modes for
  different document collections: semantic search for narrative documents,
  vector search for technical specifications, hybrid for mixed corpora. A
  naive implementation duplicates the request-handling path for each strategy.
  The registry pattern isolates each strategy in its own class while keeping
  the API layer strategy-agnostic — adding a new retrieval approach requires
  zero changes to the router, middleware, or auth layers.

NOVELTY:
  The key design decision is pushing the strategy selection into the approach
  registry rather than into the endpoint or a conditional chain. Each approach
  is fully self-contained: it owns its retrieval logic, context formatting,
  and generation call. The registry resolves the approach from a string key in
  the request overrides, so the entire system behaviour can be changed via a
  frontend config dropdown without any backend code change. The abstract base
  class enforces the `run()` contract at registration time (via isinstance
  check), catching missing method implementations at startup rather than at
  first request — a pattern borrowed from plugin systems in compiler toolchains.
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from typing import Any

logger = logging.getLogger(__name__)


# ── Base interface ─────────────────────────────────────────────────────────

class BaseApproach(ABC):
    """
    Contract that every retrieval approach must implement.

    All approaches receive the same inputs and return the same output shape,
    making them transparently substitutable from the router's perspective.
    """

    @abstractmethod
    async def run(
        self,
        messages: list[dict],
        auth_claims: dict[str, Any],
        overrides: dict[str, Any],
    ) -> dict[str, Any]:
        """
        Execute the approach and return a response dict.

        Returns:
            {
                "message":      ChatCompletionMessage,
                "search_query": str,
                "sources":      list[str],
            }
        """
        ...

    @property
    @abstractmethod
    def name(self) -> str:
        """Human-readable approach name for logging and audit."""
        ...


# ── Registry ──────────────────────────────────────────────────────────────

class ApproachRegistry:
    """
    Maps string keys to registered approach instances.

    Approaches are registered at application startup (in app.py) and resolved
    at request time from the `overrides.approach` field sent by the frontend.

    Usage:
        registry = ApproachRegistry()
        registry.register("hybrid",   ChatReadRetrieveReadApproach(...))
        registry.register("semantic", SemanticApproach(...))
        registry.register("vision",   VisionAugmentedApproach(...))

        # At request time:
        approach = registry.resolve(overrides.get("approach", "hybrid"))
        result   = await approach.run(messages, auth_claims, overrides)
    """

    def __init__(self, default_key: str = "hybrid"):
        self._registry: dict[str, BaseApproach] = {}
        self._default_key = default_key

    def register(self, key: str, approach: BaseApproach) -> None:
        """
        Register an approach under a string key.

        Raises TypeError if `approach` does not implement BaseApproach.
        Called at startup so misconfiguration is caught before serving traffic.
        """
        if not isinstance(approach, BaseApproach):
            raise TypeError(
                f"Approach '{key}' must subclass BaseApproach, "
                f"got {type(approach).__name__}."
            )
        self._registry[key] = approach
        logger.info("Registered retrieval approach: '%s' (%s)", key, approach.name)

    def resolve(self, key: str) -> BaseApproach:
        """
        Look up an approach by key. Falls back to the default if key is unknown.

        Falling back (rather than raising) prevents a bad frontend config value
        from taking down the service — the fallback is logged as a warning so
        it surfaces in monitoring without causing a 500.
        """
        approach = self._registry.get(key)
        if approach is None:
            logger.warning(
                "Unknown approach key '%s' — falling back to '%s'.",
                key, self._default_key,
            )
            approach = self._registry.get(self._default_key)
        if approach is None:
            raise RuntimeError(
                f"Default approach '{self._default_key}' is not registered. "
                "Call registry.register() before serving requests."
            )
        return approach

    def available_approaches(self) -> list[str]:
        """Return registered approach keys (for health-check endpoints)."""
        return list(self._registry.keys())


# ── Example: vision-augmented approach stub ───────────────────────────────

class VisionAugmentedApproach(BaseApproach):
    """
    Extends the hybrid approach with GPT-4V processing for image-heavy pages.

    Pages identified as image-dominant during ingestion are stored with a
    `has_images` flag. At retrieval time, flagged pages are routed through
    GPT-4V to extract visual content descriptions, which are merged into the
    retrieval context before answer generation.
    """

    def __init__(self, base_approach: BaseApproach, vision_client, image_store):
        self._base = base_approach
        self._vision = vision_client
        self._image_store = image_store

    @property
    def name(self) -> str:
        return "Vision-Augmented Hybrid RAG"

    async def run(
        self,
        messages: list[dict],
        auth_claims: dict[str, Any],
        overrides: dict[str, Any],
    ) -> dict[str, Any]:
        # Run the base hybrid retrieval first
        result = await self._base.run(messages, auth_claims, overrides)

        # Augment sources that contain images with GPT-4V descriptions
        augmented_sources = []
        for source in result.get("sources", []):
            image_refs = await self._image_store.get_images_for_page(source)
            if image_refs:
                descriptions = await self._describe_images(image_refs)
                augmented_sources.append(
                    f"{source} [visual context: {'; '.join(descriptions)}]"
                )
            else:
                augmented_sources.append(source)

        result["sources"] = augmented_sources
        return result

    async def _describe_images(self, image_refs: list[str]) -> list[str]:
        """Ask GPT-4V to describe each image and return text summaries."""
        descriptions = []
        for ref in image_refs:
            response = await self._vision.chat.completions.create(
                model="gpt-4-vision-preview",
                messages=[{
                    "role": "user",
                    "content": [
                        {"type": "text",  "text": "Describe this image concisely for use as retrieval context."},
                        {"type": "image_url", "image_url": {"url": ref}},
                    ],
                }],
                max_tokens=256,
            )
            descriptions.append(response.choices[0].message.content.strip())
        return descriptions