graphiti-core 0.12.4__py3-none-any.whl → 0.13.1__py3-none-any.whl

graphiti_core/errors.py

@@ -73,3 +73,11 @@ class EntityTypeValidationError(GraphitiError):
     def __init__(self, entity_type: str, entity_type_attribute: str):
         self.message = f'{entity_type_attribute} cannot be used as an attribute for {entity_type} as it is a protected attribute name.'
         super().__init__(self.message)
+
+
+class GroupIdValidationError(GraphitiError):
+    """Raised when a group_id contains invalid characters."""
+
+    def __init__(self, group_id: str):
+        self.message = f'group_id "{group_id}" must contain only alphanumeric characters, dashes, or underscores'
+        super().__init__(self.message)

graphiti_core/graphiti.py

@@ -29,7 +29,7 @@ from graphiti_core.driver.neo4j_driver import Neo4jDriver
 from graphiti_core.edges import EntityEdge, EpisodicEdge
 from graphiti_core.embedder import EmbedderClient, OpenAIEmbedder
 from graphiti_core.graphiti_types import GraphitiClients
-from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather, validate_group_id
 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
@@ -103,6 +103,7 @@ class Graphiti:
         cross_encoder: CrossEncoderClient | None = None,
         store_raw_episode_content: bool = True,
         graph_driver: GraphDriver | None = None,
+        max_coroutines: int | None = None,
     ):
         """
         Initialize a Graphiti instance.
@@ -121,6 +122,20 @@ class Graphiti:
         llm_client : LLMClient | None, optional
             An instance of LLMClient for natural language processing tasks.
             If not provided, a default OpenAIClient will be initialized.
+        embedder : EmbedderClient | None, optional
+            An instance of EmbedderClient for embedding tasks.
+            If not provided, a default OpenAIEmbedder will be initialized.
+        cross_encoder : CrossEncoderClient | None, optional
+            An instance of CrossEncoderClient for reranking tasks.
+            If not provided, a default OpenAIRerankerClient will be initialized.
+        store_raw_episode_content : bool, optional
+            Whether to store the raw content of episodes. Defaults to True.
+        graph_driver : GraphDriver | None, optional
+            An instance of GraphDriver for database operations.
+            If not provided, a default Neo4jDriver will be initialized.
+        max_coroutines : int | None, optional
+            The maximum number of concurrent operations allowed. Overrides SEMAPHORE_LIMIT set in the environment.
+            If not set, the Graphiti default is used.
 
         Returns
         -------
@@ -145,6 +160,7 @@ class Graphiti:
 
         self.database = DEFAULT_DATABASE
         self.store_raw_episode_content = store_raw_episode_content
+        self.max_coroutines = max_coroutines
         if llm_client:
             self.llm_client = llm_client
         else:
@@ -335,6 +351,7 @@ class Graphiti:
             now = utc_now()
 
             validate_entity_types(entity_types)
+            validate_group_id(group_id)
 
             previous_episodes = (
                 await self.retrieve_episodes(
@@ -393,6 +410,7 @@ class Graphiti:
                     group_id,
                     edge_types,
                 ),
+                max_coroutines=self.max_coroutines,
             )
 
             edges = resolve_edge_pointers(extracted_edges, uuid_map)
@@ -409,6 +427,7 @@ class Graphiti:
                 extract_attributes_from_nodes(
                     self.clients, nodes, episode, previous_episodes, entity_types
                 ),
+                max_coroutines=self.max_coroutines,
             )
 
             duplicate_of_edges = build_duplicate_of_edges(episode, now, node_duplicates)
@@ -432,7 +451,8 @@ class Graphiti:
                     *[
                         update_community(self.driver, self.llm_client, self.embedder, node)
                         for node in nodes
-                    ]
+                    ],
+                    max_coroutines=self.max_coroutines,
                 )
             end = time()
             logger.info(f'Completed add_episode in {(end - start) * 1000} ms')
@@ -484,6 +504,8 @@ class Graphiti:
             start = time()
             now = utc_now()
 
+            validate_group_id(group_id)
+
             episodes = [
                 EpisodicNode(
                     name=episode.name,
@@ -499,7 +521,10 @@ class Graphiti:
             ]
 
             # Save all the episodes
-            await semaphore_gather(*[episode.save(self.driver) for episode in episodes])
+            await semaphore_gather(
+                *[episode.save(self.driver) for episode in episodes],
+                max_coroutines=self.max_coroutines,
+            )
 
             # Get previous episode context for each episode
             episode_pairs = await retrieve_previous_episodes_bulk(self.driver, episodes)
@@ -515,16 +540,21 @@ class Graphiti:
             await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
                 *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
+                max_coroutines=self.max_coroutines,
             )
 
             # Dedupe extracted nodes, compress extracted edges
             (nodes, uuid_map), extracted_edges_timestamped = await semaphore_gather(
                 dedupe_nodes_bulk(self.driver, self.llm_client, extracted_nodes),
                 extract_edge_dates_bulk(self.llm_client, extracted_edges, episode_pairs),
+                max_coroutines=self.max_coroutines,
             )
 
             # save nodes to KG
-            await semaphore_gather(*[node.save(self.driver) for node in nodes])
+            await semaphore_gather(
+                *[node.save(self.driver) for node in nodes],
+                max_coroutines=self.max_coroutines,
+            )
 
             # re-map edge pointers so that they don't point to discard dupe nodes
             extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
@@ -536,7 +566,8 @@ class Graphiti:
 
             # save episodic edges to KG
             await semaphore_gather(
-                *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers]
+                *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers],
+                max_coroutines=self.max_coroutines,
             )
 
             # Dedupe extracted edges
@@ -548,7 +579,10 @@ class Graphiti:
             # invalidate edges
 
             # save edges to KG
-            await semaphore_gather(*[edge.save(self.driver) for edge in edges])
+            await semaphore_gather(
+                *[edge.save(self.driver) for edge in edges],
+                max_coroutines=self.max_coroutines,
+            )
 
             end = time()
             logger.info(f'Completed add_episode_bulk in {(end - start) * 1000} ms')
@@ -572,11 +606,18 @@ class Graphiti:
         )
 
         await semaphore_gather(
-            *[node.generate_name_embedding(self.embedder) for node in community_nodes]
+            *[node.generate_name_embedding(self.embedder) for node in community_nodes],
+            max_coroutines=self.max_coroutines,
         )
 
-        await semaphore_gather(*[node.save(self.driver) for node in community_nodes])
-        await semaphore_gather(*[edge.save(self.driver) for edge in community_edges])
+        await semaphore_gather(
+            *[node.save(self.driver) for node in community_nodes],
+            max_coroutines=self.max_coroutines,
+        )
+        await semaphore_gather(
+            *[edge.save(self.driver) for edge in community_edges],
+            max_coroutines=self.max_coroutines,
+        )
 
         return community_nodes
 
@@ -683,7 +724,8 @@ class Graphiti:
         episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
         edges_list = await semaphore_gather(
-            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
+            *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes],
+            max_coroutines=self.max_coroutines,
         )
 
         edges: list[EntityEdge] = [edge for lst in edges_list for edge in lst]
@@ -759,6 +801,12 @@ class Graphiti:
                 if record['episode_count'] == 1:
                     nodes_to_delete.append(node)
 
-        await semaphore_gather(*[node.delete(self.driver) for node in nodes_to_delete])
-        await semaphore_gather(*[edge.delete(self.driver) for edge in edges_to_delete])
+        await semaphore_gather(
+            *[node.delete(self.driver) for node in nodes_to_delete],
+            max_coroutines=self.max_coroutines,
+        )
+        await semaphore_gather(
+            *[edge.delete(self.driver) for edge in edges_to_delete],
+            max_coroutines=self.max_coroutines,
+        )
         await episode.delete(self.driver)

graphiti_core/helpers.py

@@ -16,6 +16,7 @@ limitations under the License.
 
 import asyncio
 import os
+import re
 from collections.abc import Coroutine
 from datetime import datetime
 
@@ -25,6 +26,8 @@ from neo4j import time as neo4j_time
 from numpy._typing import NDArray
 from typing_extensions import LiteralString
 
+from graphiti_core.errors import GroupIdValidationError
+
 load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', 'neo4j')
@@ -94,12 +97,38 @@ def normalize_l2(embedding: list[float]) -> NDArray:
 # Use this instead of asyncio.gather() to bound coroutines
 async def semaphore_gather(
     *coroutines: Coroutine,
-    max_coroutines: int = SEMAPHORE_LIMIT,
+    max_coroutines: int | None = None,
 ):
-    semaphore = asyncio.Semaphore(max_coroutines)
+    semaphore = asyncio.Semaphore(max_coroutines or SEMAPHORE_LIMIT)
 
     async def _wrap_coroutine(coroutine):
         async with semaphore:
             return await coroutine
 
     return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines))
+
+
+def validate_group_id(group_id: str) -> bool:
+    """
+    Validate that a group_id contains only ASCII alphanumeric characters, dashes, and underscores.
+
+    Args:
+        group_id: The group_id to validate
+
+    Returns:
+        True if valid, False otherwise
+
+    Raises:
+        GroupIdValidationError: If group_id contains invalid characters
+    """
+
+    # Allow empty string (default case)
+    if not group_id:
+        return True
+
+    # Check if string contains only ASCII alphanumeric characters, dashes, or underscores
+    # Pattern matches: letters (a-z, A-Z), digits (0-9), hyphens (-), and underscores (_)
+    if not re.match(r'^[a-zA-Z0-9_-]+$', group_id):
+        raise GroupIdValidationError(group_id)
+
+    return True

graphiti_core/llm_client/azure_openai_client.py

@@ -14,60 +14,64 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
 import logging
-from typing import Any
+from typing import ClassVar
 
 from openai import AsyncAzureOpenAI
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
-from ..prompts.models import Message
-from .client import LLMClient
-from .config import LLMConfig, ModelSize
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .openai_base_client import BaseOpenAIClient
 
 logger = logging.getLogger(__name__)
 
 
-class AzureOpenAILLMClient(LLMClient):
+class AzureOpenAILLMClient(BaseOpenAIClient):
     """Wrapper class for AsyncAzureOpenAI that implements the LLMClient interface."""
 
-    def __init__(self, azure_client: AsyncAzureOpenAI, config: LLMConfig | None = None):
-        super().__init__(config, cache=False)
-        self.azure_client = azure_client
+    # Class-level constants
+    MAX_RETRIES: ClassVar[int] = 2
 
-    async def _generate_response(
+    def __init__(
         self,
-        messages: list[Message],
-        response_model: type[BaseModel] | None = None,
-        max_tokens: int = 1024,
-        model_size: ModelSize = ModelSize.medium,
-    ) -> dict[str, Any]:
-        """Generate response using Azure OpenAI client."""
-        # Convert messages to OpenAI format
-        openai_messages: list[ChatCompletionMessageParam] = []
-        for message in messages:
-            message.content = self._clean_input(message.content)
-            if message.role == 'user':
-                openai_messages.append({'role': 'user', 'content': message.content})
-            elif message.role == 'system':
-                openai_messages.append({'role': 'system', 'content': message.content})
-
-        # Ensure model is a string
-        model_name = self.model if self.model else 'gpt-4o-mini'
+        azure_client: AsyncAzureOpenAI,
+        config: LLMConfig | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        super().__init__(config, cache=False, max_tokens=max_tokens)
+        self.client = azure_client
 
-        try:
-            response = await self.azure_client.chat.completions.create(
-                model=model_name,
-                messages=openai_messages,
-                temperature=float(self.temperature) if self.temperature is not None else 0.7,
-                max_tokens=max_tokens,
-                response_format={'type': 'json_object'},
-            )
-            result = response.choices[0].message.content or '{}'
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ):
+        """Create a structured completion using Azure OpenAI's beta parse API."""
+        return await self.client.beta.chat.completions.parse(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format=response_model,  # type: ignore
+        )
 
-            # Parse JSON response
-            return json.loads(result)
-        except Exception as e:
-            logger.error(f'Error in Azure OpenAI LLM response: {e}')
-            raise
+    async def _create_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel] | None = None,
+    ):
+        """Create a regular completion with JSON format using Azure OpenAI."""
+        return await self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format={'type': 'json_object'},
+        )

graphiti_core/llm_client/openai_base_client.py

@@ -0,0 +1,217 @@
+"""
+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 json
+import logging
+import typing
+from abc import abstractmethod
+from typing import Any, ClassVar
+
+import openai
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from ..prompts.models import Message
+from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
+from .errors import RateLimitError, RefusalError
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gpt-4.1-mini'
+DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
+
+
+class BaseOpenAIClient(LLMClient):
+    """
+    Base client class for OpenAI-compatible APIs (OpenAI and Azure OpenAI).
+
+    This class contains shared logic for both OpenAI and Azure OpenAI clients,
+    reducing code duplication while allowing for implementation-specific differences.
+    """
+
+    # Class-level constants
+    MAX_RETRIES: ClassVar[int] = 2
+
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        cache: bool = False,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ):
+        if cache:
+            raise NotImplementedError('Caching is not implemented for OpenAI-based clients')
+
+        if config is None:
+            config = LLMConfig()
+
+        super().__init__(config, cache)
+        self.max_tokens = max_tokens
+
+    @abstractmethod
+    async def _create_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel] | None = None,
+    ) -> Any:
+        """Create a completion using the specific client implementation."""
+        pass
+
+    @abstractmethod
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ) -> Any:
+        """Create a structured completion using the specific client implementation."""
+        pass
+
+    def _convert_messages_to_openai_format(
+        self, messages: list[Message]
+    ) -> list[ChatCompletionMessageParam]:
+        """Convert internal Message format to OpenAI ChatCompletionMessageParam format."""
+        openai_messages: list[ChatCompletionMessageParam] = []
+        for m in messages:
+            m.content = self._clean_input(m.content)
+            if m.role == 'user':
+                openai_messages.append({'role': 'user', 'content': m.content})
+            elif m.role == 'system':
+                openai_messages.append({'role': 'system', 'content': m.content})
+        return openai_messages
+
+    def _get_model_for_size(self, model_size: ModelSize) -> str:
+        """Get the appropriate model name based on the requested size."""
+        if model_size == ModelSize.small:
+            return self.small_model or DEFAULT_SMALL_MODEL
+        else:
+            return self.model or DEFAULT_MODEL
+
+    def _handle_structured_response(self, response: Any) -> dict[str, Any]:
+        """Handle structured response parsing and validation."""
+        response_object = response.choices[0].message
+
+        if response_object.parsed:
+            return response_object.parsed.model_dump()
+        elif response_object.refusal:
+            raise RefusalError(response_object.refusal)
+        else:
+            raise Exception(f'Invalid response from LLM: {response_object.model_dump()}')
+
+    def _handle_json_response(self, response: Any) -> dict[str, Any]:
+        """Handle JSON response parsing."""
+        result = response.choices[0].message.content or '{}'
+        return json.loads(result)
+
+    async def _generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        model_size: ModelSize = ModelSize.medium,
+    ) -> dict[str, Any]:
+        """Generate a response using the appropriate client implementation."""
+        openai_messages = self._convert_messages_to_openai_format(messages)
+        model = self._get_model_for_size(model_size)
+
+        try:
+            if response_model:
+                response = await self._create_structured_completion(
+                    model=model,
+                    messages=openai_messages,
+                    temperature=self.temperature,
+                    max_tokens=max_tokens or self.max_tokens,
+                    response_model=response_model,
+                )
+                return self._handle_structured_response(response)
+            else:
+                response = await self._create_completion(
+                    model=model,
+                    messages=openai_messages,
+                    temperature=self.temperature,
+                    max_tokens=max_tokens or self.max_tokens,
+                )
+                return self._handle_json_response(response)
+
+        except openai.LengthFinishReasonError as e:
+            raise Exception(f'Output length exceeded max tokens {self.max_tokens}: {e}') from e
+        except openai.RateLimitError as e:
+            raise RateLimitError from e
+        except Exception as e:
+            logger.error(f'Error in generating LLM response: {e}')
+            raise
+
+    async def generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int | None = None,
+        model_size: ModelSize = ModelSize.medium,
+    ) -> dict[str, typing.Any]:
+        """Generate a response with retry logic and error handling."""
+        if max_tokens is None:
+            max_tokens = self.max_tokens
+
+        retry_count = 0
+        last_error = None
+
+        # Add multilingual extraction instructions
+        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
+
+        while retry_count <= self.MAX_RETRIES:
+            try:
+                response = await self._generate_response(
+                    messages, response_model, max_tokens, model_size
+                )
+                return response
+            except (RateLimitError, RefusalError):
+                # These errors should not trigger retries
+                raise
+            except (openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError):
+                # Let OpenAI's client handle these retries
+                raise
+            except Exception as e:
+                last_error = e
+
+                # Don't retry if we've hit the max retries
+                if retry_count >= self.MAX_RETRIES:
+                    logger.error(f'Max retries ({self.MAX_RETRIES}) exceeded. Last error: {e}')
+                    raise
+
+                retry_count += 1
+
+                # Construct a detailed error message for the LLM
+                error_context = (
+                    f'The previous response attempt was invalid. '
+                    f'Error type: {e.__class__.__name__}. '
+                    f'Error details: {str(e)}. '
+                    f'Please try again with a valid response, ensuring the output matches '
+                    f'the expected format and constraints.'
+                )
+
+                error_message = Message(role='user', content=error_context)
+                messages.append(error_message)
+                logger.warning(
+                    f'Retrying after application error (attempt {retry_count}/{self.MAX_RETRIES}): {e}'
+                )
+
+        # If we somehow get here, raise the last error
+        raise last_error or Exception('Max retries exceeded with no specific error')

graphiti_core/llm_client/openai_client.py

@@ -14,50 +14,27 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import logging
 import typing
-from typing import ClassVar
 
-import openai
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
-from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
-from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
-from .errors import RateLimitError, RefusalError
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .openai_base_client import BaseOpenAIClient
 
-logger = logging.getLogger(__name__)
 
-DEFAULT_MODEL = 'gpt-4.1-mini'
-DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
-
-
-class OpenAIClient(LLMClient):
+class OpenAIClient(BaseOpenAIClient):
     """
     OpenAIClient is a client class for interacting with OpenAI's language models.
 
-    This class extends the LLMClient and provides methods to initialize the client,
-    get an embedder, and generate responses from the language model.
+    This class extends the BaseOpenAIClient and provides OpenAI-specific implementation
+    for creating completions.
 
     Attributes:
         client (AsyncOpenAI): The OpenAI client used to interact with the API.
-        model (str): The model name to use for generating responses.
-        temperature (float): The temperature to use for generating responses.
-        max_tokens (int): The maximum number of tokens to generate in a response.
-
-    Methods:
-        __init__(config: LLMConfig | None = None, cache: bool = False, client: typing.Any = None):
-            Initializes the OpenAIClient with the provided configuration, cache setting, and client.
-
-        _generate_response(messages: list[Message]) -> dict[str, typing.Any]:
-            Generates a response from the language model based on the provided messages.
     """
 
-    # Class-level constants
-    MAX_RETRIES: ClassVar[int] = 2
-
     def __init__(
         self,
         config: LLMConfig | None = None,
@@ -72,120 +49,47 @@ class OpenAIClient(LLMClient):
             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.
-
         """
-        # removed caching to simplify the `generate_response` override
-        if cache:
-            raise NotImplementedError('Caching is not implemented for OpenAI')
+        super().__init__(config, cache, max_tokens)
 
         if config is None:
             config = LLMConfig()
 
-        super().__init__(config, cache)
-
         if client is None:
             self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
         else:
             self.client = client
 
-        self.max_tokens = max_tokens
-
-    async def _generate_response(
+    async def _create_structured_completion(
         self,
-        messages: list[Message],
-        response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
-        model_size: ModelSize = ModelSize.medium,
-    ) -> dict[str, typing.Any]:
-        openai_messages: list[ChatCompletionMessageParam] = []
-        for m in messages:
-            m.content = self._clean_input(m.content)
-            if m.role == 'user':
-                openai_messages.append({'role': 'user', 'content': m.content})
-            elif m.role == 'system':
-                openai_messages.append({'role': 'system', 'content': m.content})
-        try:
-            if model_size == ModelSize.small:
-                model = self.small_model or DEFAULT_SMALL_MODEL
-            else:
-                model = self.model or DEFAULT_MODEL
-
-            response = await self.client.beta.chat.completions.parse(
-                model=model,
-                messages=openai_messages,
-                temperature=self.temperature,
-                max_tokens=max_tokens or self.max_tokens,
-                response_format=response_model,  # type: ignore
-            )
-
-            response_object = response.choices[0].message
-
-            if response_object.parsed:
-                return response_object.parsed.model_dump()
-            elif response_object.refusal:
-                raise RefusalError(response_object.refusal)
-            else:
-                raise Exception(f'Invalid response from LLM: {response_object.model_dump()}')
-        except openai.LengthFinishReasonError as e:
-            raise Exception(f'Output length exceeded max tokens {self.max_tokens}: {e}') from e
-        except openai.RateLimitError as e:
-            raise RateLimitError from e
-        except Exception as e:
-            logger.error(f'Error in generating LLM response: {e}')
-            raise
-
-    async def generate_response(
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+    ):
+        """Create a structured completion using OpenAI's beta parse API."""
+        return await self.client.beta.chat.completions.parse(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format=response_model,  # type: ignore
+        )
+
+    async def _create_completion(
         self,
-        messages: list[Message],
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
         response_model: type[BaseModel] | None = None,
-        max_tokens: int | None = None,
-        model_size: ModelSize = ModelSize.medium,
-    ) -> dict[str, typing.Any]:
-        if max_tokens is None:
-            max_tokens = self.max_tokens
-
-        retry_count = 0
-        last_error = None
-
-        # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
-
-        while retry_count <= self.MAX_RETRIES:
-            try:
-                response = await self._generate_response(
-                    messages, response_model, max_tokens, model_size
-                )
-                return response
-            except (RateLimitError, RefusalError):
-                # These errors should not trigger retries
-                raise
-            except (openai.APITimeoutError, openai.APIConnectionError, openai.InternalServerError):
-                # Let OpenAI's client handle these retries
-                raise
-            except Exception as e:
-                last_error = e
-
-                # Don't retry if we've hit the max retries
-                if retry_count >= self.MAX_RETRIES:
-                    logger.error(f'Max retries ({self.MAX_RETRIES}) exceeded. Last error: {e}')
-                    raise
-
-                retry_count += 1
-
-                # Construct a detailed error message for the LLM
-                error_context = (
-                    f'The previous response attempt was invalid. '
-                    f'Error type: {e.__class__.__name__}. '
-                    f'Error details: {str(e)}. '
-                    f'Please try again with a valid response, ensuring the output matches '
-                    f'the expected format and constraints.'
-                )
-
-                error_message = Message(role='user', content=error_context)
-                messages.append(error_message)
-                logger.warning(
-                    f'Retrying after application error (attempt {retry_count}/{self.MAX_RETRIES}): {e}'
-                )
-
-        # If we somehow get here, raise the last error
-        raise last_error or Exception('Max retries exceeded with no specific error')
+    ):
+        """Create a regular completion with JSON format."""
+        return await self.client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            response_format={'type': 'json_object'},
+        )

graphiti_core/models/edges/edge_db_queries.py

@@ -35,7 +35,7 @@ ENTITY_EDGE_SAVE = """
         MATCH (target:Entity {uuid: $target_uuid}) 
         MERGE (source)-[r:RELATES_TO {uuid: $uuid}]->(target)
         SET r = $edge_data
-        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $fact_embedding)
+        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $edge_data.fact_embedding)
         RETURN r.uuid AS uuid"""
 
 ENTITY_EDGE_SAVE_BULK = """

graphiti_core/prompts/dedupe_nodes.py

@@ -30,7 +30,7 @@ class NodeDuplicate(BaseModel):
     )
     name: str = Field(
         ...,
-        description='Name of the entity. Should be the most complete and descriptive name possible. Do not include any JSON formatting in the Entity name.',
+        description='Name of the entity. Should be the most complete and descriptive name of the entity. Do not include any JSON formatting in the Entity name such as {}.',
     )
     additional_duplicates: list[int] = Field(
         ...,
@@ -84,19 +84,19 @@ def node(context: dict[str, Any]) -> list[Message]:
         is a duplicate entity of one of the EXISTING ENTITIES.
         
         Entities should only be considered duplicates if they refer to the *same real-world object or concept*.
+        Semantic Equivalence: if a descriptive label in existing_entities clearly refers to a named entity in context, treat them as duplicates.
 
         Do NOT mark entities as duplicates if:
         - They are related but distinct.
         - They have similar names or purposes but refer to separate instances or concepts.
 
-        Task:
-        If the NEW ENTITY represents a duplicate entity of any entity in EXISTING ENTITIES, set duplicate_entity_id to the
-        id of the EXISTING ENTITY that is the duplicate. 
-        
-        If the NEW ENTITY is not a duplicate of any of the EXISTING ENTITIES,
-        duplicate_entity_id should be set to -1.
+         TASK:
+         1. Compare `new_entity` against each item in `existing_entities`.
+         2. If it refers to the same real‐world object or concept, collect its index.
+         3. Let `duplicate_idx` = the *first* collected index, or –1 if none.
+         4. Let `additional_duplicates` = the list of *any other* collected indices (empty list if none).
         
-        Also return the name that best describes the NEW ENTITY (whether it is the name of the NEW ENTITY, a node it
+        Also return the full name of the NEW ENTITY (whether it is the name of the NEW ENTITY, a node it
         is a duplicate of, or a combination of the two).
         """,
         ),

graphiti_core/prompts/extract_edges.py

@@ -97,8 +97,7 @@ Only extract facts that:
 - The FACT TYPES provide a list of the most important types of facts, make sure to extract facts of these types
 - The FACT TYPES are not an exhaustive list, extract all facts from the message even if they do not fit into one
     of the FACT TYPES
-- The FACT TYPES each contain their fact_type_signature which represents the entity types which that fact_type is defined for.
-    A Type of Entity in the signature represents any extracted entity (it is a generic universal type for all entities).
+- The FACT TYPES each contain their fact_type_signature which represents the source and target entity types.
 
 You may use information from the PREVIOUS MESSAGES only to disambiguate references or support continuity.
 

graphiti_core/search/search_utils.py

@@ -63,7 +63,7 @@ MAX_QUERY_LENGTH = 32
 
 def fulltext_query(query: str, group_ids: list[str] | None = None):
     group_ids_filter_list = (
-        [f"group_id-'{lucene_sanitize(g)}'" for g in group_ids] if group_ids is not None else []
+        [f'group_id:"{lucene_sanitize(g)}"' for g in group_ids] if group_ids is not None else []
     )
     group_ids_filter = ''
     for f in group_ids_filter_list:
@@ -301,12 +301,12 @@ async def edge_bfs_search(
 
     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 (n:Entity)-[r:RELATES_TO]-(m:Entity)
-                            WHERE r.uuid = rel.uuid
-                            """
+                                    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 (n:Entity)-[r:RELATES_TO]-(m:Entity)
+                                    WHERE r.uuid = rel.uuid
+                                    """
         + filter_query
         + """  
                 RETURN DISTINCT
@@ -455,10 +455,10 @@ async def node_bfs_search(
 
     query = (
         """
-                    UNWIND $bfs_origin_node_uuids AS origin_uuid
-                    MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
-                    WHERE n.group_id = origin.group_id
-                    """
+                            UNWIND $bfs_origin_node_uuids AS origin_uuid
+                            MATCH (origin:Entity|Episodic {uuid: origin_uuid})-[:RELATES_TO|MENTIONS]->{1,3}(n:Entity)
+                            WHERE n.group_id = origin.group_id
+                            """
         + filter_query
         + ENTITY_NODE_RETURN
         + """

graphiti_core/utils/maintenance/node_operations.py

@@ -310,7 +310,7 @@ async def resolve_extracted_nodes(
             else extracted_node
         )
 
-        resolved_node.name = resolution.get('name')
+        # resolved_node.name = resolution.get('name')
 
         resolved_nodes.append(resolved_node)
         uuid_map[extracted_node.uuid] = resolved_node.uuid

{graphiti_core-0.12.4.dist-info → graphiti_core-0.13.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.12.4
+Version: 0.13.1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -23,7 +23,7 @@ Requires-Dist: google-genai (>=1.8.0) ; extra == "google-genai"
 Requires-Dist: groq (>=0.2.0) ; extra == "groq"
 Requires-Dist: neo4j (>=5.26.0)
 Requires-Dist: numpy (>=1.0.0)
-Requires-Dist: openai (>=1.53.0)
+Requires-Dist: openai (>=1.91.0)
 Requires-Dist: pydantic (>=2.11.5)
 Requires-Dist: python-dotenv (>=1.0.1)
 Requires-Dist: tenacity (>=9.0.0)
@@ -331,6 +331,55 @@ graphiti = Graphiti(
 # Now you can use Graphiti with Google Gemini
 ```
 
+## Using Graphiti with Ollama (Local LLM)
+
+Graphiti supports Ollama for running local LLMs and embedding models via Ollama's OpenAI-compatible API. This is ideal for privacy-focused applications or when you want to avoid API costs.
+
+
+Install the models:
+ollama pull deepseek-r1:7b         # LLM
+ollama pull nomic-embed-text       # embeddings
+
+```python
+from graphiti_core import Graphiti
+from graphiti_core.llm_client.config import LLMConfig
+from graphiti_core.llm_client.openai_client import OpenAIClient
+from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
+from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
+
+# Configure Ollama LLM client
+llm_config = LLMConfig(
+    api_key="abc",  # Ollama doesn't require a real API key
+    model="deepseek-r1:7b",
+    small_model="deepseek-r1:7b",
+    base_url="http://localhost:11434/v1", # Ollama provides this port
+)
+
+llm_client = OpenAIClient(config=llm_config)
+
+# Initialize Graphiti with Ollama clients
+graphiti = Graphiti(
+    "bolt://localhost:7687",
+    "neo4j",
+    "password",
+    llm_client=llm_client,
+    embedder=OpenAIEmbedder(
+        config=OpenAIEmbedderConfig(
+            api_key="abc",
+            embedding_model="nomic-embed-text",
+            embedding_dim=768,
+            base_url="http://localhost:11434/v1",
+        )
+    ),
+    cross_encoder=OpenAIRerankerClient(client=llm_client, config=llm_config),
+)
+
+# Now you can use Graphiti with local Ollama models
+```
+
+Ensure Ollama is running (`ollama serve`) and that you have pulled the models you want to use.
+
+
 ## Documentation
 
 - [Guides and API documentation](https://help.getzep.com/graphiti).

{graphiti_core-0.12.4.dist-info → graphiti_core-0.13.1.dist-info}/RECORD

@@ -14,34 +14,35 @@ graphiti_core/embedder/client.py,sha256=qEpSHceL_Gc4QQPJWIOnuNLemNuR_TYA4r28t2Vl
 graphiti_core/embedder/gemini.py,sha256=7En-W46YxqC5qL3vYB5Ed-Xm0hqLxi7-LgZ95c4M7ME,3263
 graphiti_core/embedder/openai.py,sha256=bIThUoLMeGlHG2-3VikzK6JZfOHKn4PKvUMx5sHxJy8,2192
 graphiti_core/embedder/voyage.py,sha256=gQhdcz2IYPSyOcDn3w8aHToVS3KQhyZrUBm4vqr3WcE,2224
-graphiti_core/errors.py,sha256=Nib1uQx2cO_VOizupmRjpFfmuRg-hFAVqTtZAuBehR8,2405
+graphiti_core/errors.py,sha256=cH_v9TPgEPeQE6GFOHIg5TvejpUCBddGarMY2Whxbwc,2707
 graphiti_core/graph_queries.py,sha256=KfWDp8xDnPa9bcHskw8NeMpeeHBtZWBCosVdu1Iwv34,7076
-graphiti_core/graphiti.py,sha256=cCUYxwNoQiEtzbPNvTEwfYOrK9UePFpRCUeFxl1-vg0,28335
+graphiti_core/graphiti.py,sha256=6TcOq65HEtJ5gQ29YHJ6HVhBkjsoWtv_mVhlcZok2MY,30399
 graphiti_core/graphiti_types.py,sha256=rL-9bvnLobunJfXU4hkD6mAj14pofKp_wq8QsFDZwDU,1035
-graphiti_core/helpers.py,sha256=sfC1M6KefKaZll6FQcpcNvWnCN6iCvBSJksAkGTItT4,3059
+graphiti_core/helpers.py,sha256=0qmGnKxxYk27JGQbx6PlM7E6nRghUrEKBym0d3WSJY4,3875
 graphiti_core/llm_client/__init__.py,sha256=QgBWUiCeBp6YiA_xqyrDvJ9jIyy1hngH8g7FWahN3nw,776
 graphiti_core/llm_client/anthropic_client.py,sha256=392rtkH_I7yOJUlQvjoOnS8Lz14WBP8egQ3OfRH0nFs,12481
-graphiti_core/llm_client/azure_openai_client.py,sha256=B6EbNIktP9FBqiFrGunVQlego2e3C5zBAbcHI55Y-OY,2680
+graphiti_core/llm_client/azure_openai_client.py,sha256=ekERggAekbb7enes1RJqdRChf_mjaZTFXsnMbxO7azQ,2497
 graphiti_core/llm_client/client.py,sha256=v_w5TBbDJYYADCXSs2r287g5Ami2Urma-GGEbHSI_Jg,5826
 graphiti_core/llm_client/config.py,sha256=90IgSBxZE_3nWdaEONVLUznI8lytPA7ZyexQz-_c55U,2560
 graphiti_core/llm_client/errors.py,sha256=pn6brRiLW60DAUIXJYKBT6MInrS4ueuH1hNLbn_JbQo,1243
 graphiti_core/llm_client/gemini_client.py,sha256=OdRAB2bWlXAi3gRmE1xVljYJ0T7JTZC82VK71wHyZi8,7722
 graphiti_core/llm_client/groq_client.py,sha256=k7zbXHfOpb4jhvvKFsccVYTq4yGGpxmY7xzNA02N2zk,2559
-graphiti_core/llm_client/openai_client.py,sha256=lLTZkd-PxEicTBmQefGoWLGTCb4QSU2Cq3x5W4kRYXg,7412
+graphiti_core/llm_client/openai_base_client.py,sha256=gfMcKPyLrylz_ouRdoenDWXyitmgfFZ17Zthbkq3Qs4,8126
+graphiti_core/llm_client/openai_client.py,sha256=ykBK94gxzE7iXux5rvOzVNA8q0Sqzq-8njPB75XcRe8,3240
 graphiti_core/llm_client/openai_generic_client.py,sha256=WElMnPqdb1CxzYH4p2-m_9rVMr5M93-eXnc3yVxBgFg,7001
 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=W2-ljKnZOt5MlD9_M4f_823GdyTMRzW2tJX0CezaixY,2284
+graphiti_core/models/edges/edge_db_queries.py,sha256=4vSWdmE5MKoDrlHJmmr2xNhVSQ-buE1O7mCX_H0Wtfk,2294
 graphiti_core/models/nodes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/models/nodes/node_db_queries.py,sha256=AQgRGVO-GgFWfLq1G6k8s86WItwpXruy3Mj4DBli-vM,2145
 graphiti_core/nodes.py,sha256=kdJY-Ugyk6J2x70w4EF_EoFNgy7D3TMOMVSUfEth6rE,18665
 graphiti_core/prompts/__init__.py,sha256=EA-x9xUki9l8wnu2l8ek_oNf75-do5tq5hVq7Zbv8Kw,101
 graphiti_core/prompts/dedupe_edges.py,sha256=AFVC1EQ0TvNkSp0G7QZmIh3YpGg9FVXo1_sT3TlRqA8,5473
-graphiti_core/prompts/dedupe_nodes.py,sha256=i8BP6h3xJZUFDNorBbphyi0kZQXVcVh8bXjw686wKdI,7670
+graphiti_core/prompts/dedupe_nodes.py,sha256=WdSnqu6O4TkEE_z1u2CEnNH0sWgBNDl4dUx20gSp464,7852
 graphiti_core/prompts/eval.py,sha256=gnBQTmwsCl3Qvwpcm7aieVszzo6y1sMCUT8jQiKTvvE,5317
 graphiti_core/prompts/extract_edge_dates.py,sha256=3Drs3CmvP0gJN5BidWSxrNvLet3HPoTybU3BUIAoc0Y,4218
-graphiti_core/prompts/extract_edges.py,sha256=1zP-Xo70UbHm1uSJIcVHhHrXIpH9ILU6ieunLPaRdnA,6761
+graphiti_core/prompts/extract_edges.py,sha256=4C9sOqdoZCqVZbmkhNCPeNx3R80rzT8eXndySjVRsHc,6622
 graphiti_core/prompts/extract_nodes.py,sha256=bJ3Inp4tMDsXgTPHKOyAtasVHEeLjbY7wfV3LwdvCsc,9747
 graphiti_core/prompts/invalidate_edges.py,sha256=yfpcs_pyctnoM77ULPZXEtKW0oHr1MeLsJzC5yrE-o4,3547
 graphiti_core/prompts/lib.py,sha256=DCyHePM4_q-CptTpEXGO_dBv9k7xDtclEaB1dGu7EcI,4092
@@ -55,7 +56,7 @@ graphiti_core/search/search_config.py,sha256=VvKg6AB_RPhoe56DBBXHRBXHThAVJ_OLFCy
 graphiti_core/search/search_config_recipes.py,sha256=4GquRphHhJlpXQhAZOySYnCzBWYoTwxlJj44eTOavZQ,7443
 graphiti_core/search/search_filters.py,sha256=jG30nMWX03xoT9ohgyHNu_Xes8GwjIF2eTv6QaiWMqw,6466
 graphiti_core/search/search_helpers.py,sha256=G5Ceaq5Pfgx0Weelqgeylp_pUHwiBnINaUYsDbURJbE,2636
-graphiti_core/search/search_utils.py,sha256=k9KKN4sYde0Hqw9BKb5T-8q-3hInIPwq9aYoGfheq6E,34877
+graphiti_core/search/search_utils.py,sha256=q-FMbSFf7mPXWdgnKQDKQIACrHI8NwqknGDnMy4dJzs,34957
 graphiti_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/utils/bulk_utils.py,sha256=RPPTAqBRg6iR1T6g6TDpfUGvkYmTEyVNrVPz_y91f-s,16196
 graphiti_core/utils/datetime_utils.py,sha256=Ti-2tnrDFRzBsbfblzsHybsM3jaDLP4-VT2t0VhpIzU,1357
@@ -63,11 +64,11 @@ graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4w
 graphiti_core/utils/maintenance/community_operations.py,sha256=2rhRqtL9gDbjXKO4-S0nGpaWvS4ck5rFiazZiogIJao,10088
 graphiti_core/utils/maintenance/edge_operations.py,sha256=Fwu2TLmQF_9EVcA-uUlt1ZiGC6RILIfKDr9W7R4gAno,21633
 graphiti_core/utils/maintenance/graph_data_operations.py,sha256=OHuiAyP1Z7dfR90dWVQ87TJQO83P0sQihJyr4WIhOhk,5362
-graphiti_core/utils/maintenance/node_operations.py,sha256=cqSZ9CmyaH-WeZlzZEmR8Q-GlaQFR31sjw_qcObCyrw,15564
+graphiti_core/utils/maintenance/node_operations.py,sha256=-PC1N-Hf2GElhLcTTSjTaCVvQdJdNpouJPB7joGACZo,15566
 graphiti_core/utils/maintenance/temporal_operations.py,sha256=mJkw9xLB4W2BsLfC5POr0r-PHWL9SIfNj_l_xu0B5ug,3410
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/utils/ontology_utils/entity_types_utils.py,sha256=QJX5cG0GSSNF_Mm_yrldr69wjVAbN_MxLhOSznz85Hk,1279
-graphiti_core-0.12.4.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.12.4.dist-info/METADATA,sha256=PAStxqsKlvFZbwfKkpOoIfDpNGK9AEVNz1s4G-KHpzM,15641
-graphiti_core-0.12.4.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-graphiti_core-0.12.4.dist-info/RECORD,,
+graphiti_core-0.13.1.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.13.1.dist-info/METADATA,sha256=BDBOGDqgXsLJu_FWwPHaVfOVSeTYJCIw5tfdcmnl884,17229
+graphiti_core-0.13.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+graphiti_core-0.13.1.dist-info/RECORD,,