add entity type validation handling (#320)

* add entity type validation handling

* bump version

Author: prasmussen15

graphiti_core/errors.py

@@ -65,3 +65,11 @@ class SearchRerankerError(GraphitiError):
     def __init__(self, text: str):
         self.message = text
         super().__init__(self.message)
+
+
+class EntityTypeValidationError(GraphitiError):
+    """Raised when an entity type uses protected attribute names."""
+
+    def __init__(self, entity_type: str, entity_type_attribute: str):
+        self.message = f'{entity_type_attribute} cannot be used as an attribute for {entity_type} as it is a protected attribute name.'
+        super().__init__(self.message)

graphiti_core/graphiti.py

@@ -76,6 +76,7 @@
     resolve_extracted_nodes,
 )
 from graphiti_core.utils.maintenance.temporal_operations import get_edge_contradictions
+from graphiti_core.utils.ontology_utils.entity_types_utils import validate_entity_types
 
 logger = logging.getLogger(__name__)
 
@@ -319,6 +320,8 @@ async def add_episode_endpoint(episode_data: EpisodeData):
             entity_edges: list[EntityEdge] = []
             now = utc_now()
 
+            validate_entity_types(entity_types)
+
             previous_episodes = (
                 await self.retrieve_episodes(
                     reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]

graphiti_core/utils/ontology_utils/entity_types_utils.py

@@ -0,0 +1,37 @@
+"""
+Copyright 2024, Zep Software, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from pydantic import BaseModel
+
+from graphiti_core.errors import EntityTypeValidationError
+from graphiti_core.nodes import EntityNode
+
+
+def validate_entity_types(
+    entity_types: dict[str, BaseModel] | None,
+) -> bool:
+    if entity_types is None:
+        return True
+
+    entity_node_field_names = EntityNode.model_fields.keys()
+
+    for entity_type_name, entity_type_model in entity_types.items():
+        entity_type_field_names = entity_type_model.model_fields.keys()
+        for entity_type_field_name in entity_type_field_names:
+            if entity_type_field_name in entity_node_field_names:
+                raise EntityTypeValidationError(entity_type_name, entity_type_field_name)
+
+    return True

mcp_server/README.md

@@ -1,8 +1,15 @@
 # Graphiti MCP Server
 
-Graphiti is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti continuously integrates user interactions, structured and unstructured enterprise data, and external information into a coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI applications.
+Graphiti is a framework for building and querying temporally-aware knowledge graphs, specifically tailored for AI agents
+operating in dynamic environments. Unlike traditional retrieval-augmented generation (RAG) methods, Graphiti
+continuously integrates user interactions, structured and unstructured enterprise data, and external information into a
+coherent, queryable graph. The framework supports incremental data updates, efficient retrieval, and precise historical
+queries without requiring complete graph recomputation, making it suitable for developing interactive, context-aware AI
+applications.
 
-This is an experimental Model Context Protocol (MCP) server implementation for Graphiti. The MCP server exposes Graphiti's key functionality through the MCP protocol, allowing AI assistants to interact with Graphiti's knowledge graph capabilities.
+This is an experimental Model Context Protocol (MCP) server implementation for Graphiti. The MCP server exposes
+Graphiti's key functionality through the MCP protocol, allowing AI assistants to interact with Graphiti's knowledge
+graph capabilities.
 
 ## Features
 
@@ -59,7 +66,7 @@ uv run graphiti_mcp_server.py
 With options:
 
 ```bash
-uv run graphiti_mcp_server.py --model gpt-4o --transport sse
+uv run graphiti_mcp_server.py --model gpt-4o-mini --transport sse
 ```
 
 Available arguments:
@@ -72,33 +79,34 @@ Available arguments:
 
 ### Docker Deployment
 
-The Graphiti MCP server can be deployed using Docker. The Dockerfile uses `uv` for package management, ensuring consistent dependency installation.
+The Graphiti MCP server can be deployed using Docker. The Dockerfile uses `uv` for package management, ensuring
+consistent dependency installation.
 
 #### Environment Configuration
 
 Before running the Docker Compose setup, you need to configure the environment variables. You have two options:
 
 1. **Using a .env file** (recommended):
 
-   - Copy the provided `.env.example` file to create a `.env` file:
-     ```bash
-     cp .env.example .env
-     ```
-   - Edit the `.env` file to set your OpenAI API key and other configuration options:
-     ```
-     # Required for LLM operations
-     OPENAI_API_KEY=your_openai_api_key_here
-     MODEL_NAME=gpt-4o
-     # Optional: OPENAI_BASE_URL only needed for non-standard OpenAI endpoints
-     # OPENAI_BASE_URL=https://api.openai.com/v1
-     ```
-   - The Docker Compose setup is configured to use this file if it exists (it's optional)
+    - Copy the provided `.env.example` file to create a `.env` file:
+      ```bash
+      cp .env.example .env
+      ```
+    - Edit the `.env` file to set your OpenAI API key and other configuration options:
+      ```
+      # Required for LLM operations
+      OPENAI_API_KEY=your_openai_api_key_here
+      MODEL_NAME=gpt-4o-mini
+      # Optional: OPENAI_BASE_URL only needed for non-standard OpenAI endpoints
+      # OPENAI_BASE_URL=https://api.openai.com/v1
+      ```
+    - The Docker Compose setup is configured to use this file if it exists (it's optional)
 
 2. **Using environment variables directly**:
-   - You can also set the environment variables when running the Docker Compose command:
-     ```bash
-     OPENAI_API_KEY=your_key MODEL_NAME=gpt-4o docker compose up
-     ```
+    - You can also set the environment variables when running the Docker Compose command:
+      ```bash
+      OPENAI_API_KEY=your_key MODEL_NAME=gpt-4o-mini docker compose up
+      ```
 
 #### Neo4j Configuration
 
@@ -154,7 +162,7 @@ To use the Graphiti MCP server with an MCP-compatible client, configure it to co
         "NEO4J_USER": "neo4j",
         "NEO4J_PASSWORD": "demodemo",
         "OPENAI_API_KEY": "${OPENAI_API_KEY}",
-        "MODEL_NAME": "gpt-4o"
+        "MODEL_NAME": "gpt-4o-mini"
       }
     }
   }
@@ -192,7 +200,7 @@ Or start the server with uv and connect to it:
         "NEO4J_USER": "neo4j",
         "NEO4J_PASSWORD": "demodemo",
         "OPENAI_API_KEY": "${OPENAI_API_KEY}",
-        "MODEL_NAME": "gpt-4o"
+        "MODEL_NAME": "gpt-4o-mini"
       }
     }
   }
@@ -215,7 +223,8 @@ The Graphiti MCP server exposes the following tools:
 
 ## Working with JSON Data
 
-The Graphiti MCP server can process structured JSON data through the `add_episode` tool with `source="json"`. This allows you to automatically extract entities and relationships from structured data:
+The Graphiti MCP server can process structured JSON data through the `add_episode` tool with `source="json"`. This
+allows you to automatically extract entities and relationships from structured data:
 
 ```
 add_episode(
@@ -236,7 +245,8 @@ To integrate the Graphiti MCP Server with the Cursor IDE, follow these steps:
 python graphiti_mcp_server.py --transport sse --use-custom-entities --group-id <your_group_id>
 ```
 
-Hint: specify a `group_id` to retain prior graph data. If you do not specify a `group_id`, the server will create a new graph
+Hint: specify a `group_id` to retain prior graph data. If you do not specify a `group_id`, the server will create a new
+graph
 
 2. Configure Cursor to connect to the Graphiti MCP server.
 
@@ -254,7 +264,8 @@ Hint: specify a `group_id` to retain prior graph data. If you do not specify a `
 
 4. Kick off an agent session in Cursor.
 
-The integration enables AI assistants in Cursor to maintain persistent memory through Graphiti's knowledge graph capabilities.
+The integration enables AI assistants in Cursor to maintain persistent memory through Graphiti's knowledge graph
+capabilities.
 
 ## Requirements
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "graphiti-core"
-version = "0.8.7"
+version = "0.8.8"
 description = "A temporal graph building library"
 authors = [
     "Paul Paliychuk <[email protected]>",