shiftgate 0.1.0__py3-none-any.whl

shiftgate/router/embedder.py

@@ -0,0 +1,95 @@
+"""
+Text embedder backed by fastembed.
+
+Uses ``BAAI/bge-small-en-v1.5`` — a compact (33 M param) model that runs
+efficiently on CPU.  The model is downloaded once by fastembed and cached in
+``~/.cache/fastembed``.
+
+A module-level singleton (``_MODEL``) is created lazily on first use so that
+importing this module is cheap.  The model is NOT re-created between calls.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+# -------------------------------------------------------------------------
+# Default model — small, CPU-friendly, strong quality/speed trade-off.
+# -------------------------------------------------------------------------
+DEFAULT_MODEL = "BAAI/bge-small-en-v1.5"
+
+# Module-level singleton; populated on first call to `_get_model()`.
+_MODEL: Any | None = None
+
+
+def _get_model(model_name: str = DEFAULT_MODEL) -> Any:
+    """Return the fastembed TextEmbedding singleton, creating it if needed.
+
+    The model is loaded once per process.  If you need a different model,
+    call ``reset_model()`` first.
+    """
+    global _MODEL
+    if _MODEL is None:
+        try:
+            from fastembed import TextEmbedding  # type: ignore
+        except ImportError as exc:
+            raise ImportError(
+                "fastembed is required for shiftgate routing. "
+                "Install it with: pip install fastembed"
+            ) from exc
+
+        logger.info("Loading embedding model '%s' (first use — one-time download may occur)…", model_name)
+        _MODEL = TextEmbedding(model_name=model_name)
+        logger.info("Embedding model loaded.")
+    return _MODEL
+
+
+def reset_model() -> None:
+    """Force the next embed call to recreate the model singleton.
+
+    Useful in tests or when switching models at runtime.
+    """
+    global _MODEL
+    _MODEL = None
+
+
+class Embedder:
+    """Thin wrapper around the fastembed TextEmbedding model.
+
+    All embedding operations are synchronous and run on CPU.  The model
+    is shared across all ``Embedder`` instances via the module-level singleton.
+    """
+
+    def __init__(self, model_name: str = DEFAULT_MODEL) -> None:
+        self._model_name = model_name
+
+    @property
+    def _model(self) -> Any:
+        return _get_model(self._model_name)
+
+    def embed(self, text: str) -> np.ndarray:
+        """Embed a single text string.
+
+        Returns a 1-D float32 numpy array of shape ``(dim,)``.
+        The vector is **not** L2-normalised here; normalisation is done
+        where appropriate (e.g. when computing task centroids).
+        """
+        # fastembed returns a generator of numpy arrays.
+        results = list(self._model.embed([text]))
+        return np.array(results[0], dtype=np.float32)
+
+    def embed_batch(self, texts: list[str]) -> np.ndarray:
+        """Embed a list of strings.
+
+        Returns a 2-D float32 numpy array of shape ``(n, dim)`` where
+        ``n = len(texts)``.
+        """
+        if not texts:
+            raise ValueError("embed_batch received an empty list.")
+        results = list(self._model.embed(texts))
+        return np.array(results, dtype=np.float32)

shiftgate/router/matcher.py

@@ -0,0 +1,115 @@
+"""
+Cosine-similarity matcher: maps query embeddings to task clusters and adapters.
+
+This module is deliberately stateless — all context (registries, embeddings)
+is passed explicitly so the functions are easy to test in isolation.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+from shiftgate.registry.schemas import AdapterEntry, TaskCluster
+
+logger = logging.getLogger(__name__)
+
+
+def top_k_tasks(
+    query_embedding: np.ndarray,
+    task_clusters: list[TaskCluster],
+    k: int = 3,
+) -> list[tuple[TaskCluster, float]]:
+    """Return the top-K task clusters by cosine similarity to the query.
+
+    Parameters
+    ----------
+    query_embedding:
+        1-D float32 array of shape ``(dim,)``.  Need not be L2-normalised;
+        this function normalises internally.
+    task_clusters:
+        All clusters in the registry.  Clusters without a computed centroid
+        are silently skipped.
+    k:
+        Number of top clusters to return.
+
+    Returns
+    -------
+    list of ``(TaskCluster, score)`` pairs sorted by score descending.
+    """
+    eligible = [t for t in task_clusters if t.embedding_centroid is not None]
+    if not eligible:
+        raise ValueError(
+            "No task cluster has a computed embedding centroid. "
+            "Run `shiftgate init` to compute embeddings."
+        )
+
+    # Stack centroids into a matrix for vectorised dot product.
+    centroid_matrix = np.array(
+        [t.embedding_centroid for t in eligible], dtype=np.float32
+    )  # shape: (n_tasks, dim)
+
+    # L2-normalise the query vector.
+    q_norm = np.linalg.norm(query_embedding)
+    if q_norm == 0:
+        raise ValueError("Query produced a zero-norm embedding.")
+    q_unit = query_embedding / q_norm
+
+    # Cosine similarity = dot(q_unit, centroid_unit) because centroids were
+    # already L2-normalised at compute time (see task_registry.py).
+    scores = centroid_matrix @ q_unit  # shape: (n_tasks,)
+
+    # Grab top-K indices.
+    k = min(k, len(eligible))
+    top_indices = np.argsort(scores)[::-1][:k]
+
+    return [(eligible[i], float(scores[i])) for i in top_indices]
+
+
+def select_adapter(
+    top_tasks: list[tuple[TaskCluster, float]],
+    adapter_registry,  # AdapterRegistry — avoid circular import with string hint
+) -> tuple[AdapterEntry, TaskCluster, float]:
+    """Select the best adapter given the ranked task list.
+
+    Strategy:
+      1. Iterate top tasks in similarity order.
+      2. For each task, try ``preferred_adapters`` then ``fallback_adapters``.
+      3. Return the first adapter that exists in the registry.
+      4. If no registered adapter matches any task, raise ``NoAdapterError``.
+
+    Parameters
+    ----------
+    top_tasks:
+        Output of ``top_k_tasks`` — list of (TaskCluster, score) descending.
+    adapter_registry:
+        ``AdapterRegistry`` instance to look up adapter IDs.
+
+    Returns
+    -------
+    ``(AdapterEntry, TaskCluster, similarity_score)``
+    """
+    for task, score in top_tasks:
+        candidates = list(task.preferred_adapters) + list(task.fallback_adapters)
+        for adapter_id in candidates:
+            adapter = adapter_registry.get_adapter(adapter_id)
+            if adapter is not None:
+                logger.debug(
+                    "Selected adapter '%s' via task '%s' (score=%.4f)",
+                    adapter.id,
+                    task.id,
+                    score,
+                )
+                return adapter, task, score
+
+    # No adapter matched — surface a helpful error.
+    task_ids = [t.id for t, _ in top_tasks]
+    raise NoAdapterError(
+        f"No registered adapter found for tasks {task_ids}. "
+        "Add adapters with `shiftgate adapter add <hf_repo>`."
+    )
+
+
+class NoAdapterError(RuntimeError):
+    """Raised when the matcher cannot find any registered adapter for a query."""

shiftgate/router/router.py

@@ -0,0 +1,97 @@
+"""
+Main routing orchestrator: query string → RoutingTrace.
+
+This module ties together the embedder, matcher, and registries.
+It is the single function that CLI commands and the runtime backend call.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from datetime import datetime, timezone
+
+from shiftgate.registry.adapter_registry import AdapterRegistry
+from shiftgate.registry.schemas import RoutingTrace
+from shiftgate.registry.task_registry import TaskRegistry
+from shiftgate.router.embedder import Embedder
+from shiftgate.router.matcher import NoAdapterError, select_adapter, top_k_tasks
+
+logger = logging.getLogger(__name__)
+
+
+def route(
+    query: str,
+    task_registry: TaskRegistry,
+    adapter_registry: AdapterRegistry,
+    embedder: Embedder,
+    top_k: int = 3,
+) -> RoutingTrace:
+    """Route a query string to the best matching adapter.
+
+    Steps
+    -----
+    1. Embed the query with the frozen embedding model.
+    2. Compute cosine similarity against all task centroid embeddings.
+    3. Select the highest-ranked task whose preferred adapters exist.
+    4. Build and return a ``RoutingTrace``.
+
+    Parameters
+    ----------
+    query:
+        The user's natural-language query or instruction.
+    task_registry:
+        Loaded ``TaskRegistry`` with pre-computed centroids.
+    adapter_registry:
+        Loaded ``AdapterRegistry``.
+    embedder:
+        ``Embedder`` instance (wraps fastembed singleton).
+    top_k:
+        Number of top task candidates to consider when walking the fallback
+        chain.  Defaults to 3.
+
+    Returns
+    -------
+    A ``RoutingTrace`` describing the decision.  The trace is **not**
+    persisted here — call ``feedback.loop.record_trace(trace)`` separately.
+
+    Raises
+    ------
+    NoAdapterError
+        If no registered adapter matches any of the top-K tasks.
+    ValueError
+        If embeddings have not been computed (missing centroids).
+    """
+    if not task_registry.embeddings_ready():
+        raise ValueError(
+            "Task embeddings are not initialised. Run `shiftgate init` first."
+        )
+
+    # Step 1: embed the query
+    query_embedding = embedder.embed(query)
+
+    # Step 2: rank tasks by similarity
+    all_tasks = task_registry.get_all_tasks()
+    ranked = top_k_tasks(query_embedding, all_tasks, k=top_k)
+
+    # Step 3: pick the best adapter
+    adapter, matched_task, score = select_adapter(ranked, adapter_registry)
+
+    # Step 4: assemble trace
+    trace = RoutingTrace(
+        id=uuid.uuid4().hex,
+        query=query,
+        matched_task_id=matched_task.id,
+        similarity_score=score,
+        selected_adapter_id=adapter.id,
+        timestamp=datetime.now(timezone.utc).isoformat(),
+    )
+
+    logger.info(
+        "Routed '%s' → task='%s' (%.2f%%) → adapter='%s'",
+        query[:60],
+        matched_task.id,
+        score * 100,
+        adapter.id,
+    )
+    return trace

shiftgate/runtime/__init__.py

@@ -0,0 +1 @@
+"""Runtime sub-package: thin clients for Ollama and vLLM backends."""

shiftgate/runtime/backend.py

@@ -0,0 +1,289 @@
+"""
+Thin clients for local LLM inference backends.
+
+Two backends are supported:
+
+Ollama (http://localhost:11434)
+    Open-source inference server that supports LoRA adapters via custom
+    Modelfiles.  To use a LoRA with Ollama, create a Modelfile like:
+
+        FROM llama3
+        ADAPTER /path/to/adapter.safetensors
+
+    and run ``ollama create my-model -f Modelfile``.  shiftgate sets
+    ``model`` to the adapter's Ollama model name (derived from adapter.id).
+
+vLLM  (http://localhost:8000)
+    Provides an OpenAI-compatible ``/v1/chat/completions`` endpoint.
+    LoRA adapters are loaded at server start-up with ``--lora-modules``
+    and are addressed by name via the ``model`` field in the request.
+
+Both backends are auto-detected by ``BackendRouter``, which pings each
+health endpoint and delegates to whichever is available.
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+
+import httpx
+
+from shiftgate.registry.schemas import AdapterEntry
+
+logger = logging.getLogger(__name__)
+
+# Default timeouts (seconds).
+_CONNECT_TIMEOUT = 3.0
+_READ_TIMEOUT = 120.0
+
+
+class BaseBackend(ABC):
+    """Abstract base for inference backends."""
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Return True if the backend can be reached."""
+
+    @abstractmethod
+    def generate(self, prompt: str, adapter: AdapterEntry) -> str:
+        """Send ``prompt`` to the backend and return the generated text."""
+
+
+# ---------------------------------------------------------------------------
+# Ollama
+# ---------------------------------------------------------------------------
+
+class OllamaBackend(BaseBackend):
+    """Thin httpx client for the Ollama inference server.
+
+    Ollama API reference: https://github.com/ollama/ollama/blob/main/docs/api.md
+
+    LoRA adapters in Ollama
+    -----------------------
+    Ollama does not have a first-class "load this adapter onto this base model"
+    API at request time.  Instead you pre-register a composite model via a
+    Modelfile::
+
+        FROM llama3
+        ADAPTER /path/to/my-lora.safetensors
+
+        ollama create my-lora-model -f Modelfile
+
+    shiftgate uses ``adapter.id`` as the Ollama model name by convention.
+    Ensure your Ollama model names match shiftgate adapter IDs.
+    """
+
+    def __init__(self, base_url: str = "http://localhost:11434") -> None:
+        self.base_url = base_url.rstrip("/")
+
+    def is_available(self) -> bool:
+        """Return True if the Ollama server responds to a health ping."""
+        try:
+            r = httpx.get(f"{self.base_url}/api/tags", timeout=_CONNECT_TIMEOUT)
+            return r.status_code == 200
+        except Exception:
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        adapter: AdapterEntry,
+        *,
+        model_name: str | None = None,
+        stream: bool = False,
+    ) -> str:
+        """Generate text via Ollama's ``/api/generate`` endpoint.
+
+        Parameters
+        ----------
+        prompt:
+            The full prompt string.
+        adapter:
+            The selected ``AdapterEntry``; ``adapter.id`` is used as the
+            Ollama model name unless overridden by ``model_name``.
+        model_name:
+            Override the Ollama model name (useful when the Ollama model name
+            differs from the shiftgate adapter ID).
+        stream:
+            If True, Ollama streams response tokens.  This client reads the
+            full stream and returns the concatenated text.
+        """
+        model = model_name or adapter.id
+        payload = {"model": model, "prompt": prompt, "stream": stream}
+
+        logger.debug("Ollama generate: model=%s", model)
+        try:
+            r = httpx.post(
+                f"{self.base_url}/api/generate",
+                json=payload,
+                timeout=_READ_TIMEOUT,
+            )
+            r.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise BackendError(f"Ollama request failed: {exc}") from exc
+
+        data = r.json()
+        return data.get("response", "")
+
+
+# ---------------------------------------------------------------------------
+# vLLM
+# ---------------------------------------------------------------------------
+
+class VLLMBackend(BaseBackend):
+    """Thin httpx client for the vLLM OpenAI-compatible inference server.
+
+    vLLM API reference: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
+
+    LoRA adapters in vLLM
+    ---------------------
+    vLLM loads LoRA adapters at startup via the ``--lora-modules`` flag::
+
+        python -m vllm.entrypoints.openai.api_server \\
+            --model meta-llama/Meta-Llama-3-8B \\
+            --lora-modules my-lora=/path/to/adapter \\
+            --enable-lora
+
+    After that, passing ``"model": "my-lora"`` in a chat completion request
+    automatically activates the adapter.
+    """
+
+    def __init__(self, base_url: str = "http://localhost:8000") -> None:
+        self.base_url = base_url.rstrip("/")
+
+    def is_available(self) -> bool:
+        """Return True if the vLLM server responds on the health endpoint."""
+        try:
+            r = httpx.get(f"{self.base_url}/health", timeout=_CONNECT_TIMEOUT)
+            return r.status_code == 200
+        except Exception:
+            return False
+
+    def generate(
+        self,
+        prompt: str,
+        adapter: AdapterEntry,
+        *,
+        lora_name: str | None = None,
+        system_prompt: str = "You are a helpful assistant.",
+    ) -> str:
+        """Generate text via vLLM's ``/v1/chat/completions`` endpoint.
+
+        Parameters
+        ----------
+        prompt:
+            The user message content.
+        adapter:
+            The selected ``AdapterEntry``; ``adapter.id`` is used as the
+            ``model`` field (which vLLM maps to the LoRA name) unless
+            overridden by ``lora_name``.
+        lora_name:
+            Override the vLLM model/lora name.
+        system_prompt:
+            System message prepended before the user message.
+        """
+        model = lora_name or adapter.id
+        payload = {
+            "model": model,
+            "messages": [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt},
+            ],
+        }
+
+        logger.debug("vLLM generate: model=%s", model)
+        try:
+            r = httpx.post(
+                f"{self.base_url}/v1/chat/completions",
+                json=payload,
+                timeout=_READ_TIMEOUT,
+            )
+            r.raise_for_status()
+        except httpx.HTTPError as exc:
+            raise BackendError(f"vLLM request failed: {exc}") from exc
+
+        data = r.json()
+        try:
+            return data["choices"][0]["message"]["content"]
+        except (KeyError, IndexError) as exc:
+            raise BackendError(f"Unexpected vLLM response format: {data}") from exc
+
+
+# ---------------------------------------------------------------------------
+# BackendRouter — auto-detects which backend is live
+# ---------------------------------------------------------------------------
+
+class BackendRouter:
+    """Detects and delegates to whichever local backend is running.
+
+    Priority: Ollama first, then vLLM.  If neither is available, calls to
+    ``generate`` raise ``NoBackendError`` with a helpful message.
+    """
+
+    def __init__(
+        self,
+        ollama_url: str = "http://localhost:11434",
+        vllm_url: str = "http://localhost:8000",
+    ) -> None:
+        self._ollama = OllamaBackend(ollama_url)
+        self._vllm = VLLMBackend(vllm_url)
+        self._active: BaseBackend | None = None
+
+    def detect(self) -> str | None:
+        """Probe both backends and return the name of the one that responds.
+
+        Returns ``"ollama"``, ``"vllm"``, or ``None``.
+        """
+        if self._ollama.is_available():
+            self._active = self._ollama
+            return "ollama"
+        if self._vllm.is_available():
+            self._active = self._vllm
+            return "vllm"
+        self._active = None
+        return None
+
+    def generate(self, prompt: str, adapter: AdapterEntry) -> str:
+        """Route the prompt to the active backend.
+
+        If no backend was detected yet, ``detect()`` is called automatically.
+
+        Raises
+        ------
+        NoBackendError
+            If neither Ollama nor vLLM is reachable.
+        """
+        if self._active is None:
+            self.detect()
+        if self._active is None:
+            raise NoBackendError(
+                "No inference backend detected.\n"
+                "  • Start Ollama : ollama serve\n"
+                "  • Start vLLM  : python -m vllm.entrypoints.openai.api_server "
+                "--model <base_model> --enable-lora\n\n"
+                "shiftgate can route queries without a backend. "
+                "Use `shiftgate route` to see routing decisions without inference."
+            )
+        return self._active.generate(prompt, adapter)
+
+    @property
+    def active_backend_name(self) -> str | None:
+        """Return 'ollama', 'vllm', or None depending on what was detected."""
+        if self._active is self._ollama:
+            return "ollama"
+        if self._active is self._vllm:
+            return "vllm"
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Exceptions
+# ---------------------------------------------------------------------------
+
+class BackendError(RuntimeError):
+    """Raised when an inference backend returns an error or unexpected response."""
+
+
+class NoBackendError(RuntimeError):
+    """Raised when no local inference backend is reachable."""

shiftgate/utils/__init__.py

@@ -0,0 +1 @@
+"""Utils sub-package: Rich terminal UI helpers."""