⚡️ Speed up validate_prompt_input_variables() by 5% in libs/langchain/langchain/memory/summary.py

[codeflash-ai]

📄 validate_prompt_input_variables() in libs/langchain/langchain/memory/summary.py

📈 Performance went up by 5% (0.05x faster)

⏱️ Runtime went down from 10.20μs to 9.70μs

Explanation and details

(click to show)

The current code appears to be pythonically correct and quite optimal function-wise. However, the exception case can be slightly rewritten by avoiding string formatting and possibly improve the function readability.

In this alternative, the f-string is replaced by the list version and then using the ''.join() operation which tends to be faster. We also replaced the direct dictionary access with the 'get' method to avoid 'KeyError' in case the "prompt" key is not in the values dictionary. The performance improvements are minor but they slightly enhance the robustness of your code.

Correctness verification

The new optimized code was tested for correctness. The results are listed below.

✅ 0 Passed − ⚙️ Existing Unit Tests

✅ 0 Passed − 🎨 Inspired Regression Tests

✅ 12 Passed − 🌀 Generated Regression Tests

(click to show generated tests)

# imports
import pytest  # used for our unit tests
from pydantic import BaseModel, root_validator  # used for our class definition and validation
from typing import Dict, Set  # used for type annotations

# function to test

class BaseChatMemory(BaseModel):  # Mock base class
    pass

class SummarizerMixin:  # Mock mixin class
    pass
from langchain.memory.summary import ConversationSummaryMemory

# unit tests

# Mock Prompt class to simulate input_variables attribute
class MockPrompt:
    def __init__(self, input_variables: Set[str]):
        self.input_variables = input_variables

# Test case for valid input
def test_valid_input():
    valid_values = {"prompt": MockPrompt({"summary", "new_lines"})}
    assert ConversationSummaryMemory.validate_prompt_input_variables(valid_values) == valid_values

# Test case for missing keys
@pytest.mark.parametrize("missing_keys", [
    {"new_lines"},
    {"summary"},
    set()
])
def test_missing_keys(missing_keys):
    invalid_values = {"prompt": MockPrompt(missing_keys)}
    with pytest.raises(ValueError):
        ConversationSummaryMemory.validate_prompt_input_variables(invalid_values)

# Test case for extra keys
def test_extra_keys():
    extra_keys = {"summary", "new_lines", "extra_key"}
    invalid_values = {"prompt": MockPrompt(extra_keys)}
    with pytest.raises(ValueError):
        ConversationSummaryMemory.validate_prompt_input_variables(invalid_values)

# Test case for non-dictionary values
@pytest.mark.parametrize("non_dict_value", [
    None,
    "string",
    123,
    [1, 2, 3],
    MockPrompt("not a set")
])
def test_non_dictionary_values(non_dict_value):
    invalid_values = {"prompt": non_dict_value}
    with pytest.raises(AttributeError):
        ConversationSummaryMemory.validate_prompt_input_variables(invalid_values)

# Test case for type mismatch
def test_type_mismatch():
    invalid_values = {"prompt": MockPrompt("not a set")}
    with pytest.raises(ValueError):
        ConversationSummaryMemory.validate_prompt_input_variables(invalid_values)

# Test case for correct keys but incorrect types
def test_correct_keys_incorrect_types():
    # Assuming the function only checks for keys, not value types
    valid_values = {"prompt": MockPrompt({"summary": 1, "new_lines": 2})}
    assert ConversationSummaryMemory.validate_prompt_input_variables(valid_values) == valid_values

# Test case for boundary conditions
def test_boundary_conditions():
    # Assuming the function only checks for keys, not value content
    valid_values = {"prompt": MockPrompt({"summary": "", "new_lines": None})}
    assert ConversationSummaryMemory.validate_prompt_input_variables(valid_values) == valid_values

# Test case for missing "prompt" key
def test_missing_prompt_key():
    invalid_values = {"not_prompt": MockPrompt({"summary", "new_lines"})}
    with pytest.raises(KeyError):
        ConversationSummaryMemory.validate_prompt_input_variables(invalid_values)