graphiti-core 0.17.4__py3-none-any.whl → 0.25.3__py3-none-any.whl

graphiti_core/llm_client/gemini_client.py

@@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, ClassVar
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import LLMConfig, ModelSize
 from .errors import RateLimitError
 
@@ -45,7 +45,7 @@ else:
 logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = 'gemini-2.5-flash'
-DEFAULT_SMALL_MODEL = 'gemini-2.5-flash-lite-preview-06-17'
+DEFAULT_SMALL_MODEL = 'gemini-2.5-flash-lite'
 
 # Maximum output tokens for different Gemini models
 GEMINI_MODEL_MAX_TOKENS = {
@@ -53,7 +53,6 @@ GEMINI_MODEL_MAX_TOKENS = {
     'gemini-2.5-pro': 65536,
     'gemini-2.5-flash': 65536,
     'gemini-2.5-flash-lite': 64000,
-    'models/gemini-2.5-flash-lite-preview-06-17': 64000,
     # Gemini 2.0 models
     'gemini-2.0-flash': 8192,
     'gemini-2.0-flash-lite': 8192,
@@ -357,6 +356,8 @@ class GeminiClient(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 Gemini language model with retry logic and error handling.
@@ -367,62 +368,79 @@ class GeminiClient(LLMClient):
             response_model (type[BaseModel] | None): An optional Pydantic model to parse the response into.
             max_tokens (int | None): The maximum number of tokens to generate in the response.
             model_size (ModelSize): The size of the model to use (small or medium).
+            group_id (str | None): Optional partition identifier for the graph.
+            prompt_name (str | None): Optional name of the prompt for tracing.
 
         Returns:
             dict[str, typing.Any]: The response from the language model.
         """
-        retry_count = 0
-        last_error = None
-        last_output = None
-
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
-
-        while retry_count < self.MAX_RETRIES:
-            try:
-                response = await self._generate_response(
-                    messages=messages,
-                    response_model=response_model,
-                    max_tokens=max_tokens,
-                    model_size=model_size,
-                )
-                last_output = (
-                    response.get('content')
-                    if isinstance(response, dict) and 'content' in response
-                    else None
-                )
-                return response
-            except RateLimitError as e:
-                # Rate limit errors should not trigger retries (fail fast)
-                raise e
-            except Exception as e:
-                last_error = e
-
-                # Check if this is a safety block - these typically shouldn't be retried
-                error_text = str(e) or (str(e.__cause__) if e.__cause__ else '')
-                if 'safety' in error_text.lower() or 'blocked' in error_text.lower():
-                    logger.warning(f'Content blocked by safety filters: {e}')
-                    raise Exception(f'Content blocked by safety filters: {e}') from e
-
-                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 exit the loop without returning, all retries are exhausted
-        logger.error('🦀 LLM generation failed and retries are exhausted.')
-        logger.error(self._get_failed_generation_log(messages, last_output))
-        logger.error(f'Max retries ({self.MAX_RETRIES}) exceeded. Last error: {last_error}')
-        raise last_error or Exception('Max retries exceeded')
+        messages[0].content += get_extraction_language_instruction(group_id)
+
+        # Wrap entire operation in tracing span
+        with self.tracer.start_span('llm.generate') as span:
+            attributes = {
+                'llm.provider': 'gemini',
+                'model.size': model_size.value,
+                'max_tokens': max_tokens or self.max_tokens,
+            }
+            if prompt_name:
+                attributes['prompt.name'] = prompt_name
+            span.add_attributes(attributes)
+
+            retry_count = 0
+            last_error = None
+            last_output = None
+
+            while retry_count < self.MAX_RETRIES:
+                try:
+                    response = await self._generate_response(
+                        messages=messages,
+                        response_model=response_model,
+                        max_tokens=max_tokens,
+                        model_size=model_size,
+                    )
+                    last_output = (
+                        response.get('content')
+                        if isinstance(response, dict) and 'content' in response
+                        else None
+                    )
+                    return response
+                except RateLimitError as e:
+                    # Rate limit errors should not trigger retries (fail fast)
+                    span.set_status('error', str(e))
+                    raise e
+                except Exception as e:
+                    last_error = e
+
+                    # Check if this is a safety block - these typically shouldn't be retried
+                    error_text = str(e) or (str(e.__cause__) if e.__cause__ else '')
+                    if 'safety' in error_text.lower() or 'blocked' in error_text.lower():
+                        logger.warning(f'Content blocked by safety filters: {e}')
+                        span.set_status('error', str(e))
+                        raise Exception(f'Content blocked by safety filters: {e}') from e
+
+                    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 exit the loop without returning, all retries are exhausted
+            logger.error('🦀 LLM generation failed and retries are exhausted.')
+            logger.error(self._get_failed_generation_log(messages, last_output))
+            logger.error(f'Max retries ({self.MAX_RETRIES}) exceeded. Last error: {last_error}')
+            span.set_status('error', str(last_error))
+            span.record_exception(last_error) if last_error else None
+            raise last_error or Exception('Max retries exceeded')

graphiti_core/llm_client/openai_base_client.py

@@ -25,7 +25,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from ..prompts.models import Message
-from .client import MULTILINGUAL_EXTRACTION_RESPONSES, LLMClient
+from .client import LLMClient, get_extraction_language_instruction
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError, RefusalError
 
@@ -33,6 +33,8 @@ logger = logging.getLogger(__name__)
 
 DEFAULT_MODEL = 'gpt-4.1-mini'
 DEFAULT_SMALL_MODEL = 'gpt-4.1-nano'
+DEFAULT_REASONING = 'minimal'
+DEFAULT_VERBOSITY = 'low'
 
 
 class BaseOpenAIClient(LLMClient):
@@ -51,6 +53,8 @@ class BaseOpenAIClient(LLMClient):
         config: LLMConfig | None = None,
         cache: bool = False,
         max_tokens: int = DEFAULT_MAX_TOKENS,
+        reasoning: str | None = DEFAULT_REASONING,
+        verbosity: str | None = DEFAULT_VERBOSITY,
     ):
         if cache:
             raise NotImplementedError('Caching is not implemented for OpenAI-based clients')
@@ -60,6 +64,8 @@ class BaseOpenAIClient(LLMClient):
 
         super().__init__(config, cache)
         self.max_tokens = max_tokens
+        self.reasoning = reasoning
+        self.verbosity = verbosity
 
     @abstractmethod
     async def _create_completion(
@@ -81,6 +87,8 @@ class BaseOpenAIClient(LLMClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel],
+        reasoning: str | None,
+        verbosity: str | None,
     ) -> Any:
         """Create a structured completion using the specific client implementation."""
         pass
@@ -107,10 +115,10 @@ class BaseOpenAIClient(LLMClient):
 
     def _handle_structured_response(self, response: Any) -> dict[str, Any]:
         """Handle structured response parsing and validation."""
-        response_object = response.choices[0].message
+        response_object = response.output_text
 
-        if response_object.parsed:
-            return response_object.parsed.model_dump()
+        if response_object:
+            return json.loads(response_object)
         elif response_object.refusal:
             raise RefusalError(response_object.refusal)
         else:
@@ -140,6 +148,8 @@ class BaseOpenAIClient(LLMClient):
                     temperature=self.temperature,
                     max_tokens=max_tokens or self.max_tokens,
                     response_model=response_model,
+                    reasoning=self.reasoning,
+                    verbosity=self.verbosity,
                 )
                 return self._handle_structured_response(response)
             else:
@@ -155,8 +165,20 @@ class BaseOpenAIClient(LLMClient):
             raise Exception(f'Output length exceeded max tokens {self.max_tokens}: {e}') from e
         except openai.RateLimitError as e:
             raise RateLimitError from e
+        except openai.AuthenticationError as e:
+            logger.error(
+                f'OpenAI Authentication Error: {e}. Please verify your API key is correct.'
+            )
+            raise
         except Exception as e:
-            logger.error(f'Error in generating LLM response: {e}')
+            # Provide more context for connection errors
+            error_msg = str(e)
+            if 'Connection error' in error_msg or 'connection' in error_msg.lower():
+                logger.error(
+                    f'Connection error communicating with OpenAI API. Please check your network connection and API key. Error: {e}'
+                )
+            else:
+                logger.error(f'Error in generating LLM response: {e}')
             raise
 
     async def generate_response(
@@ -165,53 +187,75 @@ class BaseOpenAIClient(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 with retry logic and error handling."""
         if max_tokens is None:
             max_tokens = self.max_tokens
 
-        retry_count = 0
-        last_error = None
-
         # Add multilingual extraction instructions
-        messages[0].content += MULTILINGUAL_EXTRACTION_RESPONSES
-
-        while retry_count <= self.MAX_RETRIES:
-            try:
-                response = await self._generate_response(
-                    messages, response_model, max_tokens, model_size
-                )
-                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}')
+        messages[0].content += get_extraction_language_instruction(group_id)
+
+        # Wrap entire operation in tracing span
+        with self.tracer.start_span('llm.generate') as span:
+            attributes = {
+                'llm.provider': 'openai',
+                'model.size': model_size.value,
+                'max_tokens': max_tokens,
+            }
+            if prompt_name:
+                attributes['prompt.name'] = prompt_name
+            span.add_attributes(attributes)
+
+            retry_count = 0
+            last_error = None
+
+            while retry_count <= self.MAX_RETRIES:
+                try:
+                    response = await self._generate_response(
+                        messages, response_model, max_tokens, model_size
+                    )
+                    return response
+                except (RateLimitError, RefusalError):
+                    # These errors should not trigger retries
+                    span.set_status('error', str(last_error))
                     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')
+                except (
+                    openai.APITimeoutError,
+                    openai.APIConnectionError,
+                    openai.InternalServerError,
+                ):
+                    # Let OpenAI's client handle these retries
+                    span.set_status('error', str(last_error))
+                    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}')
+                        span.set_status('error', str(e))
+                        span.record_exception(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
+            span.set_status('error', str(last_error))
+            raise last_error or Exception('Max retries exceeded with no specific error')

graphiti_core/llm_client/openai_client.py

@@ -21,7 +21,7 @@ from openai.types.chat import ChatCompletionMessageParam
 from pydantic import BaseModel
 
 from .config import DEFAULT_MAX_TOKENS, LLMConfig
-from .openai_base_client import BaseOpenAIClient
+from .openai_base_client import DEFAULT_REASONING, DEFAULT_VERBOSITY, BaseOpenAIClient
 
 
 class OpenAIClient(BaseOpenAIClient):
@@ -41,6 +41,8 @@ class OpenAIClient(BaseOpenAIClient):
         cache: bool = False,
         client: typing.Any = None,
         max_tokens: int = DEFAULT_MAX_TOKENS,
+        reasoning: str = DEFAULT_REASONING,
+        verbosity: str = DEFAULT_VERBOSITY,
     ):
         """
         Initialize the OpenAIClient with the provided configuration, cache setting, and client.
@@ -50,7 +52,7 @@ class OpenAIClient(BaseOpenAIClient):
             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.
         """
-        super().__init__(config, cache, max_tokens)
+        super().__init__(config, cache, max_tokens, reasoning, verbosity)
 
         if config is None:
             config = LLMConfig()
@@ -67,16 +69,37 @@ class OpenAIClient(BaseOpenAIClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel],
+        reasoning: str | None = None,
+        verbosity: str | None = None,
     ):
         """Create a structured completion using OpenAI's beta parse API."""
-        return await self.client.beta.chat.completions.parse(
-            model=model,
-            messages=messages,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            response_format=response_model,  # type: ignore
+        # Reasoning models (gpt-5 family) don't support temperature
+        is_reasoning_model = (
+            model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3')
         )
 
+        request_kwargs = {
+            'model': model,
+            'input': messages,  # type: ignore
+            'max_output_tokens': max_tokens,
+            'text_format': response_model,  # type: ignore
+        }
+
+        temperature_value = temperature if not is_reasoning_model else None
+        if temperature_value is not None:
+            request_kwargs['temperature'] = temperature_value
+
+        # Only include reasoning and verbosity parameters for reasoning models
+        if is_reasoning_model and reasoning is not None:
+            request_kwargs['reasoning'] = {'effort': reasoning}  # type: ignore
+
+        if is_reasoning_model and verbosity is not None:
+            request_kwargs['text'] = {'verbosity': verbosity}  # type: ignore
+
+        response = await self.client.responses.parse(**request_kwargs)
+
+        return response
+
     async def _create_completion(
         self,
         model: str,
@@ -84,12 +107,19 @@ class OpenAIClient(BaseOpenAIClient):
         temperature: float | None,
         max_tokens: int,
         response_model: type[BaseModel] | None = None,
+        reasoning: str | None = None,
+        verbosity: str | None = None,
     ):
         """Create a regular completion with JSON format."""
+        # Reasoning models (gpt-5 family) don't support temperature
+        is_reasoning_model = (
+            model.startswith('gpt-5') or model.startswith('o1') or model.startswith('o3')
+        )
+
         return await self.client.chat.completions.create(
             model=model,
             messages=messages,
-            temperature=temperature,
+            temperature=temperature if not is_reasoning_model else None,
             max_tokens=max_tokens,
             response_format={'type': 'json_object'},
         )