graphiti-core 0.7.3__py3-none-any.whl → 0.7.5__py3-none-any.whl

graphiti_core/prompts/extract_nodes.py

@@ -31,9 +31,13 @@ class MissedEntities(BaseModel):
 
 
 class EntityClassification(BaseModel):
-    entity_classification: str = Field(
+    entities: list[str] = Field(
         ...,
-        description='Dictionary of entity classifications. Key is the entity name and value is the entity type',
+        description='List of entities',
+    )
+    entity_classifications: list[str | None] = Field(
+        ...,
+        description='List of entities classifications. The index of the classification should match the index of the entity it corresponds to.',
     )
 
 
@@ -180,7 +184,8 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
     
     Guidelines:
     1. Each entity must have exactly one type
-    2. If none of the provided entity types accurately classify an extracted node, the type should be set to None
+    2. Only use the provided ENTITY TYPES as types, do not use additional types to classify entities.
+    3. If none of the provided entity types accurately classify an extracted node, the type should be set to None
 """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/search/search_utils.py

@@ -229,8 +229,8 @@ async def edge_similarity_search(
 
     query: LiteralString = (
         """
-                                                                            MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
-                                                                            """
+                                                                                MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                                                                                """
         + group_filter_query
         + filter_query
         + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
@@ -765,8 +765,9 @@ async def node_distance_reranker(
     # rerank on shortest distance
     filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
 
-    # add back in filtered center uuid
-    filtered_uuids = [center_node_uuid] + filtered_uuids
+    # add back in filtered center uuid if it was filtered out
+    if center_node_uuid in node_uuids:
+        filtered_uuids = [center_node_uuid] + filtered_uuids
 
     return filtered_uuids
 

graphiti_core/utils/maintenance/node_operations.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import ast
 import logging
 from time import time
 
@@ -163,8 +162,9 @@ async def extract_nodes(
             prompt_library.extract_nodes.classify_nodes(node_classification_context),
             response_model=EntityClassification,
         )
-        response_string = llm_response.get('entity_classification', '{}')
-        node_classifications.update(ast.literal_eval(response_string))
+        entities = llm_response.get('entities', [])
+        entity_classifications = llm_response.get('entity_classifications', [])
+        node_classifications.update(dict(zip(entities, entity_classifications)))
 
     end = time()
     logger.debug(f'Extracted new nodes: {extracted_node_names} in {(end - start) * 1000} ms')
@@ -172,7 +172,11 @@ async def extract_nodes(
     new_nodes = []
     for name in extracted_node_names:
         entity_type = node_classifications.get(name)
-        labels = ['Entity'] if entity_type is None else ['Entity', entity_type]
+        labels = (
+            ['Entity']
+            if entity_type is None or entity_type == 'None' or entity_type == 'null'
+            else ['Entity', entity_type]
+        )
 
         new_node = EntityNode(
             name=name,

{graphiti_core-0.7.3.dist-info → graphiti_core-0.7.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.7.3
+Version: 0.7.5
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.7.3.dist-info → graphiti_core-0.7.5.dist-info}/RECORD

@@ -32,7 +32,7 @@ graphiti_core/prompts/dedupe_nodes.py,sha256=mqvNATL-4Vo33vaxUEZfOq6hXXOiL-ftY0z
 graphiti_core/prompts/eval.py,sha256=csW494kKBMvWSm2SYLIRuGgNghhwNR3YwGn3veo3g_Y,3691
 graphiti_core/prompts/extract_edge_dates.py,sha256=td2yx2wnX-nLioMa0mtla3WcRyO71_wSjemT79YZGQ0,4096
 graphiti_core/prompts/extract_edges.py,sha256=vyEdW7JAPOT_eLWUi6nRmxbvucyVoyoYX2SxXfknRUg,3467
-graphiti_core/prompts/extract_nodes.py,sha256=-01MpcVd9drtmMDIpQkkzZe8YwVhedmdbZq7UNGfo24,6651
+graphiti_core/prompts/extract_nodes.py,sha256=6WwcdmB0FcRYASqNIcSBoqIgTgj_lhU-Vgy65r7oLao,6891
 graphiti_core/prompts/invalidate_edges.py,sha256=DV2mEyIhhjc0hdKEMFLQMeG0FiUCkv_X0ctCliYjQ2c,3577
 graphiti_core/prompts/lib.py,sha256=oxhlpGEgV15VOLEZiwirxmIJBIdfzfiyL58iyzFDskE,4254
 graphiti_core/prompts/models.py,sha256=cvx_Bv5RMFUD_5IUawYrbpOKLPHogai7_bm7YXrSz84,867
@@ -44,7 +44,7 @@ graphiti_core/search/search.py,sha256=DX-tcIa0SiKI2HX-b_WdjGE74A8RLWQor4p90dJluU
 graphiti_core/search/search_config.py,sha256=UZN8jFA4pBlw2O5N1cuhVRBdTwMLR9N3Oyo6sQ4MDVw,3117
 graphiti_core/search/search_config_recipes.py,sha256=yUqiLnn9vFg39M8eVwjVKfBCL_ptGrfDMQ47m_Blb0g,6885
 graphiti_core/search/search_filters.py,sha256=4MJmCXD-blMc71xB4F9K4a72qidDCigADQ_ztdG15kc,5884
-graphiti_core/search/search_utils.py,sha256=mZqq97XEjTga3pz50l__b8qVFVPRes1dNGB8tZktJ3A,25665
+graphiti_core/search/search_utils.py,sha256=Bywp7trqP_gEwqH8I-JHuj3Mljw9P2K6_XooXI75jHI,25739
 graphiti_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/utils/bulk_utils.py,sha256=P4LKO46Yle4tBdNcQ3hDHcSQFaR8UBLfoL-z1M2Wua0,14690
 graphiti_core/utils/datetime_utils.py,sha256=Ti-2tnrDFRzBsbfblzsHybsM3jaDLP4-VT2t0VhpIzU,1357
@@ -52,10 +52,10 @@ graphiti_core/utils/maintenance/__init__.py,sha256=TRY3wWWu5kn3Oahk_KKhltrWnh0NA
 graphiti_core/utils/maintenance/community_operations.py,sha256=gIw1M5HGgc2c3TXag5ygPPpAv5WsG-yoC8Lhmfr6FMs,10011
 graphiti_core/utils/maintenance/edge_operations.py,sha256=tNw56vN586JYZMgie6RLRTiHZ680-kWzDIxW8ucL6SU,12780
 graphiti_core/utils/maintenance/graph_data_operations.py,sha256=qds9ALk9PhpQs1CNZTZGpi70mqJ93Y2KhIh9X2r8MUI,6533
-graphiti_core/utils/maintenance/node_operations.py,sha256=gihbPEBH6StLQCSd9wSu582d4Owaw3l5JLR1IBDrnVs,14137
+graphiti_core/utils/maintenance/node_operations.py,sha256=ty7_U9SNu4O7mgJgNvj-8THfoYGDDlKOo54EYHOC1wA,14292
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=RdNtubCyYhOVrvcOIq2WppHls1Q-BEjtsN8r38l-Rtc,3691
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core-0.7.3.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.7.3.dist-info/METADATA,sha256=CA2-V0qWnuIo_sQecfWTQQJ-9GLb8lMvMd5FqbUVFPQ,10541
-graphiti_core-0.7.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-graphiti_core-0.7.3.dist-info/RECORD,,
+graphiti_core-0.7.5.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.7.5.dist-info/METADATA,sha256=jJE1xBlbtBNv7FfIMpATqP6dnPTinuJwRrn8ppt2JqY,10541
+graphiti_core-0.7.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+graphiti_core-0.7.5.dist-info/RECORD,,