codekingpro/portable-devtools

1"""In-memory dictionary-backed store with optional vector search.2 3!!! example "Examples"4    Basic key-value storage:5    ```python6    from langgraph.store.memory import InMemoryStore7 8    store = InMemoryStore()9    store.put(("users", "123"), "prefs", {"theme": "dark"})10    item = store.get(("users", "123"), "prefs")11    ```12 13    Vector search using LangChain embeddings:14    ```python15    from langchain.embeddings import init_embeddings16    from langgraph.store.memory import InMemoryStore17 18    store = InMemoryStore(19        index={20            "dims": 1536,21            "embed": init_embeddings("openai:text-embedding-3-small")22        }23    )24 25    # Store documents26    store.put(("docs",), "doc1", {"text": "Python tutorial"})27    store.put(("docs",), "doc2", {"text": "TypeScript guide"})28 29    # Search by similarity30    results = store.search(("docs",), query="python programming")31    ```32 33    Vector search using OpenAI SDK directly:34    ```python35    from openai import OpenAI36    from langgraph.store.memory import InMemoryStore37 38    client = OpenAI()39 40    def embed_texts(texts: list[str]) -> list[list[float]]:41        response = client.embeddings.create(42            model="text-embedding-3-small",43            input=texts44        )45        return [e.embedding for e in response.data]46 47    store = InMemoryStore(48        index={49            "dims": 1536,50            "embed": embed_texts51        }52    )53 54    # Store documents55    store.put(("docs",), "doc1", {"text": "Python tutorial"})56    store.put(("docs",), "doc2", {"text": "TypeScript guide"})57 58    # Search by similarity59    results = store.search(("docs",), query="python programming")60    ```61 62    Async vector search using OpenAI SDK:63    ```python64    from openai import AsyncOpenAI65    from langgraph.store.memory import InMemoryStore66 67    client = AsyncOpenAI()68 69    async def aembed_texts(texts: list[str]) -> list[list[float]]:70        response = await client.embeddings.create(71            model="text-embedding-3-small",72            input=texts73        )74        return [e.embedding for e in response.data]75 76    store = InMemoryStore(77        index={78            "dims": 1536,79            "embed": aembed_texts80        }81    )82 83    # Store documents84    await store.aput(("docs",), "doc1", {"text": "Python tutorial"})85    await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})86 87    # Search by similarity88    results = await store.asearch(("docs",), query="python programming")89    ```90 91Warning:92    This store keeps all data in memory. Data is lost when the process exits.93    For persistence, use a database-backed store like PostgresStore.94 95Tip:96    For vector search, install numpy for better performance:97    ```bash98    pip install numpy99    ```100"""101 102from __future__ import annotations103 104import asyncio105import concurrent.futures as cf106import functools107import logging108from collections import defaultdict109from collections.abc import Iterable110from datetime import datetime, timezone111from importlib import util112from typing import Any113 114from langchain_core.embeddings import Embeddings115 116from langgraph.store.base import (117    BaseStore,118    GetOp,119    IndexConfig,120    Item,121    ListNamespacesOp,122    MatchCondition,123    Op,124    PutOp,125    Result,126    SearchItem,127    SearchOp,128    ensure_embeddings,129    get_text_at_path,130    tokenize_path,131)132 133logger = logging.getLogger(__name__)134 135 136class InMemoryStore(BaseStore):137    """In-memory dictionary-backed store with optional vector search.138 139    !!! example "Examples"140        Basic key-value storage:141            store = InMemoryStore()142            store.put(("users", "123"), "prefs", {"theme": "dark"})143            item = store.get(("users", "123"), "prefs")144 145        Vector search with embeddings:146            from langchain.embeddings import init_embeddings147            store = InMemoryStore(index={148                "dims": 1536,149                "embed": init_embeddings("openai:text-embedding-3-small"),150                "fields": ["text"],151            })152 153            # Store documents154            store.put(("docs",), "doc1", {"text": "Python tutorial"})155            store.put(("docs",), "doc2", {"text": "TypeScript guide"})156 157            # Search by similarity158            results = store.search(("docs",), query="python programming")159 160    Note:161        Semantic search is disabled by default. You can enable it by providing an `index` configuration162        when creating the store. Without this configuration, all `index` arguments passed to163        `put` or `aput`will have no effect.164 165    Warning:166        This store keeps all data in memory. Data is lost when the process exits.167        For persistence, use a database-backed store like PostgresStore.168 169    Tip:170        For vector search, install numpy for better performance:171        ```bash172        pip install numpy173        ```174    """175 176    __slots__ = (177        "_data",178        "_vectors",179        "index_config",180        "embeddings",181    )182 183    def __init__(self, *, index: IndexConfig | None = None) -> None:184        # Both _data and _vectors are wrapped in the In-memory API185        # Do not change their names186        self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict)187        # [ns][key][path]188        self._vectors: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (189            defaultdict(lambda: defaultdict(dict))190        )191        self.index_config = index192        if self.index_config:193            self.index_config = self.index_config.copy()194            self.embeddings: Embeddings | None = ensure_embeddings(195                self.index_config.get("embed"),196            )197            self.index_config["__tokenized_fields"] = [198                (p, tokenize_path(p)) if p != "$" else (p, p)199                for p in (self.index_config.get("fields") or ["$"])200            ]201 202        else:203            self.index_config = None204            self.embeddings = None205 206    def batch(self, ops: Iterable[Op]) -> list[Result]:207        # The batch/abatch methods are treated as internal.208        # Users should access via put/search/get/list_namespaces/etc.209        results, put_ops, search_ops = self._prepare_ops(ops)210        if search_ops:211            queryinmem_store = self._embed_search_queries(search_ops)212            self._batch_search(search_ops, queryinmem_store, results)213 214        to_embed = self._extract_texts(put_ops)215        if to_embed and self.index_config and self.embeddings:216            embeddings = self.embeddings.embed_documents(list(to_embed))217            self._insertinmem_store(to_embed, embeddings)218        self._apply_put_ops(put_ops)219        return results220 221    async def abatch(self, ops: Iterable[Op]) -> list[Result]:222        # The batch/abatch methods are treated as internal.223        # Users should access via put/search/get/list_namespaces/etc.224        results, put_ops, search_ops = self._prepare_ops(ops)225        if search_ops:226            queryinmem_store = await self._aembed_search_queries(search_ops)227            self._batch_search(search_ops, queryinmem_store, results)228 229        to_embed = self._extract_texts(put_ops)230        if to_embed and self.index_config and self.embeddings:231            embeddings = await self.embeddings.aembed_documents(list(to_embed))232            self._insertinmem_store(to_embed, embeddings)233        self._apply_put_ops(put_ops)234        return results235 236    # Helpers237 238    def _filter_items(self, op: SearchOp) -> list[tuple[Item, list[list[float]]]]:239        """Filter items by namespace and filter function, return items with their embeddings."""240        namespace_prefix = op.namespace_prefix241 242        def filter_func(item: Item) -> bool:243            if not op.filter:244                return True245 246            return all(247                _compare_values(item.value.get(key), filter_value)248                for key, filter_value in op.filter.items()249            )250 251        filtered = []252        for namespace in self._data:253            if not (254                namespace[: len(namespace_prefix)] == namespace_prefix255                if len(namespace) >= len(namespace_prefix)256                else False257            ):258                continue259 260            for key, item in self._data[namespace].items():261                if filter_func(item):262                    if op.query and (embeddings := self._vectors[namespace].get(key)):263                        filtered.append((item, list(embeddings.values())))264                    else:265                        filtered.append((item, []))266        return filtered267 268    def _embed_search_queries(269        self,270        search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],271    ) -> dict[str, list[float]]:272        queryinmem_store = {}273        if self.index_config and self.embeddings and search_ops:274            queries = {op.query for (op, _) in search_ops.values() if op.query}275 276            if queries:277                with cf.ThreadPoolExecutor() as executor:278                    futures = {279                        q: executor.submit(self.embeddings.embed_query, q)280                        for q in list(queries)281                    }282                    for query, future in futures.items():283                        queryinmem_store[query] = future.result()284 285        return queryinmem_store286 287    async def _aembed_search_queries(288        self,289        search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],290    ) -> dict[str, list[float]]:291        queryinmem_store = {}292        if self.index_config and self.embeddings and search_ops:293            queries = {op.query for (op, _) in search_ops.values() if op.query}294 295            if queries:296                coros = [self.embeddings.aembed_query(q) for q in list(queries)]297                results = await asyncio.gather(*coros)298                queryinmem_store = dict(zip(queries, results, strict=False))299 300        return queryinmem_store301 302    def _batch_search(303        self,304        ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],305        queryinmem_store: dict[str, list[float]],306        results: list[Result],307    ) -> None:308        """Perform batch similarity search for multiple queries."""309        for i, (op, candidates) in ops.items():310            if not candidates:311                results[i] = []312                continue313            if op.query and queryinmem_store:314                query_embedding = queryinmem_store[op.query]315                flat_items, flat_vectors = [], []316                scoreless = []317                for item, vectors in candidates:318                    for vector in vectors:319                        flat_items.append(item)320                        flat_vectors.append(vector)321                    if not vectors:322                        scoreless.append(item)323 324                scores = _cosine_similarity(query_embedding, flat_vectors)325                sorted_results = sorted(326                    zip(scores, flat_items, strict=False),327                    key=lambda x: x[0],328                    reverse=True,329                )330                # max pooling331                seen: set[tuple[tuple[str, ...], str]] = set()332                kept: list[tuple[float | None, Item]] = []333                for score, item in sorted_results:334                    key = (item.namespace, item.key)335                    if key in seen:336                        continue337                    ix = len(seen)338                    seen.add(key)339                    if ix >= op.offset + op.limit:340                        break341                    if ix < op.offset:342                        continue343 344                    kept.append((score, item))345                if scoreless and len(kept) < op.limit:346                    # Corner case: if we request more items than what we have embedded,347                    # fill the rest with non-scored items348                    kept.extend(349                        (None, item) for item in scoreless[: op.limit - len(kept)]350                    )351 352                results[i] = [353                    SearchItem(354                        namespace=item.namespace,355                        key=item.key,356                        value=item.value,357                        created_at=item.created_at,358                        updated_at=item.updated_at,359                        score=float(score) if score is not None else None,360                    )361                    for score, item in kept362                ]363            else:364                results[i] = [365                    SearchItem(366                        namespace=item.namespace,367                        key=item.key,368                        value=item.value,369                        created_at=item.created_at,370                        updated_at=item.updated_at,371                    )372                    for (item, _) in candidates[op.offset : op.offset + op.limit]373                ]374 375    def _prepare_ops(376        self, ops: Iterable[Op]377    ) -> tuple[378        list[Result],379        dict[tuple[tuple[str, ...], str], PutOp],380        dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],381    ]:382        results: list[Result] = []383        put_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}384        search_ops: dict[385            int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]386        ] = {}387        for i, op in enumerate(ops):388            if isinstance(op, GetOp):389                item = self._data[op.namespace].get(op.key)390                results.append(item)391            elif isinstance(op, SearchOp):392                search_ops[i] = (op, self._filter_items(op))393                results.append(None)394            elif isinstance(op, ListNamespacesOp):395                results.append(self._handle_list_namespaces(op))396            elif isinstance(op, PutOp):397                put_ops[(op.namespace, op.key)] = op398                results.append(None)399            else:400                raise ValueError(f"Unknown operation type: {type(op)}")401 402        return results, put_ops, search_ops403 404    def _apply_put_ops(self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]) -> None:405        for (namespace, key), op in put_ops.items():406            if op.value is None:407                self._data[namespace].pop(key, None)408                self._vectors[namespace].pop(key, None)409            else:410                self._data[namespace][key] = Item(411                    value=op.value,412                    key=key,413                    namespace=namespace,414                    created_at=datetime.now(timezone.utc),415                    updated_at=datetime.now(timezone.utc),416                )417 418    def _extract_texts(419        self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]420    ) -> dict[str, list[tuple[tuple[str, ...], str, str]]]:421        if put_ops and self.index_config and self.embeddings:422            to_embed = defaultdict(list)423 424            for op in put_ops.values():425                if op.value is not None and op.index is not False:426                    if op.index is None:427                        paths = self.index_config["__tokenized_fields"]428                    else:429                        paths = [(ix, tokenize_path(ix)) for ix in op.index]430                    for path, field in paths:431                        texts = get_text_at_path(op.value, field)432                        if texts:433                            if len(texts) > 1:434                                for i, text in enumerate(texts):435                                    to_embed[text].append(436                                        (op.namespace, op.key, f"{path}.{i}")437                                    )438 439                            else:440                                to_embed[texts[0]].append((op.namespace, op.key, path))441 442            return to_embed443 444        return {}445 446    def _insertinmem_store(447        self,448        to_embed: dict[str, list[tuple[tuple[str, ...], str, str]]],449        embeddings: list[list[float]],450    ) -> None:451        indices = [index for indices in to_embed.values() for index in indices]452        if len(indices) != len(embeddings):453            raise ValueError(454                f"Number of embeddings ({len(embeddings)}) does not"455                f" match number of indices ({len(indices)})"456            )457        for embedding, (ns, key, path) in zip(embeddings, indices, strict=False):458            self._vectors[ns][key][path] = embedding459 460    def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]:461        all_namespaces = list(462            self._data.keys()463        )  # Avoid collection size changing while iterating464        namespaces = all_namespaces465        if op.match_conditions:466            namespaces = [467                ns468                for ns in namespaces469                if all(_does_match(condition, ns) for condition in op.match_conditions)470            ]471 472        if op.max_depth is not None:473            namespaces = sorted({ns[: op.max_depth] for ns in namespaces})474        else:475            namespaces = sorted(namespaces)476        return namespaces[op.offset : op.offset + op.limit]477 478 479@functools.lru_cache(maxsize=1)480def _check_numpy() -> bool:481    if bool(util.find_spec("numpy")):482        return True483    logger.warning(484        "NumPy not found in the current Python environment. "485        "The InMemoryStore will use a pure Python implementation for vector operations, "486        "which may significantly impact performance, especially for large datasets or frequent searches. "487        "For optimal speed and efficiency, consider installing NumPy: "488        "pip install numpy"489    )490    return False491 492 493def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:494    """495    Compute cosine similarity between a vector X and a matrix Y.496    Lazy import numpy for efficiency.497    """498    if not Y:499        return []500    if _check_numpy():501        import numpy as np502 503        X_arr = np.array(X) if not isinstance(X, np.ndarray) else X504        Y_arr = np.array(Y) if not isinstance(Y, np.ndarray) else Y505        X_norm = np.linalg.norm(X_arr)506        Y_norm = np.linalg.norm(Y_arr, axis=1)507 508        # Avoid division by zero509        mask = Y_norm != 0510        similarities = np.zeros_like(Y_norm)511        similarities[mask] = np.dot(Y_arr[mask], X_arr) / (Y_norm[mask] * X_norm)512        return similarities.tolist()513 514    similarities = []515    for y in Y:516        dot_product = sum(a * b for a, b in zip(X, y, strict=False))517        norm1 = sum(a * a for a in X) ** 0.5518        norm2 = sum(a * a for a in y) ** 0.5519        similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0520        similarities.append(similarity)521 522    return similarities523 524 525def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool:526    """Whether a namespace key matches a match condition."""527    match_type = match_condition.match_type528    path = match_condition.path529 530    if len(key) < len(path):531        return False532 533    if match_type == "prefix":534        for k_elem, p_elem in zip(key, path, strict=False):535            if p_elem == "*":536                continue  # Wildcard matches any element537            if k_elem != p_elem:538                return False539        return True540    elif match_type == "suffix":541        for k_elem, p_elem in zip(reversed(key), reversed(path), strict=False):542            if p_elem == "*":543                continue  # Wildcard matches any element544            if k_elem != p_elem:545                return False546        return True547    else:548        raise ValueError(f"Unsupported match type: {match_type}")549 550 551def _compare_values(item_value: Any, filter_value: Any) -> bool:552    """Compare values in a JSONB-like way, handling nested objects."""553    if isinstance(filter_value, dict):554        if any(k.startswith("$") for k in filter_value):555            return all(556                _apply_operator(item_value, op_key, op_value)557                for op_key, op_value in filter_value.items()558            )559        if not isinstance(item_value, dict):560            return False561        return all(562            _compare_values(item_value.get(k), v) for k, v in filter_value.items()563        )564    elif isinstance(filter_value, (list, tuple)):565        return (566            isinstance(item_value, (list, tuple))567            and len(item_value) == len(filter_value)568            and all(569                _compare_values(iv, fv)570                for iv, fv in zip(item_value, filter_value, strict=False)571            )572        )573    else:574        return item_value == filter_value575 576 577def _apply_operator(value: Any, operator: str, op_value: Any) -> bool:578    """Apply a comparison operator, matching PostgreSQL's JSONB behavior."""579    if operator == "$eq":580        return value == op_value581    elif operator == "$gt":582        return float(value) > float(op_value)583    elif operator == "$gte":584        return float(value) >= float(op_value)585    elif operator == "$lt":586        return float(value) < float(op_value)587    elif operator == "$lte":588        return float(value) <= float(op_value)589    elif operator == "$ne":590        return value != op_value591    else:592        raise ValueError(f"Unsupported operator: {operator}")593