khazad 0.1.2__py3-none-any.whl

khazad/__init__.py

@@ -0,0 +1,109 @@
+"""Khazad — transparent semantic cache for LLM API calls.
+
+"You shall not pass" — Khazad stands between your application and
+expensive LLM API calls, turning semantically equivalent requests
+away at the bridge.
+
+Usage::
+
+    # Functional singleton API
+    import khazad
+    khazad.init(redis_url="redis://localhost:6379", threshold=0.92)
+    khazad.stop()
+
+    # Or manage the instance explicitly
+    from khazad import Khazad
+    cache = Khazad(redis_url="redis://localhost:6379", threshold=0.92)
+    cache.stop()
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Literal
+
+from khazad._models import CacheHit, CacheScope, ParsedRequest, Stats
+from khazad.khazad import Khazad
+
+__version__ = "0.1.2"
+__all__ = [
+    "CacheHit",
+    "CacheScope",
+    "Khazad",
+    "ParsedRequest",
+    "Stats",
+    "flush",
+    "get_stats",
+    "init",
+    "is_active",
+    "stop",
+]
+
+# ---------------------------------------------------------------------------
+# Module-level singleton — functional interface over a single Khazad instance
+# ---------------------------------------------------------------------------
+
+_instance: Khazad | None = None
+
+
+def init(
+    redis_url: str = "redis://localhost:6379",
+    threshold: float = 0.90,
+    ttl: int | None = None,
+    namespace: str = "khazad",
+    embedder: str = "huggingface",
+    embedding_model: str = "redis/langcache-embed-v2",
+    log_level: Literal["DEBUG", "INFO", "WARNING", "ERROR"] = "INFO",
+    hosts: list[str] | None = None,
+    cache_scope: CacheScope | Literal["model", "host"] = CacheScope.MODEL,
+) -> None:
+    """Activate the global Khazad singleton."""
+    global _instance
+
+    if _instance is not None and _instance.is_active():
+        logging.getLogger("khazad").warning("[Khazad] Already initialized, call stop() first")
+        return
+
+    _instance = Khazad(
+        redis_url=redis_url,
+        threshold=threshold,
+        ttl=ttl,
+        namespace=namespace,
+        embedder=embedder,
+        embedding_model=embedding_model,
+        log_level=log_level,
+        hosts=hosts,
+        cache_scope=cache_scope,
+    )
+
+
+def stop() -> None:
+    """Deactivate the global Khazad singleton."""
+    global _instance
+
+    if _instance is None or not _instance.is_active():
+        logging.getLogger("khazad").warning("[Khazad] Not currently active")
+        return
+
+    _instance.stop()
+    _instance = None
+
+
+def get_stats() -> dict:
+    """Return current cache performance metrics as a dictionary."""
+    if _instance is None:
+        return Stats().to_dict()
+    return _instance.get_stats().to_dict()
+
+
+def flush() -> None:
+    """Clear all cached entries from Redis."""
+    if _instance is None:
+        logging.getLogger("khazad").warning("[Khazad] Not initialized, nothing to flush")
+        return
+    _instance.flush()
+
+
+def is_active() -> bool:
+    """Return True if Khazad is currently intercepting HTTP traffic."""
+    return _instance is not None and _instance.is_active()

khazad/_models.py

@@ -0,0 +1,81 @@
+"""Domain models for parsed requests, cache hits and observability stats."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class CacheScope(str, Enum):
+    """How cache entries are partitioned within a provider host.
+
+    The provider host is always part of the scope, so a cached response is
+    never replayed across providers. This enum only controls whether the
+    **model** is also part of the key.
+
+    - :attr:`MODEL` (default) — each ``(host, model)`` pair gets its own vector
+      set, so a ``gpt-4o`` answer is never served to a ``gpt-4o-mini`` call.
+    - :attr:`HOST` — every model or deployment on the same provider shares one
+      vector set. Safe only for format-compatible pools (e.g. several Azure
+      OpenAI deployments, or treating ``gpt-4o`` and ``gpt-4o-mini`` as
+      interchangeable).
+
+    Members are also plain strings, so ``"model"`` / ``"host"`` are accepted
+    anywhere a :class:`CacheScope` is expected.
+    """
+
+    MODEL = "model"
+    HOST = "host"
+
+
+@dataclass(frozen=True, slots=True)
+class ParsedRequest:
+    """Semantic content extracted from a provider request body."""
+
+    prompt: str
+    model: str | None = None
+    stream: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class CacheHit:
+    """Result of a successful cache lookup."""
+
+    key: str
+    similarity: float
+    response_data: bytes
+    latency_ms: float
+
+
+@dataclass(slots=True)
+class Stats:
+    """Observable cache performance metrics."""
+
+    total_requests: int = 0
+    cache_hits: int = 0
+    cache_misses: int = 0
+    total_hit_similarity: float = 0.0
+
+    @property
+    def hit_rate(self) -> float:
+        """Return the cache hit ratio as a value between 0 and 1."""
+        if self.total_requests == 0:
+            return 0.0
+        return self.cache_hits / self.total_requests
+
+    @property
+    def avg_hit_similarity(self) -> float:
+        """Return the average cosine similarity of cache hits."""
+        if self.cache_hits == 0:
+            return 0.0
+        return self.total_hit_similarity / self.cache_hits
+
+    def to_dict(self) -> dict:
+        """Serialize stats to a plain dictionary."""
+        return {
+            "total_requests": self.total_requests,
+            "cache_hits": self.cache_hits,
+            "cache_misses": self.cache_misses,
+            "hit_rate": round(self.hit_rate, 4),
+            "avg_hit_similarity": round(self.avg_hit_similarity, 4),
+        }

khazad/_transport.py

@@ -0,0 +1,290 @@
+"""HTTP transport interceptor — patches httpx to route LLM traffic through the cache."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+import httpx
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Iterator
+
+    from khazad.khazad import Khazad, PreparedRequest
+
+logger = logging.getLogger("khazad")
+
+# ---------------------------------------------------------------------------
+# Patch install / uninstall
+# ---------------------------------------------------------------------------
+
+_original_async_init = None
+_original_sync_init = None
+
+
+def install(cache: Khazad) -> None:
+    """Monkey-patch httpx.Client and httpx.AsyncClient to use Khazad transports.
+
+    Safe to call repeatedly: only the *first* install captures the pristine
+    ``__init__`` references. Subsequent calls swap in a new cache without
+    losing the original — so ``uninstall()`` always restores real httpx.
+    """
+    global _original_async_init, _original_sync_init
+
+    if _original_async_init is None:
+        _original_async_init = httpx.AsyncClient.__init__
+    if _original_sync_init is None:
+        _original_sync_init = httpx.Client.__init__
+
+    original_async = _original_async_init
+    original_sync = _original_sync_init
+
+    def patched_async_init(self: httpx.AsyncClient, *args, **kwargs) -> None:
+        original_async(self, *args, **kwargs)
+        self._transport = CachedAsyncTransport(cache, self._transport)
+
+    def patched_sync_init(self: httpx.Client, *args, **kwargs) -> None:
+        original_sync(self, *args, **kwargs)
+        self._transport = CachedSyncTransport(cache, self._transport)
+
+    httpx.AsyncClient.__init__ = patched_async_init  # type: ignore[method-assign]
+    httpx.Client.__init__ = patched_sync_init  # type: ignore[method-assign]
+
+    logger.info("[Khazad] HTTP transport patches installed")
+
+
+def uninstall() -> None:
+    """Restore the original httpx transports."""
+    global _original_async_init, _original_sync_init
+
+    if _original_async_init is not None:
+        httpx.AsyncClient.__init__ = _original_async_init  # type: ignore[method-assign]
+        _original_async_init = None
+
+    if _original_sync_init is not None:
+        httpx.Client.__init__ = _original_sync_init  # type: ignore[method-assign]
+        _original_sync_init = None
+
+    logger.info("[Khazad] HTTP transport patches removed")
+
+
+# ---------------------------------------------------------------------------
+# Transport wrappers
+# ---------------------------------------------------------------------------
+
+
+class CachedSyncTransport(httpx.BaseTransport):
+    """Sync httpx transport that intercepts LLM requests for caching."""
+
+    def __init__(self, cache: Khazad, wrapped: httpx.BaseTransport) -> None:
+        self._cache = cache
+        self._wrapped = wrapped
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        cache = self._cache
+        prepared = cache.prepare(request) if cache.is_active() else None
+        if prepared is None:
+            return self._wrapped.handle_request(request)
+
+        hit = cache.lookup(prepared)
+        if hit is not None:
+            return _replay(prepared, hit)
+
+        response = self._wrapped.handle_request(request)
+        if response.status_code != 200:
+            return response
+
+        if _is_sse(response):
+            if not _can_capture(response):
+                return response
+            stream = _SyncTeeStream(
+                response.stream, lambda raw: _store_stream(cache, prepared, raw)
+            )
+            return _swap_stream(response, stream)
+
+        try:
+            response.read()
+            cache.store(prepared, response.content)
+        except Exception:
+            logger.warning("[Khazad] Failed to store response in cache", exc_info=True)
+        return response
+
+    def close(self) -> None:
+        self._wrapped.close()
+
+
+class CachedAsyncTransport(httpx.AsyncBaseTransport):
+    """Async httpx transport that intercepts LLM requests for caching."""
+
+    def __init__(self, cache: Khazad, wrapped: httpx.AsyncBaseTransport) -> None:
+        self._cache = cache
+        self._wrapped = wrapped
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        cache = self._cache
+        prepared = cache.prepare(request) if cache.is_active() else None
+        if prepared is None:
+            return await self._wrapped.handle_async_request(request)
+
+        # Embedding and Redis search are blocking — keep the event loop free.
+        loop = asyncio.get_running_loop()
+        hit = await loop.run_in_executor(None, cache.lookup, prepared)
+        if hit is not None:
+            return _replay(prepared, hit)
+
+        response = await self._wrapped.handle_async_request(request)
+        if response.status_code != 200:
+            return response
+
+        if _is_sse(response):
+            if not _can_capture(response):
+                return response
+            stream = _AsyncTeeStream(
+                response.stream, lambda raw: _store_stream(cache, prepared, raw)
+            )
+            return _swap_stream(response, stream)
+
+        try:
+            await response.aread()
+            await loop.run_in_executor(None, cache.store, prepared, response.content)
+        except Exception:
+            logger.warning("[Khazad] Failed to store response in cache", exc_info=True)
+        return response
+
+    async def aclose(self) -> None:
+        await self._wrapped.aclose()
+
+
+# ---------------------------------------------------------------------------
+# Cache hit / miss plumbing
+# ---------------------------------------------------------------------------
+
+
+def _replay(prepared: PreparedRequest, hit) -> httpx.Response:
+    """Build the response for a cache hit — JSON body or simulated SSE stream."""
+    if prepared.stream:
+        return httpx.Response(
+            status_code=200,
+            headers={"content-type": "text/event-stream"},
+            stream=_ReplayStream(prepared.parser.stream_chunks(hit.response_data)),
+        )
+    return prepared.parser.build_response(hit.response_data)
+
+
+def _store_stream(cache: Khazad, prepared: PreparedRequest, raw: bytes) -> None:
+    """Reconstruct a canonical JSON response from raw SSE bytes and cache it."""
+    body = prepared.parser.response_from_stream(raw)
+    if body:
+        cache.store(prepared, body)
+
+
+def _is_sse(response: httpx.Response) -> bool:
+    return "text/event-stream" in response.headers.get("content-type", "")
+
+
+def _can_capture(response: httpx.Response) -> bool:
+    """Raw stream bytes are only parseable when the body is not compressed."""
+    return response.headers.get("content-encoding", "identity").lower() in ("", "identity")
+
+
+def _swap_stream(
+    response: httpx.Response, stream: httpx.SyncByteStream | httpx.AsyncByteStream
+) -> httpx.Response:
+    """Clone a response, replacing its stream with a tee'd one."""
+    return httpx.Response(
+        status_code=response.status_code,
+        headers=response.headers,
+        stream=stream,
+        extensions=response.extensions,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Stream helpers
+# ---------------------------------------------------------------------------
+
+
+class _ReplayStream(httpx.SyncByteStream, httpx.AsyncByteStream):
+    """Serve synthesized SSE chunks to either a sync or an async client."""
+
+    def __init__(self, chunks: Iterator[bytes]) -> None:
+        self._chunks = chunks
+
+    def __iter__(self):
+        yield from self._chunks
+
+    async def __aiter__(self):
+        for chunk in self._chunks:
+            yield chunk
+            await asyncio.sleep(0)  # let other tasks run between frames
+
+
+class _SyncTeeStream(httpx.SyncByteStream):
+    """Pass chunks through untouched; cache the body when the stream ends.
+
+    SDKs (e.g. the OpenAI client) break their read loop on the terminal SSE
+    sentinel and call ``close()`` without driving the byte stream to EOF, so
+    the body is captured on *either* natural exhaustion or ``close()``. The
+    parser's :meth:`response_from_stream` decides whether the captured bytes
+    form a complete response — an aborted, partial stream reconstructs to
+    ``None`` and is never cached.
+    """
+
+    def __init__(self, inner: httpx.SyncByteStream, on_complete: Callable[[bytes], None]) -> None:
+        self._inner = inner
+        self._on_complete = on_complete
+        self._parts: list[bytes] = []
+        self._finished = False
+
+    def __iter__(self):
+        for chunk in self._inner:
+            self._parts.append(chunk)
+            yield chunk
+        self._finish()
+
+    def _finish(self) -> None:
+        if self._finished:
+            return
+        self._finished = True
+        try:
+            self._on_complete(b"".join(self._parts))
+        except Exception:
+            logger.warning("[Khazad] Failed to store streamed response", exc_info=True)
+
+    def close(self) -> None:
+        self._finish()
+        self._inner.close()
+
+
+class _AsyncTeeStream(httpx.AsyncByteStream):
+    """Async variant of :class:`_SyncTeeStream`; stores off the event loop."""
+
+    def __init__(self, inner: httpx.AsyncByteStream, on_complete: Callable[[bytes], None]) -> None:
+        self._inner = inner
+        self._on_complete = on_complete
+        self._parts: list[bytes] = []
+        self._finished = False
+
+    async def __aiter__(self):
+        async for chunk in self._inner:
+            self._parts.append(chunk)
+            yield chunk
+        self._finish()
+
+    def _finish(self) -> None:
+        if self._finished:
+            return
+        self._finished = True
+        body = b"".join(self._parts)
+        asyncio.get_running_loop().run_in_executor(None, self._safe_complete, body)
+
+    def _safe_complete(self, body: bytes) -> None:
+        try:
+            self._on_complete(body)
+        except Exception:
+            logger.warning("[Khazad] Failed to store streamed response", exc_info=True)
+
+    async def aclose(self) -> None:
+        self._finish()
+        await self._inner.aclose()

khazad/adapters/embedders/huggingface.py

@@ -0,0 +1,41 @@
+"""HuggingFace embedding adapter using sentence-transformers."""
+
+from __future__ import annotations
+
+import logging
+import threading
+from functools import cached_property
+
+from sentence_transformers import SentenceTransformer
+
+from khazad.ports.embedder import Embedder
+
+logger = logging.getLogger("khazad")
+
+
+class HuggingFaceEmbedder(Embedder):
+    """Embedder backed by a local sentence-transformers model."""
+
+    def __init__(self, model_name: str = "redis/langcache-embed-v2") -> None:
+        self._model_name = model_name
+        self._model: SentenceTransformer | None = None
+        self._load_lock = threading.Lock()
+
+    def _get_model(self) -> SentenceTransformer:
+        """Lazily load the model, guarding against concurrent double-loads."""
+        if self._model is None:
+            with self._load_lock:
+                if self._model is None:
+                    logger.info("[Khazad] Loading embedding model: %s", self._model_name)
+                    self._model = SentenceTransformer(self._model_name)
+        return self._model
+
+    def embed(self, text: str) -> list[float]:
+        """Generate a normalized embedding vector for the given text."""
+        vector = self._get_model().encode(text, normalize_embeddings=True)
+        return vector.tolist()
+
+    @cached_property
+    def dimension(self) -> int:
+        """Return the dimensionality of this model's embeddings."""
+        return self._get_model().get_sentence_embedding_dimension()  # type: ignore[return-value]

khazad/adapters/embedders/openai.py

@@ -0,0 +1,50 @@
+"""OpenAI embedding adapter (optional paid backend)."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from khazad.ports.embedder import Embedder
+
+if TYPE_CHECKING:
+    from openai import OpenAI
+
+logger = logging.getLogger("khazad")
+
+
+class OpenAIEmbedder(Embedder):
+    """Embedder backed by the OpenAI Embeddings API."""
+
+    def __init__(
+        self,
+        model: str = "text-embedding-3-small",
+        api_key: str | None = None,
+        dimension: int = 1536,
+    ) -> None:
+        self._model = model
+        self._api_key = api_key
+        self._dimension = dimension
+        self._client: OpenAI | None = None
+
+    def _get_client(self) -> OpenAI:
+        """Lazily create the OpenAI client; fails clearly if the extra is missing."""
+        if self._client is None:
+            try:
+                from openai import OpenAI
+            except ImportError as exc:
+                raise ImportError(
+                    "Install the 'openai' package: pip install khazad[openai-embeddings]"
+                ) from exc
+            self._client = OpenAI(api_key=self._api_key)
+        return self._client
+
+    def embed(self, text: str) -> list[float]:
+        """Generate an embedding via the OpenAI API."""
+        response = self._get_client().embeddings.create(input=text, model=self._model)
+        return response.data[0].embedding
+
+    @property
+    def dimension(self) -> int:
+        """Return the configured embedding dimension."""
+        return self._dimension