graphiti-core 0.3.15__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/edges.py

@@ -54,7 +54,7 @@ class Edge(BaseModel, ABC):
         DELETE e
         """,
             uuid=self.uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Deleted Edge: {self.uuid}')
@@ -82,7 +82,7 @@ class EpisodicEdge(Edge):
             uuid=self.uuid,
             group_id=self.group_id,
             created_at=self.created_at,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to neo4j: {self.uuid}')
@@ -102,7 +102,7 @@ class EpisodicEdge(Edge):
             e.created_at AS created_at
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_episodic_edge_from_record(record) for record in records]
@@ -125,7 +125,7 @@ class EpisodicEdge(Edge):
             e.created_at AS created_at
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_episodic_edge_from_record(record) for record in records]
@@ -148,7 +148,7 @@ class EpisodicEdge(Edge):
             e.created_at AS created_at
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_episodic_edge_from_record(record) for record in records]
@@ -202,7 +202,7 @@ class EntityEdge(Edge):
             expired_at=self.expired_at,
             valid_at=self.valid_at,
             invalid_at=self.invalid_at,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to neo4j: {self.uuid}')
@@ -229,7 +229,7 @@ class EntityEdge(Edge):
             e.invalid_at AS invalid_at
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_entity_edge_from_record(record) for record in records]
@@ -259,7 +259,7 @@ class EntityEdge(Edge):
             e.invalid_at AS invalid_at
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_entity_edge_from_record(record) for record in records]
@@ -289,7 +289,7 @@ class EntityEdge(Edge):
             e.invalid_at AS invalid_at
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_entity_edge_from_record(record) for record in records]
@@ -308,7 +308,7 @@ class CommunityEdge(Edge):
             uuid=self.uuid,
             group_id=self.group_id,
             created_at=self.created_at,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved edge to neo4j: {self.uuid}')
@@ -328,7 +328,7 @@ class CommunityEdge(Edge):
             e.created_at AS created_at
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_community_edge_from_record(record) for record in records]
@@ -349,7 +349,7 @@ class CommunityEdge(Edge):
             e.created_at AS created_at
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_community_edge_from_record(record) for record in records]
@@ -370,7 +370,7 @@ class CommunityEdge(Edge):
             e.created_at AS created_at
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         edges = [get_community_edge_from_record(record) for record in records]

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/nodes.py

@@ -90,7 +90,7 @@ class Node(BaseModel, ABC):
         DETACH DELETE n
         """,
             uuid=self.uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Deleted Node: {self.uuid}')
@@ -136,7 +136,7 @@ class EpisodicNode(Node):
             created_at=self.created_at,
             valid_at=self.valid_at,
             source=self.source.value,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved Node to neo4j: {self.uuid}')
@@ -158,7 +158,7 @@ class EpisodicNode(Node):
             e.source AS source
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         episodes = [get_episodic_node_from_record(record) for record in records]
@@ -184,7 +184,7 @@ class EpisodicNode(Node):
             e.source AS source
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         episodes = [get_episodic_node_from_record(record) for record in records]
@@ -207,7 +207,7 @@ class EpisodicNode(Node):
             e.source AS source
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         episodes = [get_episodic_node_from_record(record) for record in records]
@@ -237,7 +237,7 @@ class EntityNode(Node):
             summary=self.summary,
             name_embedding=self.name_embedding,
             created_at=self.created_at,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved Node to neo4j: {self.uuid}')
@@ -258,7 +258,7 @@ class EntityNode(Node):
             n.summary AS summary
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         nodes = [get_entity_node_from_record(record) for record in records]
@@ -282,7 +282,7 @@ class EntityNode(Node):
             n.summary AS summary
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         nodes = [get_entity_node_from_record(record) for record in records]
@@ -303,7 +303,7 @@ class EntityNode(Node):
             n.summary AS summary
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         nodes = [get_entity_node_from_record(record) for record in records]
@@ -324,7 +324,7 @@ class CommunityNode(Node):
             summary=self.summary,
             name_embedding=self.name_embedding,
             created_at=self.created_at,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         logger.debug(f'Saved Node to neo4j: {self.uuid}')
@@ -354,7 +354,7 @@ class CommunityNode(Node):
             n.summary AS summary
         """,
             uuid=uuid,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         nodes = [get_community_node_from_record(record) for record in records]
@@ -378,7 +378,7 @@ class CommunityNode(Node):
             n.summary AS summary
         """,
             uuids=uuids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         communities = [get_community_node_from_record(record) for record in records]
@@ -399,7 +399,7 @@ class CommunityNode(Node):
             n.summary AS summary
         """,
             group_ids=group_ids,
-            _database=DEFAULT_DATABASE,
+            database_=DEFAULT_DATABASE,
         )
 
         communities = [get_community_node_from_record(record) for record in records]