graphiti-core 0.12.0rc1__py3-none-any.whl → 0.24.3__py3-none-any.whl

graphiti_core/graphiti_types.py

@@ -14,18 +14,20 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from neo4j import AsyncDriver
 from pydantic import BaseModel, ConfigDict
 
 from graphiti_core.cross_encoder import CrossEncoderClient
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.llm_client import LLMClient
+from graphiti_core.tracer import Tracer
 
 
 class GraphitiClients(BaseModel):
-    driver: AsyncDriver
+    driver: GraphDriver
     llm_client: LLMClient
     embedder: EmbedderClient
     cross_encoder: CrossEncoderClient
+    tracer: Tracer
 
     model_config = ConfigDict(arbitrary_types_allowed=True)

graphiti_core/helpers.py

@@ -16,30 +16,47 @@ limitations under the License.
 
 import asyncio
 import os
+import re
 from collections.abc import Coroutine
 from datetime import datetime
+from typing import Any
 
 import numpy as np
 from dotenv import load_dotenv
 from neo4j import time as neo4j_time
 from numpy._typing import NDArray
-from typing_extensions import LiteralString
+from pydantic import BaseModel
+
+from graphiti_core.driver.driver import GraphProvider
+from graphiti_core.errors import GroupIdValidationError
 
 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 = int(os.getenv('MAX_REFLEXION_ITERATIONS', 0))
 DEFAULT_PAGE_LIMIT = 20
 
-RUNTIME_QUERY: LiteralString = (
-    'CYPHER runtime = parallel parallelRuntimeSupport=all\n' if USE_PARALLEL_RUNTIME else ''
-)
 
+def parse_db_date(input_date: neo4j_time.DateTime | str | None) -> datetime | None:
+    if isinstance(input_date, neo4j_time.DateTime):
+        return input_date.to_native()
+
+    if isinstance(input_date, str):
+        return datetime.fromisoformat(input_date)
+
+    return input_date
 
-def parse_db_date(neo_date: neo4j_time.DateTime | None) -> datetime | None:
-    return neo_date.to_native() if neo_date else None
+
+def get_default_group_id(provider: GraphProvider) -> str:
+    """
+    This function differentiates the default group id based on the database type.
+    For most databases, the default group id is an empty string, while there are database types that require a specific default group id.
+    """
+    if provider == GraphProvider.FALKORDB:
+        return '\\_'
+    else:
+        return ''
 
 
 def lucene_sanitize(query: str) -> str:
@@ -88,12 +105,72 @@ def normalize_l2(embedding: list[float]) -> NDArray:
 # 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)
+    max_coroutines: int | None = None,
+) -> list[Any]:
+    semaphore = asyncio.Semaphore(max_coroutines or SEMAPHORE_LIMIT)
 
     async def _wrap_coroutine(coroutine):
         async with semaphore:
             return await coroutine
 
     return await asyncio.gather(*(_wrap_coroutine(coroutine) for coroutine in coroutines))
+
+
+def validate_group_id(group_id: str | None) -> bool:
+    """
+    Validate that a group_id contains only ASCII alphanumeric characters, dashes, and underscores.
+
+    Args:
+        group_id: The group_id to validate
+
+    Returns:
+        True if valid, False otherwise
+
+    Raises:
+        GroupIdValidationError: If group_id contains invalid characters
+    """
+
+    # Allow empty string (default case)
+    if not group_id:
+        return True
+
+    # Check if string contains only ASCII alphanumeric characters, dashes, or underscores
+    # Pattern matches: letters (a-z, A-Z), digits (0-9), hyphens (-), and underscores (_)
+    if not re.match(r'^[a-zA-Z0-9_-]+$', group_id):
+        raise GroupIdValidationError(group_id)
+
+    return True
+
+
+def validate_excluded_entity_types(
+    excluded_entity_types: list[str] | None, entity_types: dict[str, type[BaseModel]] | None = None
+) -> bool:
+    """
+    Validate that excluded entity types are valid type names.
+
+    Args:
+        excluded_entity_types: List of entity type names to exclude
+        entity_types: Dictionary of available custom entity types
+
+    Returns:
+        True if valid
+
+    Raises:
+        ValueError: If any excluded type names are invalid
+    """
+    if not excluded_entity_types:
+        return True
+
+    # Build set of available type names
+    available_types = {'Entity'}  # Default type is always available
+    if entity_types:
+        available_types.update(entity_types.keys())
+
+    # Check for invalid type names
+    invalid_types = set(excluded_entity_types) - available_types
+    if invalid_types:
+        raise ValueError(
+            f'Invalid excluded entity types: {sorted(invalid_types)}. Available types: {sorted(available_types)}'
+        )
+
+    return True

graphiti_core/llm_client/__init__.py

@@ -1,3 +1,19 @@
+"""
+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.
+"""
+
 from .client import LLMClient
 from .config import LLMConfig
 from .errors import RateLimitError

graphiti_core/llm_client/anthropic_client.py

@@ -19,11 +19,8 @@ import logging
 import os
 import typing
 from json import JSONDecodeError
-from typing import Literal
+from typing import TYPE_CHECKING, Literal
 
-import anthropic
-from anthropic import AsyncAnthropic
-from anthropic.types import MessageParam, ToolChoiceParam, ToolUnionParam
 from pydantic import BaseModel, ValidationError
 
 from ..prompts.models import Message
@@ -31,9 +28,28 @@ from .client import LLMClient
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
+if TYPE_CHECKING:
+    import anthropic
+    from anthropic import AsyncAnthropic
+    from anthropic.types import MessageParam, ToolChoiceParam, ToolUnionParam
+else:
+    try:
+        import anthropic
+        from anthropic import AsyncAnthropic
+        from anthropic.types import MessageParam, ToolChoiceParam, ToolUnionParam
+    except ImportError:
+        raise ImportError(
+            'anthropic is required for AnthropicClient. '
+            'Install it with: pip install graphiti-core[anthropic]'
+        ) from None
+
+
 logger = logging.getLogger(__name__)
 
 AnthropicModel = Literal[
+    'claude-sonnet-4-5-latest',
+    'claude-sonnet-4-5-20250929',
+    'claude-haiku-4-5-latest',
     'claude-3-7-sonnet-latest',
     'claude-3-7-sonnet-20250219',
     'claude-3-5-haiku-latest',
@@ -49,7 +65,39 @@ AnthropicModel = Literal[
     'claude-2.0',
 ]
 
-DEFAULT_MODEL: AnthropicModel = 'claude-3-7-sonnet-latest'
+DEFAULT_MODEL: AnthropicModel = 'claude-haiku-4-5-latest'
+
+# Maximum output tokens for different Anthropic models
+# Based on official Anthropic documentation (as of 2025)
+# Note: These represent standard limits without beta headers.
+# Some models support higher limits with additional configuration (e.g., Claude 3.7 supports
+# 128K with 'anthropic-beta: output-128k-2025-02-19' header, but this is not currently implemented).
+ANTHROPIC_MODEL_MAX_TOKENS = {
+    # Claude 4.5 models - 64K tokens
+    'claude-sonnet-4-5-latest': 65536,
+    'claude-sonnet-4-5-20250929': 65536,
+    'claude-haiku-4-5-latest': 65536,
+    # Claude 3.7 models - standard 64K tokens
+    'claude-3-7-sonnet-latest': 65536,
+    'claude-3-7-sonnet-20250219': 65536,
+    # Claude 3.5 models
+    'claude-3-5-haiku-latest': 8192,
+    'claude-3-5-haiku-20241022': 8192,
+    'claude-3-5-sonnet-latest': 8192,
+    'claude-3-5-sonnet-20241022': 8192,
+    'claude-3-5-sonnet-20240620': 8192,
+    # Claude 3 models - 4K tokens
+    'claude-3-opus-latest': 4096,
+    'claude-3-opus-20240229': 4096,
+    'claude-3-sonnet-20240229': 4096,
+    'claude-3-haiku-20240307': 4096,
+    # Claude 2 models - 4K tokens
+    'claude-2.1': 4096,
+    'claude-2.0': 4096,
+}
+
+# Default max tokens for models not in the mapping
+DEFAULT_ANTHROPIC_MAX_TOKENS = 8192
 
 
 class AnthropicClient(LLMClient):
@@ -164,6 +212,45 @@ class AnthropicClient(LLMClient):
         tool_choice_cast = typing.cast(ToolChoiceParam, tool_choice)
         return tool_list_cast, tool_choice_cast
 
+    def _get_max_tokens_for_model(self, model: str) -> int:
+        """Get the maximum output tokens for a specific Anthropic model.
+
+        Args:
+            model: The model name to look up
+
+        Returns:
+            int: The maximum output tokens for the model
+        """
+        return ANTHROPIC_MODEL_MAX_TOKENS.get(model, DEFAULT_ANTHROPIC_MAX_TOKENS)
+
+    def _resolve_max_tokens(self, requested_max_tokens: int | None, model: str) -> int:
+        """
+        Resolve the maximum output tokens to use based on precedence rules.
+
+        Precedence order (highest to lowest):
+        1. Explicit max_tokens parameter passed to generate_response()
+        2. Instance max_tokens set during client initialization
+        3. Model-specific maximum tokens from ANTHROPIC_MODEL_MAX_TOKENS mapping
+        4. DEFAULT_ANTHROPIC_MAX_TOKENS as final fallback
+
+        Args:
+            requested_max_tokens: The max_tokens parameter passed to generate_response()
+            model: The model name to look up model-specific limits
+
+        Returns:
+            int: The resolved maximum tokens to use
+        """
+        # 1. Use explicit parameter if provided
+        if requested_max_tokens is not None:
+            return requested_max_tokens
+
+        # 2. Use instance max_tokens if set during initialization
+        if self.max_tokens is not None:
+            return self.max_tokens
+
+        # 3. Use model-specific maximum or return DEFAULT_ANTHROPIC_MAX_TOKENS
+        return self._get_max_tokens_for_model(model)
+
     async def _generate_response(
         self,
         messages: list[Message],
@@ -191,12 +278,9 @@ class AnthropicClient(LLMClient):
         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,
-        )
+        # Resolve max_tokens dynamically based on the model's capabilities
+        # This allows different models to use their full output capacity
+        max_creation_tokens: int = self._resolve_max_tokens(max_tokens, self.model)
 
         try:
             # Create the appropriate tool based on whether response_model is provided
@@ -252,6 +336,8 @@ class AnthropicClient(LLMClient):
         response_model: type[BaseModel] | None = None,
         max_tokens: int | None = None,
         model_size: ModelSize = ModelSize.medium,
+        group_id: str | None = None,
+        prompt_name: str | None = None,
     ) -> dict[str, typing.Any]:
         """
         Generate a response from the LLM.
@@ -272,55 +358,72 @@ class AnthropicClient(LLMClient):
         if max_tokens is None:
             max_tokens = self.max_tokens
 
-        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, model_size
-                )
-
-                # 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
+        # Wrap entire operation in tracing span
+        with self.tracer.start_span('llm.generate') as span:
+            attributes = {
+                'llm.provider': 'anthropic',
+                'model.size': model_size.value,
+                'max_tokens': max_tokens,
+            }
+            if prompt_name:
+                attributes['prompt.name'] = prompt_name
+            span.add_attributes(attributes)
+
+            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, model_size
+                    )
 
-            except (RateLimitError, RefusalError):
-                # These errors should not trigger retries
-                raise
-            except Exception as e:
-                last_error = e
+                    # 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
+                    span.set_status('error', str(last_error))
+                    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}')
+                        span.set_status('error', str(e))
+                        span.record_exception(e)
+                        raise e
 
-                if retry_count >= max_retries:
                     if isinstance(e, ValidationError):
-                        logger.error(
-                            f'Validation error after {retry_count}/{max_retries} attempts: {e}'
-                        )
+                        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:
-                        logger.error(f'Max retries ({max_retries}) exceeded. Last error: {e}')
-                    raise e
+                        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.'
+                        )
 
-                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}'
                     )
 
-                # 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')
+            # If we somehow get here, raise the last error
+            span.set_status('error', str(last_error))
+            raise last_error or Exception('Max retries exceeded with no specific error')

graphiti_core/llm_client/azure_openai_client.py

@@ -0,0 +1,115 @@
+"""
+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 logging
+from typing import ClassVar
+
+from openai import AsyncAzureOpenAI, AsyncOpenAI
+from openai.types.chat import ChatCompletionMessageParam
+from pydantic import BaseModel
+
+from .config import DEFAULT_MAX_TOKENS, LLMConfig
+from .openai_base_client import BaseOpenAIClient
+
+logger = logging.getLogger(__name__)
+
+
+class AzureOpenAILLMClient(BaseOpenAIClient):
+    """Wrapper class for Azure OpenAI that implements the LLMClient interface.
+
+    Supports both AsyncAzureOpenAI and AsyncOpenAI (with Azure v1 API endpoint).
+    """
+
+    # Class-level constants
+    MAX_RETRIES: ClassVar[int] = 2
+
+    def __init__(
+        self,
+        azure_client: AsyncAzureOpenAI | AsyncOpenAI,
+        config: LLMConfig | None = None,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+        reasoning: str | None = None,
+        verbosity: str | None = None,
+    ):
+        super().__init__(
+            config,
+            cache=False,
+            max_tokens=max_tokens,
+            reasoning=reasoning,
+            verbosity=verbosity,
+        )
+        self.client = azure_client
+
+    async def _create_structured_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel],
+        reasoning: str | None,
+        verbosity: str | None,
+    ):
+        """Create a structured completion using Azure OpenAI's responses.parse API."""
+        supports_reasoning = self._supports_reasoning_features(model)
+        request_kwargs = {
+            'model': model,
+            'input': messages,
+            'max_output_tokens': max_tokens,
+            'text_format': response_model,  # type: ignore
+        }
+
+        temperature_value = temperature if not supports_reasoning else None
+        if temperature_value is not None:
+            request_kwargs['temperature'] = temperature_value
+
+        if supports_reasoning and reasoning:
+            request_kwargs['reasoning'] = {'effort': reasoning}  # type: ignore
+
+        if supports_reasoning and verbosity:
+            request_kwargs['text'] = {'verbosity': verbosity}  # type: ignore
+
+        return await self.client.responses.parse(**request_kwargs)
+
+    async def _create_completion(
+        self,
+        model: str,
+        messages: list[ChatCompletionMessageParam],
+        temperature: float | None,
+        max_tokens: int,
+        response_model: type[BaseModel] | None = None,
+    ):
+        """Create a regular completion with JSON format using Azure OpenAI."""
+        supports_reasoning = self._supports_reasoning_features(model)
+
+        request_kwargs = {
+            'model': model,
+            'messages': messages,
+            'max_tokens': max_tokens,
+            'response_format': {'type': 'json_object'},
+        }
+
+        temperature_value = temperature if not supports_reasoning else None
+        if temperature_value is not None:
+            request_kwargs['temperature'] = temperature_value
+
+        return await self.client.chat.completions.create(**request_kwargs)
+
+    @staticmethod
+    def _supports_reasoning_features(model: str) -> bool:
+        """Return True when the Azure model supports reasoning/verbosity options."""
+        reasoning_prefixes = ('o1', 'o3', 'gpt-5')
+        return model.startswith(reasoning_prefixes)