⚡️ Speed up method LexicalRetriever.get_completion by 17%

[codeflash-ai]

📄 17% (0.17x) speedup for LexicalRetriever.get_completion in cognee/modules/retrieval/lexical_retriever.py

⏱️ Runtime : 542 nanoseconds → 463 nanoseconds (best of 7 runs)

📝 Explanation and details

The optimized code achieves a 17% runtime improvement through several micro-optimizations that reduce Python's attribute lookup overhead in the critical get_context method:

Key Performance Optimizations:

1. Local variable caching for hot loop: The main optimization extracts frequently-accessed attributes (self.chunks.items(), self.scorer, self.payloads, results.append) into local variables before the scoring loop. This eliminates repeated attribute lookups during iteration, which is significant when processing many chunks.

2. Method reference optimization: Storing results.append as append_result avoids the method lookup on each loop iteration - a classic Python performance pattern for tight loops.

3. Lambda extraction: Moving the sorting key function (lambda x: x[1]) out of the nlargest call to avoid recreating it, though this is a minor gain.

4. Early return optimization: In get_completion, reordering the conditional to check if context is not None first provides a small speedup when context is pre-supplied by avoiding the async call entirely.

Why This Works:
Python's attribute access (self.attribute) and method lookups (.method()) have overhead due to the dynamic nature of object attribute resolution. By caching these lookups in local variables before entering loops, the optimized code reduces this overhead during the most computationally intensive part - scoring all chunks against the query.

Test Case Performance:
The optimization is most effective for retrieval scenarios with many chunks to score, where the loop overhead reduction compounds. For simple cases with few chunks or when context is pre-provided, the gains are minimal but still measurable.

Note that throughput remains unchanged at 371 ops/second because the core algorithmic complexity is unchanged - this optimization reduces per-operation overhead rather than changing the fundamental processing rate.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import asyncio  # used to run async functions
# --- Function to test (EXACT COPY, UNMODIFIED) ---
from heapq import nlargest
from typing import Any, Callable, Optional

import pytest  # used for our unit tests
from cognee.modules.retrieval.lexical_retriever import LexicalRetriever

# --- Mocks and minimal stubs for dependencies ---

# Minimal NoDataError for testing
class NoDataError(Exception):
    pass

# Minimal async graph engine stub
class DummyGraphEngine:
    def __init__(self, nodes):
        self._nodes = nodes

    async def get_filtered_graph_data(self, filters):
        # Only returns the nodes as per the test setup
        return self._nodes, None

# Patch point for get_graph_engine
async def get_graph_engine():
    # This will be monkeypatched in tests
    raise NotImplementedError("Should be monkeypatched in tests")

# --- Patch get_graph_engine globally for LexicalRetriever ---
LexicalRetriever._test_graph_engine = None  # for monkeypatching

# Helper to monkeypatch get_graph_engine for tests
def patch_graph_engine(nodes):
    async def _get_graph_engine():
        return DummyGraphEngine(nodes)
    global get_graph_engine
    get_graph_engine = _get_graph_engine

# --- Minimal tokenizer/scorer for tests ---
def simple_tokenizer(text):
    # Tokenize by splitting on whitespace
    return text.strip().split()

def simple_scorer(query_tokens, chunk_tokens):
    # Score by intersection count
    return len(set(query_tokens) & set(chunk_tokens))

# --- Test cases ---

@pytest.mark.asyncio











#------------------------------------------------
import asyncio  # used to run async functions
from heapq import nlargest
from typing import Any, Callable, Optional

import pytest  # used for our unit tests
from cognee.modules.retrieval.lexical_retriever import LexicalRetriever

# --- MOCKS AND TEST UTILITIES ---

# Minimal NoDataError for testing
class NoDataError(Exception):
    pass

# Patch for get_graph_engine
class DummyGraphEngine:
    def __init__(self, nodes):
        self._nodes = nodes

    async def get_filtered_graph_data(self, filters):
        # Return nodes and dummy second value
        return self._nodes, None
def make_graph_engine_patch(nodes):
    async def get_graph_engine():
        return DummyGraphEngine(nodes)
    return get_graph_engine

# --- BASIC TEST CASES ---

@pytest.mark.asyncio


async def test_get_completion_returns_context_if_provided():
    """Test get_completion returns context if provided, does not call get_context."""
    retriever = LexicalRetriever(lambda t: t.split(), lambda q, c: 1)
    provided_context = [{"id": "already", "text": "here", "type": "DocumentChunk"}]
    # Should return provided context, not call get_context
    result = await retriever.get_completion("irrelevant", context=provided_context)

# --- EDGE TEST CASES ---

@pytest.mark.asyncio

To edit these changes git checkout codeflash/optimize-LexicalRetriever.get_completion-mh2ppgg5 and push.