openai-sqlite-cache 0.0.1__py3-none-any.whl

cached_openai/_http.py

@@ -0,0 +1,515 @@
+"""Caching HTTP-client adapters used by the OpenAI SDK."""
+
+from __future__ import annotations
+
+import asyncio
+import importlib
+import threading
+import warnings
+import weakref
+import zlib
+from typing import Any, Awaitable, Callable, Iterable, List, Optional, Tuple
+
+from ._cache import CachedResponse, SQLiteCache
+from ._fingerprint import request_fingerprint
+from ._settings import CacheSettings
+
+
+def _sdk_http_module() -> Any:
+    """Resolve the HTTP module used by this installed OpenAI SDK.
+
+    OpenAI 1.x/early 2.x use ``httpx``. Newer SDK releases can use the
+    API-compatible ``httpx2`` package, so importing the SDK's own dependency is
+    more robust than pinning this package to either implementation.
+    """
+
+    base_client = importlib.import_module("openai._base_client")
+    for name in ("httpx2", "httpx"):
+        module = getattr(base_client, name, None)
+        if module is not None:
+            return module
+    for name in ("httpx2", "httpx"):
+        try:
+            return importlib.import_module(name)
+        except ImportError:
+            continue
+    raise ImportError(
+        "openai-sqlite-cache could not locate the HTTP library used by openai"
+    )
+
+
+_http = _sdk_http_module()
+_SYNC_LOCKS = [threading.Lock() for _ in range(127)]
+_ASYNC_LOCKS: "weakref.WeakKeyDictionary[Any, List[asyncio.Lock]]" = (
+    weakref.WeakKeyDictionary()
+)
+_ASYNC_LOCKS_GUARD = threading.Lock()
+_warning_lock = threading.Lock()
+_warned_messages = set()
+
+
+def _warn_once(message: str) -> None:
+    with _warning_lock:
+        if message in _warned_messages:
+            return
+        _warned_messages.add(message)
+    warnings.warn(message, RuntimeWarning, stacklevel=3)
+
+
+def _sync_lock(cache_key: str) -> Any:
+    return _SYNC_LOCKS[int(cache_key[-8:], 16) % len(_SYNC_LOCKS)]
+
+
+def _async_lock(cache_key: str) -> asyncio.Lock:
+    loop = asyncio.get_running_loop()
+    with _ASYNC_LOCKS_GUARD:
+        locks = _ASYNC_LOCKS.get(loop)
+        if locks is None:
+            locks = [asyncio.Lock() for _ in range(127)]
+            _ASYNC_LOCKS[loop] = locks
+    return locks[int(cache_key[-8:], 16) % len(locks)]
+
+
+_DROPPED_RESPONSE_HEADERS = {
+    "connection",
+    "content-encoding",
+    "content-length",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailer",
+    "transfer-encoding",
+    "upgrade",
+}
+
+
+def _storable_headers(headers: Any, *, raw_body: bool) -> List[Tuple[str, str]]:
+    items = (
+        headers.multi_items() if hasattr(headers, "multi_items") else headers.items()
+    )
+    return [
+        (str(name), str(value))
+        for name, value in items
+        if str(name).lower() not in _DROPPED_RESPONSE_HEADERS
+        or (raw_body and str(name).lower() == "content-encoding")
+    ]
+
+
+class _SSETerminalDetector:
+    """Find an SSE done frame incrementally, including gzip/deflate streams."""
+
+    _markers = (b"data: [DONE]", b"data:[DONE]")
+
+    def __init__(self, content_encoding: str):
+        self.complete = False
+        self._tail = b""
+        self._encoding = content_encoding.strip().lower()
+        self._decoder: Any = None
+        if self._encoding in {"gzip", "x-gzip"}:
+            self._decoder = zlib.decompressobj(zlib.MAX_WBITS | 16)
+        elif self._encoding == "deflate":
+            self._decoder = zlib.decompressobj()
+
+    def feed(self, chunk: bytes) -> bool:
+        if self.complete:
+            return True
+        if self._encoding and self._decoder is None:
+            return False
+        decoded = chunk
+        if self._decoder is not None:
+            try:
+                decoded = self._decoder.decompress(chunk)
+            except zlib.error:
+                return False
+        candidate = self._tail + decoded
+        self.complete = any(marker in candidate for marker in self._markers)
+        longest_marker = max(len(marker) for marker in self._markers)
+        self._tail = candidate[-(longest_marker - 1) :]
+        return self.complete
+
+
+class _CacheAccess:
+    def __init__(self, settings: CacheSettings):
+        self.settings = settings
+        self.backend: Optional[SQLiteCache]
+        try:
+            self.backend = SQLiteCache(settings.path)
+        except Exception as error:
+            self.backend = None
+            _warn_once(
+                "openai-sqlite-cache cache could not be opened; "
+                f"caching is disabled: {error}"
+            )
+
+    def get(self, key: str) -> Optional[CachedResponse]:
+        if self.backend is None:
+            return None
+        try:
+            return self.backend.get(key, self.settings.ttl_seconds)
+        except Exception as error:
+            _warn_once(
+                "openai-sqlite-cache cache read failed; "
+                f"requesting upstream instead: {error}"
+            )
+            return None
+
+    def put(
+        self, key: str, response: Any, body: bytes, *, raw_body: bool = False
+    ) -> None:
+        if self.backend is None or not 200 <= int(response.status_code) < 300:
+            return
+        try:
+            self.backend.put(
+                key,
+                int(response.status_code),
+                _storable_headers(response.headers, raw_body=raw_body),
+                body,
+            )
+        except Exception as error:
+            _warn_once(
+                "openai-sqlite-cache cache write failed; "
+                f"response was not cached: {error}"
+            )
+
+    @staticmethod
+    def response(record: CachedResponse, request: Any) -> Any:
+        return _http.Response(
+            status_code=record.status_code,
+            headers=record.headers,
+            content=record.body,
+            request=request,
+        )
+
+
+class _CachingSyncStream(_http.SyncByteStream):
+    def __init__(
+        self,
+        stream: Any,
+        on_complete: Callable[[bytes], None],
+        release: Callable[[], None],
+        content_encoding: str,
+    ):
+        self._stream = stream
+        self._on_complete = on_complete
+        self._release = release
+        self._content_encoding = content_encoding
+        self._terminal = _SSETerminalDetector(content_encoding)
+        self._chunks = bytearray()
+        self._stored = False
+        self._released = False
+
+    def __iter__(self) -> Iterable[bytes]:
+        completed = False
+        try:
+            for chunk in self._stream:
+                self._chunks.extend(chunk)
+                if self._terminal.feed(chunk):
+                    self._store_once()
+                yield chunk
+            completed = True
+            self._store_once()
+        finally:
+            if not completed:
+                if self._terminal.complete:
+                    self._store_once()
+                if not self._stored:
+                    self._chunks.clear()
+            self._release_once()
+
+    def _store_once(self) -> None:
+        if not self._stored:
+            self._stored = True
+            self._on_complete(bytes(self._chunks))
+
+    def _release_once(self) -> None:
+        if not self._released:
+            self._released = True
+            self._release()
+
+    def close(self) -> None:
+        try:
+            self._stream.close()
+        finally:
+            self._release_once()
+
+
+class _CachingAsyncStream(_http.AsyncByteStream):
+    def __init__(
+        self,
+        stream: Any,
+        on_complete: Callable[[bytes], Awaitable[None]],
+        release: Callable[[], None],
+        content_encoding: str,
+    ):
+        self._stream = stream
+        self._on_complete = on_complete
+        self._release = release
+        self._content_encoding = content_encoding
+        self._terminal = _SSETerminalDetector(content_encoding)
+        self._chunks = bytearray()
+        self._stored = False
+        self._released = False
+
+    async def __aiter__(self) -> Any:
+        completed = False
+        try:
+            async for chunk in self._stream:
+                self._chunks.extend(chunk)
+                if self._terminal.feed(chunk):
+                    await self._store_once()
+                yield chunk
+            completed = True
+            await self._store_once()
+        finally:
+            if not completed:
+                if self._terminal.complete:
+                    await self._store_once()
+                if not self._stored:
+                    self._chunks.clear()
+            self._release_once()
+
+    async def _store_once(self) -> None:
+        if not self._stored:
+            self._stored = True
+            await self._on_complete(bytes(self._chunks))
+
+    def _release_once(self) -> None:
+        if not self._released:
+            self._released = True
+            self._release()
+
+    async def aclose(self) -> None:
+        try:
+            if self._terminal.complete:
+                await self._store_once()
+            await self._stream.aclose()
+        finally:
+            self._release_once()
+
+
+class CachingSyncClient(_http.Client):
+    """An ``httpx.Client``-compatible proxy around any SDK HTTP client."""
+
+    def __init__(self, inner: Any, settings: CacheSettings):
+        self._cached_openai_inner = inner
+        self._cached_openai_cache = _CacheAccess(settings)
+        self._cached_openai_settings = settings
+
+    @property
+    def timeout(self) -> Any:
+        return self._cached_openai_inner.timeout
+
+    @timeout.setter
+    def timeout(self, value: Any) -> None:
+        self._cached_openai_inner.timeout = value
+
+    @property
+    def is_closed(self) -> bool:
+        return bool(self._cached_openai_inner.is_closed)
+
+    def build_request(self, *args: Any, **kwargs: Any) -> Any:
+        return self._cached_openai_inner.build_request(*args, **kwargs)
+
+    def send(self, request: Any, *args: Any, **kwargs: Any) -> Any:
+        if not self._cached_openai_settings.enabled:
+            return self._cached_openai_inner.send(request, *args, **kwargs)
+        try:
+            body = request.read()
+            cache_key = request_fingerprint(
+                str(request.method), str(request.url), request.headers, body
+            )
+        except Exception as error:
+            _warn_once(
+                "openai-sqlite-cache could not fingerprint a request; "
+                f"bypassing cache: {error}"
+            )
+            return self._cached_openai_inner.send(request, *args, **kwargs)
+        if cache_key is None:
+            return self._cached_openai_inner.send(request, *args, **kwargs)
+
+        lock = _sync_lock(cache_key)
+        lock.acquire()
+        release_here = True
+        try:
+            cached = self._cached_openai_cache.get(cache_key)
+            if cached is not None:
+                try:
+                    return self._cached_openai_cache.response(cached, request)
+                except Exception as error:
+                    _warn_once(
+                        "openai-sqlite-cache could not reconstruct a cached response; "
+                        f"requesting upstream instead: {error}"
+                    )
+
+            response = self._cached_openai_inner.send(request, *args, **kwargs)
+            if not 200 <= int(response.status_code) < 300:
+                return response
+            if kwargs.get("stream", False):
+                try:
+                    buffered_content = bytes(response.content)
+                except Exception:
+                    release_here = False
+                    response.stream = _CachingSyncStream(
+                        response.stream,
+                        lambda content: self._cached_openai_cache.put(
+                            cache_key, response, content, raw_body=True
+                        ),
+                        lock.release,
+                        str(response.headers.get("content-encoding", "")),
+                    )
+                else:
+                    self._cached_openai_cache.put(cache_key, response, buffered_content)
+            else:
+                self._cached_openai_cache.put(
+                    cache_key, response, bytes(response.content)
+                )
+            return response
+        finally:
+            if release_here:
+                lock.release()
+
+    def close(self) -> None:
+        self._cached_openai_inner.close()
+
+    def __enter__(self) -> "CachingSyncClient":
+        self._cached_openai_inner.__enter__()
+        return self
+
+    def __exit__(self, *args: Any) -> Any:
+        return self._cached_openai_inner.__exit__(*args)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._cached_openai_inner, name)
+
+
+class CachingAsyncClient(_http.AsyncClient):
+    """An ``httpx.AsyncClient``-compatible proxy around any SDK HTTP client."""
+
+    def __init__(self, inner: Any, settings: CacheSettings):
+        self._cached_openai_inner = inner
+        self._cached_openai_cache = _CacheAccess(settings)
+        self._cached_openai_settings = settings
+
+    @property
+    def timeout(self) -> Any:
+        return self._cached_openai_inner.timeout
+
+    @timeout.setter
+    def timeout(self, value: Any) -> None:
+        self._cached_openai_inner.timeout = value
+
+    @property
+    def is_closed(self) -> bool:
+        return bool(self._cached_openai_inner.is_closed)
+
+    def build_request(self, *args: Any, **kwargs: Any) -> Any:
+        return self._cached_openai_inner.build_request(*args, **kwargs)
+
+    async def send(self, request: Any, *args: Any, **kwargs: Any) -> Any:
+        if not self._cached_openai_settings.enabled:
+            return await self._cached_openai_inner.send(request, *args, **kwargs)
+        try:
+            body = await request.aread()
+            cache_key = request_fingerprint(
+                str(request.method), str(request.url), request.headers, body
+            )
+        except Exception as error:
+            _warn_once(
+                "openai-sqlite-cache could not fingerprint a request; "
+                f"bypassing cache: {error}"
+            )
+            return await self._cached_openai_inner.send(request, *args, **kwargs)
+        if cache_key is None:
+            return await self._cached_openai_inner.send(request, *args, **kwargs)
+
+        lock = _async_lock(cache_key)
+        await lock.acquire()
+        release_here = True
+        try:
+            cached = await asyncio.to_thread(self._cached_openai_cache.get, cache_key)
+            if cached is not None:
+                try:
+                    return self._cached_openai_cache.response(cached, request)
+                except Exception as error:
+                    _warn_once(
+                        "openai-sqlite-cache could not reconstruct a cached response; "
+                        f"requesting upstream instead: {error}"
+                    )
+
+            response = await self._cached_openai_inner.send(request, *args, **kwargs)
+            if not 200 <= int(response.status_code) < 300:
+                return response
+            if kwargs.get("stream", False):
+                try:
+                    buffered_content = bytes(response.content)
+                except Exception:
+                    release_here = False
+
+                    async def store(content: bytes) -> None:
+                        try:
+                            await asyncio.to_thread(
+                                self._cached_openai_cache.put,
+                                cache_key,
+                                response,
+                                content,
+                                raw_body=True,
+                            )
+                        except Exception as error:
+                            _warn_once(
+                                f"openai-sqlite-cache cache write failed: {error}"
+                            )
+
+                    response.stream = _CachingAsyncStream(
+                        response.stream,
+                        store,
+                        lock.release,
+                        str(response.headers.get("content-encoding", "")),
+                    )
+                else:
+                    await asyncio.to_thread(
+                        self._cached_openai_cache.put,
+                        cache_key,
+                        response,
+                        buffered_content,
+                    )
+            else:
+                await asyncio.to_thread(
+                    self._cached_openai_cache.put,
+                    cache_key,
+                    response,
+                    bytes(response.content),
+                )
+            return response
+        finally:
+            if release_here:
+                lock.release()
+
+    async def aclose(self) -> None:
+        await self._cached_openai_inner.aclose()
+
+    async def __aenter__(self) -> "CachingAsyncClient":
+        await self._cached_openai_inner.__aenter__()
+        return self
+
+    async def __aexit__(self, *args: Any) -> Any:
+        return await self._cached_openai_inner.__aexit__(*args)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._cached_openai_inner, name)
+
+
+def wrap_sync_client(client: Any, settings: CacheSettings) -> Any:
+    if isinstance(client, CachingSyncClient):
+        if client._cached_openai_settings == settings:
+            return client
+        client = client._cached_openai_inner
+    return CachingSyncClient(client, settings)
+
+
+def wrap_async_client(client: Any, settings: CacheSettings) -> Any:
+    if isinstance(client, CachingAsyncClient):
+        if client._cached_openai_settings == settings:
+            return client
+        client = client._cached_openai_inner
+    return CachingAsyncClient(client, settings)

cached_openai/_module_alias.py

@@ -0,0 +1,86 @@
+"""Lazy aliases for public submodules from the upstream ``openai`` package."""
+
+from __future__ import annotations
+
+import importlib
+import importlib.abc
+import importlib.util
+import sys
+from types import ModuleType
+from typing import Any, Optional
+
+
+class _AliasModule(ModuleType):
+    def __init__(self, alias: str, target_name: str):
+        super().__init__(alias)
+        target = importlib.import_module(target_name)
+        self.__dict__["_alias_target"] = target
+        self.__dict__["__doc__"] = getattr(target, "__doc__", None)
+        exported = getattr(target, "__all__", None)
+        self.__dict__["__all__"] = (
+            list(exported)
+            if exported is not None
+            else [name for name in dir(target) if not name.startswith("_")]
+        )
+        self.__dict__["__file__"] = getattr(target, "__file__", None)
+        if hasattr(target, "__path__"):
+            self.__dict__["__path__"] = list(target.__path__)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.__dict__["_alias_target"], name)
+
+    def __dir__(self) -> list:
+        return sorted(set(self.__dict__) | set(dir(self.__dict__["_alias_target"])))
+
+
+class _AliasLoader(importlib.abc.Loader):
+    def __init__(self, alias: str, target: str):
+        self.alias = alias
+        self.target = target
+
+    def create_module(self, spec: Any) -> ModuleType:
+        return _AliasModule(self.alias, self.target)
+
+    def exec_module(self, module: ModuleType) -> None:
+        return None
+
+
+class _AliasFinder(importlib.abc.MetaPathFinder):
+    def __init__(self, alias_root: str, target_root: str):
+        self.alias_root = alias_root
+        self.target_root = target_root
+
+    def find_spec(
+        self, fullname: str, path: Any = None, target: Any = None
+    ) -> Optional[Any]:
+        prefix = self.alias_root + "."
+        if not fullname.startswith(prefix):
+            return None
+        target_name = self.target_root + fullname[len(self.alias_root) :]
+        target_spec = importlib.util.find_spec(target_name)
+        if target_spec is None:
+            return None
+        is_package = target_spec.submodule_search_locations is not None
+        return importlib.util.spec_from_loader(
+            fullname,
+            _AliasLoader(fullname, target_name),
+            is_package=is_package,
+        )
+
+
+def install_module_alias(alias_root: str, target_root: str) -> ModuleType:
+    """Install a lazy module tree alias while preserving exported class identity."""
+
+    existing = sys.modules.get(alias_root)
+    if isinstance(existing, _AliasModule):
+        return existing
+    proxy = _AliasModule(alias_root, target_root)
+    sys.modules[alias_root] = proxy
+    if not any(
+        isinstance(finder, _AliasFinder)
+        and finder.alias_root == alias_root
+        and finder.target_root == target_root
+        for finder in sys.meta_path
+    ):
+        sys.meta_path.insert(0, _AliasFinder(alias_root, target_root))
+    return proxy