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

graphiti_core/utils/content_chunking.py

@@ -0,0 +1,702 @@
+"""
+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 re
+
+from graphiti_core.helpers import (
+    CHUNK_DENSITY_THRESHOLD,
+    CHUNK_MIN_TOKENS,
+    CHUNK_OVERLAP_TOKENS,
+    CHUNK_TOKEN_SIZE,
+)
+from graphiti_core.nodes import EpisodeType
+
+logger = logging.getLogger(__name__)
+
+# Approximate characters per token (conservative estimate)
+CHARS_PER_TOKEN = 4
+
+
+def estimate_tokens(text: str) -> int:
+    """Estimate token count using character-based heuristic.
+
+    Uses ~4 characters per token as a conservative estimate.
+    This is faster than actual tokenization and works across all LLM providers.
+
+    Args:
+        text: The text to estimate tokens for
+
+    Returns:
+        Estimated token count
+    """
+    return len(text) // CHARS_PER_TOKEN
+
+
+def _tokens_to_chars(tokens: int) -> int:
+    """Convert token count to approximate character count."""
+    return tokens * CHARS_PER_TOKEN
+
+
+def should_chunk(content: str, episode_type: EpisodeType) -> bool:
+    """Determine whether content should be chunked based on size and entity density.
+
+    Only chunks content that is both:
+    1. Large enough to potentially cause LLM issues (>= CHUNK_MIN_TOKENS)
+    2. High entity density (many entities per token)
+
+    Short content processes fine regardless of density. This targets the specific
+    failure case of large entity-dense inputs while preserving context for
+    prose/narrative content and avoiding unnecessary chunking of small inputs.
+
+    Args:
+        content: The content to evaluate
+        episode_type: Type of episode (json, message, text)
+
+    Returns:
+        True if content is large and has high entity density
+    """
+    tokens = estimate_tokens(content)
+
+    # Short content always processes fine - no need to chunk
+    if tokens < CHUNK_MIN_TOKENS:
+        return False
+
+    return _estimate_high_density(content, episode_type, tokens)
+
+
+def _estimate_high_density(content: str, episode_type: EpisodeType, tokens: int) -> bool:
+    """Estimate whether content has high entity density.
+
+    High-density content (many entities per token) benefits from chunking.
+    Low-density content (prose, narratives) loses context when chunked.
+
+    Args:
+        content: The content to analyze
+        episode_type: Type of episode
+        tokens: Pre-computed token count
+
+    Returns:
+        True if content appears to have high entity density
+    """
+    if episode_type == EpisodeType.json:
+        return _json_likely_dense(content, tokens)
+    else:
+        return _text_likely_dense(content, tokens)
+
+
+def _json_likely_dense(content: str, tokens: int) -> bool:
+    """Estimate entity density for JSON content.
+
+    JSON is considered dense if it has many array elements or object keys,
+    as each typically represents a distinct entity or data point.
+
+    Heuristics:
+    - Array: Count elements, estimate entities per 1000 tokens
+    - Object: Count top-level keys
+
+    Args:
+        content: JSON string content
+        tokens: Token count
+
+    Returns:
+        True if JSON appears to have high entity density
+    """
+    try:
+        data = json.loads(content)
+    except json.JSONDecodeError:
+        # Invalid JSON, fall back to text heuristics
+        return _text_likely_dense(content, tokens)
+
+    if isinstance(data, list):
+        # For arrays, each element likely contains entities
+        element_count = len(data)
+        # Estimate density: elements per 1000 tokens
+        density = (element_count / tokens) * 1000 if tokens > 0 else 0
+        return density > CHUNK_DENSITY_THRESHOLD * 1000  # Scale threshold
+    elif isinstance(data, dict):
+        # For objects, count keys recursively (shallow)
+        key_count = _count_json_keys(data, max_depth=2)
+        density = (key_count / tokens) * 1000 if tokens > 0 else 0
+        return density > CHUNK_DENSITY_THRESHOLD * 1000
+    else:
+        # Scalar value, no need to chunk
+        return False
+
+
+def _count_json_keys(data: dict, max_depth: int = 2, current_depth: int = 0) -> int:
+    """Count keys in a JSON object up to a certain depth.
+
+    Args:
+        data: Dictionary to count keys in
+        max_depth: Maximum depth to traverse
+        current_depth: Current recursion depth
+
+    Returns:
+        Count of keys
+    """
+    if current_depth >= max_depth:
+        return 0
+
+    count = len(data)
+    for value in data.values():
+        if isinstance(value, dict):
+            count += _count_json_keys(value, max_depth, current_depth + 1)
+        elif isinstance(value, list):
+            for item in value:
+                if isinstance(item, dict):
+                    count += _count_json_keys(item, max_depth, current_depth + 1)
+    return count
+
+
+def _text_likely_dense(content: str, tokens: int) -> bool:
+    """Estimate entity density for text content.
+
+    Uses capitalized words as a proxy for named entities (people, places,
+    organizations, products). High ratio of capitalized words suggests
+    high entity density.
+
+    Args:
+        content: Text content
+        tokens: Token count
+
+    Returns:
+        True if text appears to have high entity density
+    """
+    if tokens == 0:
+        return False
+
+    # Split into words
+    words = content.split()
+    if not words:
+        return False
+
+    # Count capitalized words (excluding sentence starters)
+    # A word is "capitalized" if it starts with uppercase and isn't all caps
+    capitalized_count = 0
+    for i, word in enumerate(words):
+        # Skip if it's likely a sentence starter (after . ! ? or first word)
+        if i == 0:
+            continue
+        if i > 0 and words[i - 1].rstrip()[-1:] in '.!?':
+            continue
+
+        # Check if capitalized (first char upper, not all caps)
+        cleaned = word.strip('.,!?;:\'"()[]{}')
+        if cleaned and cleaned[0].isupper() and not cleaned.isupper():
+            capitalized_count += 1
+
+    # Calculate density: capitalized words per 1000 tokens
+    density = (capitalized_count / tokens) * 1000 if tokens > 0 else 0
+
+    # Text density threshold is typically lower than JSON
+    # A well-written article might have 5-10% named entities
+    return density > CHUNK_DENSITY_THRESHOLD * 500  # Half the JSON threshold
+
+
+def chunk_json_content(
+    content: str,
+    chunk_size_tokens: int | None = None,
+    overlap_tokens: int | None = None,
+) -> list[str]:
+    """Split JSON content into chunks while preserving structure.
+
+    For arrays: splits at element boundaries, keeping complete objects.
+    For objects: splits at top-level key boundaries.
+
+    Args:
+        content: JSON string to chunk
+        chunk_size_tokens: Target size per chunk in tokens (default from env)
+        overlap_tokens: Overlap between chunks in tokens (default from env)
+
+    Returns:
+        List of JSON string chunks
+    """
+    chunk_size_tokens = chunk_size_tokens or CHUNK_TOKEN_SIZE
+    overlap_tokens = overlap_tokens or CHUNK_OVERLAP_TOKENS
+
+    chunk_size_chars = _tokens_to_chars(chunk_size_tokens)
+    overlap_chars = _tokens_to_chars(overlap_tokens)
+
+    try:
+        data = json.loads(content)
+    except json.JSONDecodeError:
+        logger.warning('Failed to parse JSON, falling back to text chunking')
+        return chunk_text_content(content, chunk_size_tokens, overlap_tokens)
+
+    if isinstance(data, list):
+        return _chunk_json_array(data, chunk_size_chars, overlap_chars)
+    elif isinstance(data, dict):
+        return _chunk_json_object(data, chunk_size_chars, overlap_chars)
+    else:
+        # Scalar value, return as-is
+        return [content]
+
+
+def _chunk_json_array(
+    data: list,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Chunk a JSON array by splitting at element boundaries."""
+    if not data:
+        return ['[]']
+
+    chunks: list[str] = []
+    current_elements: list = []
+    current_size = 2  # Account for '[]'
+
+    for element in data:
+        element_json = json.dumps(element)
+        element_size = len(element_json) + 2  # Account for comma and space
+
+        # Check if adding this element would exceed chunk size
+        if current_elements and current_size + element_size > chunk_size_chars:
+            # Save current chunk
+            chunks.append(json.dumps(current_elements))
+
+            # Start new chunk with overlap (include last few elements)
+            overlap_elements = _get_overlap_elements(current_elements, overlap_chars)
+            current_elements = overlap_elements
+            current_size = len(json.dumps(current_elements)) if current_elements else 2
+
+        current_elements.append(element)
+        current_size += element_size
+
+    # Don't forget the last chunk
+    if current_elements:
+        chunks.append(json.dumps(current_elements))
+
+    return chunks if chunks else ['[]']
+
+
+def _get_overlap_elements(elements: list, overlap_chars: int) -> list:
+    """Get elements from the end of a list that fit within overlap_chars."""
+    if not elements:
+        return []
+
+    overlap_elements: list = []
+    current_size = 2  # Account for '[]'
+
+    for element in reversed(elements):
+        element_json = json.dumps(element)
+        element_size = len(element_json) + 2
+
+        if current_size + element_size > overlap_chars:
+            break
+
+        overlap_elements.insert(0, element)
+        current_size += element_size
+
+    return overlap_elements
+
+
+def _chunk_json_object(
+    data: dict,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Chunk a JSON object by splitting at top-level key boundaries."""
+    if not data:
+        return ['{}']
+
+    chunks: list[str] = []
+    current_keys: list[str] = []
+    current_dict: dict = {}
+    current_size = 2  # Account for '{}'
+
+    for key, value in data.items():
+        entry_json = json.dumps({key: value})
+        entry_size = len(entry_json)
+
+        # Check if adding this entry would exceed chunk size
+        if current_dict and current_size + entry_size > chunk_size_chars:
+            # Save current chunk
+            chunks.append(json.dumps(current_dict))
+
+            # Start new chunk with overlap (include last few keys)
+            overlap_dict = _get_overlap_dict(current_dict, current_keys, overlap_chars)
+            current_dict = overlap_dict
+            current_keys = list(overlap_dict.keys())
+            current_size = len(json.dumps(current_dict)) if current_dict else 2
+
+        current_dict[key] = value
+        current_keys.append(key)
+        current_size += entry_size
+
+    # Don't forget the last chunk
+    if current_dict:
+        chunks.append(json.dumps(current_dict))
+
+    return chunks if chunks else ['{}']
+
+
+def _get_overlap_dict(data: dict, keys: list[str], overlap_chars: int) -> dict:
+    """Get key-value pairs from the end of a dict that fit within overlap_chars."""
+    if not data or not keys:
+        return {}
+
+    overlap_dict: dict = {}
+    current_size = 2  # Account for '{}'
+
+    for key in reversed(keys):
+        if key not in data:
+            continue
+        entry_json = json.dumps({key: data[key]})
+        entry_size = len(entry_json)
+
+        if current_size + entry_size > overlap_chars:
+            break
+
+        overlap_dict[key] = data[key]
+        current_size += entry_size
+
+    # Reverse to maintain original order
+    return dict(reversed(list(overlap_dict.items())))
+
+
+def chunk_text_content(
+    content: str,
+    chunk_size_tokens: int | None = None,
+    overlap_tokens: int | None = None,
+) -> list[str]:
+    """Split text content at natural boundaries (paragraphs, sentences).
+
+    Includes overlap to capture entities at chunk boundaries.
+
+    Args:
+        content: Text to chunk
+        chunk_size_tokens: Target size per chunk in tokens (default from env)
+        overlap_tokens: Overlap between chunks in tokens (default from env)
+
+    Returns:
+        List of text chunks
+    """
+    chunk_size_tokens = chunk_size_tokens or CHUNK_TOKEN_SIZE
+    overlap_tokens = overlap_tokens or CHUNK_OVERLAP_TOKENS
+
+    chunk_size_chars = _tokens_to_chars(chunk_size_tokens)
+    overlap_chars = _tokens_to_chars(overlap_tokens)
+
+    if len(content) <= chunk_size_chars:
+        return [content]
+
+    # Split into paragraphs first
+    paragraphs = re.split(r'\n\s*\n', content)
+
+    chunks: list[str] = []
+    current_chunk: list[str] = []
+    current_size = 0
+
+    for paragraph in paragraphs:
+        paragraph = paragraph.strip()
+        if not paragraph:
+            continue
+
+        para_size = len(paragraph)
+
+        # If a single paragraph is too large, split it by sentences
+        if para_size > chunk_size_chars:
+            # First, save current chunk if any
+            if current_chunk:
+                chunks.append('\n\n'.join(current_chunk))
+                current_chunk = []
+                current_size = 0
+
+            # Split large paragraph by sentences
+            sentence_chunks = _chunk_by_sentences(paragraph, chunk_size_chars, overlap_chars)
+            chunks.extend(sentence_chunks)
+            continue
+
+        # Check if adding this paragraph would exceed chunk size
+        if current_chunk and current_size + para_size + 2 > chunk_size_chars:
+            # Save current chunk
+            chunks.append('\n\n'.join(current_chunk))
+
+            # Start new chunk with overlap
+            overlap_text = _get_overlap_text('\n\n'.join(current_chunk), overlap_chars)
+            if overlap_text:
+                current_chunk = [overlap_text]
+                current_size = len(overlap_text)
+            else:
+                current_chunk = []
+                current_size = 0
+
+        current_chunk.append(paragraph)
+        current_size += para_size + 2  # Account for '\n\n'
+
+    # Don't forget the last chunk
+    if current_chunk:
+        chunks.append('\n\n'.join(current_chunk))
+
+    return chunks if chunks else [content]
+
+
+def _chunk_by_sentences(
+    text: str,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Split text by sentence boundaries."""
+    # Split on sentence-ending punctuation followed by whitespace
+    sentence_pattern = r'(?<=[.!?])\s+'
+    sentences = re.split(sentence_pattern, text)
+
+    chunks: list[str] = []
+    current_chunk: list[str] = []
+    current_size = 0
+
+    for sentence in sentences:
+        sentence = sentence.strip()
+        if not sentence:
+            continue
+
+        sent_size = len(sentence)
+
+        # If a single sentence is too large, split it by fixed size
+        if sent_size > chunk_size_chars:
+            if current_chunk:
+                chunks.append(' '.join(current_chunk))
+                current_chunk = []
+                current_size = 0
+
+            # Split by fixed size as last resort
+            fixed_chunks = _chunk_by_size(sentence, chunk_size_chars, overlap_chars)
+            chunks.extend(fixed_chunks)
+            continue
+
+        # Check if adding this sentence would exceed chunk size
+        if current_chunk and current_size + sent_size + 1 > chunk_size_chars:
+            chunks.append(' '.join(current_chunk))
+
+            # Start new chunk with overlap
+            overlap_text = _get_overlap_text(' '.join(current_chunk), overlap_chars)
+            if overlap_text:
+                current_chunk = [overlap_text]
+                current_size = len(overlap_text)
+            else:
+                current_chunk = []
+                current_size = 0
+
+        current_chunk.append(sentence)
+        current_size += sent_size + 1
+
+    if current_chunk:
+        chunks.append(' '.join(current_chunk))
+
+    return chunks
+
+
+def _chunk_by_size(
+    text: str,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Split text by fixed character size (last resort)."""
+    chunks: list[str] = []
+    start = 0
+
+    while start < len(text):
+        end = min(start + chunk_size_chars, len(text))
+
+        # Try to break at word boundary
+        if end < len(text):
+            space_idx = text.rfind(' ', start, end)
+            if space_idx > start:
+                end = space_idx
+
+        chunks.append(text[start:end].strip())
+
+        # Move start forward, ensuring progress even if overlap >= chunk_size
+        # Always advance by at least (chunk_size - overlap) or 1 char minimum
+        min_progress = max(1, chunk_size_chars - overlap_chars)
+        start = max(start + min_progress, end - overlap_chars)
+
+    return chunks
+
+
+def _get_overlap_text(text: str, overlap_chars: int) -> str:
+    """Get the last overlap_chars characters of text, breaking at word boundary."""
+    if len(text) <= overlap_chars:
+        return text
+
+    overlap_start = len(text) - overlap_chars
+    # Find the next word boundary after overlap_start
+    space_idx = text.find(' ', overlap_start)
+    if space_idx != -1:
+        return text[space_idx + 1 :]
+    return text[overlap_start:]
+
+
+def chunk_message_content(
+    content: str,
+    chunk_size_tokens: int | None = None,
+    overlap_tokens: int | None = None,
+) -> list[str]:
+    """Split conversation content preserving message boundaries.
+
+    Never splits mid-message. Messages are identified by patterns like:
+    - "Speaker: message"
+    - JSON message arrays
+    - Newline-separated messages
+
+    Args:
+        content: Conversation content to chunk
+        chunk_size_tokens: Target size per chunk in tokens (default from env)
+        overlap_tokens: Overlap between chunks in tokens (default from env)
+
+    Returns:
+        List of conversation chunks
+    """
+    chunk_size_tokens = chunk_size_tokens or CHUNK_TOKEN_SIZE
+    overlap_tokens = overlap_tokens or CHUNK_OVERLAP_TOKENS
+
+    chunk_size_chars = _tokens_to_chars(chunk_size_tokens)
+    overlap_chars = _tokens_to_chars(overlap_tokens)
+
+    if len(content) <= chunk_size_chars:
+        return [content]
+
+    # Try to detect message format
+    # Check if it's JSON (array of message objects)
+    try:
+        data = json.loads(content)
+        if isinstance(data, list):
+            return _chunk_message_array(data, chunk_size_chars, overlap_chars)
+    except json.JSONDecodeError:
+        pass
+
+    # Try speaker pattern (e.g., "Alice: Hello")
+    speaker_pattern = r'^([A-Za-z_][A-Za-z0-9_\s]*):(.+?)(?=^[A-Za-z_][A-Za-z0-9_\s]*:|$)'
+    if re.search(speaker_pattern, content, re.MULTILINE | re.DOTALL):
+        return _chunk_speaker_messages(content, chunk_size_chars, overlap_chars)
+
+    # Fallback to line-based chunking
+    return _chunk_by_lines(content, chunk_size_chars, overlap_chars)
+
+
+def _chunk_message_array(
+    messages: list,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Chunk a JSON array of message objects."""
+    # Delegate to JSON array chunking
+    chunks = _chunk_json_array(messages, chunk_size_chars, overlap_chars)
+    return chunks
+
+
+def _chunk_speaker_messages(
+    content: str,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Chunk messages in 'Speaker: message' format."""
+    # Split on speaker patterns
+    pattern = r'(?=^[A-Za-z_][A-Za-z0-9_\s]*:)'
+    messages = re.split(pattern, content, flags=re.MULTILINE)
+    messages = [m.strip() for m in messages if m.strip()]
+
+    if not messages:
+        return [content]
+
+    chunks: list[str] = []
+    current_messages: list[str] = []
+    current_size = 0
+
+    for message in messages:
+        msg_size = len(message)
+
+        # If a single message is too large, include it as its own chunk
+        if msg_size > chunk_size_chars:
+            if current_messages:
+                chunks.append('\n'.join(current_messages))
+                current_messages = []
+                current_size = 0
+            chunks.append(message)
+            continue
+
+        if current_messages and current_size + msg_size + 1 > chunk_size_chars:
+            chunks.append('\n'.join(current_messages))
+
+            # Get overlap (last message(s) that fit)
+            overlap_messages = _get_overlap_messages(current_messages, overlap_chars)
+            current_messages = overlap_messages
+            current_size = sum(len(m) for m in current_messages) + len(current_messages) - 1
+
+        current_messages.append(message)
+        current_size += msg_size + 1
+
+    if current_messages:
+        chunks.append('\n'.join(current_messages))
+
+    return chunks if chunks else [content]
+
+
+def _get_overlap_messages(messages: list[str], overlap_chars: int) -> list[str]:
+    """Get messages from the end that fit within overlap_chars."""
+    if not messages:
+        return []
+
+    overlap: list[str] = []
+    current_size = 0
+
+    for msg in reversed(messages):
+        msg_size = len(msg) + 1
+        if current_size + msg_size > overlap_chars:
+            break
+        overlap.insert(0, msg)
+        current_size += msg_size
+
+    return overlap
+
+
+def _chunk_by_lines(
+    content: str,
+    chunk_size_chars: int,
+    overlap_chars: int,
+) -> list[str]:
+    """Chunk content by line boundaries."""
+    lines = content.split('\n')
+
+    chunks: list[str] = []
+    current_lines: list[str] = []
+    current_size = 0
+
+    for line in lines:
+        line_size = len(line) + 1
+
+        if current_lines and current_size + line_size > chunk_size_chars:
+            chunks.append('\n'.join(current_lines))
+
+            # Get overlap lines
+            overlap_text = '\n'.join(current_lines)
+            overlap = _get_overlap_text(overlap_text, overlap_chars)
+            if overlap:
+                current_lines = overlap.split('\n')
+                current_size = len(overlap)
+            else:
+                current_lines = []
+                current_size = 0
+
+        current_lines.append(line)
+        current_size += line_size
+
+    if current_lines:
+        chunks.append('\n'.join(current_lines))
+
+    return chunks if chunks else [content]

graphiti_core/utils/datetime_utils.py

@@ -40,3 +40,16 @@ def ensure_utc(dt: datetime | None) -> datetime | None:
         return dt.astimezone(timezone.utc)
 
     return dt
+
+
+def convert_datetimes_to_strings(obj):
+    if isinstance(obj, dict):
+        return {k: convert_datetimes_to_strings(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_datetimes_to_strings(item) for item in obj]
+    elif isinstance(obj, tuple):
+        return tuple(convert_datetimes_to_strings(item) for item in obj)
+    elif isinstance(obj, datetime):
+        return obj.isoformat()
+    else:
+        return obj