⚡️ Speed up function validate_entity_types by 70%

[Saga4]

Here is an optimized version of your program.

📄 70% (0.70x) speedup for validate_entity_types in graphiti_core/utils/ontology_utils/entity_types_utils.py

⏱️ Runtime : 24.4 microseconds → 14.3 microseconds (best of 31 runs)

📝 Explanation and details

Here is an optimized version of your program.

Major inefficiency

The code computed EntityNode.model_fields.keys() every call to the function, but that is invariant and need only be computed once per module. Likewise, it looked up .keys() for every entity type model, though these could be checked far more efficiently via set intersection.

Key performance improvements.

• EntityNode.model_fields.keys() evaluated once, as a set, and reused.
• Comparison per entity type uses set intersection (O(min(N, M))) instead of repeated linear search.
• Only a single scan over each entity type model's fields, not recomputed unnecessarily.

The function signature and error handling are unchanged, and all preserved comments are respected. This implementation will be significantly faster for many input types.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Dict, Type

# imports
import pytest  # used for our unit tests
from graphiti_core.utils.ontology_utils.entity_types_utils import \
    validate_entity_types
from pydantic import BaseModel


# function to test
class EntityTypeValidationError(Exception):
    def __init__(self, entity_type_name, entity_type_field_name):
        super().__init__(f"Entity type '{entity_type_name}' has invalid field '{entity_type_field_name}'.")

# Simulate EntityNode as a Pydantic model with some reserved fields
class EntityNode(BaseModel):
    id: int
    type: str
    created_at: str
from graphiti_core.utils.ontology_utils.entity_types_utils import \
    validate_entity_types

# ----------------------------------------
# unit tests
# ----------------------------------------

# Helper: Create Pydantic models on the fly for tests
def make_model(name: str, fields: Dict[str, type]) -> Type[BaseModel]:
    return type(name, (BaseModel,), fields)

# 1. Basic Test Cases

def test_none_entity_types_returns_true():
    # Should return True if entity_types is None
    codeflash_output = validate_entity_types(None) # 338ns -> 267ns (26.6% faster)

def test_empty_entity_types_returns_true():
    # Should return True if entity_types is empty dict
    codeflash_output = validate_entity_types({}) # 2.31μs -> 555ns (316% faster)
















from types import SimpleNamespace
from typing import Dict, Type

# imports
import pytest  # used for our unit tests
from graphiti_core.utils.ontology_utils.entity_types_utils import \
    validate_entity_types
from pydantic import BaseModel


# function to test
class EntityTypeValidationError(Exception):
    def __init__(self, entity_type_name, entity_type_field_name):
        self.entity_type_name = entity_type_name
        self.entity_type_field_name = entity_type_field_name
        super().__init__(f"Field '{entity_type_field_name}' in '{entity_type_name}' conflicts with EntityNode fields.")

class EntityNode(BaseModel):
    id: int
    label: str
from graphiti_core.utils.ontology_utils.entity_types_utils import \
    validate_entity_types

# --- UNIT TESTS ---

# Helper to dynamically create Pydantic models for testing
def create_model(name, fields: Dict[str, type]) -> Type[BaseModel]:
    return type(name, (BaseModel,), fields)

# 1. BASIC TEST CASES

def test_none_entity_types():
    """Should return True when entity_types is None."""
    codeflash_output = validate_entity_types(None) # 384ns -> 360ns (6.67% faster)

def test_empty_entity_types():
    """Should return True when entity_types is an empty dict."""
    codeflash_output = validate_entity_types({}) # 3.49μs -> 655ns (433% faster)











def test_entity_type_with_non_string_field_names():
    """Should handle field names that are not strings (should not occur, but test for robustness)."""
    # Pydantic enforces string field names, so this is not really possible, but we can simulate.
    # We'll skip this as Pydantic will not allow non-string field names.

# 3. LARGE SCALE TEST CASES






from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types

def test_validate_entity_types():
    validate_entity_types({})

def test_validate_entity_types_2():
    validate_entity_types(None)

To edit these changes git checkout codeflash/optimize-validate_entity_types-mdcd3wu9 and push.

Type of Change

• Bug fix
• New feature
• Performance improvement
• Documentation/Tests

Objective

For new features and performance improvements: Clearly describe the objective and rationale for this change.

Testing

• Unit tests added/updated
• Integration tests added/updated
• All existing tests pass

Breaking Changes

• This PR contains breaking changes

If this is a breaking change, describe:

• What functionality is affected
• Migration path for existing users

Checklist

• Code follows project style guidelines (make lint passes)
• Self-review completed
• Documentation updated where necessary
• No secrets or sensitive information committed

Related Issues

Closes #[issue number]

---

Important

Optimizes validate_entity_types in entity_types_utils.py for a 70% speedup by reducing redundant computations and using set operations for field name checks.

• Performance Improvement:
  • Optimizes validate_entity_types in entity_types_utils.py for a 70% speedup.
  • Computes EntityNode.model_fields.keys() once and stores as _ENTITY_NODE_FIELD_NAMES.
  • Uses set intersection to check for field name conflicts, reducing repeated linear searches.
• Behavior:
  • Function signature and error handling remain unchanged.
  • Preserves original behavior by raising EntityTypeValidationError for the first conflict found.

^{This description was created by for 2110dbe. You can customize this summary. It will automatically update as commits are pushed.}

[danielchalef]

Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.

---

I have read the CLA Document and I hereby sign the CLA

---

_{You can retrigger this bot by commenting recheck in this Pull Request. Posted by the CLA Assistant Lite bot.}

[ellipsis-dev]

Important

Looks good to me! 👍

Reviewed everything up to 2110dbe in 47 seconds. Click for details.

• Reviewed 29 lines of code in 1 files
• Skipped 0 files when reviewing.
• Skipped posting 4 draft comments. View those below.
• Modify your settings and rules to customize what types of comments Ellipsis leaves. And don't forget to react with 👍 or 👎 to teach Ellipsis.

1. graphiti_core/utils/ontology_utils/entity_types_utils.py:31

• Draft comment:
  Nice optimization! Caching the set of EntityNode field names and using set intersection improves performance significantly.
• Reason this comment was not posted:
  Confidence changes required: 0% <= threshold 80% None

2. graphiti_core/utils/ontology_utils/entity_types_utils.py:42

• Draft comment:
  Consider adding a comment or docstring explaining that _ENTITY_NODE_FIELD_NAMES is a module-level cache. This helps clarify that it won’t update if EntityNode fields change at runtime.
• Reason this comment was not posted:
  Confidence changes required: 0% <= threshold 80% None

3. graphiti_core/utils/ontology_utils/entity_types_utils.py:37

• Draft comment:
  Using next(iter(conflict_fields)) to pick the first conflicting field is acceptable, but note that sets are unordered. If the order of error reporting matters, you might want to consider sorting or clarifying this behavior in a comment.
• Reason this comment was not posted:
  Confidence changes required: 0% <= threshold 80% None

4. graphiti_core/utils/ontology_utils/entity_types_utils.py:23

• Draft comment:
  Consider adding a docstring for the validate_entity_types function to document its parameters, return value, and exception behavior. This can help future maintainers understand the function's purpose.
• Reason this comment was not posted:
  Confidence changes required: 0% <= threshold 80% None

Workflow ID: wflow_MlQH8c2hqsgh1XvX

^{You can customize by changing your verbosity settings, reacting with 👍 or 👎, replying to comments, or adding code review rules.}

[Saga4]

I have read the CLA Document and I hereby sign the CLA