graphiti-core 0.10.0__tar.gz → 0.10.2__tar.gz

{graphiti_core-0.10.0 → graphiti_core-0.10.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.10.0
+Version: 0.10.2
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk

graphiti_core-0.10.2/graphiti_core/llm_client/anthropic_client.py

@@ -0,0 +1,332 @@
+"""
+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 os
+import typing
+from json import JSONDecodeError
+from typing import Literal
+
+import anthropic
+from anthropic import AsyncAnthropic
+from anthropic.types import MessageParam, ToolChoiceParam, ToolUnionParam
+from pydantic import BaseModel, ValidationError
+
+from ..prompts.models import Message
+from .client import LLMClient
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .errors import RateLimitError, RefusalError
+
+logger = logging.getLogger(__name__)
+
+AnthropicModel = Literal[
+    'claude-3-7-sonnet-latest',
+    'claude-3-7-sonnet-20250219',
+    'claude-3-5-haiku-latest',
+    'claude-3-5-haiku-20241022',
+    'claude-3-5-sonnet-latest',
+    'claude-3-5-sonnet-20241022',
+    'claude-3-5-sonnet-20240620',
+    'claude-3-opus-latest',
+    'claude-3-opus-20240229',
+    'claude-3-sonnet-20240229',
+    'claude-3-haiku-20240307',
+    'claude-2.1',
+    'claude-2.0',
+]
+
+DEFAULT_MODEL: AnthropicModel = 'claude-3-7-sonnet-latest'
+
+
+class AnthropicClient(LLMClient):
+    """
+    A client for the Anthropic LLM.
+
+    Args:
+        config: A configuration object for the LLM.
+        cache: Whether to cache the LLM responses.
+        client: An optional client instance to use.
+        max_tokens: The maximum number of tokens to generate.
+
+    Methods:
+        generate_response: Generate a response from the LLM.
+
+    Notes:
+        - If a LLMConfig is not provided, api_key will be pulled from the ANTHROPIC_API_KEY environment
+            variable, and all default values will be used for the LLMConfig.
+
+    """
+
+    model: AnthropicModel
+
+    def __init__(
+        self,
+        config: LLMConfig | None = None,
+        cache: bool = False,
+        client: AsyncAnthropic | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> None:
+        if config is None:
+            config = LLMConfig()
+            config.api_key = os.getenv('ANTHROPIC_API_KEY')
+            config.max_tokens = max_tokens
+
+        if config.model is None:
+            config.model = DEFAULT_MODEL
+
+        super().__init__(config, cache)
+        # Explicitly set the instance model to the config model to prevent type checking errors
+        self.model = typing.cast(AnthropicModel, config.model)
+
+        if not client:
+            self.client = AsyncAnthropic(
+                api_key=config.api_key,
+                max_retries=1,
+            )
+        else:
+            self.client = client
+
+    def _extract_json_from_text(self, text: str) -> dict[str, typing.Any]:
+        """Extract JSON from text content.
+
+        A helper method to extract JSON from text content, used when tool use fails or
+        no response_model is provided.
+
+        Args:
+            text: The text to extract JSON from
+
+        Returns:
+            Extracted JSON as a dictionary
+
+        Raises:
+            ValueError: If JSON cannot be extracted or parsed
+        """
+        try:
+            json_start = text.find('{')
+            json_end = text.rfind('}') + 1
+            if json_start >= 0 and json_end > json_start:
+                json_str = text[json_start:json_end]
+                return json.loads(json_str)
+            else:
+                raise ValueError(f'Could not extract JSON from model response: {text}')
+        except (JSONDecodeError, ValueError) as e:
+            raise ValueError(f'Could not extract JSON from model response: {text}') from e
+
+    def _create_tool(
+        self, response_model: type[BaseModel] | None = None
+    ) -> tuple[list[ToolUnionParam], ToolChoiceParam]:
+        """
+        Create a tool definition based on the response_model if provided, or a generic JSON tool if not.
+
+        Args:
+            response_model: Optional Pydantic model to use for structured output.
+
+        Returns:
+            A list containing a single tool definition for use with the Anthropic API.
+        """
+        if response_model is not None:
+            # temporary debug log
+            logger.info(f'Creating tool for response_model: {response_model}')
+            # Use the response_model to define the tool
+            model_schema = response_model.model_json_schema()
+            tool_name = response_model.__name__
+            description = model_schema.get('description', f'Extract {tool_name} information')
+        else:
+            # temporary debug log
+            logger.info('Creating generic JSON output tool')
+            # Create a generic JSON output tool
+            tool_name = 'generic_json_output'
+            description = 'Output data in JSON format'
+            model_schema = {
+                'type': 'object',
+                'additionalProperties': True,
+                'description': 'Any JSON object containing the requested information',
+            }
+
+        tool = {
+            'name': tool_name,
+            'description': description,
+            'input_schema': model_schema,
+        }
+        tool_list = [tool]
+        tool_list_cast = typing.cast(list[ToolUnionParam], tool_list)
+        tool_choice = {'type': 'tool', 'name': tool_name}
+        tool_choice_cast = typing.cast(ToolChoiceParam, tool_choice)
+        return tool_list_cast, tool_choice_cast
+
+    async def _generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int | None = None,
+    ) -> dict[str, typing.Any]:
+        """
+        Generate a response from the Anthropic LLM using tool-based approach for all requests.
+
+        Args:
+            messages: List of message objects to send to the LLM.
+            response_model: Optional Pydantic model to use for structured output.
+            max_tokens: Maximum number of tokens to generate.
+
+        Returns:
+            Dictionary containing the structured response from the LLM.
+
+        Raises:
+            RateLimitError: If the rate limit is exceeded.
+            RefusalError: If the LLM refuses to respond.
+            Exception: If an error occurs during the generation process.
+        """
+        system_message = messages[0]
+        user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]]
+        user_messages_cast = typing.cast(list[MessageParam], user_messages)
+
+        # TODO: Replace hacky min finding solution after fixing hardcoded EXTRACT_EDGES_MAX_TOKENS = 16384 in
+        # edge_operations.py. Throws errors with cheaper models that lower max_tokens.
+        max_creation_tokens: int = min(
+            max_tokens if max_tokens is not None else self.config.max_tokens,
+            DEFAULT_MAX_TOKENS,
+        )
+
+        try:
+            # Create the appropriate tool based on whether response_model is provided
+            tools, tool_choice = self._create_tool(response_model)
+            # temporary debug log
+            logger.info(f'using model: {self.model} with max_tokens: {self.max_tokens}')
+            result = await self.client.messages.create(
+                system=system_message.content,
+                max_tokens=max_creation_tokens,
+                temperature=self.temperature,
+                messages=user_messages_cast,
+                model=self.model,
+                tools=tools,
+                tool_choice=tool_choice,
+            )
+
+            # Extract the tool output from the response
+            for content_item in result.content:
+                if content_item.type == 'tool_use':
+                    if isinstance(content_item.input, dict):
+                        tool_args: dict[str, typing.Any] = content_item.input
+                    else:
+                        tool_args = json.loads(str(content_item.input))
+                    return tool_args
+
+            # If we didn't get a proper tool_use response, try to extract from text
+            # logger.debug(
+            #     f'Did not get a tool_use response, trying to extract json from text. Result: {result.content}'
+            # )
+            # temporary debug log
+            logger.info(
+                f'Did not get a tool_use response, trying to extract json from text. Result: {result.content}'
+            )
+            for content_item in result.content:
+                if content_item.type == 'text':
+                    return self._extract_json_from_text(content_item.text)
+                else:
+                    raise ValueError(
+                        f'Could not extract structured data from model response: {result.content}'
+                    )
+
+            # If we get here, we couldn't parse a structured response
+            raise ValueError(
+                f'Could not extract structured data from model response: {result.content}'
+            )
+
+        except anthropic.RateLimitError as e:
+            raise RateLimitError(f'Rate limit exceeded. Please try again later. Error: {e}') from e
+        except anthropic.APIError as e:
+            # Special case for content policy violations. We convert these to RefusalError
+            # to bypass the retry mechanism, as retrying policy-violating content will always fail.
+            # This avoids wasting API calls and provides more specific error messaging to the user.
+            if 'refused to respond' in str(e).lower():
+                raise RefusalError(str(e)) from e
+            raise e
+        except Exception as e:
+            raise e
+
+    async def generate_response(
+        self,
+        messages: list[Message],
+        response_model: type[BaseModel] | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> dict[str, typing.Any]:
+        """
+        Generate a response from the LLM.
+
+        Args:
+            messages: List of message objects to send to the LLM.
+            response_model: Optional Pydantic model to use for structured output.
+            max_tokens: Maximum number of tokens to generate.
+
+        Returns:
+            Dictionary containing the structured response from the LLM.
+
+        Raises:
+            RateLimitError: If the rate limit is exceeded.
+            RefusalError: If the LLM refuses to respond.
+            Exception: If an error occurs during the generation process.
+        """
+        retry_count = 0
+        max_retries = 2
+        last_error: Exception | None = None
+
+        while retry_count <= max_retries:
+            try:
+                response = await self._generate_response(messages, response_model, max_tokens)
+
+                # If we have a response_model, attempt to validate the response
+                if response_model is not None:
+                    # Validate the response against the response_model
+                    model_instance = response_model(**response)
+                    return model_instance.model_dump()
+
+                # If no validation needed, return the response
+                return response
+
+            except (RateLimitError, RefusalError):
+                # These errors should not trigger retries
+                raise
+            except Exception as e:
+                last_error = e
+
+                if retry_count >= max_retries:
+                    if isinstance(e, ValidationError):
+                        logger.error(
+                            f'Validation error after {retry_count}/{max_retries} attempts: {e}'
+                        )
+                    else:
+                        logger.error(f'Max retries ({max_retries}) exceeded. Last error: {e}')
+                    raise e
+
+                if isinstance(e, ValidationError):
+                    response_model_cast = typing.cast(type[BaseModel], response_model)
+                    error_context = f'The previous response was invalid. Please provide a valid {response_model_cast.__name__} object. Error: {e}'
+                else:
+                    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.'
+                    )
+
+                # Common retry logic
+                retry_count += 1
+                messages.append(Message(role='user', content=error_context))
+                logger.warning(f'Retrying after error (attempt {retry_count}/{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-0.10.0 → graphiti_core-0.10.2}/graphiti_core/llm_client/errors.py

@@ -29,3 +29,11 @@ class RefusalError(Exception):
     def __init__(self, message: str):
         self.message = message
         super().__init__(self.message)
+
+
+class EmptyResponseError(Exception):
+    """Exception raised when the LLM returns an empty response."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)

{graphiti_core-0.10.0 → graphiti_core-0.10.2}/graphiti_core/nodes.py

@@ -550,12 +550,12 @@ def get_entity_node_from_record(record: Any) -> EntityNode:
         attributes=record['attributes'],
     )
 
-    del entity_node.attributes['uuid']
-    del entity_node.attributes['name']
-    del entity_node.attributes['group_id']
-    del entity_node.attributes['name_embedding']
-    del entity_node.attributes['summary']
-    del entity_node.attributes['created_at']
+    entity_node.attributes.pop('uuid', None)
+    entity_node.attributes.pop('name', None)
+    entity_node.attributes.pop('group_id', None)
+    entity_node.attributes.pop('name_embedding', None)
+    entity_node.attributes.pop('summary', None)
+    entity_node.attributes.pop('created_at', None)
 
     return entity_node
 

{graphiti_core-0.10.0 → graphiti_core-0.10.2}/graphiti_core/prompts/summarize_nodes.py

@@ -89,6 +89,7 @@ def summarize_context(context: dict[str, Any]) -> list[Message]:
         
         Guidelines:
         1. Do not hallucinate entity property values if they cannot be found in the current context.
+        2. Only use the provided messages, entity, and entity context to set attribute values.
         
         <ENTITY>
         {context['node_name']}

{graphiti_core-0.10.0 → graphiti_core-0.10.2}/graphiti_core/utils/maintenance/node_operations.py

@@ -17,6 +17,7 @@ limitations under the License.
 import logging
 from contextlib import suppress
 from time import time
+from typing import Any
 
 import pydantic
 from pydantic import BaseModel
@@ -324,16 +325,17 @@ async def resolve_extracted_node(
         else [],
     }
 
-    summary_context = {
+    summary_context: dict[str, Any] = {
         'node_name': extracted_node.name,
         'node_summary': extracted_node.summary,
         'episode_content': episode.content if episode is not None else '',
         'previous_episodes': [ep.content for ep in previous_episodes]
         if previous_episodes is not None
         else [],
-        'attributes': [],
     }
 
+    attributes: list[dict[str, str]] = []
+
     entity_type_classes: tuple[BaseModel, ...] = tuple()
     if entity_types is not None:  # type: ignore
         entity_type_classes = entity_type_classes + tuple(
@@ -344,8 +346,15 @@ async def resolve_extracted_node(
         )
 
     for entity_type in entity_type_classes:
-        for field_name in entity_type.model_fields:
-            summary_context.get('attributes', []).append(field_name)  # type: ignore
+        for field_name, field_info in entity_type.model_fields.items():
+            attributes.append(
+                {
+                    'attribute_name': field_name,
+                    'attribute_description': field_info.description or '',
+                }
+            )
+
+    summary_context['attributes'] = attributes
 
     entity_attributes_model = pydantic.create_model(  # type: ignore
         'EntityAttributes',

{graphiti_core-0.10.0 → graphiti_core-0.10.2}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "graphiti-core"
 description = "A temporal graph building library"
-version = "0.10.0"
+version = "0.10.2"
 authors = [
     { "name" = "Paul Paliychuk", "email" = "paul@getzep.com" },
     { "name" = "Preston Rasmussen", "email" = "preston@getzep.com" },

graphiti_core-0.10.0/graphiti_core/llm_client/anthropic_client.py

@@ -1,79 +0,0 @@
-"""
-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
-
-import anthropic
-from anthropic import AsyncAnthropic
-from pydantic import BaseModel
-
-from ..prompts.models import Message
-from .client import LLMClient
-from .config import LLMConfig
-from .errors import RateLimitError
-
-logger = logging.getLogger(__name__)
-
-DEFAULT_MODEL = 'claude-3-7-sonnet-latest'
-DEFAULT_MAX_TOKENS = 8192
-
-
-class AnthropicClient(LLMClient):
-    def __init__(self, config: LLMConfig | None = None, cache: bool = False):
-        if config is None:
-            config = LLMConfig(max_tokens=DEFAULT_MAX_TOKENS)
-        elif config.max_tokens is None:
-            config.max_tokens = DEFAULT_MAX_TOKENS
-        super().__init__(config, cache)
-
-        self.client = AsyncAnthropic(
-            api_key=config.api_key,
-            # we'll use tenacity to retry
-            max_retries=1,
-        )
-
-    async def _generate_response(
-        self,
-        messages: list[Message],
-        response_model: type[BaseModel] | None = None,
-        max_tokens: int = DEFAULT_MAX_TOKENS,
-    ) -> dict[str, typing.Any]:
-        system_message = messages[0]
-        user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]] + [
-            {'role': 'assistant', 'content': '{'}
-        ]
-
-        # Ensure max_tokens is not greater than config.max_tokens or DEFAULT_MAX_TOKENS
-        max_tokens = min(max_tokens, self.config.max_tokens, DEFAULT_MAX_TOKENS)
-
-        try:
-            result = await self.client.messages.create(
-                system='Only include JSON in the response. Do not include any additional text or explanation of the content.\n'
-                + system_message.content,
-                max_tokens=max_tokens,
-                temperature=self.temperature,
-                messages=user_messages,  # type: ignore
-                model=self.model or DEFAULT_MODEL,
-            )
-
-            return json.loads('{' + result.content[0].text)  # type: ignore
-        except anthropic.RateLimitError as e:
-            raise RateLimitError from e
-        except Exception as e:
-            logger.error(f'Error in generating LLM response: {e}')
-            raise