⚡️ Speed up method Langchain.list by 13%

[codeflash-ai]

📄 13% (0.13x) speedup for Langchain.list in mem0/vector_stores/langchain.py

⏱️ Runtime : 8.19 milliseconds → 7.26 milliseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 12% speedup through several key optimizations in the _parse_output method:

1. List Comprehension over For-Loop
The original code used a for-loop with .append() to build the result list for Document objects. The optimized version replaces this with a list comprehension, which is inherently faster in Python due to reduced bytecode overhead.

2. Tuple instead of List for Constants
Changed keys = ["ids", "distances", "metadatas"] to keys = ("ids", "distances", "metadatas"). Tuples have slightly better performance for iteration since they're immutable.

3. Pre-computed Length Checks
The original code performed expensive isinstance() and len() checks inside the main loop for each vector. The optimized version pre-computes these lengths once:

ids_len = len(ids) if isinstance(ids, list) and ids is not None else 0

This eliminates redundant type checking and length calculations that were happening 6000+ times in large datasets.

4. Simplified Conditional Logic
The optimized version uses direct index bounds checking (i < ids_len) instead of complex nested conditions, reducing computational overhead per iteration.

5. Cached Attribute Access
In the list() method, the optimized code caches self.client._collection in a local variable to avoid repeated attribute lookups, and uses getattr() with a default to handle missing attributes more efficiently.

These optimizations are particularly effective for large datasets, as shown in the test results where 1000-vector test cases show 23-24% speedups. The pre-computed lengths and simplified conditionals eliminate the quadratic behavior that was occurring in the original nested condition checks within the main processing loop.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 64 Passed
⏪ Replay Tests	✅ 4 Passed
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import logging
from typing import Dict, List

# imports
import pytest
from mem0.vector_stores.langchain import Langchain


# Mock OutputData for test purposes
class OutputData:
    def __init__(self, id, score, payload):
        self.id = id
        self.score = score
        self.payload = payload

    def __eq__(self, other):
        return (
            isinstance(other, OutputData)
            and self.id == other.id
            and self.score == other.score
            and self.payload == other.payload
        )

    def __repr__(self):
        return f"OutputData(id={self.id}, score={self.score}, payload={self.payload})"

# Mock Document for test purposes
class MockDocument:
    def __init__(self, id=None, metadata=None):
        self.id = id
        self.metadata = metadata or {}

# Mock client and collection for testing
class MockCollection:
    def __init__(self, get_return_value=None, raise_exception=False):
        self._get_return_value = get_return_value
        self._raise_exception = raise_exception
        self.last_where = None
        self.last_limit = None

    def get(self, where=None, limit=None):
        self.last_where = where
        self.last_limit = limit
        if self._raise_exception:
            raise RuntimeError("Simulated error in get")
        return self._get_return_value

class MockClient:
    def __init__(self, get_return_value=None, raise_exception=False):
        self._collection = MockCollection(get_return_value, raise_exception)

# ------------------- UNIT TESTS -------------------

# 1. BASIC TEST CASES

def test_list_returns_empty_when_collection_returns_empty_dict():
    """Test that list() returns an empty list when collection.get returns an empty dict."""
    client = MockClient(get_return_value={})
    langchain = Langchain(client)
    codeflash_output = langchain.list() # 1.80μs -> 1.59μs (13.1% faster)

def test_list_returns_empty_when_collection_returns_none():
    """Test that list() returns an empty list when collection.get returns None."""
    client = MockClient(get_return_value=None)
    langchain = Langchain(client)
    codeflash_output = langchain.list() # 1.32μs -> 1.31μs (0.687% faster)

def test_list_returns_single_vector_with_full_data():
    """Test that list() returns a single vector with all fields populated."""
    data = {
        "ids": [["abc123"]],
        "distances": [[0.42]],
        "metadatas": [[{"foo": "bar"}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 11.6μs -> 11.5μs (1.61% faster)
    expected = [ [OutputData(id="abc123", score=0.42, payload={"foo": "bar"})] ]

def test_list_returns_multiple_vectors():
    """Test that list() returns multiple vectors correctly."""
    data = {
        "ids": [["id1", "id2"]],
        "distances": [[0.1, 0.2]],
        "metadatas": [[{"a": 1}, {"b": 2}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 10.0μs -> 9.21μs (8.83% faster)
    expected = [[
        OutputData(id="id1", score=0.1, payload={"a": 1}),
        OutputData(id="id2", score=0.2, payload={"b": 2}),
    ]]

def test_list_with_limit_passes_limit_to_collection():
    """Test that the limit argument is passed to the collection.get method."""
    data = {
        "ids": [["id1"]],
        "distances": [[0.1]],
        "metadatas": [[{"a": 1}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(limit=5); _ = codeflash_output # 8.54μs -> 7.82μs (9.22% faster)

def test_list_with_filters_passes_filters_to_collection():
    """Test that the filters argument is passed to the collection.get method as where clause."""
    data = {
        "ids": [["id1"]],
        "distances": [[0.1]],
        "metadatas": [[{"a": 1}]],
    }
    filters = {"user_id": "u1"}
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(filters=filters); _ = codeflash_output # 8.45μs -> 7.92μs (6.69% faster)

# 2. EDGE TEST CASES

def test_list_handles_missing_fields_in_collection_data():
    """Test that missing fields in collection.get result are handled gracefully."""
    data = {
        "ids": [["id1", "id2"]],
        # 'distances' missing
        "metadatas": [[{"a": 1}, {"b": 2}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 9.82μs -> 9.29μs (5.69% faster)
    expected = [[
        OutputData(id="id1", score=None, payload={"a": 1}),
        OutputData(id="id2", score=None, payload={"b": 2}),
    ]]

def test_list_handles_fields_with_inconsistent_lengths():
    """Test that list() handles fields of different lengths safely."""
    data = {
        "ids": [["id1", "id2", "id3"]],
        "distances": [[0.1, 0.2]],
        "metadatas": [[{"a": 1}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 10.3μs -> 9.74μs (5.74% faster)
    expected = [[
        OutputData(id="id1", score=0.1, payload={"a": 1}),
        OutputData(id="id2", score=0.2, payload=None),
        OutputData(id="id3", score=None, payload=None),
    ]]

def test_list_handles_flat_lists_instead_of_nested():
    """Test that list() can handle flat lists instead of nested lists."""
    data = {
        "ids": ["id1", "id2"],
        "distances": [0.1, 0.2],
        "metadatas": [{"a": 1}, {"b": 2}],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 9.27μs -> 8.54μs (8.52% faster)
    expected = [[
        OutputData(id="id1", score=0.1, payload={"a": 1}),
        OutputData(id="id2", score=0.2, payload={"b": 2}),
    ]]

def test_list_handles_empty_lists():
    """Test that list() returns an empty list if all fields are empty lists."""
    data = {
        "ids": [[]],
        "distances": [[]],
        "metadatas": [[]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list() # 4.31μs -> 4.13μs (4.41% faster)

def test_list_handles_no_collection_attribute():
    """Test that list() returns [] if client has no _collection attribute."""
    class NoCollectionClient:
        pass
    langchain = Langchain(NoCollectionClient())
    codeflash_output = langchain.list() # 629ns -> 684ns (8.04% slower)

def test_list_handles_collection_without_get():
    """Test that list() returns [] if _collection has no get method."""
    class NoGetCollection:
        pass
    class Client:
        _collection = NoGetCollection()
    langchain = Langchain(Client())
    codeflash_output = langchain.list() # 789ns -> 820ns (3.78% slower)

def test_list_handles_collection_get_raises_exception(caplog):
    """Test that list() logs error and returns [] if collection.get raises an exception."""
    client = MockClient(raise_exception=True)
    langchain = Langchain(client)
    with caplog.at_level(logging.ERROR):
        codeflash_output = langchain.list(); result = codeflash_output # 483μs -> 483μs (0.027% faster)

def test_list_handles_document_list_input():
    """Test that list() can parse a list of Document-like objects."""
    # This is not a normal use case for list(), but _parse_output should handle it
    langchain = Langchain(MockClient(get_return_value={}))
    docs = [MockDocument(id="d1", metadata={"foo": "bar"}), MockDocument(id="d2", metadata={"x": 1})]
    result = langchain._parse_output(docs)
    expected = [
        OutputData(id="d1", score=None, payload={"foo": "bar"}),
        OutputData(id="d2", score=None, payload={"x": 1}),
    ]

def test_list_handles_ids_with_none():
    """Test that list() handles ids containing None values."""
    data = {
        "ids": [[None, "id2"]],
        "distances": [[0.1, 0.2]],
        "metadatas": [[{"a": 1}, {"b": 2}]],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 10.7μs -> 10.1μs (5.32% faster)
    expected = [[
        OutputData(id=None, score=0.1, payload={"a": 1}),
        OutputData(id="id2", score=0.2, payload={"b": 2}),
    ]]

# 3. LARGE SCALE TEST CASES

def test_list_handles_large_number_of_vectors():
    """Test that list() can handle a large number of vectors efficiently."""
    N = 1000
    data = {
        "ids": [ [f"id{i}" for i in range(N)] ],
        "distances": [ [float(i)/N for i in range(N)] ],
        "metadatas": [ [{"val": i} for i in range(N)] ],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 819μs -> 666μs (22.9% faster)
    expected = [[
        OutputData(id=f"id{i}", score=float(i)/N, payload={"val": i}) for i in range(N)
    ]]

def test_list_handles_large_flat_lists():
    """Test that list() can handle large flat lists."""
    N = 1000
    data = {
        "ids": [f"id{i}" for i in range(N)],
        "distances": [float(i)/N for i in range(N)],
        "metadatas": [{"val": i} for i in range(N)],
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 826μs -> 667μs (23.7% faster)
    expected = [[
        OutputData(id=f"id{i}", score=float(i)/N, payload={"val": i}) for i in range(N)
    ]]

def test_list_handles_large_inconsistent_lengths():
    """Test that list() handles large lists with inconsistent lengths."""
    N = 1000
    data = {
        "ids": [ [f"id{i}" for i in range(N)] ],
        "distances": [ [float(i)/N for i in range(N//2)] ],  # Only half as many distances
        "metadatas": [ [{"val": i} for i in range(N//3)] ],  # Only a third as many metadatas
    }
    client = MockClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 760μs -> 613μs (23.9% faster)
    # Expected: first N//3 have id, score, payload; next N//2-N//3 have id, score, None; rest have id, None, None
    expected = []
    for i in range(N):
        id_val = f"id{i}"
        score_val = float(i)/N if i < N//2 else None
        payload_val = {"val": i} if i < N//3 else None
        expected.append(OutputData(id=id_val, score=score_val, payload=payload_val))
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import logging
from typing import Dict, List

# imports
import pytest
from mem0.vector_stores.langchain import Langchain


# Dummy OutputData for testing (since the original is not provided)
class OutputData:
    def __init__(self, id, score, payload):
        self.id = id
        self.score = score
        self.payload = payload

    def __eq__(self, other):
        if not isinstance(other, OutputData):
            return False
        return self.id == other.id and self.score == other.score and self.payload == other.payload

    def __repr__(self):
        return f"OutputData(id={self.id!r}, score={self.score!r}, payload={self.payload!r})"


# Dummy Document class for testing
class DummyDocument:
    def __init__(self, id, metadata):
        self.id = id
        self.metadata = metadata

# Dummy VectorStore client for testing
class DummyCollection:
    def __init__(self, get_return_value):
        self._get_return_value = get_return_value
        self.get_called_with = []

    def get(self, where=None, limit=None):
        self.get_called_with.append({'where': where, 'limit': limit})
        return self._get_return_value

class DummyClient:
    def __init__(self, get_return_value):
        self._collection = DummyCollection(get_return_value)

# unit tests

# -----------------------------
# BASIC TEST CASES
# -----------------------------

def test_list_returns_empty_on_none_result():
    """Test that list returns an empty list if get returns None."""
    client = DummyClient(get_return_value=None)
    langchain = Langchain(client)
    codeflash_output = langchain.list() # 1.75μs -> 1.66μs (5.61% faster)

def test_list_returns_empty_on_empty_dict():
    """Test that list returns an empty list if get returns an empty dict."""
    client = DummyClient(get_return_value={})
    langchain = Langchain(client)
    codeflash_output = langchain.list() # 1.39μs -> 1.40μs (0.645% slower)

def test_list_single_vector():
    """Test that list returns correct OutputData for a single vector."""
    data = {
        "ids": ["id1"],
        "distances": [0.123],
        "metadatas": [{"foo": "bar"}]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 12.8μs -> 12.3μs (3.39% faster)

def test_list_multiple_vectors():
    """Test that list returns correct OutputData for multiple vectors."""
    data = {
        "ids": ["id1", "id2"],
        "distances": [0.1, 0.2],
        "metadatas": [{"foo": "bar"}, {"baz": "qux"}]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 10.2μs -> 9.43μs (8.54% faster)
    expected = [
        OutputData(id="id1", score=0.1, payload={"foo": "bar"}),
        OutputData(id="id2", score=0.2, payload={"baz": "qux"})
    ]

def test_list_with_filters_and_limit():
    """Test that filters and limit are passed to the client's get method."""
    data = {
        "ids": ["id1"],
        "distances": [0.5],
        "metadatas": [{"foo": "bar"}]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    filters = {"user_id": "u1"}
    limit = 1
    langchain.list(filters=filters, limit=limit) # 8.70μs -> 8.22μs (5.78% faster)

def test_list_with_nested_lists():
    """Test that nested lists in ids/distances/metadatas are flattened."""
    data = {
        "ids": [["id1", "id2"]],
        "distances": [[0.1, 0.2]],
        "metadatas": [[{"foo": "bar"}, {"baz": "qux"}]]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 9.49μs -> 8.76μs (8.38% faster)
    expected = [
        OutputData(id="id1", score=0.1, payload={"foo": "bar"}),
        OutputData(id="id2", score=0.2, payload={"baz": "qux"})
    ]

# -----------------------------
# EDGE TEST CASES
# -----------------------------

def test_list_missing_keys():
    """Test that missing keys are handled gracefully."""
    data = {
        "ids": ["id1"]
        # distances and metadatas missing
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 7.73μs -> 7.50μs (3.12% faster)
    expected = [OutputData(id="id1", score=None, payload=None)]

def test_list_empty_lists():
    """Test that empty lists for ids, distances, metadatas are handled."""
    data = {
        "ids": [],
        "distances": [],
        "metadatas": []
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 4.27μs -> 4.10μs (4.05% faster)

def test_list_lists_of_different_lengths():
    """Test that lists of different lengths are handled (shortest wins)."""
    data = {
        "ids": ["id1", "id2", "id3"],
        "distances": [0.1, 0.2],
        "metadatas": [{"foo": "bar"}]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 11.3μs -> 10.8μs (4.44% faster)
    expected = [
        OutputData(id="id1", score=0.1, payload={"foo": "bar"}),
        OutputData(id="id2", score=0.2, payload=None),
        OutputData(id="id3", score=None, payload=None)
    ]

def test_list_with_non_dict_result():
    """Test that non-dict result (e.g., list of Document) is handled."""
    docs = [DummyDocument("id1", {"foo": "bar"}), DummyDocument("id2", {"baz": "qux"})]
    # Patch DummyCollection.get to return a list
    class DummyCollectionWithList(DummyCollection):
        def get(self, where=None, limit=None):
            self.get_called_with.append({'where': where, 'limit': limit})
            return docs
    client = DummyClient(get_return_value=None)
    client._collection = DummyCollectionWithList(get_return_value=None)
    langchain = Langchain(client)
    # Patch _parse_output to handle Document list
    result = langchain._parse_output(docs)
    expected = [
        OutputData(id="id1", score=None, payload={"foo": "bar"}),
        OutputData(id="id2", score=None, payload={"baz": "qux"})
    ]

def test_list_handles_exceptions_and_logs(caplog):
    """Test that exceptions in get are caught and logged, and [] is returned."""
    class FailingCollection:
        def get(self, where=None, limit=None):
            raise RuntimeError("fail!")
    class FailingClient:
        def __init__(self):
            self._collection = FailingCollection()
    client = FailingClient()
    langchain = Langchain(client)
    with caplog.at_level(logging.ERROR):
        codeflash_output = langchain.list(); result = codeflash_output # 469μs -> 469μs (0.053% slower)

def test_list_with_none_values():
    """Test handling of None in ids, distances, metadatas."""
    data = {
        "ids": [None, "id2"],
        "distances": [None, 0.2],
        "metadatas": [None, {"baz": "qux"}]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 13.3μs -> 12.9μs (2.91% faster)
    expected = [
        OutputData(id=None, score=None, payload=None),
        OutputData(id="id2", score=0.2, payload={"baz": "qux"})
    ]

# -----------------------------
# LARGE SCALE TEST CASES
# -----------------------------

def test_list_large_number_of_vectors():
    """Test handling of a large number of vectors (e.g., 1000)."""
    n = 1000
    data = {
        "ids": [f"id{i}" for i in range(n)],
        "distances": [float(i) for i in range(n)],
        "metadatas": [{"val": i} for i in range(n)]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 822μs -> 665μs (23.6% faster)
    expected = [OutputData(id=f"id{i}", score=float(i), payload={"val": i}) for i in range(n)]

def test_list_large_with_missing_fields():
    """Test large scale with some missing fields."""
    n = 1000
    data = {
        "ids": [f"id{i}" for i in range(n)],
        # distances missing
        "metadatas": [{"val": i} for i in range(n)]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 787μs -> 660μs (19.3% faster)
    expected = [OutputData(id=f"id{i}", score=None, payload={"val": i}) for i in range(n)]

def test_list_large_with_nested_lists():
    """Test large scale with nested lists."""
    n = 1000
    data = {
        "ids": [[f"id{i}" for i in range(n)]],
        "distances": [[float(i) for i in range(n)]],
        "metadatas": [[{"val": i} for i in range(n)]]
    }
    client = DummyClient(get_return_value=data)
    langchain = Langchain(client)
    codeflash_output = langchain.list(); result = codeflash_output # 821μs -> 664μs (23.6% faster)
    expected = [OutputData(id=f"id{i}", score=float(i), payload={"val": i}) for i in range(n)]
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_testsvector_storestest_opensearch_py_testsvector_storestest_upstash_vector_py_testsllmstest_l__replay_test_0.py::test_mem0_vector_stores_langchain_Langchain_list	2.23ms	2.21ms	0.786%✅

To edit these changes git checkout codeflash/optimize-Langchain.list-mhl5zscf and push.