schemagate 0.1.0__py3-none-any.whl

schemagate/__init__.py

@@ -0,0 +1,29 @@
+"""schemagate -- identity-scoped schema selection for NL2SQL.
+
+    from schemagate import Catalog, Principal
+    cat = Catalog().bootstrap("postgresql://localhost/app")
+    sel = cat.select("revenue by month", principal=Principal("okta:jdoe"))
+    sel.prompt_fragment()
+"""
+from .identity import Principal, IdentityError
+from .models import Column, ForeignKey, ObjectDoc, Selection, Scored
+from .embedder import HashingEmbedder, cosine_distance, tokenize
+from .stores.memory import MemoryStore
+from .catalog import Catalog
+
+__version__ = "0.1.0"
+__all__ = ["Catalog", "Principal", "IdentityError", "ObjectDoc", "Column",
+           "ForeignKey", "Selection", "Scored", "HashingEmbedder",
+           "MemoryStore", "cosine_distance", "tokenize"]
+
+
+def __getattr__(name):
+    # Lazy, so importing schemagate never pulls in oracledb or
+    # sentence-transformers. Both raise a clear ImportError naming the extra.
+    if name == "OracleStore":
+        from .stores.oracle import OracleStore
+        return OracleStore
+    if name == "SentenceTransformerEmbedder":
+        from .embedders.hf import SentenceTransformerEmbedder
+        return SentenceTransformerEmbedder
+    raise AttributeError(f"module 'schemagate' has no attribute {name!r}")

schemagate/ai/__init__.py

@@ -0,0 +1,28 @@
+"""Optional AI features. schemagate works fully without importing this package.
+
+Nothing here is required. The default catalog uses an offline, deterministic
+embedder and no network. Import from ``schemagate.ai`` only when you want a model
+to write catalog descriptions or produce embeddings, and bring your own key.
+
+    from schemagate.ai import SchemaDescriber, AnthropicProvider
+    cat.describe(SchemaDescriber(AnthropicProvider(model="claude-sonnet-4-5")))
+"""
+from .describe import SchemaDescriber
+from .embedder import APIEmbedder
+from .providers import (
+    AnthropicProvider,
+    CallableProvider,
+    GeminiProvider,
+    OpenAIProvider,
+    Provider,
+    ProviderError,
+    auto_provider,
+    available_providers,
+)
+
+__all__ = [
+    "SchemaDescriber", "APIEmbedder",
+    "Provider", "ProviderError", "CallableProvider",
+    "AnthropicProvider", "OpenAIProvider", "GeminiProvider",
+    "auto_provider", "available_providers",
+]

schemagate/ai/describe.py

@@ -0,0 +1,170 @@
+"""AI-generated catalog descriptions.
+
+Retrieval quality is limited by how much meaning the schema text carries.
+``CUST_ORD_LN_T`` with columns ``ID``, ``QTY``, ``AMT`` tells a retriever
+almost nothing, which is why the offline embedder scores 50% on questions
+phrased in business words rather than identifier words.
+
+A ``SchemaDescriber`` asks a model to write one sentence per object saying
+what it holds and when to use it, and stores that on
+``ObjectDoc.description``. That text is already part of ``embed_text()``,
+so descriptions improve both vector and BM25 retrieval with no other change.
+
+**Only schema metadata is sent.** Object names, column names, types,
+nullability, existing comments, and foreign keys. No rows, no sample values,
+no query results, no credentials -- ``ObjectDoc`` does not carry row data,
+and ``_render`` cannot reach any. There is a test asserting this.
+
+**A human hint always wins.** ``Catalog.hint()`` is applied after
+descriptions and outranks them everywhere, so a wrong AI description can be
+corrected without regenerating anything.
+
+Descriptions cost money, so results are cached by content: an object is
+re-described only when its structure or the model changes.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import pathlib
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict, Iterable, List, Optional, Sequence
+
+from ..models import ObjectDoc
+from .providers import Provider, ProviderError
+
+_SYSTEM = (
+    "You document database schemas for a SQL-generating assistant. "
+    "Given one table or view definition, reply with a single sentence, at "
+    "most 25 words, saying what the object holds and what question it "
+    "answers. Use the business meaning, not the column list. Do not repeat "
+    "the object name. Do not speculate about data you cannot see. No "
+    "preamble, no markdown, no quotes -- just the sentence."
+)
+
+
+def _render(doc: ObjectDoc, max_columns: int = 30) -> str:
+    """The only thing ever sent to a provider. Metadata, never rows."""
+    lines = [f"{doc.kind} {doc.qname}"]
+    if doc.description:
+        lines.append(f"existing comment: {doc.description}")
+    for col in doc.columns[:max_columns]:
+        bits = f"  {col.name} {col.type}"
+        if col.pk:
+            bits += " PK"
+        if col.comment:
+            bits += f"  -- {col.comment}"
+        lines.append(bits)
+    if len(doc.columns) > max_columns:
+        lines.append(f"  ...{len(doc.columns) - max_columns} more columns")
+    for fk in doc.foreign_keys:
+        lines.append(f"  FK {','.join(fk.columns)} -> {fk.ref_table}")
+    return "\n".join(lines)
+
+
+def _fingerprint(doc: ObjectDoc, model: str) -> str:
+    """Cache key: changes when the object's structure or the model changes."""
+    payload = json.dumps({
+        "q": doc.qname, "k": doc.kind, "m": model,
+        "c": [[c.name, c.type, c.pk, c.comment] for c in doc.columns],
+        "f": [[fk.columns, fk.ref_table] for fk in doc.foreign_keys],
+        "d": doc.description,
+    }, sort_keys=True, default=str)
+    return hashlib.sha256(payload.encode("utf-8")).hexdigest()[:32]
+
+
+class SchemaDescriber:
+    """Generate one-sentence descriptions for catalog objects.
+
+    Parameters
+    ----------
+    provider
+        Any object with ``complete(system, prompt, max_tokens)``.
+    cache_path
+        JSON file for generated descriptions. Highly recommended: it makes
+        re-runs free and keeps a diffable record of what the model wrote.
+    workers
+        Parallel requests. Keep it modest; providers rate-limit.
+    strict
+        ``False`` (default) skips objects whose call failed and carries on,
+        because a partial catalog still selects. ``True`` re-raises.
+    """
+
+    def __init__(self, provider: Provider, cache_path: Optional[str] = None,
+                 max_columns: int = 30, workers: int = 4,
+                 max_tokens: int = 120, strict: bool = False):
+        self.provider = provider
+        self.cache_path = pathlib.Path(cache_path) if cache_path else None
+        self.max_columns = max_columns
+        self.workers = max(1, int(workers))
+        self.max_tokens = max_tokens
+        self.strict = strict
+        self.failures: List[str] = []
+        self._cache: Dict[str, str] = self._load_cache()
+
+    # ---------------- cache ----------------
+
+    def _load_cache(self) -> Dict[str, str]:
+        if not self.cache_path or not self.cache_path.exists():
+            return {}
+        try:
+            return json.loads(self.cache_path.read_text("utf-8"))
+        except (json.JSONDecodeError, OSError):
+            return {}   # a corrupt cache must never break cataloguing
+
+    def _save_cache(self) -> None:
+        if not self.cache_path:
+            return
+        try:
+            self.cache_path.parent.mkdir(parents=True, exist_ok=True)
+            tmp = self.cache_path.with_suffix(self.cache_path.suffix + ".tmp")
+            tmp.write_text(json.dumps(self._cache, indent=2, sort_keys=True),
+                           encoding="utf-8")
+            os.replace(tmp, self.cache_path)
+        except OSError:
+            pass    # a read-only disk is not a reason to lose the run
+
+    # ---------------- generation ----------------
+
+    def _describe_one(self, doc: ObjectDoc) -> Optional[str]:
+        key = _fingerprint(doc, getattr(self.provider, "name", "?"))
+        if key in self._cache:
+            return self._cache[key]
+        try:
+            text = self.provider.complete(_SYSTEM, _render(doc, self.max_columns),
+                                          max_tokens=self.max_tokens)
+        except ProviderError:
+            self.failures.append(doc.qname)
+            if self.strict:
+                raise
+            return None
+        text = " ".join((text or "").split()).strip().strip('"')
+        if not text:
+            self.failures.append(doc.qname)
+            return None
+        self._cache[key] = text
+        return text
+
+    def describe(self, docs: Iterable[ObjectDoc]) -> Dict[str, str]:
+        """Return ``{qname: description}``. Objects that failed are absent."""
+        docs = list(docs)
+        self.failures = []
+        if not docs:
+            return {}
+        if self.workers == 1:
+            results = [self._describe_one(d) for d in docs]
+        else:
+            with ThreadPoolExecutor(max_workers=self.workers) as pool:
+                results = list(pool.map(self._describe_one, docs))
+        self._save_cache()
+        return {d.qname: t for d, t in zip(docs, results) if t}
+
+    # what the provider would receive, for review before spending money
+    def preview(self, doc: ObjectDoc) -> str:
+        return _render(doc, self.max_columns)
+
+    def estimate_calls(self, docs: Sequence[ObjectDoc]) -> int:
+        """How many billed calls ``describe()`` would make right now."""
+        name = getattr(self.provider, "name", "?")
+        return sum(1 for d in docs if _fingerprint(d, name) not in self._cache)

schemagate/ai/embedder.py

@@ -0,0 +1,128 @@
+"""Embeddings from a hosted model.
+
+The offline ``HashingEmbedder`` matches subwords, which is strong on
+identifier-shaped questions and weak on pure paraphrase. An API embedder
+trades money, latency and a network dependency for semantic matching.
+
+Measure before adopting: on identifier-heavy schema text the offline
+embedder is often competitive, and ``tests/bench.py`` runs against your
+own schema. This is a swap, not an upgrade.
+
+    from schemagate import Catalog
+    from schemagate.ai import APIEmbedder, OpenAIProvider
+
+    provider = OpenAIProvider(model="gpt-4.1-mini",
+                              embed_model="text-embedding-3-small")
+    cat = Catalog(embedder=APIEmbedder(provider, dim=1536))
+
+Vectors are cached in memory for the life of the embedder and, if you pass
+``cache_path``, on disk -- so re-indexing an unchanged schema is free.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+import os
+import pathlib
+from typing import Dict, List, Optional, Sequence
+
+from .providers import Provider, ProviderError
+
+
+def _l2(vec: List[float]) -> List[float]:
+    n = math.sqrt(sum(v * v for v in vec))
+    return [v / n for v in vec] if n else vec
+
+
+class APIEmbedder:
+    """Embedder backed by a provider's embedding endpoint.
+
+    Parameters
+    ----------
+    provider
+        Any object with ``embed(texts) -> list[list[float]]``.
+    dim
+        Expected dimension. Checked on the first response, because a silent
+        dimension change invalidates every persisted vector.
+    batch_size
+        Texts per request.
+    cache_path
+        Optional JSON cache keyed by SHA-256 of the text plus provider name.
+    """
+
+    def __init__(self, provider: Provider, dim: int, batch_size: int = 64,
+                 cache_path: Optional[str] = None, normalize: bool = True):
+        if not hasattr(provider, "embed"):
+            raise TypeError(
+                f"{getattr(provider, 'name', provider)!r} has no embed(); "
+                "pass a provider built with an embed model")
+        self.provider = provider
+        self.dim = int(dim)
+        self.batch_size = max(1, int(batch_size))
+        self.normalize = normalize
+        self.name = f"api:{getattr(provider, 'name', 'provider')}"
+        self.cache_path = pathlib.Path(cache_path) if cache_path else None
+        self._cache: Dict[str, List[float]] = self._load()
+
+    # ---------------- cache ----------------
+
+    def _key(self, text: str) -> str:
+        seed = f"{self.name}:{self.dim}:{text}"
+        return hashlib.sha256(seed.encode("utf-8")).hexdigest()[:32]
+
+    def _load(self) -> Dict[str, List[float]]:
+        if not self.cache_path or not self.cache_path.exists():
+            return {}
+        try:
+            return json.loads(self.cache_path.read_text("utf-8"))
+        except (json.JSONDecodeError, OSError):
+            return {}
+
+    def _save(self) -> None:
+        if not self.cache_path:
+            return
+        try:
+            self.cache_path.parent.mkdir(parents=True, exist_ok=True)
+            tmp = self.cache_path.with_suffix(self.cache_path.suffix + ".tmp")
+            tmp.write_text(json.dumps(self._cache), encoding="utf-8")
+            os.replace(tmp, self.cache_path)
+        except OSError:
+            pass
+
+    # ---------------- embedding ----------------
+
+    def embed(self, texts: Sequence[str]) -> List[List[float]]:
+        texts = list(texts)
+        out: List[Optional[List[float]]] = [None] * len(texts)
+        todo: List[int] = []
+        for i, text in enumerate(texts):
+            hit = self._cache.get(self._key(text))
+            if hit is not None:
+                out[i] = list(hit)
+            else:
+                todo.append(i)
+
+        for start in range(0, len(todo), self.batch_size):
+            chunk = todo[start:start + self.batch_size]
+            vectors = self.provider.embed([texts[i] for i in chunk])
+            if len(vectors) != len(chunk):
+                raise ProviderError(
+                    f"{self.name} returned {len(vectors)} vectors for "
+                    f"{len(chunk)} inputs")
+            for i, vector in zip(chunk, vectors):
+                vector = [float(v) for v in vector]
+                if len(vector) != self.dim:
+                    raise ProviderError(
+                        f"{self.name} returned {len(vector)}-dim vectors but "
+                        f"dim={self.dim} was declared; a changed embedding "
+                        "model invalidates every persisted vector, so this "
+                        "is refused rather than mixed")
+                if self.normalize:
+                    vector = _l2(vector)
+                self._cache[self._key(texts[i])] = vector
+                out[i] = vector
+
+        if todo:
+            self._save()
+        return [v for v in out if v is not None]

schemagate/ai/providers.py

@@ -0,0 +1,268 @@
+"""Adapters for hosted AI models. All optional; none is ever required.
+
+schemagate works with no AI provider at all -- the default embedder is offline and
+deterministic. A provider buys you two things: better catalog descriptions
+(``schemagate.ai.SchemaDescriber``) and semantic embeddings
+(``schemagate.ai.APIEmbedder``). Both are opt-in and both degrade to the offline
+path if the provider is unavailable.
+
+Bring your own key. Nothing here reads a key from anywhere except the
+environment variable of the provider you chose, or the value you pass.
+
+``model`` is a required argument on every provider. That is deliberate:
+model identifiers change often, and a library that hardcodes one eventually
+ships a default that 404s for everyone. Pass the model you actually have
+access to.
+
+    from schemagate.ai import AnthropicProvider, SchemaDescriber
+
+    provider = AnthropicProvider(model="claude-sonnet-4-5")
+    cat.describe(SchemaDescriber(provider))
+
+If your provider is not one of the three below, or its SDK changes shape,
+use ``CallableProvider`` and keep control of the call yourself::
+
+    CallableProvider(lambda system, prompt: my_llm(system, prompt))
+"""
+from __future__ import annotations
+
+import os
+from typing import Any, Callable, List, Optional, Protocol, Sequence, runtime_checkable
+
+
+@runtime_checkable
+class Provider(Protocol):
+    """What schemagate needs from a model. Implement either method, or both."""
+
+    name: str
+
+    def complete(self, system: str, prompt: str, max_tokens: int = 1024) -> str: ...
+
+
+class ProviderError(RuntimeError):
+    """A provider call failed. Callers decide whether that is fatal."""
+
+
+# --------------------------------------------------------------------------
+# Bring your own callable -- the escape hatch that never breaks on an SDK bump
+# --------------------------------------------------------------------------
+
+class CallableProvider:
+    """Wrap any function ``f(system, prompt) -> str``.
+
+    Use this for a provider schemagate does not ship, a gateway, a local model,
+    or when an SDK changes and you do not want to wait for a release.
+    """
+
+    def __init__(self, fn: Callable[[str, str], str], name: str = "callable",
+                 embed_fn: Optional[Callable[[Sequence[str]], List[List[float]]]] = None):
+        self._fn = fn
+        self._embed_fn = embed_fn
+        self.name = name
+
+    def complete(self, system: str, prompt: str, max_tokens: int = 1024) -> str:
+        return self._fn(system, prompt)
+
+    def embed(self, texts: Sequence[str]) -> List[List[float]]:
+        if self._embed_fn is None:
+            raise ProviderError(f"{self.name} was built without an embed_fn")
+        return self._embed_fn(texts)
+
+
+# --------------------------------------------------------------------------
+# Anthropic (Claude)
+# --------------------------------------------------------------------------
+
+class AnthropicProvider:
+    """Claude via the official ``anthropic`` SDK.
+
+    ``pip install anthropic``. Key from ``api_key=`` or ``ANTHROPIC_API_KEY``.
+    """
+
+    env_var = "ANTHROPIC_API_KEY"
+
+    def __init__(self, model: str, api_key: Optional[str] = None,
+                 client: Any = None, timeout: float = 60.0):
+        if not model:
+            raise ValueError("model is required, e.g. model='claude-sonnet-4-5'")
+        self.model = model
+        self.name = f"anthropic:{model}"
+        if client is not None:
+            self._client = client
+            return
+        try:
+            import anthropic
+        except ImportError as e:
+            raise ImportError("pip install 'schemagate[anthropic]' to use "
+                              "AnthropicProvider") from e
+        key = api_key or os.environ.get(self.env_var)
+        if not key:
+            raise ValueError(f"no API key: pass api_key= or set {self.env_var}")
+        self._client = anthropic.Anthropic(api_key=key, timeout=timeout)
+
+    def complete(self, system: str, prompt: str, max_tokens: int = 1024) -> str:
+        try:
+            msg = self._client.messages.create(
+                model=self.model, max_tokens=max_tokens, system=system,
+                messages=[{"role": "user", "content": prompt}],
+            )
+        except Exception as e:
+            raise ProviderError(f"{self.name}: {e}") from e
+        return "".join(b.text for b in msg.content if getattr(b, "type", "") == "text")
+
+
+# --------------------------------------------------------------------------
+# OpenAI (GPT) -- also covers Azure OpenAI and any OpenAI-compatible gateway
+# --------------------------------------------------------------------------
+
+class OpenAIProvider:
+    """GPT via the official ``openai`` SDK.
+
+    ``pip install openai``. Key from ``api_key=`` or ``OPENAI_API_KEY``.
+    Point ``base_url`` at any OpenAI-compatible endpoint (Azure, vLLM,
+    OpenRouter, Ollama) to use this adapter with something else.
+    """
+
+    env_var = "OPENAI_API_KEY"
+
+    def __init__(self, model: str, api_key: Optional[str] = None,
+                 client: Any = None, base_url: Optional[str] = None,
+                 embed_model: Optional[str] = None, timeout: float = 60.0):
+        if not model:
+            raise ValueError("model is required, e.g. model='gpt-4.1-mini'")
+        self.model = model
+        self.embed_model = embed_model
+        self.name = f"openai:{model}"
+        if client is not None:
+            self._client = client
+            return
+        try:
+            import openai
+        except ImportError as e:
+            raise ImportError("pip install 'schemagate[openai]' to use "
+                              "OpenAIProvider") from e
+        key = api_key or os.environ.get(self.env_var)
+        if not key:
+            raise ValueError(f"no API key: pass api_key= or set {self.env_var}")
+        self._client = openai.OpenAI(api_key=key, base_url=base_url,
+                                     timeout=timeout)
+
+    def complete(self, system: str, prompt: str, max_tokens: int = 1024) -> str:
+        try:
+            resp = self._client.chat.completions.create(
+                model=self.model, max_tokens=max_tokens,
+                messages=[{"role": "system", "content": system},
+                          {"role": "user", "content": prompt}],
+            )
+        except Exception as e:
+            raise ProviderError(f"{self.name}: {e}") from e
+        return resp.choices[0].message.content or ""
+
+    def embed(self, texts: Sequence[str]) -> List[List[float]]:
+        if not self.embed_model:
+            raise ProviderError(
+                "pass embed_model= to use OpenAIProvider for embeddings, "
+                "e.g. embed_model='text-embedding-3-small'")
+        try:
+            resp = self._client.embeddings.create(model=self.embed_model,
+                                                  input=list(texts))
+        except Exception as e:
+            raise ProviderError(f"{self.name}: {e}") from e
+        return [list(d.embedding) for d in resp.data]
+
+
+# --------------------------------------------------------------------------
+# Google (Gemini)
+# --------------------------------------------------------------------------
+
+class GeminiProvider:
+    """Gemini via the ``google-genai`` SDK.
+
+    ``pip install google-genai``. Key from ``api_key=``, ``GEMINI_API_KEY``
+    or ``GOOGLE_API_KEY``.
+    """
+
+    env_var = "GEMINI_API_KEY"
+    alt_env_var = "GOOGLE_API_KEY"
+
+    def __init__(self, model: str, api_key: Optional[str] = None,
+                 client: Any = None, embed_model: Optional[str] = None):
+        if not model:
+            raise ValueError("model is required, e.g. model='gemini-2.5-flash'")
+        self.model = model
+        self.embed_model = embed_model
+        self.name = f"gemini:{model}"
+        if client is not None:
+            self._client = client
+            return
+        try:
+            from google import genai
+        except ImportError as e:
+            raise ImportError("pip install 'schemagate[gemini]' to use "
+                              "GeminiProvider") from e
+        key = (api_key or os.environ.get(self.env_var)
+               or os.environ.get(self.alt_env_var))
+        if not key:
+            raise ValueError(
+                f"no API key: pass api_key= or set {self.env_var}")
+        self._client = genai.Client(api_key=key)
+
+    def complete(self, system: str, prompt: str, max_tokens: int = 1024) -> str:
+        try:
+            resp = self._client.models.generate_content(
+                model=self.model, contents=f"{system}\n\n{prompt}")
+        except Exception as e:
+            raise ProviderError(f"{self.name}: {e}") from e
+        return resp.text or ""
+
+    def embed(self, texts: Sequence[str]) -> List[List[float]]:
+        if not self.embed_model:
+            raise ProviderError(
+                "pass embed_model= to use GeminiProvider for embeddings")
+        try:
+            resp = self._client.models.embed_content(model=self.embed_model,
+                                                     contents=list(texts))
+        except Exception as e:
+            raise ProviderError(f"{self.name}: {e}") from e
+        return [list(e.values) for e in resp.embeddings]
+
+
+# --------------------------------------------------------------------------
+# Detection
+# --------------------------------------------------------------------------
+
+#: checked in order; first key present wins
+_AUTO_ORDER = [
+    ("ANTHROPIC_API_KEY", AnthropicProvider),
+    ("OPENAI_API_KEY", OpenAIProvider),
+    ("GEMINI_API_KEY", GeminiProvider),
+    ("GOOGLE_API_KEY", GeminiProvider),
+]
+
+
+def available_providers(env: Optional[dict] = None) -> List[str]:
+    """Names of providers whose API key is present. Never returns a key."""
+    env = os.environ if env is None else env
+    seen, out = set(), []
+    for var, cls in _AUTO_ORDER:
+        if env.get(var) and cls.__name__ not in seen:
+            seen.add(cls.__name__)
+            out.append(cls.__name__)
+    return out
+
+
+def auto_provider(model: str, env: Optional[dict] = None, **kwargs):
+    """Build a provider from whichever API key is in the environment.
+
+    Convenience only. Construct the provider directly when you care which
+    one you get -- this picks by key presence, not by capability.
+    """
+    env = os.environ if env is None else env
+    for var, cls in _AUTO_ORDER:
+        if env.get(var):
+            return cls(model=model, api_key=env[var], **kwargs)
+    raise ValueError(
+        "no provider API key found; set one of "
+        + ", ".join(v for v, _ in _AUTO_ORDER)
+        + " or construct a provider directly. schemagate works without one."
+    )