⚡️ Speed up method JSONalyzeQueryEngine._get_prompts by 64%

[codeflash-ai]

📄 64% (0.64x) speedup for JSONalyzeQueryEngine._get_prompts in llama-index-core/llama_index/core/query_engine/jsonalyze_query_engine.py

⏱️ Runtime : 37.0 microseconds → 22.5 microseconds (best of 250 runs)

📝 Explanation and details

The optimization achieves a 64% speedup by eliminating redundant dictionary creation in the frequently-called _get_prompts() method through two key changes:

What was optimized:

1. Pre-computed prompt dictionary caching: The original code recreated a dictionary with prompt objects on every _get_prompts() call. The optimized version pre-computes this dictionary once during __init__ and stores it in self._prompts_dict, then simply returns the cached reference.

2. Added __slots__: Defined __slots__ to reduce memory overhead and speed up attribute access, which is beneficial since this class likely gets instantiated frequently in query processing pipelines.

Why this leads to speedup:
The line profiler shows the original _get_prompts() spent significant time on dictionary construction (115.3μs total) with three separate operations: creating the dict literal, accessing self._jsonalyze_prompt, and accessing self._response_synthesis_prompt. The optimized version reduces this to a single attribute lookup (43.4μs total), eliminating ~62% of the work.

Performance characteristics:

• Best for: Workloads where _get_prompts() is called repeatedly with the same engine instance, which appears to be the common case based on test results showing consistent 50-95% improvements across all scenarios
• Scalability: Performance gains are consistent regardless of prompt size, custom vs default prompts, or engine configuration complexity
• Memory trade-off: Slight increase in per-instance memory (one extra dict) for significant CPU savings on repeated calls

The cached approach is particularly effective because prompt configurations are immutable after engine initialization, making this a safe and highly beneficial optimization for query engine performance.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from llama_index.core.query_engine.jsonalyze_query_engine import
JSONalyzeQueryEngine

--- Minimal stubs for dependencies (to allow import-free testability) ---

class BasePromptTemplate:
def init(self, template: str, prompt_type=None):
self.template = template
self.prompt_type = prompt_type

def __eq__(self, other):
    # For testing, compare on template and prompt_type only
    return (
        isinstance(other, BasePromptTemplate)
        and self.template == other.template
        and self.prompt_type == other.prompt_type
    )

class PromptType:
SQL_RESPONSE_SYNTHESIS = "SQL_RESPONSE_SYNTHESIS"

class PromptTemplate(BasePromptTemplate):
pass

DEFAULT_JSONALYZE_PROMPT = BasePromptTemplate("DEFAULT_JSONALYZE_PROMPT")
DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL = (
"Given a query, synthesize a response based on SQL query results"
" to satisfy the query. Only include details that are relevant to"
" the query. If you don't know the answer, then say that.\n"
"SQL Query: {sql_query}\n"
"Table Schema: {table_schema}\n"
"SQL Response: {sql_response}\n"
"Query: {query_str}\n"
"Response: "
)
DEFAULT_RESPONSE_SYNTHESIS_PROMPT = PromptTemplate(
DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL,
prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS,
)

class DummyCallbackManager:
pass

class DummyServiceContext:
def init(self, llm=None, callback_manager=None):
self.llm = llm
self.callback_manager = callback_manager

class DummyLLM:
pass

class DummySQLParser:
pass

def callback_manager_from_settings_or_context(settings, context):
return DummyCallbackManager()

def llm_from_settings_or_context(settings, context):
return DummyLLM()

class BaseQueryEngine:
def init(self, callback_manager=None):
self.callback_manager = callback_manager
from llama_index.core.query_engine.jsonalyze_query_engine import
JSONalyzeQueryEngine

--- Unit tests ---

1. Basic Test Cases

def test_get_prompts_returns_default_prompts():
"""Test that _get_prompts returns default prompts when none are provided."""
engine = JSONalyzeQueryEngine(list_of_dict=[])
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 499ns -> 262ns (90.5% faster)

def test_get_prompts_with_custom_prompts():
"""Test that _get_prompts returns custom prompts if provided."""
custom_jsonalyze_prompt = BasePromptTemplate("CUSTOM_JSONALYZE")
custom_response_prompt = BasePromptTemplate("CUSTOM_RESPONSE")
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze_prompt,
response_synthesis_prompt=custom_response_prompt,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 497ns -> 258ns (92.6% faster)

def test_get_prompts_with_mixed_prompts():
"""Test that _get_prompts returns a mix of default and custom prompts."""
custom_jsonalyze_prompt = BasePromptTemplate("CUSTOM_JSONALYZE")
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze_prompt,
# response_synthesis_prompt not set
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 472ns -> 283ns (66.8% faster)

def test_get_prompts_with_service_context():
"""Test _get_prompts works even if service_context is provided."""
context = DummyServiceContext()
engine = JSONalyzeQueryEngine(list_of_dict=[], service_context=context)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 504ns -> 260ns (93.8% faster)

2. Edge Test Cases

def test_get_prompts_with_empty_list_of_dict():
"""Test _get_prompts when list_of_dict is empty."""
engine = JSONalyzeQueryEngine(list_of_dict=[])
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 495ns -> 259ns (91.1% faster)

def test_get_prompts_with_large_prompt_texts():
"""Test _get_prompts with very large prompt templates."""
big_text = "X" * 1000
custom_jsonalyze_prompt = BasePromptTemplate(big_text)
custom_response_prompt = BasePromptTemplate(big_text[::-1])
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze_prompt,
response_synthesis_prompt=custom_response_prompt,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 484ns -> 256ns (89.1% faster)

def test_get_prompts_with_non_string_prompt_templates():
"""Test _get_prompts with prompt templates that are not strings (edge case)."""
class WeirdPrompt(BasePromptTemplate):
def init(self, data):
self.data = data
def eq(self, other):
return isinstance(other, WeirdPrompt) and self.data == other.data
weird_prompt = WeirdPrompt(data={"foo": 123})
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=weird_prompt,
response_synthesis_prompt=weird_prompt,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 510ns -> 260ns (96.2% faster)

def test_get_prompts_with_none_prompts():
"""Test _get_prompts with both prompts explicitly set to None."""
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=None,
response_synthesis_prompt=None,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 490ns -> 273ns (79.5% faster)

def test_get_prompts_with_unusual_table_name():
"""Test get_prompts with a strange table name (should not affect prompts)."""
engine = JSONalyzeQueryEngine(
list_of_dict=[],
table_name="!@#$%^&*()+=-[]{}|;:',.<>/?`~"
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 490ns -> 252ns (94.4% faster)

3. Large Scale Test Cases

def test_get_prompts_with_large_list_of_dict():
"""Test _get_prompts with a large list_of_dict (should not affect prompts)."""
big_list = [{"a": i, "b": str(i)} for i in range(1000)]
engine = JSONalyzeQueryEngine(list_of_dict=big_list)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 469ns -> 272ns (72.4% faster)

def test_get_prompts_with_many_custom_prompts():
"""Test get_prompts with many different custom prompt objects (stress test)."""
for i in range(100):
custom_jsonalyze_prompt = BasePromptTemplate(f"jsonalyze{i}")
custom_response_prompt = BasePromptTemplate(f"response_{i}")
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze_prompt,
response_synthesis_prompt=custom_response_prompt,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 21.6μs -> 13.9μs (54.9% faster)

def test_get_prompts_scalability_with_long_strings():
"""Test get_prompts with extremely long string prompts (performance/robustness)."""
long_prompt = "PROMPT" + ("x" * 900)
custom_jsonalyze_prompt = BasePromptTemplate(long_prompt)
custom_response_prompt = BasePromptTemplate(long_prompt[::-1])
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze_prompt,
response_synthesis_prompt=custom_response_prompt,
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 438ns -> 227ns (93.0% faster)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import pytest
from llama_index.core.query_engine.jsonalyze_query_engine import
JSONalyzeQueryEngine

Function and class definitions (as provided above, slightly simplified for test context)

class BasePromptTemplate:
def init(self, template_str, prompt_type=None):
self.template_str = template_str
self.prompt_type = prompt_type

def __eq__(self, other):
    return (
        isinstance(other, BasePromptTemplate)
        and self.template_str == other.template_str
        and self.prompt_type == other.prompt_type
    )

class PromptType:
SQL_RESPONSE_SYNTHESIS = "SQL_RESPONSE_SYNTHESIS"
JSONALYZE = "JSONALYZE"

class PromptTemplate(BasePromptTemplate):
pass

DEFAULT_JSONALYZE_PROMPT = PromptTemplate("Default JSONalyze prompt", prompt_type=PromptType.JSONALYZE)
DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL = (
"Given a query, synthesize a response based on SQL query results"
" to satisfy the query. Only include details that are relevant to"
" the query. If you don't know the answer, then say that.\n"
"SQL Query: {sql_query}\n"
"Table Schema: {table_schema}\n"
"SQL Response: {sql_response}\n"
"Query: {query_str}\n"
"Response: "
)
DEFAULT_RESPONSE_SYNTHESIS_PROMPT = PromptTemplate(
DEFAULT_RESPONSE_SYNTHESIS_PROMPT_TMPL,
prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS,
)
DEFAULT_TABLE_NAME = "items"

class DummySQLParser:
pass
from llama_index.core.query_engine.jsonalyze_query_engine import
JSONalyzeQueryEngine

-----------------------

UNIT TESTS START HERE

-----------------------

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

def test_get_prompts_returns_default_prompts():
"""Test that default prompts are returned when no custom prompts are provided."""
engine = JSONalyzeQueryEngine(list_of_dict=[])
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 428ns -> 273ns (56.8% faster)

def test_get_prompts_returns_custom_jsonalyze_prompt():
"""Test that a custom jsonalyze_prompt is returned if provided."""
custom_prompt = PromptTemplate("Custom JSONalyze prompt", prompt_type=PromptType.JSONALYZE)
engine = JSONalyzeQueryEngine(list_of_dict=[], jsonalyze_prompt=custom_prompt)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 439ns -> 256ns (71.5% faster)

def test_get_prompts_returns_custom_response_synthesis_prompt():
"""Test that a custom response_synthesis_prompt is returned if provided."""
custom_response_prompt = PromptTemplate("Custom synthesis prompt", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(list_of_dict=[], response_synthesis_prompt=custom_response_prompt)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 431ns -> 253ns (70.4% faster)

def test_get_prompts_returns_both_custom_prompts():
"""Test that both prompts are returned if both are provided."""
custom_jsonalyze = PromptTemplate("J", prompt_type=PromptType.JSONALYZE)
custom_response = PromptTemplate("R", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze,
response_synthesis_prompt=custom_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 417ns -> 259ns (61.0% faster)

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

def test_get_prompts_with_empty_prompt_templates():
"""Test engine with empty template strings for prompts."""
empty_jsonalyze = PromptTemplate("", prompt_type=PromptType.JSONALYZE)
empty_response = PromptTemplate("", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=empty_jsonalyze,
response_synthesis_prompt=empty_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 470ns -> 223ns (111% faster)

def test_get_prompts_with_none_list_of_dict():
"""Test engine with None for list_of_dict (should still work for _get_prompts)."""
engine = JSONalyzeQueryEngine(list_of_dict=None)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 505ns -> 271ns (86.3% faster)

def test_get_prompts_with_non_prompt_template_objects():
"""Test engine with objects that are not PromptTemplate but have similar attributes."""
class FakePrompt:
def init(self, template_str, prompt_type):
self.template_str = template_str
self.prompt_type = prompt_type
def eq(self, other):
return (
hasattr(other, "template_str") and
self.template_str == other.template_str and
getattr(other, "prompt_type", None) == self.prompt_type
)
fake_jsonalyze = FakePrompt("Fake", PromptType.JSONALYZE)
fake_response = FakePrompt("FakeResp", PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=fake_jsonalyze,
response_synthesis_prompt=fake_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 457ns -> 261ns (75.1% faster)

def test_get_prompts_with_unusual_prompt_type():
"""Test engine with a custom prompt type."""
unusual_prompt = PromptTemplate("Unusual", prompt_type="UNUSUAL_TYPE")
engine = JSONalyzeQueryEngine(list_of_dict=[], jsonalyze_prompt=unusual_prompt)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 490ns -> 271ns (80.8% faster)

def test_get_prompts_with_long_prompt_strings():
"""Test engine with very long prompt strings."""
long_str = "A" * 1000
long_jsonalyze = PromptTemplate(long_str, prompt_type=PromptType.JSONALYZE)
long_response = PromptTemplate(long_str, prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=long_jsonalyze,
response_synthesis_prompt=long_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 480ns -> 250ns (92.0% faster)

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

def test_get_prompts_with_large_list_of_dict():
"""Test engine with a large list_of_dict (should not affect _get_prompts)."""
large_list = [{"a": i, "b": str(i)} for i in range(1000)]
engine = JSONalyzeQueryEngine(list_of_dict=large_list)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 504ns -> 254ns (98.4% faster)

def test_get_prompts_with_many_custom_prompts():
"""Test engine with many different custom prompt objects (should always return the one provided)."""
for i in range(10): # limit to 10 for test speed
custom_jsonalyze = PromptTemplate(f"Custom{i}", prompt_type=PromptType.JSONALYZE)
custom_response = PromptTemplate(f"Resp{i}", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze,
response_synthesis_prompt=custom_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 2.68μs -> 1.51μs (76.6% faster)

def test_get_prompts_with_large_prompt_objects():
"""Test engine with large prompt objects (long template_str and many attributes)."""
class LargePromptTemplate(BasePromptTemplate):
def init(self, template_str, prompt_type):
super().init(template_str, prompt_type)
self.extra = "x" * 500
self.more = [i for i in range(100)]
large_jsonalyze = LargePromptTemplate("J"*500, PromptType.JSONALYZE)
large_response = LargePromptTemplate("R"*500, PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=large_jsonalyze,
response_synthesis_prompt=large_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 440ns -> 251ns (75.3% faster)

----------- NEGATIVE TEST CASES ------------

def test_get_prompts_with_missing_attributes():
"""Test engine with prompt objects missing expected attributes should raise AttributeError when accessed."""
class IncompletePrompt:
def init(self):
pass
incomplete_jsonalyze = IncompletePrompt()
incomplete_response = IncompletePrompt()
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=incomplete_jsonalyze,
response_synthesis_prompt=incomplete_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 455ns -> 262ns (73.7% faster)
# Should not raise error on return, but accessing .template_str should fail
with pytest.raises(AttributeError):
_ = prompts["jsonalyze_prompt"].template_str
with pytest.raises(AttributeError):
_ = prompts["response_synthesis_prompt"].template_str

----------- DETERMINISM TEST CASE ------------

def test_get_prompts_is_deterministic():
"""Test that repeated calls return the same prompt objects."""
custom_jsonalyze = PromptTemplate("Deterministic", prompt_type=PromptType.JSONALYZE)
custom_response = PromptTemplate("DeterministicResp", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze,
response_synthesis_prompt=custom_response
)
codeflash_output = engine._get_prompts(); prompts1 = codeflash_output # 410ns -> 259ns (58.3% faster)
codeflash_output = engine._get_prompts(); prompts2 = codeflash_output # 272ns -> 163ns (66.9% faster)

----------- TYPE TEST CASE ------------

def test_get_prompts_returns_dict_with_expected_keys():
"""Test that the returned dict has exactly the expected keys."""
engine = JSONalyzeQueryEngine(list_of_dict=[])
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 454ns -> 262ns (73.3% faster)

----------- IMMUTABILITY TEST CASE ------------

def test_get_prompts_does_not_modify_prompts():
"""Test that modifying the returned dict does not affect the engine's internal state."""
custom_jsonalyze = PromptTemplate("Immut", prompt_type=PromptType.JSONALYZE)
custom_response = PromptTemplate("ImmutResp", prompt_type=PromptType.SQL_RESPONSE_SYNTHESIS)
engine = JSONalyzeQueryEngine(
list_of_dict=[],
jsonalyze_prompt=custom_jsonalyze,
response_synthesis_prompt=custom_response
)
codeflash_output = engine._get_prompts(); prompts = codeflash_output # 416ns -> 253ns (64.4% faster)
prompts["jsonalyze_prompt"] = "changed"
# The next call should still return the original prompt object
codeflash_output = engine._get_prompts(); new_prompts = codeflash_output # 279ns -> 147ns (89.8% faster)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-JSONalyzeQueryEngine._get_prompts-mhvami6y and push.