graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/cross_encoder/bge_reranker_client.py

@@ -15,8 +15,18 @@ limitations under the License.
 """
 
 import asyncio
-
-from sentence_transformers import CrossEncoder
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sentence_transformers import CrossEncoder
+else:
+    try:
+        from sentence_transformers import CrossEncoder
+    except ImportError:
+        raise ImportError(
+            'sentence-transformers is required for BGERerankerClient. '
+            'Install it with: pip install graphiti-core[sentence-transformers]'
+        ) from None
 
 from graphiti_core.cross_encoder.client import CrossEncoderClient
 

graphiti_core/cross_encoder/gemini_reranker_client.py

@@ -0,0 +1,161 @@
+"""
+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 logging
+import re
+from typing import TYPE_CHECKING
+
+from ..helpers import semaphore_gather
+from ..llm_client import LLMConfig, RateLimitError
+from .client import CrossEncoderClient
+
+if TYPE_CHECKING:
+    from google import genai
+    from google.genai import types
+else:
+    try:
+        from google import genai
+        from google.genai import types
+    except ImportError:
+        raise ImportError(
+            'google-genai is required for GeminiRerankerClient. '
+            'Install it with: pip install graphiti-core[google-genai]'
+        ) from None
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gemini-2.5-flash-lite'
+
+
+class GeminiRerankerClient(CrossEncoderClient):
+    """
+    Google Gemini Reranker Client
+    """
+
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        client: 'genai.Client | None' = None,
+    ):
+        """
+        Initialize the GeminiRerankerClient with the provided configuration and client.
+
+        The Gemini Developer API does not yet support logprobs. Unlike the OpenAI reranker,
+        this reranker uses the Gemini API to perform direct relevance scoring of passages.
+        Each passage is scored individually on a 0-100 scale.
+
+        Args:
+            config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
+            client (genai.Client | None): An optional async client instance to use. If not provided, a new genai.Client is created.
+        """
+        if config is None:
+            config = LLMConfig()
+
+        self.config = config
+        if client is None:
+            self.client = genai.Client(api_key=config.api_key)
+        else:
+            self.client = client
+
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
+        """
+        Rank passages based on their relevance to the query using direct scoring.
+
+        Each passage is scored individually on a 0-100 scale, then normalized to [0,1].
+        """
+        if len(passages) <= 1:
+            return [(passage, 1.0) for passage in passages]
+
+        # Generate scoring prompts for each passage
+        scoring_prompts = []
+        for passage in passages:
+            prompt = f"""Rate how well this passage answers or relates to the query. Use a scale from 0 to 100.
+
+Query: {query}
+
+Passage: {passage}
+
+Provide only a number between 0 and 100 (no explanation, just the number):"""
+
+            scoring_prompts.append(
+                [
+                    types.Content(
+                        role='user',
+                        parts=[types.Part.from_text(text=prompt)],
+                    ),
+                ]
+            )
+
+        try:
+            # Execute all scoring requests concurrently - O(n) API calls
+            responses = await semaphore_gather(
+                *[
+                    self.client.aio.models.generate_content(
+                        model=self.config.model or DEFAULT_MODEL,
+                        contents=prompt_messages,  # type: ignore
+                        config=types.GenerateContentConfig(
+                            system_instruction='You are an expert at rating passage relevance. Respond with only a number from 0-100.',
+                            temperature=0.0,
+                            max_output_tokens=3,
+                        ),
+                    )
+                    for prompt_messages in scoring_prompts
+                ]
+            )
+
+            # Extract scores and create results
+            results = []
+            for passage, response in zip(passages, responses, strict=True):
+                try:
+                    if hasattr(response, 'text') and response.text:
+                        # Extract numeric score from response
+                        score_text = response.text.strip()
+                        # Handle cases where model might return non-numeric text
+                        score_match = re.search(r'\b(\d{1,3})\b', score_text)
+                        if score_match:
+                            score = float(score_match.group(1))
+                            # Normalize to [0, 1] range and clamp to valid range
+                            normalized_score = max(0.0, min(1.0, score / 100.0))
+                            results.append((passage, normalized_score))
+                        else:
+                            logger.warning(
+                                f'Could not extract numeric score from response: {score_text}'
+                            )
+                            results.append((passage, 0.0))
+                    else:
+                        logger.warning('Empty response from Gemini for passage scoring')
+                        results.append((passage, 0.0))
+                except (ValueError, AttributeError) as e:
+                    logger.warning(f'Error parsing score from Gemini response: {e}')
+                    results.append((passage, 0.0))
+
+            # Sort by score in descending order (highest relevance first)
+            results.sort(reverse=True, key=lambda x: x[1])
+            return results
+
+        except Exception as e:
+            # Check if it's a rate limit error based on Gemini API error codes
+            error_message = str(e).lower()
+            if (
+                'rate limit' in error_message
+                or 'quota' in error_message
+                or 'resource_exhausted' in error_message
+                or '429' in str(e)
+            ):
+                raise RateLimitError from e
+
+            logger.error(f'Error in generating LLM response: {e}')
+            raise

graphiti_core/cross_encoder/openai_reranker_client.py

@@ -22,7 +22,7 @@ import openai
 from openai import AsyncAzureOpenAI, AsyncOpenAI
 
 from ..helpers import semaphore_gather
-from ..llm_client import LLMConfig, RateLimitError
+from ..llm_client import LLMConfig, OpenAIClient, RateLimitError
 from ..prompts import Message
 from .client import CrossEncoderClient
 
@@ -35,7 +35,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
     def __init__(
         self,
         config: LLMConfig | None = None,
-        client: AsyncOpenAI | AsyncAzureOpenAI | None = None,
+        client: AsyncOpenAI | AsyncAzureOpenAI | OpenAIClient | None = None,
     ):
         """
         Initialize the OpenAIRerankerClient with the provided configuration and client.
@@ -45,7 +45,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
 
         Args:
             config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
-            client (AsyncOpenAI | AsyncAzureOpenAI | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
+            client (AsyncOpenAI | AsyncAzureOpenAI | OpenAIClient | None): An optional async client instance to use. If not provided, a new AsyncOpenAI client is created.
         """
         if config is None:
             config = LLMConfig()
@@ -53,6 +53,8 @@ class OpenAIRerankerClient(CrossEncoderClient):
         self.config = config
         if client is None:
             self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
+        elif isinstance(client, OpenAIClient):
+            self.client = client.client
         else:
             self.client = client
 
@@ -82,7 +84,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
             responses = await semaphore_gather(
                 *[
                     self.client.chat.completions.create(
-                        model=DEFAULT_MODEL,
+                        model=self.config.model or DEFAULT_MODEL,
                         messages=openai_messages,
                         temperature=0,
                         max_tokens=1,
@@ -106,7 +108,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
                 if len(top_logprobs) == 0:
                     continue
                 norm_logprobs = np.exp(top_logprobs[0].logprob)
-                if bool(top_logprobs[0].token):
+                if top_logprobs[0].token.strip().split(' ')[0].lower() == 'true':
                     scores.append(norm_logprobs)
                 else:
                     scores.append(1 - norm_logprobs)

graphiti_core/decorators.py

@@ -0,0 +1,110 @@
+"""
+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 functools
+import inspect
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.helpers import semaphore_gather
+from graphiti_core.search.search_config import SearchResults
+
+F = TypeVar('F', bound=Callable[..., Awaitable[Any]])
+
+
+def handle_multiple_group_ids(func: F) -> F:
+    """
+    Decorator for FalkorDB methods that need to handle multiple group_ids.
+    Runs the function for each group_id separately and merges results.
+    """
+
+    @functools.wraps(func)
+    async def wrapper(self, *args, **kwargs):
+        group_ids_func_pos = get_parameter_position(func, 'group_ids')
+        group_ids_pos = (
+            group_ids_func_pos - 1 if group_ids_func_pos is not None else None
+        )  # Adjust for zero-based index
+        group_ids = kwargs.get('group_ids')
+
+        # If not in kwargs and position exists, get from args
+        if group_ids is None and group_ids_pos is not None and len(args) > group_ids_pos:
+            group_ids = args[group_ids_pos]
+
+        # Only handle FalkorDB with multiple group_ids
+        if (
+            hasattr(self, 'clients')
+            and hasattr(self.clients, 'driver')
+            and self.clients.driver.provider == GraphProvider.FALKORDB
+            and group_ids
+            and len(group_ids) > 1
+        ):
+            # Execute for each group_id concurrently
+            driver = self.clients.driver
+
+            async def execute_for_group(gid: str):
+                # Remove group_ids from args if it was passed positionally
+                filtered_args = list(args)
+                if group_ids_pos is not None and len(args) > group_ids_pos:
+                    filtered_args.pop(group_ids_pos)
+
+                return await func(
+                    self,
+                    *filtered_args,
+                    **{**kwargs, 'group_ids': [gid], 'driver': driver.clone(database=gid)},
+                )
+
+            results = await semaphore_gather(
+                *[execute_for_group(gid) for gid in group_ids],
+                max_coroutines=getattr(self, 'max_coroutines', None),
+            )
+
+            # Merge results based on type
+            if isinstance(results[0], SearchResults):
+                return SearchResults.merge(results)
+            elif isinstance(results[0], list):
+                return [item for result in results for item in result]
+            elif isinstance(results[0], tuple):
+                # Handle tuple outputs (like build_communities returning (nodes, edges))
+                merged_tuple = []
+                for i in range(len(results[0])):
+                    component_results = [result[i] for result in results]
+                    if isinstance(component_results[0], list):
+                        merged_tuple.append(
+                            [item for component in component_results for item in component]
+                        )
+                    else:
+                        merged_tuple.append(component_results)
+                return tuple(merged_tuple)
+            else:
+                return results
+
+        # Normal execution
+        return await func(self, *args, **kwargs)
+
+    return wrapper  # type: ignore
+
+
+def get_parameter_position(func: Callable, param_name: str) -> int | None:
+    """
+    Returns the positional index of a parameter in the function signature.
+    If the parameter is not found, returns None.
+    """
+    sig = inspect.signature(func)
+    for idx, (name, _param) in enumerate(sig.parameters.items()):
+        if name == param_name:
+            return idx
+    return None

graphiti_core/driver/__init__.py

@@ -0,0 +1,19 @@
+"""
+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 neo4j import Neo4jDriver
+
+__all__ = ['Neo4jDriver']

graphiti_core/driver/driver.py

@@ -0,0 +1,124 @@
+"""
+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 copy
+import logging
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from enum import Enum
+from typing import Any
+
+from dotenv import load_dotenv
+
+from graphiti_core.driver.graph_operations.graph_operations import GraphOperationsInterface
+from graphiti_core.driver.search_interface.search_interface import SearchInterface
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SIZE = 10
+
+load_dotenv()
+
+ENTITY_INDEX_NAME = os.environ.get('ENTITY_INDEX_NAME', 'entities')
+EPISODE_INDEX_NAME = os.environ.get('EPISODE_INDEX_NAME', 'episodes')
+COMMUNITY_INDEX_NAME = os.environ.get('COMMUNITY_INDEX_NAME', 'communities')
+ENTITY_EDGE_INDEX_NAME = os.environ.get('ENTITY_EDGE_INDEX_NAME', 'entity_edges')
+
+
+class GraphProvider(Enum):
+    NEO4J = 'neo4j'
+    FALKORDB = 'falkordb'
+    KUZU = 'kuzu'
+    NEPTUNE = 'neptune'
+
+
+class GraphDriverSession(ABC):
+    provider: GraphProvider
+
+    async def __aenter__(self):
+        return self
+
+    @abstractmethod
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Falkor, but method must exist
+        pass
+
+    @abstractmethod
+    async def run(self, query: str, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def execute_write(self, func, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class GraphDriver(ABC):
+    provider: GraphProvider
+    fulltext_syntax: str = (
+        ''  # Neo4j (default) syntax does not require a prefix for fulltext queries
+    )
+    _database: str
+    default_group_id: str = ''
+    search_interface: SearchInterface | None = None
+    graph_operations_interface: GraphOperationsInterface | None = None
+
+    @abstractmethod
+    def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def session(self, database: str | None = None) -> GraphDriverSession:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def delete_all_indexes(self) -> Coroutine:
+        raise NotImplementedError()
+
+    def with_database(self, database: str) -> 'GraphDriver':
+        """
+        Returns a shallow copy of this driver with a different default database.
+        Reuses the same connection (e.g. FalkorDB, Neo4j).
+        """
+        cloned = copy.copy(self)
+        cloned._database = database
+
+        return cloned
+
+    @abstractmethod
+    async def build_indices_and_constraints(self, delete_existing: bool = False):
+        raise NotImplementedError()
+
+    def clone(self, database: str) -> 'GraphDriver':
+        """Clone the driver with a different database or graph name."""
+        return self
+
+    def build_fulltext_query(
+        self, query: str, group_ids: list[str] | None = None, max_query_length: int = 128
+    ) -> str:
+        """
+        Specific fulltext query builder for database providers.
+        Only implemented by providers that need custom fulltext query building.
+        """
+        raise NotImplementedError(f'build_fulltext_query not implemented for {self.provider}')