graphiti-core 0.5.0rc4__py3-none-any.whl → 0.5.1__py3-none-any.whl

graphiti_core/cross_encoder/openai_reranker_client.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from typing import Any
 
@@ -22,6 +21,7 @@ import openai
 from openai import AsyncOpenAI
 from pydantic import BaseModel
 
+from ..helpers import semaphore_gather
 from ..llm_client import LLMConfig, RateLimitError
 from ..prompts import Message
 from .client import CrossEncoderClient
@@ -75,7 +75,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
             for passage in passages
         ]
         try:
-            responses = await asyncio.gather(
+            responses = await semaphore_gather(
                 *[
                     self.client.chat.completions.create(
                         model=DEFAULT_MODEL,

graphiti_core/graphiti.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from datetime import datetime
 from time import time
@@ -27,7 +26,7 @@ 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.helpers import DEFAULT_DATABASE, semaphore_gather
 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
@@ -340,13 +339,13 @@ class Graphiti:
 
             # Calculate Embeddings
 
-            await asyncio.gather(
+            await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes]
             )
 
             # Find relevant nodes already in the graph
             existing_nodes_lists: list[list[EntityNode]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[get_relevant_nodes(self.driver, [node]) for node in extracted_nodes]
                 )
             )
@@ -354,7 +353,7 @@ class Graphiti:
             # Resolve extracted nodes with nodes already in the graph and extract facts
             logger.debug(f'Extracted nodes: {[(n.name, n.uuid) for n in extracted_nodes]}')
 
-            (mentioned_nodes, uuid_map), extracted_edges = await asyncio.gather(
+            (mentioned_nodes, uuid_map), extracted_edges = await semaphore_gather(
                 resolve_extracted_nodes(
                     self.llm_client,
                     extracted_nodes,
@@ -374,7 +373,7 @@ class Graphiti:
             )
 
             # calculate embeddings
-            await asyncio.gather(
+            await semaphore_gather(
                 *[
                     edge.generate_embedding(self.embedder)
                     for edge in extracted_edges_with_resolved_pointers
@@ -383,7 +382,7 @@ class Graphiti:
 
             # Resolve extracted edges with related edges already in the graph
             related_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -404,7 +403,7 @@ class Graphiti:
             )
 
             existing_source_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -419,7 +418,7 @@ class Graphiti:
             )
 
             existing_target_edges_list: list[list[EntityEdge]] = list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         get_relevant_edges(
                             self.driver,
@@ -468,7 +467,7 @@ class Graphiti:
 
             # Update any communities
             if update_communities:
-                await asyncio.gather(
+                await semaphore_gather(
                     *[
                         update_community(self.driver, self.llm_client, self.embedder, node)
                         for node in nodes
@@ -538,7 +537,7 @@ class Graphiti:
             ]
 
             # Save all the episodes
-            await asyncio.gather(*[episode.save(self.driver) for episode in episodes])
+            await semaphore_gather(*[episode.save(self.driver) for episode in episodes])
 
             # Get previous episode context for each episode
             episode_pairs = await retrieve_previous_episodes_bulk(self.driver, episodes)
@@ -551,19 +550,19 @@ class Graphiti:
             ) = await extract_nodes_and_edges_bulk(self.llm_client, episode_pairs)
 
             # Generate embeddings
-            await asyncio.gather(
+            await semaphore_gather(
                 *[node.generate_name_embedding(self.embedder) for node in extracted_nodes],
                 *[edge.generate_embedding(self.embedder) for edge in extracted_edges],
             )
 
             # Dedupe extracted nodes, compress extracted edges
-            (nodes, uuid_map), extracted_edges_timestamped = await asyncio.gather(
+            (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),
             )
 
             # save nodes to KG
-            await asyncio.gather(*[node.save(self.driver) for node in nodes])
+            await semaphore_gather(*[node.save(self.driver) for node in nodes])
 
             # re-map edge pointers so that they don't point to discard dupe nodes
             extracted_edges_with_resolved_pointers: list[EntityEdge] = resolve_edge_pointers(
@@ -574,7 +573,7 @@ class Graphiti:
             )
 
             # save episodic edges to KG
-            await asyncio.gather(
+            await semaphore_gather(
                 *[edge.save(self.driver) for edge in episodic_edges_with_resolved_pointers]
             )
 
@@ -587,7 +586,7 @@ class Graphiti:
             # invalidate edges
 
             # save edges to KG
-            await asyncio.gather(*[edge.save(self.driver) for edge in edges])
+            await semaphore_gather(*[edge.save(self.driver) for edge in edges])
 
             end = time()
             logger.info(f'Completed add_episode_bulk in {(end - start) * 1000} ms')
@@ -610,12 +609,12 @@ class Graphiti:
             self.driver, self.llm_client, group_ids
         )
 
-        await asyncio.gather(
+        await semaphore_gather(
             *[node.generate_name_embedding(self.embedder) for node in community_nodes]
         )
 
-        await asyncio.gather(*[node.save(self.driver) for node in community_nodes])
-        await asyncio.gather(*[edge.save(self.driver) for edge in community_edges])
+        await semaphore_gather(*[node.save(self.driver) for node in community_nodes])
+        await semaphore_gather(*[edge.save(self.driver) for edge in community_edges])
 
         return community_nodes
 
@@ -698,7 +697,7 @@ class Graphiti:
     async def get_episode_mentions(self, episode_uuids: list[str]) -> SearchResults:
         episodes = await EpisodicNode.get_by_uuids(self.driver, episode_uuids)
 
-        edges_list = await asyncio.gather(
+        edges_list = await semaphore_gather(
             *[EntityEdge.get_by_uuids(self.driver, episode.entity_edges) for episode in episodes]
         )
 

graphiti_core/helpers.py

@@ -14,7 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+import asyncio
 import os
+from collections.abc import Coroutine
 from datetime import datetime
 
 import numpy as np
@@ -25,6 +27,7 @@ load_dotenv()
 
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
+SEMAPHORE_LIMIT = int(os.getenv('SEMAPHORE_LIMIT', 20))
 MAX_REFLEXION_ITERATIONS = 2
 DEFAULT_PAGE_LIMIT = 20
 
@@ -80,3 +83,14 @@ def normalize_l2(embedding: list[float]) -> list[float]:
     else:
         norm = np.linalg.norm(embedding_array, 2, axis=1, keepdims=True)
         return (np.where(norm == 0, embedding_array, embedding_array / norm)).tolist()
+
+
+# Use this instead of asyncio.gather() to bound coroutines
+async def semaphore_gather(*coroutines: Coroutine, max_coroutines: int = SEMAPHORE_LIMIT):
+    semaphore = asyncio.Semaphore(max_coroutines)
+
+    async def _wrap_coroutine(coroutine):
+        async with semaphore:
+            return await coroutine
+
+    return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines))

graphiti_core/llm_client/client.py

@@ -56,6 +56,28 @@ class LLMClient(ABC):
         self.cache_enabled = cache
         self.cache_dir = Cache(DEFAULT_CACHE_DIR)  # Create a cache directory
 
+    def _clean_input(self, input: str) -> str:
+        """Clean input string of invalid unicode and control characters.
+
+        Args:
+            input: Raw input string to be cleaned
+
+        Returns:
+            Cleaned string safe for LLM processing
+        """
+        # Clean any invalid Unicode
+        cleaned = input.encode('utf-8', errors='ignore').decode('utf-8')
+
+        # Remove zero-width characters and other invisible unicode
+        zero_width = '\u200b\u200c\u200d\ufeff\u2060'
+        for char in zero_width:
+            cleaned = cleaned.replace(char, '')
+
+        # Remove control characters except newlines, returns, and tabs
+        cleaned = ''.join(char for char in cleaned if ord(char) >= 32 or char in '\n\r\t')
+
+        return cleaned
+
     @retry(
         stop=stop_after_attempt(4),
         wait=wait_random_exponential(multiplier=10, min=5, max=120),
@@ -106,6 +128,9 @@ class LLMClient(ABC):
                 logger.debug(f'Cache hit for {cache_key}')
                 return cached_response
 
+        for message in messages:
+            message.content = self._clean_input(message.content)
+
         response = await self._generate_response_with_retry(messages, response_model)
 
         if self.cache_enabled:

graphiti_core/llm_client/openai_client.py

@@ -88,6 +88,7 @@ class OpenAIClient(LLMClient):
     ) -> 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':

graphiti_core/llm_client/openai_generic_client.py

@@ -0,0 +1,163 @@
+"""
+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 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 LLMClient
+from .config import LLMConfig
+from .errors import RateLimitError, RefusalError
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = 'gpt-4o-mini'
+
+
+class OpenAIGenericClient(LLMClient):
+    """
+    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.
+
+    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, cache: bool = False, client: typing.Any = 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.
+
+        """
+        # removed caching to simplify the `generate_response` override
+        if cache:
+            raise NotImplementedError('Caching is not implemented for OpenAI')
+
+        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
+
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> 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:
+            response = await self.client.chat.completions.create(
+                model=self.model or DEFAULT_MODEL,
+                messages=openai_messages,
+                temperature=self.temperature,
+                max_tokens=self.max_tokens,
+                response_format={'type': 'json_object'},
+            )
+            result = response.choices[0].message.content or ''
+            return json.loads(result)
+        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
+    ) -> dict[str, typing.Any]:
+        retry_count = 0
+        last_error = None
+
+        if response_model is not None:
+            serialized_model = json.dumps(response_model.model_json_schema())
+            messages[
+                -1
+            ].content += (
+                f'\n\nRespond with a JSON object in the following format:\n\n{serialized_model}'
+            )
+
+        while retry_count <= self.MAX_RETRIES:
+            try:
+                response = await self._generate_response(messages, response_model)
+                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/search/search.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
@@ -25,6 +24,7 @@ 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
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.nodes import CommunityNode, EntityNode
 from graphiti_core.search.search_config import (
     DEFAULT_SEARCH_LIMIT,
@@ -78,7 +78,7 @@ async def search(
 
     # if group_ids is empty, set it to None
     group_ids = group_ids if group_ids else None
-    edges, nodes, communities = await asyncio.gather(
+    edges, nodes, communities = await semaphore_gather(
         edge_search(
             driver,
             cross_encoder,
@@ -141,7 +141,7 @@ async def edge_search(
         return []
 
     search_results: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 edge_fulltext_search(driver, query, group_ids, 2 * limit),
                 edge_similarity_search(
@@ -226,7 +226,7 @@ async def node_search(
         return []
 
     search_results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 node_fulltext_search(driver, query, group_ids, 2 * limit),
                 node_similarity_search(
@@ -295,7 +295,7 @@ async def community_search(
         return []
 
     search_results: list[list[CommunityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 community_fulltext_search(driver, query, group_ids, 2 * limit),
                 community_similarity_search(

graphiti_core/search/search_utils.py

@@ -14,10 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from collections import defaultdict
 from time import time
+from typing import Any
 
 import numpy as np
 from neo4j import AsyncDriver, Query
@@ -29,6 +29,7 @@ from graphiti_core.helpers import (
     USE_PARALLEL_RUNTIME,
     lucene_sanitize,
     normalize_l2,
+    semaphore_gather,
 )
 from graphiti_core.nodes import (
     CommunityNode,
@@ -191,12 +192,27 @@ async def edge_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
-    query: LiteralString = """
-                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 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
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE r.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+        query_params['source_node_uuid'] = source_node_uuid
+        query_params['target_node_uuid'] = target_node_uuid
+
+        if source_node_uuid is not None:
+            group_filter_query += '\nAND (n.uuid IN [$source_uuid, $target_uuid])'
+
+        if target_node_uuid is not None:
+            group_filter_query += '\nAND (m.uuid IN [$source_uuid, $target_uuid])'
+
+    query: LiteralString = (
+        """
+                        MATCH (n:Entity)-[r:RELATES_TO]->(m:Entity)
+                        """
+        + group_filter_query
+        + """\nWITH DISTINCT r, vector.similarity.cosine(r.fact_embedding, $search_vector) AS score
                 WHERE score > $min_score
                 RETURN
                     r.uuid AS uuid,
@@ -214,9 +230,11 @@ async def edge_similarity_search(
                 ORDER BY score DESC
                 LIMIT $limit
         """
+    )
 
     records, _, _ = await driver.execute_query(
         runtime_query + query,
+        query_params,
         search_vector=search_vector,
         source_uuid=source_node_uuid,
         target_uuid=target_node_uuid,
@@ -325,11 +343,20 @@ async def node_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE n.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
             MATCH (n:Entity)
-            WHERE $group_ids IS NULL OR n.group_id IN $group_ids
+            """
+        + group_filter_query
+        + """
             WITH n, vector.similarity.cosine(n.name_embedding, $search_vector) AS score
             WHERE score > $min_score
             RETURN
@@ -342,6 +369,7 @@ async def node_similarity_search(
             ORDER BY score DESC
             LIMIT $limit
             """,
+        query_params,
         search_vector=search_vector,
         group_ids=group_ids,
         limit=limit,
@@ -436,11 +464,20 @@ async def community_similarity_search(
         'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
     )
 
+    query_params: dict[str, Any] = {}
+
+    group_filter_query: LiteralString = ''
+    if group_ids is not None:
+        group_filter_query += 'WHERE comm.group_id IN $group_ids'
+        query_params['group_ids'] = group_ids
+
     records, _, _ = await driver.execute_query(
         runtime_query
         + """
            MATCH (comm:Community)
-           WHERE ($group_ids IS NULL OR comm.group_id IN $group_ids)
+           """
+        + group_filter_query
+        + """
            WITH comm, vector.similarity.cosine(comm.name_embedding, $search_vector) AS score
            WHERE score > $min_score
            RETURN
@@ -512,7 +549,7 @@ async def hybrid_node_search(
 
     start = time()
     results: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[node_fulltext_search(driver, q, group_ids, 2 * limit) for q in queries],
             *[node_similarity_search(driver, e, group_ids, 2 * limit) for e in embeddings],
         )
@@ -582,7 +619,7 @@ async def get_relevant_edges(
     relevant_edges: list[EntityEdge] = []
     relevant_edge_uuids = set()
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             edge_similarity_search(
                 driver,

graphiti_core/utils/bulk_utils.py

@@ -14,7 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 import typing
 from collections import defaultdict
@@ -26,6 +25,7 @@ from numpy import dot, sqrt
 from pydantic import BaseModel
 
 from graphiti_core.edges import Edge, EntityEdge, EpisodicEdge
+from graphiti_core.helpers import semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.models.edges.edge_db_queries import (
     ENTITY_EDGE_SAVE_BULK,
@@ -71,7 +71,7 @@ class RawEpisode(BaseModel):
 async def retrieve_previous_episodes_bulk(
     driver: AsyncDriver, episodes: list[EpisodicNode]
 ) -> list[tuple[EpisodicNode, list[EpisodicNode]]]:
-    previous_episodes_list = await asyncio.gather(
+    previous_episodes_list = await semaphore_gather(
         *[
             retrieve_episodes(
                 driver, episode.valid_at, last_n=EPISODE_WINDOW_LEN, group_ids=[episode.group_id]
@@ -118,7 +118,7 @@ async def add_nodes_and_edges_bulk_tx(
 async def extract_nodes_and_edges_bulk(
     llm_client: LLMClient, episode_tuples: list[tuple[EpisodicNode, list[EpisodicNode]]]
 ) -> tuple[list[EntityNode], list[EntityEdge], list[EpisodicEdge]]:
-    extracted_nodes_bulk = await asyncio.gather(
+    extracted_nodes_bulk = await semaphore_gather(
         *[
             extract_nodes(llm_client, episode, previous_episodes)
             for episode, previous_episodes in episode_tuples
@@ -130,7 +130,7 @@ async def extract_nodes_and_edges_bulk(
         [episode[1] for episode in episode_tuples],
     )
 
-    extracted_edges_bulk = await asyncio.gather(
+    extracted_edges_bulk = await semaphore_gather(
         *[
             extract_edges(
                 llm_client,
@@ -171,13 +171,13 @@ async def dedupe_nodes_bulk(
     node_chunks = [nodes[i : i + CHUNK_SIZE] for i in range(0, len(nodes), CHUNK_SIZE)]
 
     existing_nodes_chunks: list[list[EntityNode]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_nodes(driver, node_chunk) for node_chunk in node_chunks]
         )
     )
 
     results: list[tuple[list[EntityNode], dict[str, str]]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_nodes(llm_client, node_chunk, existing_nodes_chunks[i])
                 for i, node_chunk in enumerate(node_chunks)
@@ -205,13 +205,13 @@ async def dedupe_edges_bulk(
     ]
 
     relevant_edges_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[get_relevant_edges(driver, edge_chunk, None, None) for edge_chunk in edge_chunks]
         )
     )
 
     resolved_edge_chunks: list[list[EntityEdge]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 dedupe_extracted_edges(llm_client, edge_chunk, relevant_edges_chunks[i])
                 for i, edge_chunk in enumerate(edge_chunks)
@@ -292,7 +292,9 @@ async def compress_nodes(
             # add both nodes to the shortest chunk
             node_chunks[-1].extend([n, m])
 
-    results = await asyncio.gather(*[dedupe_node_list(llm_client, chunk) for chunk in node_chunks])
+    results = await semaphore_gather(
+        *[dedupe_node_list(llm_client, chunk) for chunk in node_chunks]
+    )
 
     extended_map = dict(uuid_map)
     compressed_nodes: list[EntityNode] = []
@@ -315,7 +317,9 @@ async def compress_edges(llm_client: LLMClient, edges: list[EntityEdge]) -> list
     # We build a map of the edges based on their source and target nodes.
     edge_chunks = chunk_edges_by_nodes(edges)
 
-    results = await asyncio.gather(*[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks])
+    results = await semaphore_gather(
+        *[dedupe_edge_list(llm_client, chunk) for chunk in edge_chunks]
+    )
 
     compressed_edges: list[EntityEdge] = []
     for edge_chunk in results:
@@ -368,7 +372,7 @@ async def extract_edge_dates_bulk(
         episode.uuid: (episode, previous_episodes) for episode, previous_episodes in episode_pairs
     }
 
-    results = await asyncio.gather(
+    results = await semaphore_gather(
         *[
             extract_edge_dates(
                 llm_client,

graphiti_core/utils/maintenance/community_operations.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 
 from graphiti_core.edges import CommunityEdge
 from graphiti_core.embedder import EmbedderClient
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import (
     CommunityNode,
@@ -71,7 +71,7 @@ async def get_community_clusters(
 
         community_clusters.extend(
             list(
-                await asyncio.gather(
+                await semaphore_gather(
                     *[EntityNode.get_by_uuids(driver, cluster) for cluster in cluster_uuids]
                 )
             )
@@ -164,7 +164,7 @@ async def build_community(
             odd_one_out = summaries.pop()
             length -= 1
         new_summaries: list[str] = list(
-            await asyncio.gather(
+            await semaphore_gather(
                 *[
                     summarize_pair(llm_client, (str(left_summary), str(right_summary)))
                     for left_summary, right_summary in zip(
@@ -207,7 +207,9 @@ async def build_communities(
             return await build_community(llm_client, cluster)
 
     communities: list[tuple[CommunityNode, list[CommunityEdge]]] = list(
-        await asyncio.gather(*[limited_build_community(cluster) for cluster in community_clusters])
+        await semaphore_gather(
+            *[limited_build_community(cluster) for cluster in community_clusters]
+        )
     )
 
     community_nodes: list[CommunityNode] = []

graphiti_core/utils/maintenance/edge_operations.py

@@ -14,13 +14,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from datetime import datetime
 from time import time
 
 from graphiti_core.edges import CommunityEdge, EntityEdge, EpisodicEdge
-from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS
+from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
 from graphiti_core.prompts import prompt_library
@@ -199,7 +198,7 @@ async def resolve_extracted_edges(
 ) -> tuple[list[EntityEdge], list[EntityEdge]]:
     # resolve edges with related edges in the graph, extract temporal information, and find invalidation candidates
     results: list[tuple[EntityEdge, list[EntityEdge]]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 resolve_extracted_edge(
                     llm_client,
@@ -266,7 +265,7 @@ async def resolve_extracted_edge(
     current_episode: EpisodicNode,
     previous_episodes: list[EpisodicNode],
 ) -> tuple[EntityEdge, list[EntityEdge]]:
-    resolved_edge, (valid_at, invalid_at), invalidation_candidates = await asyncio.gather(
+    resolved_edge, (valid_at, invalid_at), invalidation_candidates = await semaphore_gather(
         dedupe_extracted_edge(llm_client, extracted_edge, related_edges),
         extract_edge_dates(llm_client, extracted_edge, current_episode, previous_episodes),
         get_edge_contradictions(llm_client, extracted_edge, existing_edges),

graphiti_core/utils/maintenance/graph_data_operations.py

@@ -14,14 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from datetime import datetime, timezone
 
 from neo4j import AsyncDriver
 from typing_extensions import LiteralString
 
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, semaphore_gather
 from graphiti_core.nodes import EpisodeType, EpisodicNode
 
 EPISODE_WINDOW_LEN = 3
@@ -38,7 +37,7 @@ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bo
             database_=DEFAULT_DATABASE,
         )
         index_names = [record['name'] for record in records]
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 driver.execute_query(
                     """DROP INDEX $name""",
@@ -82,7 +81,7 @@ async def build_indices_and_constraints(driver: AsyncDriver, delete_existing: bo
 
     index_queries: list[LiteralString] = range_indices + fulltext_indices
 
-    await asyncio.gather(
+    await semaphore_gather(
         *[
             driver.execute_query(
                 query,

graphiti_core/utils/maintenance/node_operations.py

@@ -14,11 +14,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import asyncio
 import logging
 from time import time
 
-from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS
+from graphiti_core.helpers import MAX_REFLEXION_ITERATIONS, semaphore_gather
 from graphiti_core.llm_client import LLMClient
 from graphiti_core.nodes import EntityNode, EpisodeType, EpisodicNode
 from graphiti_core.prompts import prompt_library
@@ -223,7 +222,7 @@ async def resolve_extracted_nodes(
     uuid_map: dict[str, str] = {}
     resolved_nodes: list[EntityNode] = []
     results: list[tuple[EntityNode, dict[str, str]]] = list(
-        await asyncio.gather(
+        await semaphore_gather(
             *[
                 resolve_extracted_node(
                     llm_client, extracted_node, existing_nodes, episode, previous_episodes
@@ -275,7 +274,7 @@ async def resolve_extracted_node(
         else [],
     }
 
-    llm_response, node_summary_response = await asyncio.gather(
+    llm_response, node_summary_response = await semaphore_gather(
         llm_client.generate_response(
             prompt_library.dedupe_nodes.node(context), response_model=NodeDuplicate
         ),

graphiti_core/utils/maintenance/temporal_operations.py

@@ -55,7 +55,7 @@ async def extract_edge_dates(
         try:
             valid_at_datetime = ensure_utc(datetime.fromisoformat(valid_at.replace('Z', '+00:00')))
         except ValueError as e:
-            logger.error(f'Error parsing valid_at date: {e}. Input: {valid_at}')
+            logger.warning(f'WARNING: Error parsing valid_at date: {e}. Input: {valid_at}')
 
     if invalid_at:
         try:
@@ -63,7 +63,7 @@ async def extract_edge_dates(
                 datetime.fromisoformat(invalid_at.replace('Z', '+00:00'))
             )
         except ValueError as e:
-            logger.error(f'Error parsing invalid_at date: {e}. Input: {invalid_at}')
+            logger.warning(f'WARNING: Error parsing invalid_at date: {e}. Input: {invalid_at}')
 
     return valid_at_datetime, invalid_at_datetime
 

{graphiti_core-0.5.0rc4.dist-info → graphiti_core-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: graphiti-core
-Version: 0.5.0rc4
+Version: 0.5.1
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

{graphiti_core-0.5.0rc4.dist-info → graphiti_core-0.5.1.dist-info}/RECORD

@@ -2,22 +2,23 @@ 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=xgXZqB_qoaWQPjnmuf1ne38YPyhhvApySKcQDaHc9R4,1435
 graphiti_core/cross_encoder/client.py,sha256=KLsbfWKOEaAV3adFe3XZlAeb-gje9_sVKCVZTaJP3ac,1441
-graphiti_core/cross_encoder/openai_reranker_client.py,sha256=F0S9ksusyxdlcRp4yRJuCyAEg-YqgdCwXjZbH8L-Xxo,4063
+graphiti_core/cross_encoder/openai_reranker_client.py,sha256=DECh36QdkM4mE5qgg158ss6G4si2DjgOTqrzGOz5C9s,4089
 graphiti_core/edges.py,sha256=A9tlOtSIVYy_OOKfOk5fAmZ13g8euuAe76im6nKJs0o,14766
 graphiti_core/embedder/__init__.py,sha256=eWd-0sPxflnYXLoWNT9sxwCIFun5JNO9Fk4E-ZXXf8Y,164
 graphiti_core/embedder/client.py,sha256=HKIlpPLnzFT81jurPkry6z8F8nxfZVfejdcfxHVUSFU,995
 graphiti_core/embedder/openai.py,sha256=FzEM9rtSDK1wTb4iYKjNjjdFf8BEBTDxG2vM_E-5W-8,1621
 graphiti_core/embedder/voyage.py,sha256=7kqrLG75J3Q6cdA2Nlx1JSYtpk2141ckdl3OtDDw0vU,1882
 graphiti_core/errors.py,sha256=ddHrHGQxhwkVAtSph4AV84UoOlgwZufMczXPwB7uqPo,1795
-graphiti_core/graphiti.py,sha256=XzgnEbcnRvYXeS1Ul0bngCe0Z3-coZIoVNuAaxcwcvM,27000
-graphiti_core/helpers.py,sha256=DYS8bR7QdtZoZZu5nrJgSAeixYB2hQm-vehHQFnqI7M,2375
+graphiti_core/graphiti.py,sha256=qS3q9FmygplN2H7JL0jyO0VGvAc8IOl89xyipyJc84o,27039
+graphiti_core/helpers.py,sha256=pb0bt2C2NPuGLFf86juS0UJLfyyXa2ju-42apU6xtWA,2886
 graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
 graphiti_core/llm_client/anthropic_client.py,sha256=4hU_PXObkdiT_gUNj9G-Cj6vHFPQ0QpolwaRnJwDiL4,2497
-graphiti_core/llm_client/client.py,sha256=rt0Ic0RiKNbleHCkMYaqI6SUH_K9SkbJtdv7TypDD0A,3937
+graphiti_core/llm_client/client.py,sha256=3eidjHWxTQvCHTUOztSm2SN4-VQhgO0BghLwnHK6u7Y,4801
 graphiti_core/llm_client/config.py,sha256=VwtvD0B7TNqE6Cl-rvH5v-bAfmjMLhEUuFmHSPt10EI,2339
 graphiti_core/llm_client/errors.py,sha256=Vk0mj2SgNDg8E8p7m1UyUaerqLPNLCDKPVsMEnOSBdQ,1028
 graphiti_core/llm_client/groq_client.py,sha256=A4TcbBGXyF5Br5Ggm7qnvso76L1ERO4JoCA2HlzDEyI,2421
-graphiti_core/llm_client/openai_client.py,sha256=Nz8X-4huWC3nKG9azEOSFsRYpq5bLL3S7YRuvY-hYUU,6502
+graphiti_core/llm_client/openai_client.py,sha256=RGq2emM1CioOGFTNfulT3nNdMAKHTXlI7kNp7vQJe9c,6555
+graphiti_core/llm_client/openai_generic_client.py,sha256=SkuQs3jlJAbeFc5MMPUdGcnvACeavd1FAZoj5p0TAI4,6441
 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
@@ -39,21 +40,21 @@ graphiti_core/prompts/prompt_helpers.py,sha256=-9TABwIcIQUVHcNANx6wIZd-FT2DgYKyG
 graphiti_core/prompts/summarize_nodes.py,sha256=5J_IONG7fHYiQZWnCaUyw7w2zunEaN7V89nEluRP-qY,3461
 graphiti_core/py.typed,sha256=vlmmzQOt7bmeQl9L3XJP4W6Ry0iiELepnOrinKz5KQg,79
 graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/search/search.py,sha256=5r7HpIR4cQvwH24YN_DJB6LVNN9hizviE0QXGM6swz8,11984
+graphiti_core/search/search.py,sha256=_-N8ztFv9_OpfNci-AwImhgw8NOLj5FDV_7Cv5GSiRo,12028
 graphiti_core/search/search_config.py,sha256=UZN8jFA4pBlw2O5N1cuhVRBdTwMLR9N3Oyo6sQ4MDVw,3117
 graphiti_core/search/search_config_recipes.py,sha256=yUqiLnn9vFg39M8eVwjVKfBCL_ptGrfDMQ47m_Blb0g,6885
-graphiti_core/search/search_utils.py,sha256=oDXnSlwdrjD2p2izR3_Kv_i5cZdu6mOdUsiFe_QcAN8,22546
+graphiti_core/search/search_utils.py,sha256=lL6CsFFn65tEVlfBLBZu4ynZ9cyBW5SWqc0ss9mlF0Y,23485
 graphiti_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core/utils/bulk_utils.py,sha256=8yov8bPbsQpKxy-kgOhrutic_Gm2ehPvgeAJ0G54NKs,14026
+graphiti_core/utils/bulk_utils.py,sha256=FYal4tSspGVohNsnDoyW_YjMiscySuYPuQLPSwVCy24,14110
 graphiti_core/utils/datetime_utils.py,sha256=Ti-2tnrDFRzBsbfblzsHybsM3jaDLP4-VT2t0VhpIzU,1357
 graphiti_core/utils/maintenance/__init__.py,sha256=TRY3wWWu5kn3Oahk_KKhltrWnh0NACw0FskjqF6OtlA,314
-graphiti_core/utils/maintenance/community_operations.py,sha256=wHla4w6WaCcCFw7tl6Ce9CLZ4odlp6jXXzkB6IZBLCo,9965
-graphiti_core/utils/maintenance/edge_operations.py,sha256=DTVwu12yWrlNmcS8zfQ4730QMRukFNiUiGWzMa1zydU,12673
-graphiti_core/utils/maintenance/graph_data_operations.py,sha256=w66_SLlvPapuG91YGGfR3bG2sM6cJ2XPHIaxM0slAdE,6526
-graphiti_core/utils/maintenance/node_operations.py,sha256=rwQUNG8t7wEI_LBZIhA5vjPJl5MYOzcpEOKviqTEycA,12029
-graphiti_core/utils/maintenance/temporal_operations.py,sha256=MWg1HErEPVqfLNz-lO0qbuuBT-xw1f2Gq5YoBzrWDfA,3669
+graphiti_core/utils/maintenance/community_operations.py,sha256=gIw1M5HGgc2c3TXag5ygPPpAv5WsG-yoC8Lhmfr6FMs,10011
+graphiti_core/utils/maintenance/edge_operations.py,sha256=8Ajq6GS9xH3LRHHw4UAiCeYsWmqJAPYdwfJ_HBGYq2Q,12680
+graphiti_core/utils/maintenance/graph_data_operations.py,sha256=qds9ALk9PhpQs1CNZTZGpi70mqJ93Y2KhIh9X2r8MUI,6533
+graphiti_core/utils/maintenance/node_operations.py,sha256=A-6H2ohqcGJRA_sg_-0m_AA7syiP_gVBsyY7VTTbfuA,12036
+graphiti_core/utils/maintenance/temporal_operations.py,sha256=RdNtubCyYhOVrvcOIq2WppHls1Q-BEjtsN8r38l-Rtc,3691
 graphiti_core/utils/maintenance/utils.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-graphiti_core-0.5.0rc4.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.5.0rc4.dist-info/METADATA,sha256=lImN-I3lfkdtNbE6P1IAcRmk7E5_vjWFXf97a-IMTq8,10061
-graphiti_core-0.5.0rc4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-graphiti_core-0.5.0rc4.dist-info/RECORD,,
+graphiti_core-0.5.1.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.5.1.dist-info/METADATA,sha256=UedUsekIenAFbGcLsXBcDi_MZV1qWRdMPt5IQHXQLEI,10058
+graphiti_core-0.5.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+graphiti_core-0.5.1.dist-info/RECORD,,