graphiti-core 0.4.2__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

graphiti_core/cross_encoder/bge_reranker_client.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 import asyncio
-from typing import List, Tuple
 
 from sentence_transformers import CrossEncoder
 
@@ -26,7 +25,7 @@ class BGERerankerClient(CrossEncoderClient):
     def __init__(self):
         self.model = CrossEncoder('BAAI/bge-reranker-v2-m3')
 
-    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
         if not passages:
             return []
 

graphiti_core/cross_encoder/client.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 from abc import ABC, abstractmethod
-from typing import List, Tuple
 
 
 class CrossEncoderClient(ABC):
@@ -26,16 +25,16 @@ class CrossEncoderClient(ABC):
     """
 
     @abstractmethod
-    async def rank(self, query: str, passages: List[str]) -> List[Tuple[str, float]]:
+    async def rank(self, query: str, passages: list[str]) -> list[tuple[str, float]]:
         """
         Rank the given passages based on their relevance to the query.
 
         Args:
             query (str): The query string.
-            passages (List[str]): A list of passages to rank.
+            passages (list[str]): A list of passages to rank.
 
         Returns:
-            List[Tuple[str, float]]: A list of tuples containing the passage and its score,
+            list[tuple[str, float]]: A list of tuples containing the passage and its score,
                                      sorted in descending order of relevance.
         """
         pass

graphiti_core/edges.py

@@ -23,10 +23,11 @@ from uuid import uuid4
 
 from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
+from typing_extensions import LiteralString
 
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
-from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
+from graphiti_core.helpers import DEFAULT_DATABASE, DEFAULT_PAGE_LIMIT, parse_db_date
 from graphiti_core.models.edges.edge_db_queries import (
     COMMUNITY_EDGE_SAVE,
     ENTITY_EDGE_SAVE,
@@ -50,7 +51,7 @@ class Edge(BaseModel, ABC):
     async def delete(self, driver: AsyncDriver):
         result = await driver.execute_query(
             """
-        MATCH (n)-[e {uuid: $uuid}]->(m)
+        MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
         DELETE e
         """,
             uuid=self.uuid,
@@ -137,19 +138,34 @@ class EpisodicEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid As uuid,
             e.group_id AS group_id,
             n.uuid AS source_node_uuid, 
             m.uuid AS target_node_uuid, 
             e.created_at AS created_at
+        ORDER BY e.uuid DESC 
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -274,11 +290,22 @@ class EntityEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity)-[e:RELATES_TO]->(m:Entity)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid AS uuid,
             n.uuid AS source_node_uuid,
@@ -292,8 +319,12 @@ class EntityEdge(Edge):
             e.expired_at AS expired_at,
             e.valid_at AS valid_at,
             e.invalid_at AS invalid_at
+        ORDER BY e.uuid DESC 
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -365,19 +396,34 @@ class CommunityEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
         WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             e.uuid As uuid,
             e.group_id AS group_id,
             n.uuid AS source_node_uuid, 
             m.uuid AS target_node_uuid, 
             e.created_at AS created_at
+        ORDER BY e.uuid DESC
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )

graphiti_core/embedder/client.py

@@ -15,7 +15,7 @@ limitations under the License.
 """
 
 from abc import ABC, abstractmethod
-from typing import Iterable, List, Literal
+from collections.abc import Iterable
 
 from pydantic import BaseModel, Field
 
@@ -23,12 +23,12 @@ EMBEDDING_DIM = 1024
 
 
 class EmbedderConfig(BaseModel):
-    embedding_dim: Literal[1024] = Field(default=EMBEDDING_DIM, frozen=True)
+    embedding_dim: int = Field(default=EMBEDDING_DIM, frozen=True)
 
 
 class EmbedderClient(ABC):
     @abstractmethod
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         pass

graphiti_core/embedder/openai.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Iterable, List
+from collections.abc import Iterable
 
 from openai import AsyncOpenAI
 from openai.types import EmbeddingModel
@@ -42,7 +42,7 @@ class OpenAIEmbedder(EmbedderClient):
         self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
 
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         result = await self.client.embeddings.create(
             input=input_data, model=self.config.embedding_model

graphiti_core/embedder/voyage.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from typing import Iterable, List
+from collections.abc import Iterable
 
 import voyageai  # type: ignore
 from pydantic import Field
@@ -41,11 +41,11 @@ class VoyageAIEmbedder(EmbedderClient):
         self.client = voyageai.AsyncClient(api_key=config.api_key)
 
     async def create(
-        self, input_data: str | List[str] | Iterable[int] | Iterable[Iterable[int]]
+        self, input_data: str | list[str] | Iterable[int] | Iterable[Iterable[int]]
     ) -> list[float]:
         if isinstance(input_data, str):
             input_list = [input_data]
-        elif isinstance(input_data, List):
+        elif isinstance(input_data, list):
             input_list = [str(i) for i in input_data if i]
         else:
             input_list = [str(i) for i in input_data if i is not None]

graphiti_core/graphiti.py

@@ -318,17 +318,21 @@ class Graphiti:
             previous_episodes = await self.retrieve_episodes(
                 reference_time, last_n=RELEVANT_SCHEMA_LIMIT, group_ids=[group_id]
             )
-            episode = EpisodicNode(
-                name=name,
-                group_id=group_id,
-                labels=[],
-                source=source,
-                content=episode_body,
-                source_description=source_description,
-                created_at=now,
-                valid_at=reference_time,
+
+            episode = (
+                await EpisodicNode.get_by_uuid(self.driver, uuid)
+                if uuid is not None
+                else EpisodicNode(
+                    name=name,
+                    group_id=group_id,
+                    labels=[],
+                    source=source,
+                    content=episode_body,
+                    source_description=source_description,
+                    created_at=now,
+                    valid_at=reference_time,
+                )
             )
-            episode.uuid = uuid if uuid is not None else episode.uuid
 
             # Extract entities as nodes
 

graphiti_core/helpers.py

@@ -26,6 +26,7 @@ load_dotenv()
 DEFAULT_DATABASE = os.getenv('DEFAULT_DATABASE', None)
 USE_PARALLEL_RUNTIME = bool(os.getenv('USE_PARALLEL_RUNTIME', False))
 MAX_REFLEXION_ITERATIONS = 2
+DEFAULT_PAGE_LIMIT = 20
 
 
 def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:

graphiti_core/llm_client/anthropic_client.py

@@ -20,6 +20,7 @@ import typing
 
 import anthropic
 from anthropic import AsyncAnthropic
+from pydantic import BaseModel
 
 from ..prompts.models import Message
 from .client import LLMClient
@@ -46,7 +47,9 @@ class AnthropicClient(LLMClient):
             max_retries=1,
         )
 
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         system_message = messages[0]
         user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]] + [
             {'role': 'assistant', 'content': '{'}

graphiti_core/llm_client/client.py

@@ -22,6 +22,7 @@ from abc import ABC, abstractmethod
 
 import httpx
 from diskcache import Cache
+from pydantic import BaseModel
 from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
 
 from ..prompts.models import Message
@@ -66,14 +67,18 @@ class LLMClient(ABC):
         else None,
         reraise=True,
     )
-    async def _generate_response_with_retry(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response_with_retry(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         try:
-            return await self._generate_response(messages)
+            return await self._generate_response(messages, response_model)
         except (httpx.HTTPStatusError, RateLimitError) as e:
             raise e
 
     @abstractmethod
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         pass
 
     def _get_cache_key(self, messages: list[Message]) -> str:
@@ -82,7 +87,17 @@ class LLMClient(ABC):
         key_str = f'{self.model}:{message_str}'
         return hashlib.md5(key_str.encode()).hexdigest()
 
-    async def generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
+        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}'
+            )
+
         if self.cache_enabled:
             cache_key = self._get_cache_key(messages)
 
@@ -91,7 +106,7 @@ class LLMClient(ABC):
                 logger.debug(f'Cache hit for {cache_key}')
                 return cached_response
 
-        response = await self._generate_response_with_retry(messages)
+        response = await self._generate_response_with_retry(messages, response_model)
 
         if self.cache_enabled:
             self.cache_dir.set(cache_key, response)

graphiti_core/llm_client/errors.py

@@ -21,3 +21,11 @@ class RateLimitError(Exception):
     def __init__(self, message='Rate limit exceeded. Please try again later.'):
         self.message = message
         super().__init__(self.message)
+
+
+class RefusalError(Exception):
+    """Exception raised when the LLM refuses to generate a response."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)

graphiti_core/llm_client/groq_client.py

@@ -21,6 +21,7 @@ import typing
 import groq
 from groq import AsyncGroq
 from groq.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
 
 from ..prompts.models import Message
 from .client import LLMClient
@@ -43,7 +44,9 @@ class GroqClient(LLMClient):
 
         self.client = AsyncGroq(api_key=config.api_key)
 
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    async def _generate_response(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
         msgs: list[ChatCompletionMessageParam] = []
         for m in messages:
             if m.role == 'user':

graphiti_core/llm_client/openai_client.py

@@ -14,18 +14,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import json
 import logging
 import typing
 
 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
+from .errors import RateLimitError, RefusalError
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +65,10 @@ class OpenAIClient(LLMClient):
             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()
 
@@ -75,7 +79,9 @@ class OpenAIClient(LLMClient):
         else:
             self.client = client
 
-    async def _generate_response(self, messages: list[Message]) -> dict[str, typing.Any]:
+    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:
             if m.role == 'user':
@@ -83,17 +89,33 @@ class OpenAIClient(LLMClient):
             elif m.role == 'system':
                 openai_messages.append({'role': 'system', 'content': m.content})
         try:
-            response = await self.client.chat.completions.create(
+            response = await self.client.beta.chat.completions.parse(
                 model=self.model or DEFAULT_MODEL,
                 messages=openai_messages,
                 temperature=self.temperature,
                 max_tokens=self.max_tokens,
-                response_format={'type': 'json_object'},
+                response_format=response_model,  # type: ignore
             )
-            result = response.choices[0].message.content or ''
-            return json.loads(result)
+
+            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('No response from LLM')
+        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
+    ) -> dict[str, typing.Any]:
+        response = await self._generate_response(messages, response_model)
+
+        return response

graphiti_core/nodes.py

@@ -24,10 +24,11 @@ from uuid import uuid4
 
 from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
+from typing_extensions import LiteralString
 
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import NodeNotFoundError
-from graphiti_core.helpers import DEFAULT_DATABASE
+from graphiti_core.helpers import DEFAULT_DATABASE, DEFAULT_PAGE_LIMIT
 from graphiti_core.models.nodes.node_db_queries import (
     COMMUNITY_NODE_SAVE,
     ENTITY_NODE_SAVE,
@@ -207,10 +208,21 @@ class EpisodicNode(Node):
         return episodes
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (e:Episodic) WHERE e.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
             RETURN DISTINCT
             e.content AS content,
             e.created_at AS created_at,
@@ -220,8 +232,12 @@ class EpisodicNode(Node):
             e.group_id AS group_id,
             e.source_description AS source_description,
             e.source AS source
+        ORDER BY e.uuid DESC
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -308,10 +324,21 @@ class EntityNode(Node):
         return nodes
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity) WHERE n.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             n.uuid As uuid, 
             n.name AS name,
@@ -319,8 +346,12 @@ class EntityNode(Node):
             n.group_id AS group_id,
             n.created_at AS created_at, 
             n.summary AS summary
+        ORDER BY n.uuid DESC
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -407,10 +438,21 @@ class CommunityNode(Node):
         return communities
 
     @classmethod
-    async def get_by_group_ids(cls, driver: AsyncDriver, group_ids: list[str]):
+    async def get_by_group_ids(
+        cls,
+        driver: AsyncDriver,
+        group_ids: list[str],
+        limit: int = DEFAULT_PAGE_LIMIT,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community) WHERE n.group_id IN $group_ids
+        """
+            + cursor_query
+            + """
         RETURN
             n.uuid As uuid, 
             n.name AS name,
@@ -418,8 +460,12 @@ class CommunityNode(Node):
             n.group_id AS group_id,
             n.created_at AS created_at, 
             n.summary AS summary
+        ORDER BY n.uuid DESC
+        LIMIT $limit
         """,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )

graphiti_core/prompts/dedupe_edges.py

@@ -15,11 +15,30 @@ limitations under the License.
 """
 
 import json
-from typing import Any, Protocol, TypedDict
+from typing import Any, Optional, Protocol, TypedDict
+
+from pydantic import BaseModel, Field
 
 from .models import Message, PromptFunction, PromptVersion
 
 
+class EdgeDuplicate(BaseModel):
+    is_duplicate: bool = Field(..., description='true or false')
+    uuid: Optional[str] = Field(
+        None,
+        description="uuid of the existing edge like '5d643020624c42fa9de13f97b1b3fa39' or null",
+    )
+
+
+class UniqueFact(BaseModel):
+    uuid: str = Field(..., description='unique identifier of the fact')
+    fact: str = Field(..., description='fact of a unique edge')
+
+
+class UniqueFacts(BaseModel):
+    unique_facts: list[UniqueFact]
+
+
 class Prompt(Protocol):
     edge: PromptVersion
     edge_list: PromptVersion
@@ -56,12 +75,6 @@ def edge(context: dict[str, Any]) -> list[Message]:
 
         Guidelines:
         1. The facts do not need to be completely identical to be duplicates, they just need to express the same information.
-
-        Respond with a JSON object in the following format:
-            {{
-                "is_duplicate": true or false,
-                "uuid": uuid of the existing edge like "5d643020624c42fa9de13f97b1b3fa39" or null,
-            }}
         """,
         ),
     ]
@@ -90,16 +103,6 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
         3. Facts will often discuss the same or similar relation between identical entities
         4. The final list should have only unique facts. If 3 facts are all duplicates of each other, only one of their
             facts should be in the response
-
-        Respond with a JSON object in the following format:
-        {{
-            "unique_facts": [
-                {{
-                    "uuid": "unique identifier of the fact",
-                    "fact": "fact of a unique edge"
-                }}
-            ]
-        }}
         """,
         ),
     ]

graphiti_core/prompts/dedupe_nodes.py

@@ -15,11 +15,25 @@ limitations under the License.
 """
 
 import json
-from typing import Any, Protocol, TypedDict
+from typing import Any, Optional, Protocol, TypedDict
+
+from pydantic import BaseModel, Field
 
 from .models import Message, PromptFunction, PromptVersion
 
 
+class NodeDuplicate(BaseModel):
+    is_duplicate: bool = Field(..., description='true or false')
+    uuid: Optional[str] = Field(
+        None,
+        description="uuid of the existing node like '5d643020624c42fa9de13f97b1b3fa39' or null",
+    )
+    name: str = Field(
+        ...,
+        description="Updated name of the new node (use the best name between the new node's name, an existing duplicate name, or a combination of both)",
+    )
+
+
 class Prompt(Protocol):
     node: PromptVersion
     node_list: PromptVersion