auditkit 1.0.0__py3-none-any.whl

auditkit/model/lexsi.py

@@ -0,0 +1,155 @@
+"""Lexsi platform model backend with optional MLflow auto-logging."""
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from . import (
+    Model, Request, Result_, Generated, resolve_params, DEFAULT_TEMPERATURE,
+    reject_generation_kwargs, split_session_kwargs,
+)
+from ..errors import ExtraNotInstalled
+
+# The Lexsi platform's real accepted parameter set beyond temperature/
+# max_tokens isn't documented in this codebase -- kept deliberately
+# conservative (only what was already hardcoded here) rather than guessing
+# at unverified field names for a proprietary API.
+_KEY_MAP = {
+    "temperature": "temperature",
+    "max_tokens": "max_tokens",
+}
+
+
+class LexsiModel(Model):
+    """Call the Lexsi platform API with optional MLflow tracking.
+
+    Parameters
+    ----------
+    model
+        Model name on the Lexsi platform.
+    api_key
+        Lexsi API key. Falls back to ``LEXSI_API_KEY`` env var.
+    api_base
+        Base URL for the Lexsi API.
+    mlflow_tracking_uri
+        If set, auto-log every run to MLflow.
+    mlflow_experiment
+        MLflow experiment name (defaults to "auditkit-lexsi").
+    """
+
+    name = "lexsi"
+
+    def __init__(
+        self,
+        model: str = "lexsi-3.5",
+        api_key: str | None = None,
+        api_base: str = "https://api.lexsi.ai/v1",
+        mlflow_tracking_uri: str | None = None,
+        mlflow_experiment: str | None = None,
+        name: str = "lexsi",
+        **kwargs: Any,
+    ) -> None:
+        reject_generation_kwargs(kwargs, "LexsiModel")
+        self._model_name = model
+        self._api_key = api_key or os.environ.get("LEXSI_API_KEY")
+        self._api_base = api_base.rstrip("/")
+        self._mlflow_tracking_uri = mlflow_tracking_uri
+        self._mlflow_experiment = mlflow_experiment or "auditkit-lexsi"
+        # verify=/proxies=/cert=/timeout= are real requests.Session concerns,
+        # not API request-body fields -- split them out so e.g.
+        # LexsiModel(verify=False) actually disables TLS verification
+        # instead of silently becoming an extra, ignored JSON field.
+        self._session_kwargs, self._extra_kwargs = split_session_kwargs(kwargs)
+        self._session = None
+        self.name = name
+        self._total_requests = 0
+        self._total_latency_ms = 0.0
+
+    def _ensure_session(self) -> None:
+        if self._session is not None:
+            return
+        try:
+            import requests
+        except ImportError:
+            raise ExtraNotInstalled("requests", "pip install requests")
+        self._session = requests.Session()
+        self._session.headers.update({
+            "Authorization": f"Bearer {self._api_key}",
+            "Content-Type": "application/json",
+        })
+        for attr in ("verify", "proxies", "cert"):
+            if attr in self._session_kwargs:
+                setattr(self._session, attr, self._session_kwargs[attr])
+
+    def generate(self, requests: list[Request]) -> list[Result_]:
+        self._ensure_session()
+        if self._api_key is None:
+            raise ValueError("Lexsi API key required (pass api_key= or set LEXSI_API_KEY)")
+
+        mlflow_active = False
+        if self._mlflow_tracking_uri:
+            try:
+                import mlflow
+                mlflow.set_tracking_uri(self._mlflow_tracking_uri)
+                mlflow.set_experiment(self._mlflow_experiment)
+                mlflow_active = True
+            except ImportError:
+                pass
+
+        results = []
+        run_metrics = {}
+        defaults = {"temperature": DEFAULT_TEMPERATURE, "max_tokens": 1024}
+        last_gen_kwargs = dict(defaults)
+
+        for r in requests:
+            prompt = r.prompt if isinstance(r.prompt, str) else str(r.prompt)
+            last_gen_kwargs = resolve_params(r, defaults, _KEY_MAP)
+            import time
+            t0 = time.monotonic()
+            resp = self._session.post(
+                f"{self._api_base}/completions",
+                json={
+                    "model": self._model_name,
+                    "prompt": prompt,
+                    **last_gen_kwargs,
+                    **self._extra_kwargs,
+                },
+                timeout=self._session_kwargs.get("timeout", 120),
+            )
+            latency_ms = (time.monotonic() - t0) * 1000
+            self._total_requests += 1
+            self._total_latency_ms += latency_ms
+            try:
+                resp.raise_for_status()
+            except Exception as e:
+                # requests' default HTTPError message drops the response
+                # body, which is where the real API actually explains what's
+                # wrong. Surface it instead of a bare "400 Client Error".
+                raise type(e)(f"{e} -- response body: {resp.text}") from None
+            data = resp.json()
+            text = data.get("choices", [{}])[0].get("text", "")
+            results.append(Result_(completions=[Generated(text=text)], latency_ms=latency_ms))
+
+        if mlflow_active:
+            import mlflow
+            with mlflow.start_run(run_name=f"lexsi-batch-{self._total_requests}"):
+                mlflow.log_param("model", self._model_name)
+                mlflow.log_param("temperature", last_gen_kwargs.get("temperature", DEFAULT_TEMPERATURE))
+                mlflow.log_param("max_tokens", last_gen_kwargs.get("max_tokens", 1024))
+                mlflow.log_param("batch_size", len(requests))
+                mlflow.log_metric("total_requests", self._total_requests)
+                mlflow.log_metric("avg_latency_ms", self._total_latency_ms / max(self._total_requests, 1))
+                if run_metrics:
+                    for k, v in run_metrics.items():
+                        mlflow.log_metric(k, v)
+
+        return results
+
+    def log_mlflow_params(self, **params: Any) -> None:
+        """Manually log parameters to the active MLflow run."""
+        try:
+            import mlflow
+            for k, v in params.items():
+                mlflow.log_param(k, v)
+        except ImportError:
+            pass

auditkit/model/litellm_gen.py

@@ -0,0 +1,65 @@
+"""LiteLLM unified API backend (optional extra: auditkit[litellm])."""
+from __future__ import annotations
+
+from typing import Any
+
+from . import (
+    Model, Request, Result_, Generated, resolve_params, resolve_messages,
+    DEFAULT_TEMPERATURE, reject_generation_kwargs,
+)
+from ..errors import ExtraNotInstalled
+
+# LiteLLM normalizes most providers to OpenAI-style kwarg names.
+_KEY_MAP = {
+    "temperature": "temperature",
+    "top_p": "top_p",
+    "max_tokens": "max_tokens",
+    "stop_sequences": "stop",
+    "presence_penalty": "presence_penalty",
+    "frequency_penalty": "frequency_penalty",
+    "num_completions": "n",
+    "seed": "seed",
+}
+
+
+class LiteLLMModel(Model):
+    """Call any LLM through the LiteLLM unified interface."""
+
+    name = "litellm"
+    threadsafe = True
+
+    def __init__(self, model: str = "gpt-4o-mini",
+                 name: str = "litellm", **kwargs: Any) -> None:
+        reject_generation_kwargs(kwargs, "LiteLLMModel")
+        self._model_name = model
+        self._extra_kwargs = kwargs
+        self.name = name
+
+    def generate(self, requests: list[Request]) -> list[Result_]:
+        try:
+            from litellm import completion
+        except ImportError:
+            raise ExtraNotInstalled("litellm", "pip install auditkit[litellm]")
+        results = []
+        for r in requests:
+            defaults = {"temperature": DEFAULT_TEMPERATURE, "max_tokens": 1024}
+            gen_kwargs = resolve_params(r, defaults, _KEY_MAP)
+            resp = completion(
+                model=self._model_name,
+                messages=resolve_messages(r),
+                **gen_kwargs,
+                **self._extra_kwargs,
+            )
+            text = resp.choices[0].message.content or ""
+            usage = {}
+            resp_usage = getattr(resp, "usage", None)
+            if resp_usage is not None:
+                # LiteLLM normalizes every provider's response to this
+                # OpenAI-shaped usage object regardless of backend.
+                usage = {
+                    "prompt_tokens": resp_usage.prompt_tokens,
+                    "completion_tokens": resp_usage.completion_tokens,
+                    "total_tokens": resp_usage.total_tokens,
+                }
+            results.append(Result_(completions=[Generated(text=text)], usage=usage))
+        return results

auditkit/model/openai.py

@@ -0,0 +1,90 @@
+"""OpenAI model backend (optional extra: auditkit[openai])."""
+from __future__ import annotations
+
+from typing import Any
+
+from . import (
+    Model, Request, Result_, Generated, resolve_params, resolve_messages,
+    DEFAULT_TEMPERATURE, reject_generation_kwargs,
+)
+from ..errors import ExtraNotInstalled
+
+# AuditKit-internal RunConfig key -> OpenAI chat.completions.create() kwarg name.
+# Only these keys are ever pulled from Request.params; timeout/max_retries/
+# top_k/best_of are either Runner-only or unsupported by this API, so they're
+# deliberately absent and can never reach the real API call.
+_KEY_MAP = {
+    "temperature": "temperature",
+    "top_p": "top_p",
+    "max_tokens": "max_tokens",
+    "stop_sequences": "stop",
+    "presence_penalty": "presence_penalty",
+    "frequency_penalty": "frequency_penalty",
+    "num_completions": "n",
+    "seed": "seed",
+}
+
+# Real openai.OpenAI() CLIENT-construction kwargs -- must reach the client
+# constructor in _ensure_client(), not chat.completions.create(). Previously
+# blindly spread into **self._extra_kwargs on the per-request call instead,
+# where max_retries/default_headers/default_query/organization/project
+# don't exist at all and raise TypeError (timeout happens to be valid on
+# both, so it alone never surfaced this). Anything else in **kwargs is
+# assumed to be a real chat.completions.create()-level kwarg and keeps
+# going there, unaffected. api_key/base_url are already handled via their
+# own explicit constructor params, not this set.
+_CLIENT_KWARGS = frozenset({
+    "timeout", "max_retries", "default_headers", "default_query",
+    "organization", "project",
+})
+
+
+class OpenAIModel(Model):
+    """Call the OpenAI chat completions API."""
+
+    name = "openai"
+
+    def __init__(self, model: str = "gpt-4o-mini", api_key: str | None = None,
+                 api_base: str | None = None, name: str = "openai", **kwargs: Any) -> None:
+        reject_generation_kwargs(kwargs, "OpenAIModel")
+        self._model_name = model
+        self._extra_kwargs = kwargs
+        self._client = None
+        self._api_key = api_key
+        self._api_base = api_base
+        self.name = name
+
+    def _ensure_client(self) -> None:
+        if self._client is not None:
+            return
+        try:
+            import openai
+        except ImportError:
+            raise ExtraNotInstalled("openai", "pip install auditkit[openai]")
+        client_kwargs = {k: v for k, v in self._extra_kwargs.items() if k in _CLIENT_KWARGS}
+        self._client = openai.OpenAI(api_key=self._api_key, base_url=self._api_base, **client_kwargs)
+
+    def generate(self, requests: list[Request]) -> list[Result_]:
+        self._ensure_client()
+        results = []
+        for r in requests:
+            defaults = {"temperature": DEFAULT_TEMPERATURE, "max_tokens": 1024}
+            gen_kwargs = resolve_params(r, defaults, _KEY_MAP)
+            call_kwargs = {k: v for k, v in self._extra_kwargs.items() if k not in _CLIENT_KWARGS}
+            resp = self._client.chat.completions.create(
+                model=self._model_name,
+                messages=resolve_messages(r),
+                **gen_kwargs,
+                **call_kwargs,
+            )
+            text = resp.choices[0].message.content or ""
+            usage = {}
+            resp_usage = getattr(resp, "usage", None)
+            if resp_usage is not None:
+                usage = {
+                    "prompt_tokens": resp_usage.prompt_tokens,
+                    "completion_tokens": resp_usage.completion_tokens,
+                    "total_tokens": resp_usage.total_tokens,
+                }
+            results.append(Result_(completions=[Generated(text=text)], usage=usage))
+        return results

auditkit/model/openrouter_gen.py

@@ -0,0 +1,152 @@
+"""OpenRouter model backend (optional extra: auditkit[requests]).
+
+OpenRouter's API is a documented OpenAI-compatible endpoint
+(https://openrouter.ai/api/v1, standard Bearer auth, /chat/completions
+request/response shape) -- this backend is a thin, pre-configured wrapper so
+you don't have to pass api_base=/api_key= through the generic api: backend
+by hand every time. Same shape as GroqModel; OpenRouter itself is a router
+in front of many providers (OpenAI, Anthropic, Meta, Mistral, Google, ...),
+selected via a "provider/model" name (e.g. "openai/gpt-4o",
+"anthropic/claude-3.5-sonnet", "meta-llama/llama-3.1-70b-instruct") rather
+than one fixed model family.
+"""
+from __future__ import annotations
+
+import os
+import time
+from typing import Any
+
+from . import (
+    Model, Request, Result_, Generated, resolve_params, resolve_messages,
+    DEFAULT_TEMPERATURE, reject_generation_kwargs, split_session_kwargs,
+)
+from ..errors import ExtraNotInstalled
+
+_KEY_MAP = {
+    "temperature": "temperature",
+    "top_p": "top_p",
+    "max_tokens": "max_tokens",
+    "stop_sequences": "stop",
+    "presence_penalty": "presence_penalty",
+    "frequency_penalty": "frequency_penalty",
+    "num_completions": "n",
+    "seed": "seed",
+}
+
+
+class OpenRouterModel(Model):
+    """Call the OpenRouter chat completions API.
+
+    Parameters
+    ----------
+    model
+        Model name on OpenRouter, in ``"provider/model"`` form (e.g.
+        ``"openai/gpt-4o"``, ``"anthropic/claude-3.5-sonnet"``,
+        ``"meta-llama/llama-3.1-70b-instruct"``). The provider prefix is
+        OpenRouter's own routing syntax, unrelated to AuditKit's
+        ``"openrouter:"`` spec prefix -- ``"openrouter:openai/gpt-4o"``
+        parses correctly since only the text before the *first* colon is
+        AuditKit's own prefix.
+    api_key
+        OpenRouter API key. Falls back to the ``OPENROUTER_API_KEY`` env var.
+    api_base
+        Base URL (default: ``https://openrouter.ai/api/v1``).
+    site_url, app_name
+        Optional -- OpenRouter's own leaderboard/analytics attribution
+        headers (``HTTP-Referer``/``X-Title``). Not required for the API
+        to work; harmless to omit.
+    """
+
+    name = "openrouter"
+    threadsafe = False
+
+    def __init__(
+        self,
+        model: str = "openai/gpt-4o-mini",
+        api_key: str | None = None,
+        api_base: str = "https://openrouter.ai/api/v1",
+        site_url: str | None = None,
+        app_name: str | None = None,
+        name: str = "openrouter",
+        **kwargs: Any,
+    ) -> None:
+        reject_generation_kwargs(kwargs, "OpenRouterModel")
+        self._model_name = model
+        self._api_base = api_base.rstrip("/")
+        self._api_key = api_key or os.environ.get("OPENROUTER_API_KEY")
+        self._site_url = site_url
+        self._app_name = app_name
+        # verify=/proxies=/cert=/timeout= are real requests.Session concerns,
+        # not API request-body fields -- split them out so e.g.
+        # OpenRouterModel(verify=False) actually disables TLS verification
+        # instead of silently becoming an extra, ignored JSON field.
+        self._session_kwargs, self._extra_kwargs = split_session_kwargs(kwargs)
+        self._session = None
+        self.name = name
+
+    def _ensure_session(self) -> None:
+        if self._session is not None:
+            return
+        try:
+            import requests
+        except ImportError:
+            raise ExtraNotInstalled("requests", "pip install auditkit[requests]")
+        self._session = requests.Session()
+        headers = {
+            "Authorization": f"Bearer {self._api_key}" if self._api_key else "",
+            "Content-Type": "application/json",
+        }
+        if self._site_url:
+            headers["HTTP-Referer"] = self._site_url
+        if self._app_name:
+            headers["X-Title"] = self._app_name
+        self._session.headers.update(headers)
+        for attr in ("verify", "proxies", "cert"):
+            if attr in self._session_kwargs:
+                setattr(self._session, attr, self._session_kwargs[attr])
+
+    def generate(self, requests: list[Request]) -> list[Result_]:
+        self._ensure_session()
+        results = []
+        for r in requests:
+            defaults = {"temperature": DEFAULT_TEMPERATURE, "max_tokens": 1024}
+            gen_kwargs = resolve_params(r, defaults, _KEY_MAP)
+            body = {
+                "model": self._model_name,
+                "messages": resolve_messages(r),
+                **gen_kwargs,
+                **self._extra_kwargs,
+            }
+            timeout = self._session_kwargs.get("timeout", 120)
+            t0 = time.monotonic()
+            resp = self._session.post(f"{self._api_base}/chat/completions", json=body, timeout=timeout)
+            try:
+                resp.raise_for_status()
+            except Exception as e:
+                # requests' default HTTPError message drops the response body,
+                # which is where OpenRouter (and the underlying provider it
+                # routed to) actually explains what's wrong -- an invalid
+                # model slug, an unsupported param for that specific model,
+                # rate limiting, etc. Surface it instead of a bare "400
+                # Client Error: Bad Request".
+                raise type(e)(f"{e} -- response body: {resp.text}") from None
+            latency_ms = (time.monotonic() - t0) * 1000
+            data = resp.json()
+            message = data.get("choices", [{}])[0].get("message", {}) or {}
+            # Reasoning models (e.g. some OpenRouter-routed models) can return
+            # an explicit `"content": null` -- not a missing key -- when
+            # max_tokens is hit while the model is still inside its internal
+            # `reasoning` field. dict.get(key, "") only falls back to "" when
+            # the key is *absent*, so a present-but-null value passes through
+            # as None otherwise; `or ""` catches both cases.
+            text = message.get("content") or ""
+            raw_usage = data.get("usage") or {}
+            usage = {}
+            if raw_usage:
+                usage = {
+                    "prompt_tokens": raw_usage.get("prompt_tokens", 0),
+                    "completion_tokens": raw_usage.get("completion_tokens", 0),
+                    "total_tokens": raw_usage.get("total_tokens", 0),
+                }
+            results.append(Result_(completions=[Generated(text=text)], usage=usage, latency_ms=latency_ms))
+        return results