plotaviz 0.1.0__py3-none-any.whl

plotaviz/core/llm/__init__.py

@@ -0,0 +1,211 @@
+"""The LLM layer — optional intelligence on top of a fully working local engine.
+
+PlotaViz is **offline-first**. The rules + scoring engine answers instantly with no provider
+configured, and it is the fallback whenever an LLM call fails, times out, or returns something
+that does not validate. A model is never on the critical path.
+
+The layer does two jobs:
+
+1. **A second opinion on selection** when the local scores are close enough that the ranking is
+   effectively a coin flip.
+2. **Natural language → chart spec**, which is the only way to get from "revenue by region over
+   time" to a rendered chart.
+
+Both return a *validated chart spec*. No code is ever generated, returned, or executed.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from ..errors import LLMError, ProviderNotConfigured
+from ..profiler import DatasetProfile
+from ..spec import ChartSpec
+from .anthropic import AnthropicProvider
+from .base import (
+    KEYRING_SERVICE,
+    LLMResult,
+    Provider,
+    build_prompt,
+    delete_api_key,
+    extract_json,
+    get_api_key,
+    redact,
+    sample_rows_for,
+    set_api_key,
+)
+from .gemini import GeminiProvider
+from .ollama import OllamaProvider
+from .openai import OpenAIProvider
+
+logger = logging.getLogger(__name__)
+
+#: Every provider PlotaViz can use, keyed by the name stored in settings and the keyring.
+PROVIDERS: dict[str, type[Provider]] = {
+    AnthropicProvider.name: AnthropicProvider,
+    OpenAIProvider.name: OpenAIProvider,
+    GeminiProvider.name: GeminiProvider,
+    OllamaProvider.name: OllamaProvider,
+}
+
+__all__ = [
+    "KEYRING_SERVICE",
+    "PROVIDERS",
+    "AnthropicProvider",
+    "GeminiProvider",
+    "LLMAssistant",
+    "LLMResult",
+    "OllamaProvider",
+    "OpenAIProvider",
+    "Provider",
+    "build_prompt",
+    "delete_api_key",
+    "extract_json",
+    "get_api_key",
+    "get_provider",
+    "redact",
+    "sample_rows_for",
+    "set_api_key",
+]
+
+
+def get_provider(name: str, **kwargs: Any) -> Provider:
+    """Instantiate a provider by name.
+
+    Raises:
+        ProviderNotConfigured: If ``name`` is not a known provider.
+    """
+    cls = PROVIDERS.get(name.lower().strip())
+    if cls is None:
+        raise ProviderNotConfigured(
+            f"Unknown LLM provider {name!r}.",
+            hint=f"Available providers: {', '.join(PROVIDERS)}.",
+        )
+    return cls(**kwargs)
+
+
+class LLMAssistant:
+    """Optional LLM help, with the local engine always underneath.
+
+    Args:
+        provider: A configured provider, or ``None`` for local-only operation.
+        consent: Whether the user has agreed to send schema and statistics to a remote service.
+            Required before any provider that leaves the machine is used. Ollama does not need it.
+
+    Example:
+        >>> assistant = LLMAssistant(get_provider("ollama"), consent=True)
+        >>> question = "revenue by region over time"
+        >>> result = assistant.from_question(profile, question)  # doctest: +SKIP
+    """
+
+    def __init__(self, provider: Provider | None = None, *, consent: bool = False) -> None:
+        self.provider = provider
+        self.consent = consent
+
+    # ------------------------------------------------------------------ state
+
+    @property
+    def available(self) -> bool:
+        """Whether a request could be made right now — configured, and consented to if remote."""
+        if self.provider is None:
+            return False
+        if self.provider.sends_data_off_machine and not self.consent:
+            return False
+        return self.provider.is_configured()
+
+    def consent_prompt(self) -> str | None:
+        """The text to show before the first remote request, or ``None`` if none is needed."""
+        if self.provider is None or not self.provider.sends_data_off_machine:
+            return None
+        return (
+            f"PlotaViz will send this dataset's column names, summary statistics, and up to "
+            f"{5} sample rows to {self.provider.name}. The dataset itself is never uploaded.\n\n"
+            "To keep everything on this machine, use the Ollama provider instead."
+        )
+
+    # ------------------------------------------------------------------ jobs
+
+    def from_question(
+        self,
+        profile: DatasetProfile,
+        question: str,
+        *,
+        sample_rows: list[dict[str, Any]] | None = None,
+    ) -> LLMResult:
+        """Turn a natural-language request into a validated chart spec.
+
+        Raises:
+            LLMError: If no provider is available, or the request fails. There is no local
+                fallback for this job — a free-text question genuinely needs a model — so the
+                caller should show the error and leave the current chart alone.
+        """
+        if self.provider is None:
+            raise ProviderNotConfigured(
+                "The natural-language query bar needs an LLM provider.",
+                hint=(
+                    "Configure one in Settings → LLM provider. For a fully local setup, install "
+                    "Ollama and pull a model. The chart recommendations work without any of this."
+                ),
+            )
+        if self.provider.sends_data_off_machine and not self.consent:
+            raise LLMError(
+                "Sending data to a remote provider has not been approved for this session.",
+                hint=(
+                    "Accept the consent prompt in Settings, or switch to the local Ollama provider."
+                ),
+            )
+
+        prompt_rows = sample_rows or []
+        result = self.provider.suggest_chart(profile, question=question)
+        logger.debug(
+            "NL query answered by %s/%s with %d sample rows",
+            result.provider,
+            result.model,
+            len(prompt_rows),
+        )
+        return result
+
+    def refine(
+        self,
+        profile: DatasetProfile,
+        candidates: list[ChartSpec],
+        *,
+        fallback: ChartSpec | None = None,
+    ) -> ChartSpec:
+        """Ask for a second opinion on an ambiguous ranking, falling back silently.
+
+        This is the failure-tolerant path: a timeout, a bad key, a malformed reply, or no
+        provider at all all end the same way — with the local engine's answer. The user's chart
+        appears either way.
+
+        Args:
+            profile: The dataset profile.
+            candidates: The local engine's ranked candidates.
+            fallback: Spec to return on any failure. Defaults to the top candidate.
+
+        Returns:
+            Either the model's validated spec or the fallback.
+        """
+        default = fallback or (candidates[0] if candidates else None)
+        if default is None:
+            raise LLMError("There are no candidate charts to refine.")
+        if not self.available:
+            return default
+
+        try:
+            assert self.provider is not None  # narrowed by self.available
+            result = self.provider.suggest_chart(profile, candidates=candidates)
+        except LLMError as exc:
+            logger.info("LLM refinement failed, using the local ranking: %s", redact(str(exc)))
+            return default
+        except Exception as exc:
+            logger.warning("Unexpected LLM error, using the local ranking: %s", redact(str(exc)))
+            return default
+
+        spec = result.spec
+        if not spec.why:
+            spec.why = f"Suggested by {result.provider} ({result.model})."
+        else:
+            spec.why = f"{spec.why} (Suggested by {result.provider}.)"
+        return spec

plotaviz/core/llm/anthropic.py

@@ -0,0 +1,61 @@
+"""Anthropic (Claude) provider."""
+
+from __future__ import annotations
+
+from ..errors import LLMError
+from .base import Provider, redact
+
+
+class AnthropicProvider(Provider):
+    """Talks to the Anthropic Messages API.
+
+    Requires the ``anthropic`` package, which is an optional extra:
+    ``pip install "plotaviz[llm]"``.
+    """
+
+    name = "anthropic"
+    default_model = "claude-sonnet-4-5"
+    available_models = (
+        "claude-opus-4-5",
+        "claude-sonnet-4-5",
+        "claude-haiku-4-5",
+    )
+
+    def complete(self, system: str, user: str, *, timeout: float) -> str:
+        """Send one message and return the text of the reply."""
+        try:
+            import anthropic
+        except ImportError as exc:
+            raise LLMError(
+                "The Anthropic provider needs the anthropic package.",
+                hint='Install it with: pip install "plotaviz[llm]"',
+            ) from exc
+
+        client = anthropic.Anthropic(api_key=self.require_key(), timeout=timeout)
+        try:
+            message = client.messages.create(
+                model=self.model,
+                max_tokens=1024,
+                system=system,
+                messages=[{"role": "user", "content": user}],
+            )
+        except anthropic.AuthenticationError as exc:
+            raise LLMError(
+                "Anthropic rejected the API key.",
+                hint="Check the key in Settings → LLM provider.",
+            ) from exc
+        except anthropic.RateLimitError as exc:
+            raise LLMError(
+                "Anthropic is rate limiting this key.", hint="Wait a moment and try again."
+            ) from exc
+        except anthropic.APITimeoutError as exc:
+            raise LLMError(
+                f"Anthropic did not respond within {timeout:.0f} seconds.",
+                hint="PlotaViz will use its own rules engine instead.",
+            ) from exc
+        except anthropic.APIError as exc:
+            raise LLMError("The Anthropic request failed.", hint=redact(str(exc))) from exc
+
+        return "".join(
+            block.text for block in message.content if getattr(block, "type", "") == "text"
+        )

plotaviz/core/llm/base.py

@@ -0,0 +1,373 @@
+"""Provider abstraction and the prompt/validation machinery around it.
+
+A provider's entire job is to turn two strings into one string. Everything else — building the
+prompt, deciding what data is safe to send, extracting JSON from a chatty reply, validating the
+returned spec against the real schema — happens here, once, so every provider behaves identically
+and a new one is about thirty lines.
+
+Two rules are structural, not stylistic:
+
+**The model returns a chart spec, never code.** PlotaViz does not execute anything a model
+produces. The reply is parsed as JSON and validated against the dataframe's actual columns; an
+unknown column is a rejection with a clear message, not a traceback.
+
+**The model sees schema, statistics, and a few sample rows.** Never the full dataset. This is a
+privacy boundary and a token-cost boundary at the same time.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+from ..errors import LLMError, ProviderNotConfigured
+from ..profiler import DatasetProfile
+from ..spec import CHART_TYPES, ChartSpec
+
+#: Service name used for keyring entries. One entry per provider.
+KEYRING_SERVICE = "plotaviz"
+
+#: Default per-request timeout in seconds. LLM latency must never hang the UI.
+DEFAULT_TIMEOUT = 30.0
+
+#: Sample rows included in a prompt. Small on purpose.
+SAMPLE_ROWS = 5
+
+_SYSTEM_PROMPT = f"""You are the chart-selection assistant inside PlotaViz, a desktop data \
+visualization tool.
+
+You receive a dataset's schema and summary statistics — never the dataset itself. You reply with \
+a single JSON object describing which chart to draw. You never write code, never write prose \
+outside the JSON, and never invent column names.
+
+Reply with exactly this shape:
+
+{{
+  "chart": one of {list(CHART_TYPES)},
+  "x": column name or null,
+  "y": column name or null,
+  "color": column name or null,
+  "agg": one of ["sum","mean","median","min","max","count","nunique"] or null,
+  "title": short human title or null,
+  "options": {{}},
+  "why": one or two sentences explaining the choice in plain language
+}}
+
+Rules:
+- Every column name you use MUST appear in the provided schema, spelled exactly.
+- Use "agg" whenever y would otherwise have several values per x.
+- Prefer readable charts: do not put a 500-category column on a bar axis.
+- Columns marked "is_identifier": true are keys, not measures. Do not plot them.
+- Respond with the JSON object alone. No markdown fences, no commentary.
+"""
+
+
+@dataclass
+class LLMResult:
+    """What the LLM layer returns to the app.
+
+    Attributes:
+        spec: The validated chart spec.
+        narrative: Optional insight text the model offered alongside the choice.
+        provider: Which provider answered.
+        model: Which model answered.
+        raw: The raw reply, kept for the "show me what it said" affordance.
+    """
+
+    spec: ChartSpec
+    narrative: str = ""
+    provider: str = ""
+    model: str = ""
+    raw: str = ""
+
+
+class Provider(ABC):
+    """Base class for an LLM provider.
+
+    Subclasses implement :meth:`complete` and declare :attr:`name` and :attr:`default_model`.
+    Everything else is inherited.
+
+    Args:
+        model: Model identifier. Falls back to :attr:`default_model`.
+        api_key: Explicit key. When omitted, the key is read from the OS keyring.
+        timeout: Per-request timeout in seconds.
+    """
+
+    #: Identifier used in settings and keyring entries.
+    name: str = "base"
+
+    #: Model used when the user has not chosen one.
+    default_model: str = ""
+
+    #: Models offered in the settings dialog.
+    available_models: tuple[str, ...] = ()
+
+    #: Whether this provider sends data off the machine. False for local runtimes.
+    sends_data_off_machine: bool = True
+
+    def __init__(
+        self,
+        model: str | None = None,
+        api_key: str | None = None,
+        timeout: float = DEFAULT_TIMEOUT,
+    ) -> None:
+        self.model = model or self.default_model
+        self._api_key = api_key
+        self.timeout = timeout
+
+    # ------------------------------------------------------------------ credentials
+
+    @property
+    def api_key(self) -> str | None:
+        """The API key, from the constructor or the OS keyring.
+
+        Keys are never read from environment variables or config files by default — the keyring
+        is the one storage location, so there is one place to audit.
+        """
+        if self._api_key:
+            return self._api_key
+        return get_api_key(self.name)
+
+    def require_key(self) -> str:
+        """Return the API key or explain how to set one.
+
+        Raises:
+            ProviderNotConfigured: If no key is stored.
+        """
+        key = self.api_key
+        if not key:
+            raise ProviderNotConfigured(
+                f"No API key is stored for {self.name}.",
+                hint="Add one in Settings → LLM provider. Keys are kept in your OS keychain.",
+            )
+        return key
+
+    def is_configured(self) -> bool:
+        """Whether this provider could make a request right now."""
+        return bool(self.api_key)
+
+    # ------------------------------------------------------------------ contract
+
+    @abstractmethod
+    def complete(self, system: str, user: str, *, timeout: float) -> str:
+        """Send one prompt and return the raw text reply.
+
+        Args:
+            system: System prompt.
+            user: User message.
+            timeout: Seconds to wait before giving up.
+
+        Raises:
+            LLMError: On any transport, auth, or API failure.
+        """
+
+    # ------------------------------------------------------------------ tasks
+
+    def suggest_chart(
+        self,
+        profile: DatasetProfile,
+        *,
+        question: str | None = None,
+        candidates: list[ChartSpec] | None = None,
+    ) -> LLMResult:
+        """Ask the model for a chart spec.
+
+        Serves both jobs the LLM layer exists for: with ``question`` it is natural-language →
+        spec; without one it is a second opinion on an ambiguous ranking.
+
+        Args:
+            profile: The dataset profile. Only its schema summary is sent.
+            question: The user's natural-language request, if any.
+            candidates: The local engine's top candidates, sent as context so the model can agree
+                with a good ranking rather than inventing something worse.
+
+        Returns:
+            A validated :class:`LLMResult`.
+
+        Raises:
+            LLMError: If the request fails or the reply is not a usable spec. Callers are
+                expected to catch this and fall back to the local engine.
+        """
+        prompt = build_prompt(profile, question=question, candidates=candidates)
+        raw = self.complete(_SYSTEM_PROMPT, prompt, timeout=self.timeout)
+        payload = extract_json(raw)
+        spec = ChartSpec.from_dict(payload)
+        spec.source = "llm"
+        spec.validate(list(profile.columns))
+        return LLMResult(
+            spec=spec,
+            narrative=str(payload.get("narrative") or payload.get("insight") or ""),
+            provider=self.name,
+            model=self.model,
+            raw=raw,
+        )
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}(model={self.model!r})"
+
+
+# ---------------------------------------------------------------------------- prompts
+
+
+def build_prompt(
+    profile: DatasetProfile,
+    *,
+    question: str | None = None,
+    candidates: list[ChartSpec] | None = None,
+    sample_rows: list[dict[str, Any]] | None = None,
+) -> str:
+    """Assemble the user message.
+
+    Contains the schema, per-column statistics, at most :data:`SAMPLE_ROWS` example rows, and the
+    local engine's candidates. Deliberately does not contain the dataset.
+    """
+    summary = profile.schema_summary()
+    parts = [
+        "Dataset schema and summary statistics:",
+        json.dumps(summary, indent=2, default=str),
+    ]
+
+    if sample_rows:
+        parts += [
+            "",
+            f"A sample of {min(len(sample_rows), SAMPLE_ROWS)} rows (illustrative only):",
+            json.dumps(sample_rows[:SAMPLE_ROWS], indent=2, default=str),
+        ]
+
+    if candidates:
+        ranked = [
+            {"chart": c.chart, "x": c.x, "y": c.y, "color": c.color, "score": round(c.score, 3)}
+            for c in candidates[:5]
+        ]
+        parts += [
+            "",
+            "PlotaViz's own rules engine ranked these candidates. Their scores were close, "
+            "which is why you are being asked:",
+            json.dumps(ranked, indent=2),
+        ]
+
+    if question:
+        parts += [
+            "",
+            f"The user asked for: {question!r}",
+            "",
+            "Return the chart spec that answers it.",
+        ]
+    else:
+        parts += ["", "Return the single chart spec that best reveals what is in this dataset."]
+
+    return "\n".join(parts)
+
+
+def sample_rows_for(df: Any, n: int = SAMPLE_ROWS) -> list[dict[str, Any]]:
+    """Take a few example rows in a JSON-safe form.
+
+    Kept tiny by design — enough for the model to see value formats, not enough to constitute
+    sending someone's data anywhere.
+    """
+    try:
+        head = df.head(n)
+        return json.loads(head.to_json(orient="records", date_format="iso")) or []
+    except Exception:
+        return []
+
+
+def extract_json(text: str) -> dict[str, Any]:
+    """Pull the JSON object out of a model reply.
+
+    Models add markdown fences and preambles no matter how firmly they are told not to, so this
+    tolerates both rather than failing on formatting.
+
+    Raises:
+        LLMError: If no JSON object can be found or parsed.
+    """
+    if not text or not text.strip():
+        raise LLMError("The model returned an empty response.")
+
+    cleaned = text.strip()
+    fenced = re.search(r"```(?:json)?\s*(.+?)\s*```", cleaned, re.DOTALL)
+    if fenced:
+        cleaned = fenced.group(1).strip()
+
+    try:
+        payload = json.loads(cleaned)
+    except json.JSONDecodeError:
+        start, end = cleaned.find("{"), cleaned.rfind("}")
+        if start == -1 or end <= start:
+            raise LLMError(
+                "The model's reply did not contain a chart specification.",
+                hint="Try rephrasing, or use the ranked recommendations from the local engine.",
+            ) from None
+        try:
+            payload = json.loads(cleaned[start : end + 1])
+        except json.JSONDecodeError as exc:
+            raise LLMError("The model's reply was not valid JSON.", hint=str(exc)) from exc
+
+    if not isinstance(payload, dict):
+        raise LLMError("The model returned JSON, but not a chart specification object.")
+    return payload
+
+
+# ---------------------------------------------------------------------------- keyring
+
+
+def get_api_key(provider: str) -> str | None:
+    """Read a provider's API key from the OS keyring.
+
+    Returns ``None`` when nothing is stored or no keyring backend is available — a missing key is
+    an ordinary state, since PlotaViz works fully offline.
+    """
+    try:
+        import keyring
+
+        return keyring.get_password(KEYRING_SERVICE, provider)
+    except Exception:
+        return None
+
+
+def set_api_key(provider: str, key: str) -> None:
+    """Store a provider's API key in the OS keyring.
+
+    Raises:
+        LLMError: If no keyring backend is available. PlotaViz does not fall back to writing the
+            key to a file — that is exactly the failure mode the keyring exists to prevent.
+    """
+    try:
+        import keyring
+
+        keyring.set_password(KEYRING_SERVICE, provider, key)
+    except Exception as exc:
+        raise LLMError(
+            "Could not store the API key in your system keychain.",
+            hint=(
+                "PlotaViz will not write keys to disk in plain text. On Linux, install a Secret "
+                f"Service backend such as gnome-keyring or kwallet. Underlying error: {exc}"
+            ),
+        ) from exc
+
+
+def delete_api_key(provider: str) -> None:
+    """Remove a provider's stored API key. Missing entries are not an error."""
+    try:
+        import keyring
+
+        keyring.delete_password(KEYRING_SERVICE, provider)
+    except Exception:
+        pass
+
+
+def redact(text: str) -> str:
+    """Mask anything that looks like an API key before it reaches a log or an error dialog."""
+    patterns = (
+        r"sk-ant-[A-Za-z0-9\-_]{8,}",
+        r"sk-[A-Za-z0-9]{16,}",
+        r"AIza[0-9A-Za-z\-_]{20,}",
+        r"Bearer\s+[A-Za-z0-9\-._~+/]{16,}",
+    )
+    out = text
+    for pattern in patterns:
+        out = re.sub(pattern, "[redacted]", out)
+    return out

plotaviz/core/llm/gemini.py

@@ -0,0 +1,51 @@
+"""Google Gemini provider."""
+
+from __future__ import annotations
+
+from ..errors import LLMError
+from .base import Provider, redact
+
+
+class GeminiProvider(Provider):
+    """Talks to the Google Gemini API.
+
+    Requires the ``google-genai`` package: ``pip install "plotaviz[llm]"``.
+    """
+
+    name = "gemini"
+    default_model = "gemini-2.0-flash"
+    available_models = ("gemini-2.5-pro", "gemini-2.0-flash", "gemini-2.0-flash-lite")
+
+    def complete(self, system: str, user: str, *, timeout: float) -> str:
+        """Send one message and return the text of the reply."""
+        try:
+            from google import genai
+            from google.genai import types
+        except ImportError as exc:
+            raise LLMError(
+                "The Gemini provider needs the google-genai package.",
+                hint='Install it with: pip install "plotaviz[llm]"',
+            ) from exc
+
+        client = genai.Client(api_key=self.require_key())
+        try:
+            response = client.models.generate_content(
+                model=self.model,
+                contents=user,
+                config=types.GenerateContentConfig(
+                    system_instruction=system,
+                    response_mime_type="application/json",
+                    max_output_tokens=1024,
+                    http_options=types.HttpOptions(timeout=int(timeout * 1000)),
+                ),
+            )
+        except Exception as exc:
+            message = str(exc)
+            if "API key" in message or "PERMISSION_DENIED" in message:
+                raise LLMError(
+                    "Google rejected the API key.",
+                    hint="Check the key in Settings → LLM provider.",
+                ) from exc
+            raise LLMError("The Gemini request failed.", hint=redact(message)) from exc
+
+        return response.text or ""