Exception Handling Patterns

from __future__ import annotations

from typing import Any, ClassVar, Dict, Optional


class RagkitError(Exception):
    """Base class for every ragkit error. Catch this at application boundaries."""

    retryable: ClassVar[bool] = False
    code: ClassVar[str] = "ragkit_error"

    def __init__(self, message: str, **context: Any) -> None:
        super().__init__(message)
        self.message = message
        self.context = context

    def __str__(self) -> str:
        extra = ", ".join(f"{k}={v!r}" for k, v in self.context.items() if v is not None)
        return f"{self.message} ({extra})" if extra else self.message


class ConfigurationError(RagkitError):
    code = "configuration"


class InputValidationError(RagkitError, ValueError):
    code = "invalid_input"


class ProviderError(RagkitError):
    code = "provider"

    def __init__(self, message: str, *, provider: str, model: Optional[str] = None, **context: Any) -> None:
        super().__init__(message, provider=provider, model=model, **context)
        self.provider = provider
        self.model = model


class RateLimitError(ProviderError):
    retryable = True
    code = "rate_limited"

    def __init__(self, message: str, *, retry_after: Optional[float] = None, **context: Any) -> None:
        super().__init__(message, retry_after=retry_after, **context)
        self.retry_after = retry_after


class ContextLengthError(ProviderError):
    code = "context_length_exceeded"

    def __init__(self, message: str, *, tokens: int, limit: int, **context: Any) -> None:
        super().__init__(message, tokens=tokens, limit=limit, **context)
        self.tokens = tokens
        self.limit = limit


HTTP_STATUS: Dict[type, int] = {RagkitError: 500, InputValidationError: 422, ProviderError: 502,
                                RateLimitError: 429, ContextLengthError: 413}

def build_prompt(template: str, variables: dict) -> str:
    missing = sorted(set(PLACEHOLDER.findall(template)) - variables.keys())
    if missing:
        raise InputValidationError(f"prompt template is missing variables: {', '.join(missing)}",
                                   template_id=variables.get("template_id"))
    return template.format(**variables)

except ProviderError:
    metrics.increment("provider_errors")
    raise                       # same exception, original traceback

except ProviderError as err:
    raise err                   # works, but adds this line to the traceback: prefer bare raise

# edge
class QdrantStore:
    async def search(self, vector: list, top_k: int) -> list:
        try:
            return await self._client.search(self.collection, query_vector=vector, limit=top_k)
        except (httpx.TimeoutException, httpx.ConnectError) as err:
            raise VectorStoreError("qdrant unreachable", store="qdrant", collection=self.collection) from err


# policy
async def retrieve(query: str) -> list:
    vector = await embedder.embed_one(query)
    try:
        return await store.search(vector, top_k=8)
    except VectorStoreError:
        logger.warning("vector search unavailable; falling back to BM25", exc_info=True)
        return bm25.search(query, top_k=8)


# boundary: see section 9

def load_model_config(path: str) -> ModelConfig:
    try:
        fh = open(path, encoding="utf-8")
    except FileNotFoundError:
        raise ConfigurationError(f"model config not found: {path}") from None
    try:
        raw = json.load(fh)
    except json.JSONDecodeError as err:
        raise ConfigurationError(f"{path}:{err.lineno}:{err.colno}: {err.msg}") from err
    else:
        return ModelConfig.model_validate(raw)   # errors here are NOT caught by the JSON handler
    finally:
        fh.close()

results, failures = [], []
for doc in docs:
    try:
        chunks = chunker.split(doc.text)
    except ValueError as err:
        failures.append((doc.id, err))
        logger.warning("skipping doc %s: %s", doc.id, err)
        continue
    results.extend(chunks)
if failures and len(failures) == len(docs):
    raise IngestionError("every document failed to chunk") from failures[0][1]

try:
    payload = json.loads(text)
except json.JSONDecodeError as err:
    raise StructuredOutputError("model did not return JSON", raw_output=text[:500]) from err

for i, batch in enumerate(batched(texts, 64)):
    try:
        vectors.extend(await embedder.embed(batch))
    except ProviderError as err:
        err.add_note(f"batch {i} ({len(batch)} texts, first id={batch[0].id})")
        raise

def add_note(err: BaseException, note: str) -> None:
    if hasattr(err, "add_note"):
        err.add_note(note)
    else:
        err.__notes__ = [*getattr(err, "__notes__", []), note]  # shown by the 3.11+ traceback printer only

from contextlib import contextmanager


@contextmanager
def error_context(**fields):
    try:
        yield
    except Exception as err:
        add_note(err, ", ".join(f"{k}={v!r}" for k, v in fields.items()))
        raise


with error_context(step="rerank", model=RERANKER, query_id=qid):
    ranked = reranker.rank(query, candidates)

async with httpx.AsyncClient(timeout=30) as http, tracer.start_as_current_span("rag.answer") as span:
    ...

from contextlib import ExitStack

with ExitStack() as stack:
    shards = [stack.enter_context(open(p, "rb")) for p in shard_paths]
    index = build_index(shards)   # every opened shard is closed, even if a later open() fails

def close_all(clients) -> None:
    errors = []
    for c in clients:
        try:
            c.close()
        except Exception as err:      # keep closing the rest
            errors.append(err)
    if errors:
        raise ExceptionGroup("failed to close clients", errors)  # 3.11+

logger = logging.getLogger(__name__)

try:
    answer = await rag.answer(question)
except ProviderError:
    logger.exception("answer failed", extra={"request_id": rid, "model": settings.model})
    raise

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse

app = FastAPI()


@app.exception_handler(RagkitError)
async def handle_ragkit_error(request: Request, exc: RagkitError) -> JSONResponse:
    status = next((HTTP_STATUS[c] for c in type(exc).__mro__ if c in HTTP_STATUS), 500)
    headers = {}
    if getattr(exc, "retry_after", None):
        headers["Retry-After"] = str(math.ceil(exc.retry_after))
    if status >= 500:
        logger.error("request failed", exc_info=exc, extra={"path": request.url.path})
        body = {"error": exc.code, "message": "upstream service error"}
    else:
        body = {"error": exc.code, "message": exc.message, **exc.context}
    return JSONResponse(status_code=status, content=body, headers=headers)


@app.exception_handler(Exception)
async def handle_unexpected(request: Request, exc: Exception) -> JSONResponse:
    logger.error("unhandled error", exc_info=exc, extra={"path": request.url.path})
    return JSONResponse(status_code=500, content={"error": "internal_error"})

def main(argv: Optional[list] = None) -> int:
    args = parse_args(argv)
    try:
        run(args)
    except RagkitError as err:
        print(f"error: {err}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("interrupted", file=sys.stderr)
        return 130
    return 0


if __name__ == "__main__":
    sys.exit(main())

async def consume(queue) -> None:
    async for msg in queue:
        try:
            await handle(msg)
        except Exception as err:              # boundary: one bad message must not stop the worker
            logger.exception("message failed", extra={"msg_id": msg.id, "attempt": msg.attempt})
            if is_retryable(err) and msg.attempt < 5:
                await msg.nack(requeue=True)
            else:
                await msg.dead_letter(reason=type(err).__name__)
        else:
            await msg.ack()

from pydantic import BaseModel, Field, ValidationError


class GenerationParams(BaseModel):
    temperature: float = Field(0.2, ge=0.0, le=2.0)
    max_tokens: int = Field(512, gt=0, le=8192)
    top_p: float = Field(1.0, gt=0.0, le=1.0)


try:
    params = GenerationParams.model_validate(user_payload)
except ValidationError as err:
    raise InputValidationError("invalid generation parameters", errors=err.errors()) from err

def merge_scores(dense: np.ndarray, sparse: np.ndarray) -> np.ndarray:
    if dense.shape != sparse.shape:
        raise ValueError(f"shape mismatch: dense {dense.shape} vs sparse {sparse.shape}")
    fused = 0.7 * dense + 0.3 * sparse
    assert fused.shape == dense.shape  # noqa: EXC009 - invariant: never false unless this code is wrong
    return fused

import warnings


def chat(prompt: str, *, model: str = "small", engine: Optional[str] = None) -> str:
    if engine is not None:
        warnings.warn("`engine` is deprecated, use `model`", DeprecationWarning, stacklevel=2)
        model = engine
    ...

from dataclasses import dataclass
from typing import Generic, TypeVar, Union

T = TypeVar("T")


@dataclass(frozen=True)
class Ok(Generic[T]):
    value: T


@dataclass(frozen=True)
class Err:
    error: Exception


Result = Union[Ok[T], Err]


def embed_each(texts: list) -> list:
    out = []
    for text in texts:
        try:
            out.append(Ok(embedder.embed_one(text)))
        except ProviderError as err:
            out.append(Err(err))
    return out


for text, res in zip(texts, embed_each(texts)):
    if isinstance(res, Ok):
        index.add(text, res.value)
    else:
        logger.warning("embedding failed for %r: %s", text[:40], res.error)

match res:
    case Ok(value=vec):
        index.add(text, vec)
    case Err(error=RateLimitError() as err):
        retry_queue.put((text, err.retry_after))
    case Err(error=err):
        dead_letter.append((text, err))

class ContextLengthError(ProviderError):
    def __init__(self, message: str, *, tokens: int, limit: int, **context):
        super().__init__(message, tokens=tokens, limit=limit, **context)
        self.tokens, self.limit = tokens, limit

    def __reduce__(self):
        return (_rebuild, (type(self), self.message, self.context))


def _rebuild(cls, message, context):
    return cls(message, **context)