graphiti-core 0.21.0rc13__py3-none-any.whl → 0.22.0__py3-none-any.whl

graphiti_core/graphiti_types.py

@@ -20,6 +20,7 @@ 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):
@@ -27,5 +28,6 @@ class GraphitiClients(BaseModel):
     llm_client: LLMClient
     embedder: EmbedderClient
     cross_encoder: CrossEncoderClient
+    tracer: Tracer
 
     model_config = ConfigDict(arbitrary_types_allowed=True)

graphiti_core/llm_client/anthropic_client.py

@@ -265,6 +265,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.
@@ -285,55 +287,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/client.py

@@ -26,6 +26,7 @@ from pydantic import BaseModel
 from tenacity import retry, retry_if_exception, stop_after_attempt, wait_random_exponential
 
 from ..prompts.models import Message
+from ..tracer import NoOpTracer, Tracer
 from .config import DEFAULT_MAX_TOKENS, LLMConfig, ModelSize
 from .errors import RateLimitError
 
@@ -33,12 +34,16 @@ DEFAULT_TEMPERATURE = 0
 DEFAULT_CACHE_DIR = './llm_cache'
 
 
-def get_extraction_language_instruction() -> str:
+def get_extraction_language_instruction(group_id: str | None = None) -> str:
     """Returns instruction for language extraction behavior.
 
     Override this function to customize language extraction:
     - Return empty string to disable multilingual instructions
     - Return custom instructions for specific language requirements
+    - Use group_id to provide different instructions per group/partition
+
+    Args:
+        group_id: Optional partition identifier for the graph
 
     Returns:
         str: Language instruction to append to system messages
@@ -70,11 +75,16 @@ class LLMClient(ABC):
         self.max_tokens = config.max_tokens
         self.cache_enabled = cache
         self.cache_dir = None
+        self.tracer: Tracer = NoOpTracer()
 
         # Only create the cache directory if caching is enabled
         if self.cache_enabled:
             self.cache_dir = Cache(DEFAULT_CACHE_DIR)
 
+    def set_tracer(self, tracer: Tracer) -> None:
+        """Set the tracer for this LLM client."""
+        self.tracer = tracer
+
     def _clean_input(self, input: str) -> str:
         """Clean input string of invalid unicode and control characters.
 
@@ -142,6 +152,8 @@ class LLMClient(ABC):
         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]:
         if max_tokens is None:
             max_tokens = self.max_tokens
@@ -155,28 +167,64 @@ class LLMClient(ABC):
             )
 
         # Add multilingual extraction instructions
-        messages[0].content += get_extraction_language_instruction()
-
-        if self.cache_enabled and self.cache_dir is not None:
-            cache_key = self._get_cache_key(messages)
-
-            cached_response = self.cache_dir.get(cache_key)
-            if cached_response is not None:
-                logger.debug(f'Cache hit for {cache_key}')
-                return cached_response
+        messages[0].content += get_extraction_language_instruction(group_id)
 
         for message in messages:
             message.content = self._clean_input(message.content)
 
-        response = await self._generate_response_with_retry(
-            messages, response_model, max_tokens, model_size
-        )
-
-        if self.cache_enabled and self.cache_dir is not None:
-            cache_key = self._get_cache_key(messages)
-            self.cache_dir.set(cache_key, response)
-
-        return response
+        # Wrap entire operation in tracing span
+        with self.tracer.start_span('llm.generate') as span:
+            attributes = {
+                'llm.provider': self._get_provider_type(),
+                'model.size': model_size.value,
+                'max_tokens': max_tokens,
+                'cache.enabled': self.cache_enabled,
+            }
+            if prompt_name:
+                attributes['prompt.name'] = prompt_name
+            span.add_attributes(attributes)
+
+            # Check cache first
+            if self.cache_enabled and self.cache_dir is not None:
+                cache_key = self._get_cache_key(messages)
+                cached_response = self.cache_dir.get(cache_key)
+                if cached_response is not None:
+                    logger.debug(f'Cache hit for {cache_key}')
+                    span.add_attributes({'cache.hit': True})
+                    return cached_response
+
+            span.add_attributes({'cache.hit': False})
+
+            # Execute LLM call
+            try:
+                response = await self._generate_response_with_retry(
+                    messages, response_model, max_tokens, model_size
+                )
+            except Exception as e:
+                span.set_status('error', str(e))
+                span.record_exception(e)
+                raise
+
+            # Cache response if enabled
+            if self.cache_enabled and self.cache_dir is not None:
+                cache_key = self._get_cache_key(messages)
+                self.cache_dir.set(cache_key, response)
+
+            return response
+
+    def _get_provider_type(self) -> str:
+        """Get provider type from class name."""
+        class_name = self.__class__.__name__.lower()
+        if 'openai' in class_name:
+            return 'openai'
+        elif 'anthropic' in class_name:
+            return 'anthropic'
+        elif 'gemini' in class_name:
+            return 'gemini'
+        elif 'groq' in class_name:
+            return 'groq'
+        else:
+            return 'unknown'
 
     def _get_failed_generation_log(self, messages: list[Message], output: str | None) -> str:
         """

graphiti_core/llm_client/gemini_client.py

@@ -357,6 +357,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 +369,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 += get_extraction_language_instruction()
-
-        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

@@ -175,53 +175,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 += get_extraction_language_instruction()
-
-        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')