graphiti-core 0.10.1__py3-none-any.whl → 0.10.2__py3-none-any.whl

graphiti_core/llm_client/anthropic_client.py

@@ -16,64 +16,317 @@ 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 pydantic import BaseModel
+from anthropic.types import MessageParam, ToolChoiceParam, ToolUnionParam
+from pydantic import BaseModel, ValidationError
 
 from ..prompts.models import Message
 from .client import LLMClient
-from .config import LLMConfig
-from .errors import RateLimitError
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .errors import RateLimitError, RefusalError
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MODEL = 'claude-3-7-sonnet-latest'
-DEFAULT_MAX_TOKENS = 8192
+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):
-    def __init__(self, config: LLMConfig | None = None, cache: bool = False):
+    """
+    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(max_tokens=DEFAULT_MAX_TOKENS)
-        elif config.max_tokens is None:
-            config.max_tokens = DEFAULT_MAX_TOKENS
+            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)
 
-        self.client = AsyncAnthropic(
-            api_key=config.api_key,
-            # we'll use tenacity to retry
-            max_retries=1,
-        )
+        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 = DEFAULT_MAX_TOKENS,
+        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:]] + [
-            {'role': 'assistant', 'content': '{'}
-        ]
+        user_messages = [{'role': m.role, 'content': m.content} for m in messages[1:]]
+        user_messages_cast = typing.cast(list[MessageParam], user_messages)
 
-        # 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)
+        # 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='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,
+                system=system_message.content,
+                max_tokens=max_creation_tokens,
                 temperature=self.temperature,
-                messages=user_messages,  # type: ignore
-                model=self.model or DEFAULT_MODEL,
+                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}'
             )
 
-            return json.loads('{' + result.content[0].text)  # type: ignore
         except anthropic.RateLimitError as e:
-            raise RateLimitError from 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:
-            logger.error(f'Error in generating LLM response: {e}')
-            raise
+            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/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/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/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.1.dist-info → graphiti_core-0.10.2.dist-info}/METADATA

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

{graphiti_core-0.10.1.dist-info → graphiti_core-0.10.2.dist-info}/RECORD

@@ -13,10 +13,10 @@ graphiti_core/errors.py,sha256=Nib1uQx2cO_VOizupmRjpFfmuRg-hFAVqTtZAuBehR8,2405
 graphiti_core/graphiti.py,sha256=mdIkciD4o3fjD5PapshKvNv0F6eE76rXZCAB00vcDG8,30807
 graphiti_core/helpers.py,sha256=7BQzUBFmoBDA2OIDdFtoN4W-vXOhPRIsF0uDb7PsNi0,2913
 graphiti_core/llm_client/__init__.py,sha256=PA80TSMeX-sUXITXEAxMDEt3gtfZgcJrGJUcyds1mSo,207
-graphiti_core/llm_client/anthropic_client.py,sha256=dTM8rKhk9TZAU4O-0jFMivOwJvWM-gHpp5gLmuJHiGQ,2723
+graphiti_core/llm_client/anthropic_client.py,sha256=qVLtRdlYksdl221lIBv7saOmuJtgG5p3af0HUsVxseM,12926
 graphiti_core/llm_client/client.py,sha256=jLyrn--opI6_ekCgmCaHLvC9XRhU25HYeoKVIivPHgQ,5507
 graphiti_core/llm_client/config.py,sha256=JO-biZwGwakTKGKyNQQqt9fjqyXxtqP-h0sB4wsJ2Kk,2339
-graphiti_core/llm_client/errors.py,sha256=Vk0mj2SgNDg8E8p7m1UyUaerqLPNLCDKPVsMEnOSBdQ,1028
+graphiti_core/llm_client/errors.py,sha256=pn6brRiLW60DAUIXJYKBT6MInrS4ueuH1hNLbn_JbQo,1243
 graphiti_core/llm_client/gemini_client.py,sha256=uibmwppDgkEv60FsIhS-oakuafTUFgpu3qr5Kdcbhz4,7321
 graphiti_core/llm_client/groq_client.py,sha256=EesX0_iFOIvvIc3ql6Xa8EOm0dbGJ_o3VpqyDM83mKg,2498
 graphiti_core/llm_client/openai_client.py,sha256=ZqkBa127i5gCIA2UoiXuPDM-357xFK_uwKlgafDhVyQ,6989
@@ -39,7 +39,7 @@ graphiti_core/prompts/invalidate_edges.py,sha256=DV2mEyIhhjc0hdKEMFLQMeG0FiUCkv_
 graphiti_core/prompts/lib.py,sha256=DCyHePM4_q-CptTpEXGO_dBv9k7xDtclEaB1dGu7EcI,4092
 graphiti_core/prompts/models.py,sha256=NgxdbPHJpBEcpbXovKyScgpBc73Q-GIW-CBDlBtDjto,894
 graphiti_core/prompts/prompt_helpers.py,sha256=-9TABwIcIQUVHcNANx6wIZd-FT2DgYKyGTfx4IGYq2I,64
-graphiti_core/prompts/summarize_nodes.py,sha256=PeA1Taov5KBNNBKgrCPeF1tLg4_SMgT-Ilz2P6xbx-M,4051
+graphiti_core/prompts/summarize_nodes.py,sha256=CDXeWCbv34BcHn6I9lkYN-YqQv98bSt0cOZMnVho2Lk,4146
 graphiti_core/py.typed,sha256=vlmmzQOt7bmeQl9L3XJP4W6Ry0iiELepnOrinKz5KQg,79
 graphiti_core/search/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 graphiti_core/search/search.py,sha256=3D460AUKNn6llVkTghpautP925VDKwaQXXutntZTvRg,15745
@@ -55,11 +55,11 @@ graphiti_core/utils/maintenance/__init__.py,sha256=vW4H1KyapTl-OOz578uZABYcpND4w
 graphiti_core/utils/maintenance/community_operations.py,sha256=pQUv0pfInC1Pho7C4BN8gC3_bks7wRAZpJn2bmw6gT8,10008
 graphiti_core/utils/maintenance/edge_operations.py,sha256=9i0PBgaW3dLPTLmx-9j1W86Rb4sPc1bG4Y3TjPn07Gg,12794
 graphiti_core/utils/maintenance/graph_data_operations.py,sha256=F3dkXg63nJU7OFLzaiT67a2_kFiWSU9Vr59iIHgsXQs,7030
-graphiti_core/utils/maintenance/node_operations.py,sha256=WhZQixx05dAFFQAd5KTXJ8Gc8YCahCAAW9uPCucjyBw,15556
+graphiti_core/utils/maintenance/node_operations.py,sha256=7TLBBJ5m-3_UK53_UNm0hQJh3iiJP0tIVG7Sd6oOioE,15801
 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/utils/ontology_utils/entity_types_utils.py,sha256=QJX5cG0GSSNF_Mm_yrldr69wjVAbN_MxLhOSznz85Hk,1279
-graphiti_core-0.10.1.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
-graphiti_core-0.10.1.dist-info/METADATA,sha256=fAqTBXkTPCDh-bw0798LlCmjprM3ugqejB-HopOGR4o,14659
-graphiti_core-0.10.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-graphiti_core-0.10.1.dist-info/RECORD,,
+graphiti_core-0.10.2.dist-info/LICENSE,sha256=KCUwCyDXuVEgmDWkozHyniRyWjnWUWjkuDHfU6o3JlA,11325
+graphiti_core-0.10.2.dist-info/METADATA,sha256=TDLWSn0xZGDl8QKnWqaw9tw_aKRRGTNLqPECCSLb-j8,14659
+graphiti_core-0.10.2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+graphiti_core-0.10.2.dist-info/RECORD,,