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

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,19 @@ 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
+from .errors import RateLimitError, RefusalError
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +54,9 @@ class OpenAIClient(LLMClient):
             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
     ):
@@ -65,6 +69,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,25 +83,81 @@ 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:
+            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(
+            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(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(
+        self, messages: list[Message], response_model: type[BaseModel] | None = None
+    ) -> dict[str, typing.Any]:
+        retry_count = 0
+        last_error = None
+
+        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/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/nodes.py

@@ -16,7 +16,7 @@ limitations under the License.
 
 import logging
 from abc import ABC, abstractmethod
-from datetime import datetime, timezone
+from datetime import datetime
 from enum import Enum
 from time import time
 from typing import Any
@@ -24,6 +24,7 @@ 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
@@ -33,6 +34,7 @@ from graphiti_core.models.nodes.node_db_queries import (
     ENTITY_NODE_SAVE,
     EPISODIC_NODE_SAVE,
 )
+from graphiti_core.utils.datetime_utils import utc_now
 
 logger = logging.getLogger(__name__)
 
@@ -78,7 +80,7 @@ class Node(BaseModel, ABC):
     name: str = Field(description='name of the node')
     group_id: str = Field(description='partition of the graph')
     labels: list[str] = Field(default_factory=list)
-    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    created_at: datetime = Field(default_factory=lambda: utc_now())
 
     @abstractmethod
     async def save(self, driver: AsyncDriver): ...
@@ -207,10 +209,22 @@ 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 | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND e.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None 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 +234,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_query,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -308,10 +326,22 @@ 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 | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None 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 +349,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_query,
             group_ids=group_ids,
+            created_at=created_at,
+            limit=limit,
             database_=DEFAULT_DATABASE,
             routing_='r',
         )
@@ -407,10 +441,22 @@ 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 | None = None,
+        created_at: datetime | None = None,
+    ):
+        cursor_query: LiteralString = 'AND n.created_at < $created_at' if created_at else ''
+        limit_query: LiteralString = 'LIMIT $limit' if limit is not None 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 +464,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_query,
             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

graphiti_core/prompts/eval.py

@@ -17,9 +17,26 @@ limitations under the License.
 import json
 from typing import Any, Protocol, TypedDict
 
+from pydantic import BaseModel, Field
+
 from .models import Message, PromptFunction, PromptVersion
 
 
+class QueryExpansion(BaseModel):
+    query: str = Field(..., description='query optimized for database search')
+
+
+class QAResponse(BaseModel):
+    ANSWER: str = Field(..., description='how Alice would answer the question')
+
+
+class EvalResponse(BaseModel):
+    is_correct: bool = Field(..., description='boolean if the answer is correct or incorrect')
+    reasoning: str = Field(
+        ..., description='why you determined the response was correct or incorrect'
+    )
+
+
 class Prompt(Protocol):
     qa_prompt: PromptVersion
     eval_prompt: PromptVersion
@@ -41,10 +58,6 @@ def query_expansion(context: dict[str, Any]) -> list[Message]:
     <QUESTION>
     {json.dumps(context['query'])}
     </QUESTION>
-    respond with a JSON object in the following format:
-    {{
-        "query": "query optimized for database search"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),
@@ -67,10 +80,6 @@ def qa_prompt(context: dict[str, Any]) -> list[Message]:
     <QUESTION>
     {context['query']}
     </QUESTION>
-    respond with a JSON object in the following format:
-    {{
-        "ANSWER": "how Alice would answer the question"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),
@@ -96,12 +105,6 @@ def eval_prompt(context: dict[str, Any]) -> list[Message]:
     <RESPONSE>
     {context['response']}
     </RESPONSE>
-    
-    respond with a JSON object in the following format:
-    {{
-        "is_correct": "boolean if the answer is correct or incorrect"
-        "reasoning": "why you determined the response was correct or incorrect"
-    }}
     """
     return [
         Message(role='system', content=sys_prompt),

graphiti_core/prompts/extract_edge_dates.py

@@ -14,11 +14,24 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-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 EdgeDates(BaseModel):
+    valid_at: Optional[str] = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact became true or was established. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.',
+    )
+    invalid_at: Optional[str] = Field(
+        None,
+        description='The date and time when the relationship described by the edge fact stopped being true or ended. YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null.',
+    )
+
+
 class Prompt(Protocol):
     v1: PromptVersion
 
@@ -60,7 +73,7 @@ def v1(context: dict[str, Any]) -> list[Message]:
             Analyze the conversation and determine if there are dates that are part of the edge fact. Only set dates if they explicitly relate to the formation or alteration of the relationship itself.
 
             Guidelines:
-            1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SSZ) for datetimes.
+            1. Use ISO 8601 format (YYYY-MM-DDTHH:MM:SS.SSSSSSZ) for datetimes.
             2. Use the reference timestamp as the current time when determining the valid_at and invalid_at dates.
             3. If the fact is written in the present tense, use the Reference Timestamp for the valid_at date
             4. If no temporal information is found that establishes or changes the relationship, leave the fields as null.
@@ -69,11 +82,6 @@ def v1(context: dict[str, Any]) -> list[Message]:
             7. If only a date is mentioned without a specific time, use 00:00:00 (midnight) for that date.
             8. If only year is mentioned, use January 1st of that year at 00:00:00.
             9. Always include the time zone offset (use Z for UTC if no specific time zone is mentioned).
-            Respond with a JSON object:
-            {{
-                "valid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
-                "invalid_at": "YYYY-MM-DDTHH:MM:SS.SSSSSSZ or null",
-            }}
             """,
         ),
     ]