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

graphiti_core/cross_encoder/bge_reranker_client.py

@@ -0,0 +1,45 @@
+"""
+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.
+"""
+
+import asyncio
+from typing import List, Tuple
+
+from sentence_transformers import CrossEncoder
+
+from graphiti_core.cross_encoder.client import CrossEncoderClient
+
+
+class BGERerankerClient(CrossEncoderClient):
+    def __init__(self):
+        self.model = CrossEncoder('BAAI/bge-reranker-v2-m3')
+
+    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+        if not passages:
+            return []
+
+        input_pairs = [[query, passage] for passage in passages]
+
+        # Run the synchronous predict method in an executor
+        loop = asyncio.get_running_loop()
+        scores = await loop.run_in_executor(None, self.model.predict, input_pairs)
+
+        ranked_passages = sorted(
+            [(passage, float(score)) for passage, score in zip(passages, scores)],
+            key=lambda x: x[1],
+            reverse=True,
+        )
+
+        return ranked_passages

graphiti_core/cross_encoder/client.py

@@ -0,0 +1,41 @@
+"""
+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 abc import ABC, abstractmethod
+from typing import List, Tuple
+
+
+class CrossEncoderClient(ABC):
+    """
+    CrossEncoderClient is an abstract base class that defines the interface
+    for cross-encoder models used for ranking passages based on their relevance to a query.
+    It allows for different implementations of cross-encoder models to be used interchangeably.
+    """
+
+    @abstractmethod
+    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+        """
+        Rank the given passages based on their relevance to the query.
+
+        Args:
+            query (str): The query string.
+            passages (List[str]): A list of passages to rank.
+
+        Returns:
+            List[Tuple[str, float]]: A list of tuples containing the passage and its score,
+                                     sorted in descending order of relevance.
+        """
+        pass

graphiti_core/cross_encoder/openai_reranker_client.py

@@ -0,0 +1,113 @@
+"""
+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.
+"""
+
+import asyncio
+import logging
+from typing import Any
+
+import openai
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+from ..llm_client import LLMConfig, RateLimitError
+from ..prompts import Message
+from .client import CrossEncoderClient
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gpt-4o-mini'
+
+
+class BooleanClassifier(BaseModel):
+    isTrue: bool
+
+
+class OpenAIRerankerClient(CrossEncoderClient):
+    def __init__(self, config: LLMConfig | None = None):
+        """
+        Initialize the OpenAIClient with the provided configuration, cache setting, and client.
+
+        Args:
+            config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
+            cache (bool): Whether to use caching for responses. Defaults to False.
+            client (Any | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
+
+        """
+        if config is None:
+            config = LLMConfig()
+
+        self.config = config
+        self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
+        openai_messages_list: Any = [
+            [
+                Message(
+                    role='system',
+                    content='You are an expert tasked with determining whether the passage is relevant to the query',
+                ),
+                Message(
+                    role='user',
+                    content=f"""
+                           Respond with "True" if PASSAGE is relevant to QUERY and "False" otherwise. 
+                           <PASSAGE>
+                           {query}
+                           </PASSAGE>
+                           {passage}
+                           <QUERY>
+                           </QUERY>
+                           """,
+                ),
+            ]
+            for passage in passages
+        ]
+        try:
+            responses = await asyncio.gather(
+                *[
+                    self.client.chat.completions.create(
+                        model=DEFAULT_MODEL,
+                        messages=openai_messages,
+                        temperature=0,
+                        max_tokens=1,
+                        logit_bias={'6432': 1, '7983': 1},
+                        logprobs=True,
+                        top_logprobs=2,
+                    )
+                    for openai_messages in openai_messages_list
+                ]
+            )
+
+            responses_top_logprobs = [
+                response.choices[0].logprobs.content[0].top_logprobs
+                if response.choices[0].logprobs is not None
+                and response.choices[0].logprobs.content is not None
+                else []
+                for response in responses
+            ]
+            scores: list[float] = []
+            for top_logprobs in responses_top_logprobs:
+                for logprob in top_logprobs:
+                    if bool(logprob.token):
+                        scores.append(logprob.logprob)
+
+            results = [(passage, score) for passage, score in zip(passages, scores)]
+            results.sort(reverse=True, key=lambda x: x[1])
+            return results
+        except openai.RateLimitError as e:
+            raise RateLimitError from e
+        except Exception as e:
+            logger.error(f'Error in generating LLM response: {e}')
+            raise

graphiti_core/graphiti.py

@@ -23,8 +23,11 @@ from dotenv import load_dotenv
 from neo4j import AsyncGraphDatabase
 from pydantic import BaseModel
 
+from graphiti_core.cross_encoder.client import CrossEncoderClient
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
+from graphiti_core.helpers import DEFAULT_DATABASE
 from graphiti_core.llm_client import LLMClient, OpenAIClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.search.search import SearchConfig, search
@@ -92,6 +95,7 @@ class Graphiti:
         password: str,
         llm_client: LLMClient | None = None,
         embedder: EmbedderClient | None = None,
+        cross_encoder: CrossEncoderClient | None = None,
         store_raw_episode_content: bool = True,
     ):
         """
@@ -131,7 +135,7 @@ class Graphiti:
         Graphiti if you're using the default OpenAIClient.
         """
         self.driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
-        self.database = 'neo4j'
+        self.database = DEFAULT_DATABASE
         self.store_raw_episode_content = store_raw_episode_content
         if llm_client:
             self.llm_client = llm_client
@@ -141,6 +145,10 @@ class Graphiti:
             self.embedder = embedder
         else:
             self.embedder = OpenAIEmbedder()
+        if cross_encoder:
+            self.cross_encoder = cross_encoder
+        else:
+            self.cross_encoder = OpenAIRerankerClient()
 
     async def close(self):
         """
@@ -648,6 +656,7 @@ class Graphiti:
             await search(
                 self.driver,
                 self.embedder,
+                self.cross_encoder,
                 query,
                 group_ids,
                 search_config,
@@ -663,8 +672,18 @@ class Graphiti:
         config: SearchConfig,
         group_ids: list[str] | None = None,
         center_node_uuid: str | None = None,
+        bfs_origin_node_uuids: list[str] | None = None,
     ) -> SearchResults:
-        return await search(self.driver, self.embedder, query, group_ids, config, center_node_uuid)
+        return await search(
+            self.driver,
+            self.embedder,
+            self.cross_encoder,
+            query,
+            group_ids,
+            config,
+            center_node_uuid,
+            bfs_origin_node_uuids,
+        )
 
     async def get_nodes_by_query(
         self,
@@ -716,7 +735,13 @@ class Graphiti:
 
         nodes = (
             await search(
-                self.driver, self.embedder, query, group_ids, search_config, center_node_uuid
+                self.driver,
+                self.embedder,
+                self.cross_encoder,
+                query,
+                group_ids,
+                search_config,
+                center_node_uuid,
             )
         ).nodes
         return nodes

graphiti_core/search/search.py

@@ -21,6 +21,7 @@ from time import time
 
 from neo4j import AsyncDriver
 
+from graphiti_core.cross_encoder.client import CrossEncoderClient
 from graphiti_core.edges import EntityEdge
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import SearchRerankerError
@@ -39,6 +40,7 @@ from graphiti_core.search.search_config import (
 from graphiti_core.search.search_utils import (
     community_fulltext_search,
     community_similarity_search,
+    edge_bfs_search,
     edge_fulltext_search,
     edge_similarity_search,
     episode_mentions_reranker,
@@ -55,40 +57,49 @@ logger = logging.getLogger(__name__)
 async def search(
     driver: AsyncDriver,
     embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
     group_ids: list[str] | None,
     config: SearchConfig,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
 ) -> SearchResults:
     start = time()
-    query = query.replace('\n', ' ')
+    query_vector = await embedder.create(input=[query.replace('\n', ' ')])
+
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
     edges, nodes, communities = await asyncio.gather(
         edge_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.edge_config,
             center_node_uuid,
+            bfs_origin_node_uuids,
             config.limit,
         ),
         node_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.node_config,
             center_node_uuid,
+            bfs_origin_node_uuids,
             config.limit,
         ),
         community_search(
             driver,
-            embedder,
+            cross_encoder,
             query,
+            query_vector,
             group_ids,
             config.community_config,
+            bfs_origin_node_uuids,
             config.limit,
         ),
     )
@@ -99,27 +110,27 @@ async def search(
         communities=communities,
     )
 
-    end = time()
+    latency = (time() - start) * 1000
 
-    logger.info(f'search returned context for query {query} in {(end - start) * 1000} ms')
+    logger.debug(f'search returned context for query {query} in {latency} ms')
 
     return results
 
 
 async def edge_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: EdgeSearchConfig | None,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[EntityEdge]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[EntityEdge]] = list(
         await asyncio.gather(
             *[
@@ -127,6 +138,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),
             ]
         )
     )
@@ -147,6 +159,10 @@ async def edge_search(
         reranked_uuids = maximal_marginal_relevance(
             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}
+        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:
         if center_node_uuid is None:
             raise SearchRerankerError('No center node provided for Node Distance reranker')
@@ -177,18 +193,18 @@ async def edge_search(
 
 async def node_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: NodeSearchConfig | None,
     center_node_uuid: str | None = None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[EntityNode]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[EntityNode]] = list(
         await asyncio.gather(
             *[
@@ -215,6 +231,12 @@ async def node_search(
         reranked_uuids = maximal_marginal_relevance(
             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
+        }
+        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:
         reranked_uuids = await episode_mentions_reranker(driver, search_result_uuids)
     elif config.reranker == NodeReranker.node_distance:
@@ -231,17 +253,17 @@ async def node_search(
 
 async def community_search(
     driver: AsyncDriver,
-    embedder: EmbedderClient,
+    cross_encoder: CrossEncoderClient,
     query: str,
+    query_vector: list[float],
     group_ids: list[str] | None,
     config: CommunitySearchConfig | None,
+    bfs_origin_node_uuids: list[str] | None = None,
     limit=DEFAULT_SEARCH_LIMIT,
 ) -> list[CommunityNode]:
     if config is None:
         return []
 
-    query_vector = await embedder.create(input=[query])
-
     search_results: list[list[CommunityNode]] = list(
         await asyncio.gather(
             *[
@@ -273,6 +295,12 @@ async def community_search(
         reranked_uuids = maximal_marginal_relevance(
             query_vector, search_result_uuids_and_vectors, config.mmr_lambda
         )
+    elif config.reranker == CommunityReranker.cross_encoder:
+        summary_to_uuid_map = {
+            node.summary: node.uuid for result in search_results for node in result
+        }
+        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]
 
     reranked_communities = [community_uuid_map[uuid] for uuid in reranked_uuids]
 

graphiti_core/search/search_config.py

@@ -20,7 +20,11 @@ from pydantic import BaseModel, Field
 
 from graphiti_core.edges import EntityEdge
 from graphiti_core.nodes import CommunityNode, EntityNode
-from graphiti_core.search.search_utils import DEFAULT_MIN_SCORE, DEFAULT_MMR_LAMBDA
+from graphiti_core.search.search_utils import (
+    DEFAULT_MIN_SCORE,
+    DEFAULT_MMR_LAMBDA,
+    MAX_SEARCH_DEPTH,
+)
 
 DEFAULT_SEARCH_LIMIT = 10
 
@@ -28,11 +32,13 @@ DEFAULT_SEARCH_LIMIT = 10
 class EdgeSearchMethod(Enum):
     cosine_similarity = 'cosine_similarity'
     bm25 = 'bm25'
+    bfs = 'breadth_first_search'
 
 
 class NodeSearchMethod(Enum):
     cosine_similarity = 'cosine_similarity'
     bm25 = 'bm25'
+    bfs = 'breadth_first_search'
 
 
 class CommunitySearchMethod(Enum):
@@ -45,6 +51,7 @@ class EdgeReranker(Enum):
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class NodeReranker(Enum):
@@ -52,11 +59,13 @@ class NodeReranker(Enum):
     node_distance = 'node_distance'
     episode_mentions = 'episode_mentions'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class CommunityReranker(Enum):
     rrf = 'reciprocal_rank_fusion'
     mmr = 'mmr'
+    cross_encoder = 'cross_encoder'
 
 
 class EdgeSearchConfig(BaseModel):
@@ -64,6 +73,7 @@ class EdgeSearchConfig(BaseModel):
     reranker: EdgeReranker = Field(default=EdgeReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class NodeSearchConfig(BaseModel):
@@ -71,6 +81,7 @@ class NodeSearchConfig(BaseModel):
     reranker: NodeReranker = Field(default=NodeReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class CommunitySearchConfig(BaseModel):
@@ -78,6 +89,7 @@ class CommunitySearchConfig(BaseModel):
     reranker: CommunityReranker = Field(default=CommunityReranker.rrf)
     sim_min_score: float = Field(default=DEFAULT_MIN_SCORE)
     mmr_lambda: float = Field(default=DEFAULT_MMR_LAMBDA)
+    bfs_max_depth: int = Field(default=MAX_SEARCH_DEPTH)
 
 
 class SearchConfig(BaseModel):

graphiti_core/search/search_config_recipes.py

@@ -48,14 +48,41 @@ COMBINED_HYBRID_SEARCH_MMR = SearchConfig(
     edge_config=EdgeSearchConfig(
         search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
         reranker=EdgeReranker.mmr,
+        mmr_lambda=1,
     ),
     node_config=NodeSearchConfig(
         search_methods=[NodeSearchMethod.bm25, NodeSearchMethod.cosine_similarity],
         reranker=NodeReranker.mmr,
+        mmr_lambda=1,
     ),
     community_config=CommunitySearchConfig(
         search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
         reranker=CommunityReranker.mmr,
+        mmr_lambda=1,
+    ),
+)
+
+# Performs a full-text search, similarity search, and bfs with cross_encoder reranking over edges, nodes, and communities
+COMBINED_HYBRID_SEARCH_CROSS_ENCODER = SearchConfig(
+    edge_config=EdgeSearchConfig(
+        search_methods=[
+            EdgeSearchMethod.bm25,
+            EdgeSearchMethod.cosine_similarity,
+            EdgeSearchMethod.bfs,
+        ],
+        reranker=EdgeReranker.cross_encoder,
+    ),
+    node_config=NodeSearchConfig(
+        search_methods=[
+            NodeSearchMethod.bm25,
+            NodeSearchMethod.cosine_similarity,
+            NodeSearchMethod.bfs,
+        ],
+        reranker=NodeReranker.cross_encoder,
+    ),
+    community_config=CommunitySearchConfig(
+        search_methods=[CommunitySearchMethod.bm25, CommunitySearchMethod.cosine_similarity],
+        reranker=CommunityReranker.cross_encoder,
     ),
 )
 
@@ -81,7 +108,6 @@ EDGE_HYBRID_SEARCH_NODE_DISTANCE = SearchConfig(
         search_methods=[EdgeSearchMethod.bm25, EdgeSearchMethod.cosine_similarity],
         reranker=EdgeReranker.node_distance,
     ),
-    limit=30,
 )
 
 # performs a hybrid search over edges with episode mention reranking

graphiti_core/search/search_utils.py

@@ -37,6 +37,7 @@ logger = logging.getLogger(__name__)
 RELEVANT_SCHEMA_LIMIT = 3
 DEFAULT_MIN_SCORE = 0.6
 DEFAULT_MMR_LAMBDA = 0.5
+MAX_SEARCH_DEPTH = 3
 MAX_QUERY_LENGTH = 128
 
 
@@ -79,21 +80,21 @@ async def get_mentioned_nodes(
     driver: AsyncDriver, episodes: list[EpisodicNode]
 ) -> list[EntityNode]:
     episode_uuids = [episode.uuid for episode in episodes]
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-            MATCH (episode:Episodic)-[:MENTIONS]->(n:Entity) WHERE episode.uuid IN $uuids
-            RETURN DISTINCT
-                n.uuid As uuid, 
-                n.group_id AS group_id,
-                n.name AS name,
-                n.name_embedding AS name_embedding,
-                n.created_at AS created_at, 
-                n.summary AS summary
-            """,
-            {'uuids': episode_uuids},
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+        MATCH (episode:Episodic)-[:MENTIONS]->(n:Entity) WHERE episode.uuid IN $uuids
+        RETURN DISTINCT
+            n.uuid As uuid, 
+            n.group_id AS group_id,
+            n.name AS name,
+            n.name_embedding AS name_embedding,
+            n.created_at AS created_at, 
+            n.summary AS summary
+        """,
+        uuids=episode_uuids,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
 
     nodes = [get_entity_node_from_record(record) for record in records]
 
@@ -104,21 +105,21 @@ async def get_communities_by_nodes(
     driver: AsyncDriver, nodes: list[EntityNode]
 ) -> list[CommunityNode]:
     node_uuids = [node.uuid for node in nodes]
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-        MATCH (c:Community)-[:HAS_MEMBER]->(n:Entity) WHERE n.uuid IN $uuids
-        RETURN DISTINCT
-            c.uuid As uuid, 
-            c.group_id AS group_id,
-            c.name AS name,
-            c.name_embedding AS name_embedding
-            c.created_at AS created_at, 
-            c.summary AS summary
-        """,
-            {'uuids': node_uuids},
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+    MATCH (c:Community)-[:HAS_MEMBER]->(n:Entity) WHERE n.uuid IN $uuids
+    RETURN DISTINCT
+        c.uuid As uuid, 
+        c.group_id AS group_id,
+        c.name AS name,
+        c.name_embedding AS name_embedding
+        c.created_at AS created_at, 
+        c.summary AS summary
+    """,
+        uuids=node_uuids,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
 
     communities = [get_community_node_from_record(record) for record in records]
 
@@ -141,8 +142,10 @@ async def edge_fulltext_search(
     cypher_query = Query("""
               CALL db.index.fulltext.queryRelationships("edge_name_and_fact", $query) 
               YIELD relationship AS rel, score
-              MATCH (n:Entity)-[r {uuid: rel.uuid}]-(m:Entity)
-              RETURN 
+              MATCH (n:Entity)-[r {uuid: rel.uuid}]->(m:Entity)
+              WHERE ($source_uuid IS NULL OR n.uuid IN [$source_uuid, $target_uuid])
+              AND ($target_uuid IS NULL OR m.uuid IN [$source_uuid, $target_uuid])
+              RETURN
                     r.uuid AS uuid,
                     r.group_id AS group_id,
                     n.uuid AS source_node_uuid,
@@ -158,18 +161,16 @@ async def edge_fulltext_search(
                 ORDER BY score DESC LIMIT $limit
                 """)
 
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            cypher_query,
-            {
-                'query': fuzzy_query,
-                'source_uuid': source_node_uuid,
-                'target_uuid': target_node_uuid,
-                'group_ids': group_ids,
-                'limit': limit,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        cypher_query,
+        query=fuzzy_query,
+        source_uuid=source_node_uuid,
+        target_uuid=target_node_uuid,
+        group_ids=group_ids,
+        limit=limit,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
 
     edges = [get_entity_edge_from_record(record) for record in records]
 
@@ -188,17 +189,17 @@ async def edge_similarity_search(
     # vector similarity search over embedded facts
     query = Query("""
                 CYPHER runtime = parallel parallelRuntimeSupport=all
-                MATCH (n:Entity)-[r:RELATES_TO]-(m:Entity)
+                MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
                 WHERE ($group_ids IS NULL OR r.group_id IN $group_ids)
-                AND ($source_uuid IS NULL OR n.uuid = $source_uuid)
-                AND ($target_uuid IS NULL OR m.uuid = $target_uuid)
-                WITH n, r, m, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
+                AND ($source_uuid IS NULL OR n.uuid IN [$source_uuid, $target_uuid])
+                AND ($target_uuid IS NULL OR m.uuid IN [$source_uuid, $target_uuid])
+                WITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
                 WHERE score > $min_score
                 RETURN
                     r.uuid AS uuid,
                     r.group_id AS group_id,
-                    n.uuid AS source_node_uuid,
-                    m.uuid AS target_node_uuid,
+                    startNode(r).uuid AS source_node_uuid,
+                    endNode(r).uuid AS target_node_uuid,
                     r.created_at AS created_at,
                     r.name AS name,
                     r.fact AS fact,
@@ -211,19 +212,59 @@ async def edge_similarity_search(
                 LIMIT $limit
         """)
 
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            query,
-            {
-                'search_vector': search_vector,
-                'source_uuid': source_node_uuid,
-                'target_uuid': target_node_uuid,
-                'group_ids': group_ids,
-                'limit': limit,
-                'min_score': min_score,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        query,
+        search_vector=search_vector,
+        source_uuid=source_node_uuid,
+        target_uuid=target_node_uuid,
+        group_ids=group_ids,
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+
+    edges = [get_entity_edge_from_record(record) for record in records]
+
+    return edges
+
+
+async def edge_bfs_search(
+    driver: AsyncDriver,
+    bfs_origin_node_uuids: list[str] | None,
+    bfs_max_depth: int,
+) -> list[EntityEdge]:
+    # vector similarity search over embedded facts
+    if bfs_origin_node_uuids is None:
+        return []
+
+    query = Query("""
+                UNWIND $bfs_origin_node_uuids AS origin_uuid
+                MATCH path = (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
+                UNWIND relationships(path) AS rel
+                MATCH ()-[r:RELATES_TO {uuid: rel.uuid}]-()
+                RETURN DISTINCT
+                    r.uuid AS uuid,
+                    r.group_id AS group_id,
+                    startNode(r).uuid AS source_node_uuid,
+                    endNode(r).uuid AS target_node_uuid,
+                    r.created_at AS created_at,
+                    r.name AS name,
+                    r.fact AS fact,
+                    r.fact_embedding AS fact_embedding,
+                    r.episodes AS episodes,
+                    r.expired_at AS expired_at,
+                    r.valid_at AS valid_at,
+                    r.invalid_at AS invalid_at
+        """)
+
+    records, _, _ = await driver.execute_query(
+        query,
+        bfs_origin_node_uuids=bfs_origin_node_uuids,
+        depth=bfs_max_depth,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
 
     edges = [get_entity_edge_from_record(record) for record in records]
 
@@ -241,28 +282,26 @@ async def node_fulltext_search(
     if fuzzy_query == '':
         return []
 
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-    CALL db.index.fulltext.queryNodes("node_name_and_summary", $query) 
-    YIELD node AS n, score
-    RETURN
-        n.uuid AS uuid,
-        n.group_id AS group_id, 
-        n.name AS name, 
-        n.name_embedding AS name_embedding,
-        n.created_at AS created_at, 
-        n.summary AS summary
-    ORDER BY score DESC
-    LIMIT $limit
-    """,
-            {
-                'query': fuzzy_query,
-                'group_ids': group_ids,
-                'limit': limit,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+        CALL db.index.fulltext.queryNodes("node_name_and_summary", $query) 
+        YIELD node AS n, score
+        RETURN
+            n.uuid AS uuid,
+            n.group_id AS group_id, 
+            n.name AS name, 
+            n.name_embedding AS name_embedding,
+            n.created_at AS created_at, 
+            n.summary AS summary
+        ORDER BY score DESC
+        LIMIT $limit
+        """,
+        query=fuzzy_query,
+        group_ids=group_ids,
+        limit=limit,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
     nodes = [get_entity_node_from_record(record) for record in records]
 
     return nodes
@@ -276,32 +315,62 @@ async def node_similarity_search(
     min_score: float = DEFAULT_MIN_SCORE,
 ) -> list[EntityNode]:
     # vector similarity search over entity names
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-                CYPHER runtime = parallel parallelRuntimeSupport=all
-                MATCH (n:Entity)
-                WHERE $group_ids IS NULL OR n.group_id IN $group_ids
-                WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
-                WHERE score > $min_score
-                RETURN
-                    n.uuid As uuid,
-                    n.group_id AS group_id,
-                    n.name AS name, 
-                    n.name_embedding AS name_embedding,
-                    n.created_at AS created_at, 
-                    n.summary AS summary
-                ORDER BY score DESC
-                LIMIT $limit
-                """,
-            {
-                'search_vector': search_vector,
-                'group_ids': group_ids,
-                'limit': limit,
-                'min_score': min_score,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+            CYPHER runtime = parallel parallelRuntimeSupport=all
+            MATCH (n:Entity)
+            WHERE $group_ids IS NULL OR n.group_id IN $group_ids
+            WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
+            WHERE score > $min_score
+            RETURN
+                n.uuid As uuid,
+                n.group_id AS group_id,
+                n.name AS name, 
+                n.name_embedding AS name_embedding,
+                n.created_at AS created_at, 
+                n.summary AS summary
+            ORDER BY score DESC
+            LIMIT $limit
+            """,
+        search_vector=search_vector,
+        group_ids=group_ids,
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
+    nodes = [get_entity_node_from_record(record) for record in records]
+
+    return nodes
+
+
+async def node_bfs_search(
+    driver: AsyncDriver,
+    bfs_origin_node_uuids: list[str] | None,
+    bfs_max_depth: int,
+) -> list[EntityNode]:
+    # vector similarity search over entity names
+    if bfs_origin_node_uuids is None:
+        return []
+
+    records, _, _ = await driver.execute_query(
+        """
+            UNWIND $bfs_origin_node_uuids AS origin_uuid
+            MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
+            RETURN DISTINCT
+                n.uuid As uuid,
+                n.group_id AS group_id,
+                n.name AS name, 
+                n.name_embedding AS name_embedding,
+                n.created_at AS created_at, 
+                n.summary AS summary
+            LIMIT $limit
+            """,
+        bfs_origin_node_uuids=bfs_origin_node_uuids,
+        depth=bfs_max_depth,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
     nodes = [get_entity_node_from_record(record) for record in records]
 
     return nodes
@@ -318,28 +387,26 @@ async def community_fulltext_search(
     if fuzzy_query == '':
         return []
 
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-    CALL db.index.fulltext.queryNodes("community_name", $query) 
-    YIELD node AS comm, score
-    RETURN
-        comm.uuid AS uuid,
-        comm.group_id AS group_id, 
-        comm.name AS name, 
-        comm.name_embedding AS name_embedding,
-        comm.created_at AS created_at, 
-        comm.summary AS summary
-    ORDER BY score DESC
-    LIMIT $limit
-    """,
-            {
-                'query': fuzzy_query,
-                'group_ids': group_ids,
-                'limit': limit,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+        CALL db.index.fulltext.queryNodes("community_name", $query) 
+        YIELD node AS comm, score
+        RETURN
+            comm.uuid AS uuid,
+            comm.group_id AS group_id, 
+            comm.name AS name, 
+            comm.name_embedding AS name_embedding,
+            comm.created_at AS created_at, 
+            comm.summary AS summary
+        ORDER BY score DESC
+        LIMIT $limit
+        """,
+        query=fuzzy_query,
+        group_ids=group_ids,
+        limit=limit,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -353,32 +420,30 @@ async def community_similarity_search(
     min_score=DEFAULT_MIN_SCORE,
 ) -> list[CommunityNode]:
     # vector similarity search over entity names
-    async with driver.session(database=DEFAULT_DATABASE) as session:
-        result = await session.run(
-            """
-                CYPHER runtime = parallel parallelRuntimeSupport=all
-                MATCH (comm:Community)
-                WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
-                WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
-                WHERE score > $min_score
-                RETURN
-                    comm.uuid As uuid,
-                    comm.group_id AS group_id,
-                    comm.name AS name, 
-                    comm.name_embedding AS name_embedding,
-                    comm.created_at AS created_at, 
-                    comm.summary AS summary
-                ORDER BY score DESC
-                LIMIT $limit
-                """,
-            {
-                'search_vector': search_vector,
-                'group_ids': group_ids,
-                'limit': limit,
-                'min_score': min_score,
-            },
-        )
-        records = [record async for record in result]
+    records, _, _ = await driver.execute_query(
+        """
+            CYPHER runtime = parallel parallelRuntimeSupport=all
+            MATCH (comm:Community)
+            WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
+            WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
+            WHERE score > $min_score
+            RETURN
+                comm.uuid As uuid,
+                comm.group_id AS group_id,
+                comm.name AS name, 
+                comm.name_embedding AS name_embedding,
+                comm.created_at AS created_at, 
+                comm.summary AS summary
+            ORDER BY score DESC
+            LIMIT $limit
+            """,
+        search_vector=search_vector,
+        group_ids=group_ids,
+        limit=limit,
+        min_score=min_score,
+        database_=DEFAULT_DATABASE,
+        routing_='r',
+    )
     communities = [get_community_node_from_record(record) for record in records]
 
     return communities
@@ -554,32 +619,27 @@ async def node_distance_reranker(
     driver: AsyncDriver, node_uuids: list[str], center_node_uuid: str
 ) -> list[str]:
     # filter out node_uuid center node node uuid
-    filtered_uuids = list(filter(lambda uuid: uuid != center_node_uuid, node_uuids))
+    filtered_uuids = list(filter(lambda node_uuid: node_uuid != center_node_uuid, node_uuids))
     scores: dict[str, float] = {}
 
     # Find the shortest path to center node
     query = Query("""
-        MATCH p = SHORTEST 1 (center:Entity {uuid: $center_uuid})-[:RELATES_TO]-+(n:Entity {uuid: $node_uuid})
-        RETURN length(p) AS score
+        UNWIND $node_uuids AS node_uuid
+        MATCH p = SHORTEST 1 (center:Entity {uuid: $center_uuid})-[:RELATES_TO]-+(n:Entity {uuid: node_uuid})
+        RETURN length(p) AS score, node_uuid AS uuid
         """)
 
-    path_results = await asyncio.gather(
-        *[
-            driver.execute_query(
-                query,
-                node_uuid=uuid,
-                center_uuid=center_node_uuid,
-                database_=DEFAULT_DATABASE,
-            )
-            for uuid in filtered_uuids
-        ]
+    path_results, _, _ = await driver.execute_query(
+        query,
+        node_uuids=filtered_uuids,
+        center_uuid=center_node_uuid,
+        database_=DEFAULT_DATABASE,
     )
 
-    for uuid, result in zip(filtered_uuids, path_results):
-        records = result[0]
-        record = records[0] if len(records) > 0 else None
-        distance: float = record['score'] if record is not None else float('inf')
-        scores[uuid] = distance
+    for result in path_results:
+        uuid = result['uuid']
+        score = result['score'] if 'score' in result else float('inf')
+        scores[uuid] = score
 
     # rerank on shortest distance
     filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])
@@ -596,25 +656,20 @@ async def episode_mentions_reranker(driver: AsyncDriver, node_uuids: list[list[s
     scores: dict[str, float] = {}
 
     # Find the shortest path to center node
-    query = Query("""  
-        MATCH (episode:Episodic)-[r:MENTIONS]->(n:Entity {uuid: $node_uuid})
-        RETURN count(*) AS score
+    query = Query("""
+        UNWIND $node_uuids AS node_uuid 
+        MATCH (episode:Episodic)-[r:MENTIONS]->(n:Entity {uuid: node_uuid})
+        RETURN count(*) AS score, n.uuid AS uuid
         """)
 
-    result_scores = await asyncio.gather(
-        *[
-            driver.execute_query(
-                query,
-                node_uuid=uuid,
-                database_=DEFAULT_DATABASE,
-            )
-            for uuid in sorted_uuids
-        ]
+    results, _, _ = await driver.execute_query(
+        query,
+        node_uuids=sorted_uuids,
+        database_=DEFAULT_DATABASE,
     )
 
-    for uuid, result in zip(sorted_uuids, result_scores):
-        record = result[0][0]
-        scores[uuid] = record['score']
+    for result in results:
+        scores[result['uuid']] = result['score']
 
     # rerank on shortest distance
     sorted_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.3.16
+Version: 0.3.17
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -17,7 +17,6 @@ Requires-Dist: numpy (>=1.0.0)
 Requires-Dist: openai (>=1.50.2,<2.0.0)
 Requires-Dist: pydantic (>=2.8.2,<3.0.0)
 Requires-Dist: tenacity (<9.0.0)
-Requires-Dist: voyageai (>=0.2.3,<0.3.0)
 Description-Content-Type: text/markdown
 
 <div align="center">

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

@@ -1,11 +1,15 @@
 graphiti_core/__init__.py,sha256=e5SWFkRiaUwfprYIeIgVIh7JDedNiloZvd3roU-0aDY,55
+graphiti_core/cross_encoder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+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/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/errors.py,sha256=ddHrHGQxhwkVAtSph4AV84UoOlgwZufMczXPwB7uqPo,1795
-graphiti_core/graphiti.py,sha256=BBYuSDgGj8FZKm6ldNntn8Dv7jFccFSZK1_kTDZNUQE,26945
+graphiti_core/graphiti.py,sha256=c9Rh777TrHYffPF6qvFAfm-m-PA4kD8a3ZW_ShsZGxE,27714
 graphiti_core/helpers.py,sha256=kqC2TD8Auwty4sG7KH4BuRMX413oTChGaAT_XUt9ZjU,2108
 graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
 graphiti_core/llm_client/anthropic_client.py,sha256=4l2PbCjIoeRr7UJ2DUh2grYLTtE2vNaWlo72IIRQDeI,2405
@@ -34,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=2YvUjOWnPYVa2DvZAqOuKbcdCxjX5bSynpQICqFqqGU,9040
-graphiti_core/search/search_config.py,sha256=dWcanEmMoL42RHF-jcZO9C2G9BdqjkI9w-5xe9Wd2Xg,2737
-graphiti_core/search/search_config_recipes.py,sha256=FpASZLdyMdTSwY4ISHrjRUnFKVCego7Wd3j5RPN-ris,4907
-graphiti_core/search/search_utils.py,sha256=4OChWhtJXAtiOUeyZ3AoEWROY5JKJzEi-TzhlkZZfoo,21020
+graphiti_core/search/search.py,sha256=jhYIb7ERK-zGT9N0M8wXjIdFMyU0pUKzpnHmmBxdG3A,10710
+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/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
@@ -47,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.16.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.3.16.dist-info/METADATA,sha256=CXX1YrYZICJQzNzufvTOCMIecS-i8I9ZnpGiuRiRhio,9437
-graphiti_core-0.3.16.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-graphiti_core-0.3.16.dist-info/RECORD,,
+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,,