ragaudit-cli 0.3.0__py3-none-any.whl

ragaudit/llm/factory.py

@@ -0,0 +1,70 @@
+"""Builds the LLM client for v2 commands, and enforces --max-llm-calls budgets.
+
+Provider is `mock` (the offline MockLLM) or `openai` (the existing
+OpenAI-compatible LLMClient). It comes from the `--llm` flag, else the
+RAGAUDIT_LLM environment variable, else `openai`, which keeps v1's
+"configure with RAGAUDIT_API_*" behaviour as the default.
+"""
+
+import os
+from typing import Protocol
+
+from ragaudit.llm.client import LLMClient
+from ragaudit.llm.mock import MockLLM
+
+PROVIDER_ENV = "RAGAUDIT_LLM"
+PROVIDERS = ("openai", "mock")
+
+
+class Completer(Protocol):
+    model: str
+
+    def complete(self, system: str, user: str) -> str: ...
+
+
+class LLMBudgetExceeded(Exception):
+    """Raised when a call would exceed --max-llm-calls. Callers stop cleanly and save progress."""
+
+
+def resolve_provider(provider: str | None = None) -> str:
+    resolved = (provider or os.environ.get(PROVIDER_ENV) or "openai").lower()
+    if resolved not in PROVIDERS:
+        raise ValueError(f"unknown LLM provider {resolved!r}; expected one of {', '.join(PROVIDERS)}")
+    return resolved
+
+
+def make_client(provider: str | None = None, model: str | None = None, mock_model: str = "mock-llm") -> Completer:
+    """An LLM client for the given provider. Real clients read RAGAUDIT_API_BASE/_API_KEY (and _MODEL if no model)."""
+    if resolve_provider(provider) == "mock":
+        return MockLLM(model=model or mock_model)
+    return LLMClient(model=model)
+
+
+class CallBudget:
+    """A shared counter of LLM calls with an optional cap."""
+
+    def __init__(self, max_calls: int | None = None) -> None:
+        self.max_calls = max_calls
+        self.used = 0
+
+    @property
+    def remaining(self) -> int | None:
+        return None if self.max_calls is None else max(self.max_calls - self.used, 0)
+
+    def spend(self) -> None:
+        if self.max_calls is not None and self.used >= self.max_calls:
+            raise LLMBudgetExceeded(f"LLM call budget of {self.max_calls} reached")
+        self.used += 1
+
+
+class BudgetedCompleter:
+    """Wraps a Completer so every call is counted against a shared CallBudget."""
+
+    def __init__(self, inner: Completer, budget: CallBudget) -> None:
+        self._inner = inner
+        self.budget = budget
+        self.model = inner.model
+
+    def complete(self, system: str, user: str) -> str:
+        self.budget.spend()
+        return self._inner.complete(system, user)

ragaudit/llm/json_output.py

@@ -0,0 +1,72 @@
+"""Pull one JSON object out of an LLM reply, tolerating the ways real models get it slightly wrong.
+
+Every v2 prompt asks for strict JSON, but local models often wrap it in a
+markdown fence, add a sentence before or after it, put a raw newline inside
+a string, or drop the final closing brace. Each of those used to throw away
+an otherwise good answer. This tries, in order:
+
+1. the reply as is (inside a markdown fence, if there is one);
+2. the span from the first "{" to the last "}";
+3. the span from the first "{" with any unclosed brackets closed.
+
+Strings are parsed with strict=False, so raw control characters inside a
+string are accepted (callers that need one-line text check that
+themselves). Anything still unparseable returns None.
+"""
+
+import json
+import re
+
+_FENCE = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL)
+
+
+def _close_brackets(text: str) -> str | None:
+    """Append the closers for any "{" / "[" left open, or None if the text ends inside a string."""
+    stack: list[str] = []
+    in_string = escaped = False
+    for char in text:
+        if in_string:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == '"':
+                in_string = False
+        elif char == '"':
+            in_string = True
+        elif char in "{[":
+            stack.append("}" if char == "{" else "]")
+        elif char in "}]":
+            if not stack or stack.pop() != char:
+                return None
+    if in_string:
+        return None
+    return text + "".join(reversed(stack))
+
+
+def _candidates(text: str) -> list[str]:
+    candidates = [text]
+    start, end = text.find("{"), text.rfind("}")
+    if start != -1 and end > start:
+        candidates.append(text[start : end + 1])
+    if start != -1:
+        closed = _close_brackets(text[start:].rstrip())
+        if closed is not None:
+            candidates.append(closed)
+    return candidates
+
+
+def parse_json_object(raw: str) -> dict | None:
+    """The first JSON object found in raw, or None."""
+    text = raw.strip()
+    fenced = _FENCE.search(text)
+    if fenced:
+        text = fenced.group(1)
+    for candidate in _candidates(text):
+        try:
+            data = json.loads(candidate, strict=False)
+        except json.JSONDecodeError:
+            continue
+        if isinstance(data, dict):
+            return data
+    return None