codekingpro/portable-devtools

1"""Create a key-value store for any langchain serializable object."""2 3from collections.abc import Callable4from typing import Any5 6from langchain_core.documents import Document7from langchain_core.load import Serializable, dumps, loads8from langchain_core.stores import BaseStore, ByteStore9 10from langchain_classic.storage.encoder_backed import EncoderBackedStore11 12 13def _dump_as_bytes(obj: Serializable) -> bytes:14    """Return a bytes representation of a `Document`."""15    return dumps(obj).encode("utf-8")16 17 18def _dump_document_as_bytes(obj: Any) -> bytes:19    """Return a bytes representation of a `Document`."""20    if not isinstance(obj, Document):21        msg = "Expected a Document instance"22        raise TypeError(msg)23    return dumps(obj).encode("utf-8")24 25 26def _load_document_from_bytes(serialized: bytes) -> Document:27    """Return a document from a bytes representation."""28    obj = loads(serialized.decode("utf-8"), allowed_objects=[Document])29    if not isinstance(obj, Document):30        msg = f"Expected a Document instance. Got {type(obj)}"31        raise TypeError(msg)32    return obj33 34 35def _load_from_bytes(serialized: bytes) -> Serializable:36    """Return a `Serializable` from a bytes representation."""37    # The default allowlist (`'core'`) is unsafe with untrusted input - a38    # tampered byte payload can reconstruct any core class with39    # attacker-controlled kwargs (custom `base_url`, headers, model name,40    # etc.). The byte store backing this loader must be treated as a trust41    # boundary - see the danger note on `create_lc_store`. If the store can42    # be written to by anyone you do not already trust, use43    # `create_kv_docstore` instead.44    return loads(serialized.decode("utf-8"))45 46 47def _identity(x: str) -> str:48    """Return the same object."""49    return x50 51 52# PUBLIC API53 54 55def create_lc_store(56    store: ByteStore,57    *,58    key_encoder: Callable[[str], str] | None = None,59) -> BaseStore[str, Serializable]:60    """Create a store for LangChain serializable objects from a bytes store.61 62    !!! danger "Treat the underlying byte store as a trust boundary"63 64        Reads from this store are deserialized with65        `langchain_core.load.loads`, which instantiates Python objects from66        the stored payload. The same threat model applies: a payload can67        carry constructor kwargs (custom `base_url`, headers, model name,68        etc.) that get applied during `__init__`, so the bytes are69        effectively executable configuration rather than plain data.70 71        **Never back this store with anything an attacker can write to** —72        for example a shared cache that other tenants can populate, an73        S3 bucket without strict write controls, or a Redis instance reused74        across trust boundaries. A single tampered value will instantiate75        attacker-controlled classes the next time the store is read.76 77        If you cannot guarantee the store is write-restricted to your own78        process, use `create_kv_docstore` instead — it pins79        `allowed_objects=[Document]` so a tampered value can at worst80        produce a `Document`, never a chat model or LLM with a redirected81        endpoint.82 83    Args:84        store: A bytes store to use as the underlying store.85        key_encoder: A function to encode keys; if `None` uses identity function.86 87    Returns:88        A key-value store for `Document` objects.89    """90    return EncoderBackedStore(91        store,92        key_encoder or _identity,93        _dump_as_bytes,94        _load_from_bytes,95    )96 97 98def create_kv_docstore(99    store: ByteStore,100    *,101    key_encoder: Callable[[str], str] | None = None,102) -> BaseStore[str, Document]:103    """Create a store for langchain `Document` objects from a bytes store.104 105    This store does run time type checking to ensure that the values are106    `Document` objects.107 108    Args:109        store: A bytes store to use as the underlying store.110        key_encoder: A function to encode keys; if `None`, uses identity function.111 112    Returns:113        A key-value store for `Document` objects.114    """115    return EncoderBackedStore(116        store,117        key_encoder or _identity,118        _dump_document_as_bytes,119        _load_document_from_bytes,120    )121