shiftgate 0.1.0__py3-none-any.whl

shiftgate/feedback/loop.py

@@ -0,0 +1,182 @@
+"""
+Feedback loop: persist routing traces and compute adapter acceptance scores.
+
+Traces are stored as newline-delimited JSON in ``~/.shiftgate/traces.jsonl``.
+Each line is a serialised ``RoutingTrace``.  This format is append-only and
+easy to stream-process without loading the entire file into memory.
+
+Workflow
+--------
+1. After every ``shiftgate route`` / ``shiftgate run``, call ``record_trace``.
+2. User runs ``shiftgate feedback accept`` or ``shiftgate feedback reject``.
+3. Call ``mark_accepted(trace_id, accepted)`` to annotate the trace.
+4. ``compute_adapter_scores()`` aggregates acceptance rates per adapter.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from shiftgate.registry.schemas import RoutingTrace
+
+logger = logging.getLogger(__name__)
+
+_SHIFTGATE_DIR = Path.home() / ".shiftgate"
+_TRACES_PATH = _SHIFTGATE_DIR / "traces.jsonl"
+
+# How many recent traces to scan when ``mark_accepted`` searches by trace ID.
+_RECENT_SCAN_LIMIT = 200
+
+
+def record_trace(trace: RoutingTrace) -> None:
+    """Append a ``RoutingTrace`` as a JSON line to the traces log.
+
+    The file is created (along with its parent directory) on first write.
+    """
+    _SHIFTGATE_DIR.mkdir(parents=True, exist_ok=True)
+    line = trace.model_dump_json()
+    with _TRACES_PATH.open("a", encoding="utf-8") as fh:
+        fh.write(line + "\n")
+    logger.debug("Trace %s recorded.", trace.id)
+
+
+def get_last_trace() -> RoutingTrace | None:
+    """Return the most recently recorded trace, or None if no traces exist."""
+    if not _TRACES_PATH.exists():
+        return None
+    last_line: str | None = None
+    with _TRACES_PATH.open("r", encoding="utf-8") as fh:
+        for line in fh:
+            line = line.strip()
+            if line:
+                last_line = line
+    if last_line is None:
+        return None
+    return RoutingTrace.model_validate_json(last_line)
+
+
+def mark_accepted(trace_id: str, accepted: bool) -> bool:
+    """Set the ``accepted`` field on a specific trace.
+
+    Rewrites the last ``_RECENT_SCAN_LIMIT`` lines of the traces file in-place
+    (only those lines, prepending unchanged older lines).  Trades slight memory
+    use for simplicity.
+
+    Parameters
+    ----------
+    trace_id:
+        The ``RoutingTrace.id`` hex string to update.
+    accepted:
+        True = good routing decision, False = bad routing decision.
+
+    Returns
+    -------
+    True if the trace was found and updated; False if not found.
+    """
+    if not _TRACES_PATH.exists():
+        logger.warning("No traces file found at %s.", _TRACES_PATH)
+        return False
+
+    lines = _TRACES_PATH.read_text(encoding="utf-8").splitlines()
+    updated = False
+
+    for i in range(len(lines) - 1, max(-1, len(lines) - _RECENT_SCAN_LIMIT - 1), -1):
+        line = lines[i].strip()
+        if not line:
+            continue
+        try:
+            data = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if data.get("id") == trace_id:
+            data["accepted"] = accepted
+            lines[i] = json.dumps(data, ensure_ascii=False)
+            updated = True
+            break
+
+    if updated:
+        _TRACES_PATH.write_text("\n".join(lines) + "\n", encoding="utf-8")
+        logger.debug("Trace %s marked accepted=%s.", trace_id, accepted)
+
+    return updated
+
+
+def mark_last_accepted(accepted: bool) -> RoutingTrace | None:
+    """Convenience: mark the most recent trace as accepted/rejected.
+
+    Returns the updated trace, or None if no traces exist.
+    """
+    trace = get_last_trace()
+    if trace is None:
+        return None
+    mark_accepted(trace.id, accepted)
+    trace.accepted = accepted
+    return trace
+
+
+def load_all_traces() -> list[RoutingTrace]:
+    """Load all traces from disk into memory.
+
+    For large files prefer streaming with ``iter_traces()`` instead.
+    """
+    return list(iter_traces())
+
+
+def iter_traces():
+    """Yield ``RoutingTrace`` objects one at a time from the traces file."""
+    if not _TRACES_PATH.exists():
+        return
+    with _TRACES_PATH.open("r", encoding="utf-8") as fh:
+        for line in fh:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                yield RoutingTrace.model_validate_json(line)
+            except Exception as exc:
+                logger.warning("Skipping malformed trace line: %s", exc)
+
+
+def compute_adapter_scores() -> dict[str, float]:
+    """Compute the acceptance rate for each adapter across all rated traces.
+
+    Returns
+    -------
+    A dict mapping ``adapter_id`` → acceptance rate (0.0 – 1.0).
+    Only adapters with at least one rated trace are included.
+    Adapters with a 0 % acceptance rate are included with score 0.0.
+    """
+    totals: dict[str, int] = {}
+    accepted_counts: dict[str, int] = {}
+
+    for trace in iter_traces():
+        if trace.accepted is None:
+            continue
+        aid = trace.selected_adapter_id
+        totals[aid] = totals.get(aid, 0) + 1
+        if trace.accepted:
+            accepted_counts[aid] = accepted_counts.get(aid, 0) + 1
+
+    return {
+        aid: accepted_counts.get(aid, 0) / total
+        for aid, total in totals.items()
+    }
+
+
+def get_trace_stats() -> dict[str, int]:
+    """Return summary statistics about the traces file.
+
+    Keys: ``total``, ``accepted``, ``rejected``, ``unrated``.
+    """
+    stats = {"total": 0, "accepted": 0, "rejected": 0, "unrated": 0}
+    for trace in iter_traces():
+        stats["total"] += 1
+        if trace.accepted is True:
+            stats["accepted"] += 1
+        elif trace.accepted is False:
+            stats["rejected"] += 1
+        else:
+            stats["unrated"] += 1
+    return stats

shiftgate/registry/__init__.py

@@ -0,0 +1 @@
+"""Registry sub-package: adapter catalog, task clusters, and Pydantic schemas."""

shiftgate/registry/adapter_registry.py

@@ -0,0 +1,162 @@
+"""
+Adapter registry: load, persist, and manage AdapterEntry definitions.
+
+The registry reads from (in priority order):
+  1. ``~/.shiftgate/adapters.json``  — user-edited / previously saved
+  2. ``<package>/../../data/default_adapters.json``  — bundled defaults (empty list)
+
+Adapters can be added by passing a HuggingFace repo ID string or a full
+``AdapterEntry`` object.  When a bare HF repo ID is provided, metadata is
+fetched from the Hub to fill in the entry automatically.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from shiftgate.registry.schemas import AdapterEntry
+
+logger = logging.getLogger(__name__)
+
+_SHIFTGATE_DIR = Path.home() / ".shiftgate"
+_USER_ADAPTERS_PATH = _SHIFTGATE_DIR / "adapters.json"
+_DEFAULT_ADAPTERS_PATH = Path(__file__).parent.parent.parent / "data" / "default_adapters.json"
+
+
+class AdapterRegistry:
+    """In-memory store for AdapterEntry objects, backed by a JSON file.
+
+    Usage::
+
+        registry = AdapterRegistry.load()
+        registry.add_adapter(AdapterEntry(id="my-lora", ...))
+        registry.save()
+    """
+
+    def __init__(self, adapters: list[AdapterEntry], source_path: Path) -> None:
+        self._adapters: dict[str, AdapterEntry] = {a.id: a for a in adapters}
+        self._source_path = source_path
+
+    # ------------------------------------------------------------------
+    # Factory / persistence
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def load(cls) -> "AdapterRegistry":
+        """Load the adapter registry from disk.
+
+        Prefers ``~/.shiftgate/adapters.json``.  Falls back to the bundled
+        ``data/default_adapters.json`` (which ships as an empty list).
+        """
+        if _USER_ADAPTERS_PATH.exists():
+            source = _USER_ADAPTERS_PATH
+        elif _DEFAULT_ADAPTERS_PATH.exists():
+            source = _DEFAULT_ADAPTERS_PATH
+        else:
+            logger.warning("No adapter registry found; starting empty.")
+            return cls([], source_path=_USER_ADAPTERS_PATH)
+
+        logger.debug("Loading adapter registry from %s", source)
+        raw = json.loads(source.read_text(encoding="utf-8"))
+        adapters = [AdapterEntry.model_validate(a) for a in raw]
+        return cls(adapters, source_path=source)
+
+    def save(self) -> None:
+        """Persist the current registry to ``~/.shiftgate/adapters.json``."""
+        _SHIFTGATE_DIR.mkdir(parents=True, exist_ok=True)
+        data = [a.model_dump() for a in self._adapters.values()]
+        _USER_ADAPTERS_PATH.write_text(
+            json.dumps(data, indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+        logger.debug("Adapter registry saved to %s", _USER_ADAPTERS_PATH)
+
+    # ------------------------------------------------------------------
+    # CRUD
+    # ------------------------------------------------------------------
+
+    def get_adapter(self, adapter_id: str) -> AdapterEntry | None:
+        """Return an adapter by ID, or None if not found."""
+        return self._adapters.get(adapter_id)
+
+    def list_adapters(self) -> list[AdapterEntry]:
+        """Return all registered adapters."""
+        return list(self._adapters.values())
+
+    def add_adapter(self, adapter: AdapterEntry | str, **kwargs: object) -> AdapterEntry:
+        """Add or replace an adapter in the registry.
+
+        Parameters
+        ----------
+        adapter:
+            Either a fully-constructed ``AdapterEntry`` or a HuggingFace
+            repo ID string (e.g. ``"username/my-lora-adapter"``).  When a
+            string is provided the repo ID is used as ``hf_repo`` and a
+            best-effort ID slug is derived from it.  Extra keyword arguments
+            (``tags``, ``base_model``, ``description``) override auto-derived
+            values.
+        """
+        if isinstance(adapter, str):
+            adapter = _adapter_from_hf_repo(adapter, **kwargs)
+
+        self._adapters[adapter.id] = adapter
+        logger.debug("Adapter '%s' added to registry.", adapter.id)
+        return adapter
+
+    def remove_adapter(self, adapter_id: str) -> bool:
+        """Remove an adapter by ID. Returns True if it existed."""
+        if adapter_id in self._adapters:
+            del self._adapters[adapter_id]
+            return True
+        return False
+
+    def __len__(self) -> int:
+        return len(self._adapters)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _adapter_from_hf_repo(hf_repo: str, **kwargs: object) -> AdapterEntry:
+    """Construct a minimal AdapterEntry from a HuggingFace repo ID.
+
+    Tries to pull card metadata from the Hub.  If that fails (offline, private
+    repo, etc.) it builds a stub entry from the repo ID alone.
+
+    Extra ``kwargs`` are merged after auto-detection and override any
+    auto-derived fields (``tags``, ``base_model``, ``description``).
+    """
+    # Derive a clean ID slug from the repo path (e.g. "org/my-lora" → "my-lora")
+    slug = hf_repo.split("/")[-1].lower().replace("_", "-")
+
+    entry_data: dict = {
+        "id": slug,
+        "name": slug.replace("-", " ").title(),
+        "base_model": kwargs.pop("base_model", "unknown"),
+        "task_tags": kwargs.pop("tags", []),
+        "description": kwargs.pop("description", f"Imported from {hf_repo}"),
+        "hf_repo": hf_repo,
+    }
+    entry_data.update(kwargs)
+
+    # Attempt to enrich from HuggingFace Hub metadata.
+    try:
+        from huggingface_hub import hf_hub_download, model_info  # type: ignore
+
+        info = model_info(hf_repo)
+        if info.card_data:
+            card = info.card_data
+            if hasattr(card, "base_model") and card.base_model:
+                base = card.base_model
+                entry_data["base_model"] = base[0] if isinstance(base, list) else base
+            if hasattr(card, "tags") and card.tags and not entry_data["task_tags"]:
+                entry_data["task_tags"] = list(card.tags)[:8]
+        if info.id:
+            entry_data["name"] = info.id.split("/")[-1]
+    except Exception as exc:
+        logger.debug("Could not fetch HF metadata for '%s': %s", hf_repo, exc)
+
+    return AdapterEntry.model_validate(entry_data)

shiftgate/registry/schemas.py

@@ -0,0 +1,115 @@
+"""
+Pydantic v2 schemas for shiftgate's core data model.
+
+Three top-level types:
+  - AdapterEntry  : a LoRA adapter (or fine-tuned model) in the registry
+  - TaskCluster   : a group of semantically related tasks with example queries
+  - RoutingTrace  : one routing decision, optionally annotated with user feedback
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class AdapterEntry(BaseModel):
+    """A LoRA adapter registered with shiftgate.
+
+    Adapters can live on HuggingFace (``hf_repo``) or locally (``local_path``).
+    At least one of the two must be set for inference to work, though the
+    registry itself does not enforce this so adapters can be catalogued before
+    they are downloaded.
+    """
+
+    id: str = Field(
+        description="Unique slug, e.g. 'python-lora-llama3'. Used as a stable reference key."
+    )
+    name: str = Field(description="Human-readable display name.")
+    base_model: str = Field(
+        description="The base model this adapter was trained on, e.g. 'meta-llama/Meta-Llama-3-8B'."
+    )
+    task_tags: list[str] = Field(
+        default_factory=list,
+        description="Free-form tags describing the adapter's specialisation, e.g. ['code', 'python'].",
+    )
+    description: str = Field(default="", description="Short prose description of the adapter's purpose.")
+    hf_repo: str | None = Field(
+        default=None,
+        description="HuggingFace Hub repository ID, e.g. 'username/my-lora-adapter'.",
+    )
+    local_path: str | None = Field(
+        default=None,
+        description="Absolute path to a local .safetensors file or adapter directory.",
+    )
+    benchmark_score: float | None = Field(
+        default=None,
+        description="Optional benchmark score (0–1) reported by the adapter author.",
+    )
+    context_length: int = Field(
+        default=4096,
+        description="Maximum context window in tokens.",
+    )
+    memory_mb: int | None = Field(
+        default=None,
+        description="Approximate VRAM/RAM usage in MB when the adapter is loaded.",
+    )
+
+
+class TaskCluster(BaseModel):
+    """A cluster of semantically related tasks used for routing.
+
+    During ``shiftgate init``, the ``validation_examples`` are embedded and
+    averaged to produce ``embedding_centroid``.  At routing time the query
+    embedding is compared against every cluster's centroid.
+    """
+
+    id: str = Field(
+        description="Unique slug, e.g. 'code_python'. Used as a stable routing key."
+    )
+    name: str = Field(description="Human-readable cluster name, e.g. 'Python Code Generation'.")
+    description: str = Field(description="Short description of what tasks belong here.")
+    validation_examples: list[str] = Field(
+        description="3–10 representative query strings used to compute the centroid embedding.",
+    )
+    embedding_centroid: list[float] | None = Field(
+        default=None,
+        description="Pre-computed mean embedding of the validation_examples. Populated by init.",
+    )
+    preferred_adapters: list[str] = Field(
+        default_factory=list,
+        description="Adapter IDs in priority order. The first available adapter is selected.",
+    )
+    fallback_adapters: list[str] = Field(
+        default_factory=list,
+        description="Adapter IDs to try when none of the preferred adapters are available.",
+    )
+
+
+class RoutingTrace(BaseModel):
+    """A single routing decision recorded for observability and feedback.
+
+    Traces are appended as JSON lines to ``~/.shiftgate/traces.jsonl``.
+    The ``accepted`` field starts as ``None`` and is filled in via
+    ``shiftgate feedback accept/reject``.
+    """
+
+    id: str = Field(
+        description="Unique trace ID (UUID4 hex string) for targeted feedback updates."
+    )
+    query: str = Field(description="The original user query that triggered this routing decision.")
+    matched_task_id: str = Field(description="ID of the TaskCluster that won the similarity match.")
+    similarity_score: float = Field(
+        description="Cosine similarity between the query embedding and the winning centroid (0–1)."
+    )
+    selected_adapter_id: str = Field(description="ID of the adapter that was selected for inference.")
+    accepted: bool | None = Field(
+        default=None,
+        description="User feedback: True = good routing, False = bad routing, None = not yet rated.",
+    )
+    latency_ms: float | None = Field(
+        default=None,
+        description="End-to-end inference latency in milliseconds (None if only routing, no run).",
+    )
+    timestamp: str = Field(
+        description="ISO-8601 UTC timestamp of when this trace was created."
+    )

shiftgate/registry/task_registry.py

@@ -0,0 +1,186 @@
+"""
+Task registry: load, persist, and manage TaskCluster definitions.
+
+The registry reads from (in priority order):
+  1. ``~/.shiftgate/tasks.json``  — user-edited / previously saved
+  2. ``<package>/../../data/default_tasks.json``  — bundled defaults
+
+On first run (``shiftgate init``) the ``compute_embeddings`` method is called
+to populate ``embedding_centroid`` for every cluster and cache them to
+``~/.shiftgate/embeddings_cache.npy`` so subsequent startups are instant.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from shiftgate.registry.schemas import TaskCluster
+
+if TYPE_CHECKING:
+    from shiftgate.router.embedder import Embedder
+
+logger = logging.getLogger(__name__)
+
+# Canonical locations
+_SHIFTGATE_DIR = Path.home() / ".shiftgate"
+_USER_TASKS_PATH = _SHIFTGATE_DIR / "tasks.json"
+_CACHE_PATH = _SHIFTGATE_DIR / "embeddings_cache.npy"
+
+# Path to the bundled default tasks, resolved relative to this file's location.
+_DEFAULT_TASKS_PATH = Path(__file__).parent.parent.parent / "data" / "default_tasks.json"
+
+
+class TaskRegistry:
+    """In-memory store for TaskCluster objects, backed by a JSON file.
+
+    Usage::
+
+        registry = TaskRegistry.load()
+        registry.compute_embeddings(embedder)
+        registry.save()
+    """
+
+    def __init__(self, tasks: list[TaskCluster], source_path: Path) -> None:
+        self._tasks: dict[str, TaskCluster] = {t.id: t for t in tasks}
+        self._source_path = source_path
+
+    # ------------------------------------------------------------------
+    # Factory / persistence
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def load(cls) -> "TaskRegistry":
+        """Load the task registry from disk.
+
+        Prefers the user's ``~/.shiftgate/tasks.json`` and falls back to the
+        bundled ``data/default_tasks.json`` if the user file does not exist.
+        """
+        if _USER_TASKS_PATH.exists():
+            source = _USER_TASKS_PATH
+        elif _DEFAULT_TASKS_PATH.exists():
+            source = _DEFAULT_TASKS_PATH
+        else:
+            raise FileNotFoundError(
+                f"No task registry found. Expected one of:\n"
+                f"  {_USER_TASKS_PATH}\n"
+                f"  {_DEFAULT_TASKS_PATH}\n"
+                "Run `shiftgate init` to set up the default registry."
+            )
+
+        logger.debug("Loading task registry from %s", source)
+        raw = json.loads(source.read_text(encoding="utf-8"))
+        tasks = [TaskCluster.model_validate(t) for t in raw]
+        instance = cls(tasks, source_path=source)
+
+        # Eagerly restore cached centroids so ``compute_embeddings`` can be
+        # skipped on normal runs (not first init).
+        instance._restore_cache()
+        return instance
+
+    def save(self) -> None:
+        """Persist the current registry to ``~/.shiftgate/tasks.json``."""
+        _SHIFTGATE_DIR.mkdir(parents=True, exist_ok=True)
+        data = [t.model_dump() for t in self._tasks.values()]
+        _USER_TASKS_PATH.write_text(
+            json.dumps(data, indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+        logger.debug("Task registry saved to %s", _USER_TASKS_PATH)
+
+    # ------------------------------------------------------------------
+    # Embedding management
+    # ------------------------------------------------------------------
+
+    def compute_embeddings(self, embedder: "Embedder") -> None:
+        """Compute and store the centroid embedding for every task cluster.
+
+        For each cluster the validation examples are embedded individually and
+        then averaged (L2-normalised mean) to form a single centroid vector.
+        The results are written back into each ``TaskCluster.embedding_centroid``
+        field **and** saved to ``~/.shiftgate/embeddings_cache.npy`` as a
+        (n_tasks × dim) float32 array for fast loading on future runs.
+        """
+        task_list = list(self._tasks.values())
+        logger.info("Computing embeddings for %d task clusters…", len(task_list))
+
+        for task in task_list:
+            all_examples = task.validation_examples
+            embeddings = embedder.embed_batch(all_examples)  # shape: (n, dim)
+            centroid = embeddings.mean(axis=0)
+            # L2-normalise so cosine similarity reduces to dot product later.
+            norm = np.linalg.norm(centroid)
+            if norm > 0:
+                centroid = centroid / norm
+            task.embedding_centroid = centroid.tolist()
+
+        # Persist centroids as a numpy array indexed by task order.
+        self._save_cache(task_list)
+        logger.info("Embeddings computed and cached.")
+
+    def _save_cache(self, task_list: list[TaskCluster]) -> None:
+        """Write centroids to the numpy cache file."""
+        _SHIFTGATE_DIR.mkdir(parents=True, exist_ok=True)
+        centroids = [t.embedding_centroid for t in task_list if t.embedding_centroid]
+        if centroids:
+            arr = np.array(centroids, dtype=np.float32)
+            np.save(_CACHE_PATH, arr)
+            logger.debug("Centroid cache written to %s (%s)", _CACHE_PATH, arr.shape)
+
+    def _restore_cache(self) -> None:
+        """Re-populate ``embedding_centroid`` from the numpy cache if available.
+
+        This avoids a full re-embedding on every startup. The cache is keyed
+        positionally — tasks must stay in the same order between runs, which is
+        true as long as the registry JSON is not manually reordered.
+        """
+        if not _CACHE_PATH.exists():
+            return
+        try:
+            arr = np.load(_CACHE_PATH)
+            task_list = list(self._tasks.values())
+            for i, task in enumerate(task_list):
+                if i < len(arr):
+                    task.embedding_centroid = arr[i].tolist()
+            logger.debug("Restored centroids from cache (%d tasks)", len(task_list))
+        except Exception as exc:
+            logger.warning("Could not restore embedding cache (%s). Re-run `shiftgate init`.", exc)
+
+    def embeddings_ready(self) -> bool:
+        """Return True if all task clusters have a computed centroid."""
+        return all(t.embedding_centroid is not None for t in self._tasks.values())
+
+    # ------------------------------------------------------------------
+    # CRUD
+    # ------------------------------------------------------------------
+
+    def get_all_tasks(self) -> list[TaskCluster]:
+        """Return all registered task clusters."""
+        return list(self._tasks.values())
+
+    def get_task(self, task_id: str) -> TaskCluster | None:
+        """Return a single task cluster by ID, or None if not found."""
+        return self._tasks.get(task_id)
+
+    def add_task(self, task: TaskCluster) -> None:
+        """Add or replace a task cluster in the registry.
+
+        If a task with the same ID already exists it is silently overwritten.
+        Call ``save()`` afterwards to persist the change.
+        """
+        self._tasks[task.id] = task
+        logger.debug("Task '%s' added to registry.", task.id)
+
+    def remove_task(self, task_id: str) -> bool:
+        """Remove a task by ID. Returns True if it existed."""
+        if task_id in self._tasks:
+            del self._tasks[task_id]
+            return True
+        return False
+
+    def __len__(self) -> int:
+        return len(self._tasks)

shiftgate/router/__init__.py

@@ -0,0 +1 @@
+"""Router sub-package: embedding, cosine matching, and routing logic."""