graphiti-core 0.3.17__py3-none-any.whl → 0.3.19__py3-none-any.whl

graphiti_core/edges.py

@@ -180,7 +180,7 @@ class EntityEdge(Edge):
         start = time()
 
         text = self.fact.replace('\n', ' ')
-        self.fact_embedding = await embedder.create(input=[text])
+        self.fact_embedding = await embedder.create(input_data=[text])
 
         end = time()
         logger.debug(f'embedded {text} in {end - start} ms')

graphiti_core/embedder/client.py

@@ -29,6 +29,6 @@ class EmbedderConfig(BaseModel):
 class EmbedderClient(ABC):
     @abstractmethod
     async def create(
-        self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         pass

graphiti_core/embedder/openai.py

@@ -42,7 +42,9 @@ class OpenAIEmbedder(EmbedderClient):
         self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 
     async def create(
-        self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
-        result = await self.client.embeddings.create(input=input, model=self.config.embedding_model)
+        result = await self.client.embeddings.create(
+            input=input_data, model=self.config.embedding_model
+        )
         return result.data[0].embedding[: self.config.embedding_dim]

graphiti_core/embedder/voyage.py

@@ -41,7 +41,18 @@ class VoyageAIEmbedder(EmbedderClient):
         self.client = voyageai.AsyncClient(api_key=config.api_key)
 
     async def create(
-        self, input: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
-        result = await self.client.embed(input, model=self.config.embedding_model)
+        if isinstance(input_data, str):
+            input_list = [input_data]
+        elif isinstance(input_data, List):
+            input_list = [str(i) for i in input_data if i]
+        else:
+            input_list = [str(i) for i in input_data if i is not None]
+
+        input_list = [i for i in input_list if i]
+        if len(input_list) == 0:
+            return []
+
+        result = await self.client.embed(input_list, model=self.config.embedding_model)
         return result.embeddings[0][: self.config.embedding_dim]

graphiti_core/llm_client/utils.py

@@ -26,7 +26,7 @@ async def generate_embedding(embedder: EmbedderClient, text: str):
     start = time()
 
     text = text.replace('\n', ' ')
-    embedding = await embedder.create(input=[text])
+    embedding = await embedder.create(input_data=[text])
 
     end = time()
     logger.debug(f'embedded text of length {len(text)} in {end - start} ms')

graphiti_core/nodes.py

@@ -222,7 +222,7 @@ class EntityNode(Node):
     async def generate_name_embedding(self, embedder: EmbedderClient):
         start = time()
         text = self.name.replace('\n', ' ')
-        self.name_embedding = await embedder.create(input=[text])
+        self.name_embedding = await embedder.create(input_data=[text])
         end = time()
         logger.debug(f'embedded {text} in {end - start} ms')
 
@@ -334,7 +334,7 @@ class CommunityNode(Node):
     async def generate_name_embedding(self, embedder: EmbedderClient):
         start = time()
         text = self.name.replace('\n', ' ')
-        self.name_embedding = await embedder.create(input=[text])
+        self.name_embedding = await embedder.create(input_data=[text])
         end = time()
         logger.debug(f'embedded {text} in {end - start} ms')
 

graphiti_core/search/search.py

@@ -45,6 +45,7 @@ from graphiti_core.search.search_utils import (
     edge_similarity_search,
     episode_mentions_reranker,
     maximal_marginal_relevance,
+    node_bfs_search,
     node_distance_reranker,
     node_fulltext_search,
     node_similarity_search,
@@ -65,7 +66,7 @@ async def search(
     bfs_origin_node_uuids: list[str] | None = None,
 ) -> SearchResults:
     start = time()
-    query_vector = await embedder.create(input=[query.replace('\n', ' ')])
+    query_vector = await embedder.create(input_data=[query.replace('\n', ' ')])
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
@@ -138,7 +139,7 @@ async def edge_search(
                 edge_similarity_search(
                     driver, query_vector, None, None, group_ids, 2 * limit, config.sim_min_score
                 ),
-                edge_bfs_search(driver, bfs_origin_node_uuids, config.bfs_max_depth),
+                edge_bfs_search(driver, bfs_origin_node_uuids, config.bfs_max_depth, 2 * limit),
             ]
         )
     )
@@ -160,7 +161,12 @@ async def edge_search(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
     elif config.reranker == EdgeReranker.cross_encoder:
-        fact_to_uuid_map = {edge.fact: edge.uuid for result in search_results for edge in result}
+        search_result_uuids = [[edge.uuid for edge in result] for result in search_results]
+
+        rrf_result_uuids = rrf(search_result_uuids)
+        rrf_edges = [edge_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
+
+        fact_to_uuid_map = {edge.fact: edge.uuid for edge in rrf_edges}
         reranked_facts = await cross_encoder.rank(query, list(fact_to_uuid_map.keys()))
         reranked_uuids = [fact_to_uuid_map[fact] for fact, _ in reranked_facts]
     elif config.reranker == EdgeReranker.node_distance:
@@ -212,6 +218,7 @@ async def node_search(
                 node_similarity_search(
                     driver, query_vector, group_ids, 2 * limit, config.sim_min_score
                 ),
+                node_bfs_search(driver, bfs_origin_node_uuids, config.bfs_max_depth, 2 * limit),
             ]
         )
     )
@@ -232,9 +239,12 @@ async def node_search(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
     elif config.reranker == NodeReranker.cross_encoder:
-        summary_to_uuid_map = {
-            node.summary: node.uuid for result in search_results for node in result
-        }
+        # use rrf as a preliminary reranker
+        rrf_result_uuids = rrf(search_result_uuids)
+        rrf_results = [node_uuid_map[uuid] for uuid in rrf_result_uuids][:limit]
+
+        summary_to_uuid_map = {node.summary: node.uuid for node in rrf_results}
+
         reranked_summaries = await cross_encoder.rank(query, list(summary_to_uuid_map.keys()))
         reranked_uuids = [summary_to_uuid_map[fact] for fact, _ in reranked_summaries]
     elif config.reranker == NodeReranker.episode_mentions:

graphiti_core/search/search_utils.py

@@ -233,6 +233,7 @@ async def edge_bfs_search(
     driver: AsyncDriver,
     bfs_origin_node_uuids: list[str] | None,
     bfs_max_depth: int,
+    limit: int,
 ) -> list[EntityEdge]:
     # vector similarity search over embedded facts
     if bfs_origin_node_uuids is None:
@@ -256,12 +257,14 @@ async def edge_bfs_search(
                     r.expired_at AS expired_at,
                     r.valid_at AS valid_at,
                     r.invalid_at AS invalid_at
+                LIMIT $limit
         """)
 
     records, _, _ = await driver.execute_query(
         query,
         bfs_origin_node_uuids=bfs_origin_node_uuids,
         depth=bfs_max_depth,
+        limit=limit,
         database_=DEFAULT_DATABASE,
         routing_='r',
     )
@@ -348,6 +351,7 @@ async def node_bfs_search(
     driver: AsyncDriver,
     bfs_origin_node_uuids: list[str] | None,
     bfs_max_depth: int,
+    limit: int,
 ) -> list[EntityNode]:
     # vector similarity search over entity names
     if bfs_origin_node_uuids is None:
@@ -368,6 +372,7 @@ async def node_bfs_search(
             """,
         bfs_origin_node_uuids=bfs_origin_node_uuids,
         depth=bfs_max_depth,
+        limit=limit,
         database_=DEFAULT_DATABASE,
         routing_='r',
     )
@@ -690,4 +695,4 @@ def maximal_marginal_relevance(
 
     candidates_with_mmr.sort(reverse=True, key=lambda c: c[1])
 
-    return [candidate[0] for candidate in candidates_with_mmr]
+    return list(set([candidate[0] for candidate in candidates_with_mmr]))

{graphiti_core-0.3.17.dist-info → graphiti_core-0.3.19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.17
+Version: 0.3.19
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -14,7 +14,7 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: diskcache (>=5.6.3,<6.0.0)
 Requires-Dist: neo4j (>=5.23.0,<6.0.0)
 Requires-Dist: numpy (>=1.0.0)
-Requires-Dist: openai (>=1.50.2,<2.0.0)
+Requires-Dist: openai (>=1.52.2,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: tenacity (<9.0.0)
 Description-Content-Type: text/markdown

{graphiti_core-0.3.17.dist-info → graphiti_core-0.3.19.dist-info}/RECORD

@@ -3,11 +3,11 @@ graphiti_core/cross_encoder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 graphiti_core/cross_encoder/bge_reranker_client.py,sha256=jsXBUHfFpGsNASHaRnfz1_miQ3x070DdU8QS4J3DciI,1466
 graphiti_core/cross_encoder/client.py,sha256=PyFYYsALQAD9wu0gL5uquPsulmaBZ0AZkJmLq2DFA-c,1472
 graphiti_core/cross_encoder/openai_reranker_client.py,sha256=ij1E1Y5G9GNP3h3h8nSUF-ZJrQ921B54uudZUsCUaDc,4063
-graphiti_core/edges.py,sha256=KgH1f-nwexEX3PCRaQHPqbD033EeiKo_s39mqZn43zk,13082
+graphiti_core/edges.py,sha256=wKXmNXtu1deT5MfcqUNGHQlBLgwt3MJ_3j35ibNbpi8,13087
 graphiti_core/embedder/__init__.py,sha256=eWd-0sPxflnYXLoWNT9sxwCIFun5JNO9Fk4E-ZXXf8Y,164
-graphiti_core/embedder/client.py,sha256=Sd9CyYXaqRazdOH8opKackrTx-y9y-T54M78XTVMzxs,1006
-graphiti_core/embedder/openai.py,sha256=28cl4qQCQeu6EGxVVPw3lPesA-Z_Cpvuhozyc1jdqVg,1586
-graphiti_core/embedder/voyage.py,sha256=pGrSquGnSiYl4nXGnutbdWchtYgZb0Fi_yW3c90dPlI,1497
+graphiti_core/embedder/client.py,sha256=gVr_xdN-d0UQfeE4Nyoa4pL3M8UxllNH4eFqmarxwys,1011
+graphiti_core/embedder/openai.py,sha256=yYUYPymx_lBlxDTGrlc03yNhPFyGG-etM2sszRK2G2U,1618
+graphiti_core/embedder/voyage.py,sha256=_eGFI5_NjNG8z7qG3jTWCdE7sAs1Yb8fiSZSJlQLD9o,1879
 graphiti_core/errors.py,sha256=ddHrHGQxhwkVAtSph4AV84UoOlgwZufMczXPwB7uqPo,1795
 graphiti_core/graphiti.py,sha256=c9Rh777TrHYffPF6qvFAfm-m-PA4kD8a3ZW_ShsZGxE,27714
 graphiti_core/helpers.py,sha256=kqC2TD8Auwty4sG7KH4BuRMX413oTChGaAT_XUt9ZjU,2108
@@ -18,13 +18,13 @@ graphiti_core/llm_client/config.py,sha256=VwtvD0B7TNqE6Cl-rvH5v-bAfmjMLhEUuFmHSP
 graphiti_core/llm_client/errors.py,sha256=-qlWwv1X-UjfsFIiNl-7yJIYvPwi7z8srVRfX4-s6uk,814
 graphiti_core/llm_client/groq_client.py,sha256=5uGWeQ903EuNxuRiaeH-_J1U2Le_b7Q1UGV_K8bQAiw,2329
 graphiti_core/llm_client/openai_client.py,sha256=xLkbpusRVFRK0zPr3kOqY31HK_XCXrpO5rqUSpcEqEU,3825
-graphiti_core/llm_client/utils.py,sha256=Ms-QhA5X9rps7NBdJeQZUgQLD3vaZRWPiTlhJa6BjXM,995
+graphiti_core/llm_client/utils.py,sha256=zKpxXEbKa369m4W7RDEf-m56kH46V1Mx3RowcWZEWWs,1000
 graphiti_core/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/edges/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/edges/edge_db_queries.py,sha256=gUC5AZeXvobsIzG8Zqxj4Sa-j0PKV4CuEPHWupFTXQs,1145
 graphiti_core/models/nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/nodes/node_db_queries.py,sha256=HxZIjVLqVRnG7OJHeZv7t3KCEWFTWvdNOBJyLSN-1Ts,918
-graphiti_core/nodes.py,sha256=ZFid7j_affdcfdHFIpzphUE0cPwg-YvC31OQuAZWJWQ,13834
+graphiti_core/nodes.py,sha256=xlObDAC6Osp6L0wwqREuzBlIFlOzuUs9lN3ZpIv4seM,13844
 graphiti_core/prompts/__init__.py,sha256=EA-x9xUki9l8wnu2l8ek_oNf75-do5tq5hVq7Zbv8Kw,101
 graphiti_core/prompts/dedupe_edges.py,sha256=DUNHdIudj50FAjkla4nc68tSFSD2yjmYHBw-Bb7ph20,6529
 graphiti_core/prompts/dedupe_nodes.py,sha256=BZ9S-PB9SSGjc5Oo8ivdgA6rZx3OGOFhKtwrBlQ0bm0,7269
@@ -38,10 +38,10 @@ graphiti_core/prompts/models.py,sha256=cvx_Bv5RMFUD_5IUawYrbpOKLPHogai7_bm7YXrSz
 graphiti_core/prompts/summarize_nodes.py,sha256=FLuZpGTABgcxuIDkx_IKH115nHEw0rIaFhcGlWveAMc,2357
 graphiti_core/py.typed,sha256=vlmmzQOt7bmeQl9L3XJP4W6Ry0iiELepnOrinKz5KQg,79
 graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/search/search.py,sha256=jhYIb7ERK-zGT9N0M8wXjIdFMyU0pUKzpnHmmBxdG3A,10710
+graphiti_core/search/search.py,sha256=Fz45sXToQ6KN_y2k94kg9UyUZeqsS14jfaGfQ5KGifQ,11176
 graphiti_core/search/search_config.py,sha256=UZN8jFA4pBlw2O5N1cuhVRBdTwMLR9N3Oyo6sQ4MDVw,3117
 graphiti_core/search/search_config_recipes.py,sha256=20jS7veJExDnXA-ovJSUJfyDHKt7GW-nng-eoiT7ATA,5810
-graphiti_core/search/search_utils.py,sha256=gakgUHtWQ1hqtbtFXbZn59P0-v12PcBut72UiTL6ivs,22437
+graphiti_core/search/search_utils.py,sha256=l8BR4GOo-A2eIXx4ybC18n6t6CeerN_9KQbYzCB6ix0,22551
 graphiti_core/utils/__init__.py,sha256=cJAcMnBZdHBQmWrZdU1PQ1YmaL75bhVUkyVpIPuOyns,260
 graphiti_core/utils/bulk_utils.py,sha256=JtoYTZPCigPa3n2E43Oe7QhFZRTA_QKNGy1jVgklHag,12614
 graphiti_core/utils/maintenance/__init__.py,sha256=TRY3wWWu5kn3Oahk_KKhltrWnh0NACw0FskjqF6OtlA,314
@@ -51,7 +51,7 @@ graphiti_core/utils/maintenance/graph_data_operations.py,sha256=w66_SLlvPapuG91Y
 graphiti_core/utils/maintenance/node_operations.py,sha256=h5nlRojbXOGJs-alpv6z6WnZ1UCixVGlAQYBQUqz8Bs,9030
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=MvaRLWrBlDeYw8CQrKish1xbYcY5ovpfdqA2hSX7v5k,3367
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core-0.3.17.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.3.17.dist-info/METADATA,sha256=bguHVSkpKofqo_zp-R195rg64568tOKJB0iaqyeKUHg,9396
-graphiti_core-0.3.17.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-graphiti_core-0.3.17.dist-info/RECORD,,
+graphiti_core-0.3.19.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.3.19.dist-info/METADATA,sha256=uwulLozcE0fov_hyYb3rIWUtuvcy6BObIuOcfng0aaQ,9396
+graphiti_core-0.3.19.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+graphiti_core-0.3.19.dist-info/RECORD,,