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

graphiti_core/llm_client/openai_generic_client.py

@@ -120,13 +120,12 @@ class OpenAIGenericClient(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]:
         if max_tokens is None:
             max_tokens = self.max_tokens
 
-        retry_count = 0
-        last_error = None
-
         if response_model is not None:
             serialized_model = json.dumps(response_model.model_json_schema())
             messages[
@@ -136,44 +135,67 @@ class OpenAIGenericClient(LLMClient):
             )
 
         # 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=max_tokens, model_size=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=max_tokens, model_size=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/models/edges/edge_db_queries.py

@@ -68,6 +68,7 @@ def get_entity_edge_save_query(provider: GraphProvider, has_aoss: bool = False)
                 MATCH (target:Entity {uuid: $edge_data.target_uuid})
                 MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
                 SET e = $edge_data
+                SET e.fact_embedding = vecf32($edge_data.fact_embedding)
                 RETURN e.uuid AS uuid
             """
         case GraphProvider.NEPTUNE:

graphiti_core/models/nodes/node_db_queries.py

@@ -133,6 +133,7 @@ def get_entity_node_save_query(provider: GraphProvider, labels: str, has_aoss: b
                 MERGE (n:Entity {{uuid: $entity_data.uuid}})
                 SET n:{labels}
                 SET n = $entity_data
+                SET n.name_embedding = vecf32($entity_data.name_embedding)
                 RETURN n.uuid AS uuid
             """
         case GraphProvider.KUZU:

graphiti_core/nodes.py

@@ -27,10 +27,6 @@ from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
 from graphiti_core.driver.driver import (
-    COMMUNITY_INDEX_NAME,
-    ENTITY_EDGE_INDEX_NAME,
-    ENTITY_INDEX_NAME,
-    EPISODE_INDEX_NAME,
     GraphDriver,
     GraphProvider,
 )
@@ -99,6 +95,9 @@ class Node(BaseModel, ABC):
     async def save(self, driver: GraphDriver): ...
 
     async def delete(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete(self, driver)
+
         match driver.provider:
             case GraphProvider.NEO4J:
                 records, _, _ = await driver.execute_query(
@@ -113,27 +112,6 @@ class Node(BaseModel, ABC):
                     uuid=self.uuid,
                 )
 
-                edge_uuids: list[str] = records[0].get('edge_uuids', []) if records else []
-
-                if driver.aoss_client:
-                    # Delete the node from OpenSearch indices
-                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
-                        await driver.aoss_client.delete(
-                            index=index,
-                            id=self.uuid,
-                            params={'routing': self.group_id},
-                        )
-
-                    # Bulk delete the detached edges
-                    if edge_uuids:
-                        actions = []
-                        for eid in edge_uuids:
-                            actions.append(
-                                {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
-                            )
-
-                        await driver.aoss_client.bulk(body=actions)
-
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -181,14 +159,18 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_group_id(cls, driver: GraphDriver, group_id: str, batch_size: int = 100):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete_by_group_id(
+                cls, driver, group_id, batch_size
+            )
+
         match driver.provider:
             case GraphProvider.NEO4J:
                 async with driver.session() as session:
                     await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community {group_id: $group_id})
-                        CALL {
-                            WITH n
+                        CALL (n) {
                             DETACH DELETE n
                         } IN TRANSACTIONS OF $batch_size ROWS
                         """,
@@ -196,31 +178,6 @@ class Node(BaseModel, ABC):
                         batch_size=batch_size,
                     )
 
-                if driver.aoss_client:
-                    await driver.aoss_client.delete_by_query(
-                        index=EPISODE_INDEX_NAME,
-                        body={'query': {'term': {'group_id': group_id}}},
-                        params={'routing': group_id},
-                    )
-
-                    await driver.aoss_client.delete_by_query(
-                        index=ENTITY_INDEX_NAME,
-                        body={'query': {'term': {'group_id': group_id}}},
-                        params={'routing': group_id},
-                    )
-
-                    await driver.aoss_client.delete_by_query(
-                        index=COMMUNITY_INDEX_NAME,
-                        body={'query': {'term': {'group_id': group_id}}},
-                        params={'routing': group_id},
-                    )
-
-                    await driver.aoss_client.delete_by_query(
-                        index=ENTITY_EDGE_INDEX_NAME,
-                        body={'query': {'term': {'group_id': group_id}}},
-                        params={'routing': group_id},
-                    )
-
             case GraphProvider.KUZU:
                 for label in ['Episodic', 'Community']:
                     await driver.execute_query(
@@ -258,6 +215,11 @@ class Node(BaseModel, ABC):
 
     @classmethod
     async def delete_by_uuids(cls, driver: GraphDriver, uuids: list[str], batch_size: int = 100):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_delete_by_uuids(
+                cls, driver, uuids, group_id=None, batch_size=batch_size
+            )
+
         match driver.provider:
             case GraphProvider.FALKORDB:
                 for label in ['Entity', 'Episodic', 'Community']:
@@ -300,7 +262,7 @@ class Node(BaseModel, ABC):
             case _:  # Neo4J, Neptune
                 async with driver.session() as session:
                     # Collect all edge UUIDs before deleting nodes
-                    result = await session.run(
+                    await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community)
                         WHERE n.uuid IN $uuids
@@ -310,18 +272,12 @@ class Node(BaseModel, ABC):
                         uuids=uuids,
                     )
 
-                    record = await result.single()
-                    edge_uuids: list[str] = (
-                        record['edge_uuids'] if record and record['edge_uuids'] else []
-                    )
-
                     # Now delete the nodes in batches
                     await session.run(
                         """
                         MATCH (n:Entity|Episodic|Community)
                         WHERE n.uuid IN $uuids
-                        CALL {
-                            WITH n
+                        CALL (n) {
                             DETACH DELETE n
                         } IN TRANSACTIONS OF $batch_size ROWS
                         """,
@@ -329,20 +285,6 @@ class Node(BaseModel, ABC):
                         batch_size=batch_size,
                     )
 
-                if driver.aoss_client:
-                    for index in (EPISODE_INDEX_NAME, ENTITY_INDEX_NAME, COMMUNITY_INDEX_NAME):
-                        await driver.aoss_client.delete_by_query(
-                            index=index,
-                            body={'query': {'terms': {'uuid': uuids}}},
-                        )
-
-                    if edge_uuids:
-                        actions = [
-                            {'delete': {'_index': ENTITY_EDGE_INDEX_NAME, '_id': eid}}
-                            for eid in edge_uuids
-                        ]
-                        await driver.aoss_client.bulk(body=actions)
-
     @classmethod
     async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
@@ -363,6 +305,9 @@ class EpisodicNode(Node):
     )
 
     async def save(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.episodic_node_save(self, driver)
+
         episode_args = {
             'uuid': self.uuid,
             'name': self.name,
@@ -375,12 +320,6 @@ class EpisodicNode(Node):
             'source': self.source.value,
         }
 
-        if driver.aoss_client:
-            await driver.save_to_aoss(  # pyright: ignore reportAttributeAccessIssue
-                'episodes',
-                [episode_args],
-            )
-
         result = await driver.execute_query(
             get_episode_node_save_query(driver.provider), **episode_args
         )
@@ -510,26 +449,14 @@ class EntityNode(Node):
         return self.name_embedding
 
     async def load_name_embedding(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_load_embeddings(self, driver)
+
         if driver.provider == GraphProvider.NEPTUNE:
             query: LiteralString = """
                 MATCH (n:Entity {uuid: $uuid})
                 RETURN [x IN split(n.name_embedding, ",") | toFloat(x)] as name_embedding
             """
-        elif driver.aoss_client:
-            resp = await driver.aoss_client.search(
-                body={
-                    'query': {'multi_match': {'query': self.uuid, 'fields': ['uuid']}},
-                    'size': 1,
-                },
-                index=ENTITY_INDEX_NAME,
-                params={'routing': self.group_id},
-            )
-
-            if resp['hits']['hits']:
-                self.name_embedding = resp['hits']['hits'][0]['_source']['name_embedding']
-                return
-            else:
-                raise NodeNotFoundError(self.uuid)
 
         else:
             query: LiteralString = """
@@ -548,6 +475,9 @@ class EntityNode(Node):
         self.name_embedding = records[0]['name_embedding']
 
     async def save(self, driver: GraphDriver):
+        if driver.graph_operations_interface:
+            return await driver.graph_operations_interface.node_save(self, driver)
+
         entity_data: dict[str, Any] = {
             'uuid': self.uuid,
             'name': self.name,
@@ -568,11 +498,8 @@ class EntityNode(Node):
             entity_data.update(self.attributes or {})
             labels = ':'.join(self.labels + ['Entity'])
 
-            if driver.aoss_client:
-                await driver.save_to_aoss(ENTITY_INDEX_NAME, [entity_data])  # pyright: ignore reportAttributeAccessIssue
-
             result = await driver.execute_query(
-                get_entity_node_save_query(driver.provider, labels, bool(driver.aoss_client)),
+                get_entity_node_save_query(driver.provider, labels),
                 entity_data=entity_data,
             )
 

graphiti_core/prompts/dedupe_edges.py

@@ -67,13 +67,13 @@ def edge(context: dict[str, Any]) -> list[Message]:
         Given the following context, determine whether the New Edge represents any of the edges in the list of Existing Edges.
 
         <EXISTING EDGES>
-        {to_prompt_json(context['related_edges'], indent=2)}
+        {to_prompt_json(context['related_edges'])}
         </EXISTING EDGES>
 
         <NEW EDGE>
-        {to_prompt_json(context['extracted_edges'], indent=2)}
+        {to_prompt_json(context['extracted_edges'])}
         </NEW EDGE>
-        
+
         Task:
         If the New Edges represents the same factual information as any edge in Existing Edges, return the id of the duplicate fact
             as part of the list of duplicate_facts.
@@ -98,7 +98,7 @@ def edge_list(context: dict[str, Any]) -> list[Message]:
         Given the following context, find all of the duplicates in a list of facts:
 
         Facts:
-        {to_prompt_json(context['edges'], indent=2)}
+        {to_prompt_json(context['edges'])}
 
         Task:
         If any facts in Facts is a duplicate of another fact, return a new fact with one of their uuid's.

graphiti_core/prompts/dedupe_nodes.py

@@ -64,20 +64,20 @@ def node(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         <PREVIOUS MESSAGES>
-        {to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+        {to_prompt_json([ep for ep in context['previous_episodes']])}
         </PREVIOUS MESSAGES>
         <CURRENT MESSAGE>
         {context['episode_content']}
         </CURRENT MESSAGE>
         <NEW ENTITY>
-        {to_prompt_json(context['extracted_node'], indent=2)}
+        {to_prompt_json(context['extracted_node'])}
         </NEW ENTITY>
         <ENTITY TYPE DESCRIPTION>
-        {to_prompt_json(context['entity_type_description'], indent=2)}
+        {to_prompt_json(context['entity_type_description'])}
         </ENTITY TYPE DESCRIPTION>
 
         <EXISTING ENTITIES>
-        {to_prompt_json(context['existing_nodes'], indent=2)}
+        {to_prompt_json(context['existing_nodes'])}
         </EXISTING ENTITIES>
         
         Given the above EXISTING ENTITIES and their attributes, MESSAGE, and PREVIOUS MESSAGES; Determine if the NEW ENTITY extracted from the conversation
@@ -125,13 +125,13 @@ def nodes(context: dict[str, Any]) -> list[Message]:
             role='user',
             content=f"""
         <PREVIOUS MESSAGES>
-        {to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+        {to_prompt_json([ep for ep in context['previous_episodes']])}
         </PREVIOUS MESSAGES>
         <CURRENT MESSAGE>
         {context['episode_content']}
         </CURRENT MESSAGE>
-        
-        
+
+
         Each of the following ENTITIES were extracted from the CURRENT MESSAGE.
         Each entity in ENTITIES is represented as a JSON object with the following structure:
         {{
@@ -142,11 +142,11 @@ def nodes(context: dict[str, Any]) -> list[Message]:
         }}
 
         <ENTITIES>
-        {to_prompt_json(context['extracted_nodes'], indent=2)}
+        {to_prompt_json(context['extracted_nodes'])}
         </ENTITIES>
 
         <EXISTING ENTITIES>
-        {to_prompt_json(context['existing_nodes'], indent=2)}
+        {to_prompt_json(context['existing_nodes'])}
         </EXISTING ENTITIES>
 
         Each entry in EXISTING ENTITIES is an object with the following structure:
@@ -197,7 +197,7 @@ def node_list(context: dict[str, Any]) -> list[Message]:
         Given the following context, deduplicate a list of nodes:
 
         Nodes:
-        {to_prompt_json(context['nodes'], indent=2)}
+        {to_prompt_json(context['nodes'])}
 
         Task:
         1. Group nodes together such that all duplicate nodes are in the same list of uuids

graphiti_core/prompts/extract_edges.py

@@ -80,7 +80,7 @@ def edge(context: dict[str, Any]) -> list[Message]:
 </FACT TYPES>
 
 <PREVIOUS_MESSAGES>
-{to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS_MESSAGES>
 
 <CURRENT_MESSAGE>
@@ -88,7 +88,7 @@ def edge(context: dict[str, Any]) -> list[Message]:
 </CURRENT_MESSAGE>
 
 <ENTITIES>
-{to_prompt_json(context['nodes'], indent=2)}
+{to_prompt_json(context['nodes'])}
 </ENTITIES>
 
 <REFERENCE_TIME>
@@ -141,7 +141,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
 
     user_prompt = f"""
 <PREVIOUS MESSAGES>
-{to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS MESSAGES>
 <CURRENT MESSAGE>
 {context['episode_content']}
@@ -175,7 +175,7 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
             content=f"""
 
         <MESSAGE>
-        {to_prompt_json(context['episode_content'], indent=2)}
+        {to_prompt_json(context['episode_content'])}
         </MESSAGE>
         <REFERENCE TIME>
         {context['reference_time']}

graphiti_core/prompts/extract_nodes.py

@@ -18,8 +18,11 @@ from typing import Any, Protocol, TypedDict
 
 from pydantic import BaseModel, Field
 
+from graphiti_core.utils.text_utils import MAX_SUMMARY_CHARS
+
 from .models import Message, PromptFunction, PromptVersion
 from .prompt_helpers import to_prompt_json
+from .snippets import summary_instructions
 
 
 class ExtractedEntity(BaseModel):
@@ -42,7 +45,8 @@ class EntityClassificationTriple(BaseModel):
     uuid: str = Field(description='UUID of the entity')
     name: str = Field(description='Name of the entity')
     entity_type: str | None = Field(
-        default=None, description='Type of the entity. Must be one of the provided types or None'
+        default=None,
+        description='Type of the entity. Must be one of the provided types or None',
     )
 
 
@@ -55,7 +59,7 @@ class EntityClassification(BaseModel):
 class EntitySummary(BaseModel):
     summary: str = Field(
         ...,
-        description='Summary containing the important information about the entity. Under 250 words',
+        description=f'Summary containing the important information about the entity. Under {MAX_SUMMARY_CHARS} characters.',
     )
 
 
@@ -89,7 +93,7 @@ def extract_message(context: dict[str, Any]) -> list[Message]:
 </ENTITY TYPES>
 
 <PREVIOUS MESSAGES>
-{to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS MESSAGES>
 
 <CURRENT MESSAGE>
@@ -197,7 +201,7 @@ def reflexion(context: dict[str, Any]) -> list[Message]:
 
     user_prompt = f"""
 <PREVIOUS MESSAGES>
-{to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+{to_prompt_json([ep for ep in context['previous_episodes']])}
 </PREVIOUS MESSAGES>
 <CURRENT MESSAGE>
 {context['episode_content']}
@@ -221,22 +225,22 @@ def classify_nodes(context: dict[str, Any]) -> list[Message]:
 
     user_prompt = f"""
     <PREVIOUS MESSAGES>
-    {to_prompt_json([ep for ep in context['previous_episodes']], indent=2)}
+    {to_prompt_json([ep for ep in context['previous_episodes']])}
     </PREVIOUS MESSAGES>
     <CURRENT MESSAGE>
     {context['episode_content']}
     </CURRENT MESSAGE>
-    
+
     <EXTRACTED ENTITIES>
     {context['extracted_entities']}
     </EXTRACTED ENTITIES>
-    
+
     <ENTITY TYPES>
     {context['entity_types']}
     </ENTITY TYPES>
-    
+
     Given the above conversation, extracted entities, and provided entity types and their descriptions, classify the extracted entities.
-    
+
     Guidelines:
     1. Each entity must have exactly one type
     2. Only use the provided ENTITY TYPES as types, do not use additional types to classify entities.
@@ -257,19 +261,18 @@ def extract_attributes(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
-
-        <MESSAGES>
-        {to_prompt_json(context['previous_episodes'], indent=2)}
-        {to_prompt_json(context['episode_content'], indent=2)}
-        </MESSAGES>
-
-        Given the above MESSAGES and the following ENTITY, update any of its attributes based on the information provided
+        Given the MESSAGES and the following ENTITY, update any of its attributes based on the information provided
         in MESSAGES. Use the provided attribute descriptions to better understand how each attribute should be determined.
 
         Guidelines:
         1. Do not hallucinate entity property values if they cannot be found in the current context.
         2. Only use the provided MESSAGES and ENTITY to set attribute values.
-        
+
+        <MESSAGES>
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
+        </MESSAGES>
+
         <ENTITY>
         {context['node']}
         </ENTITY>
@@ -287,21 +290,16 @@ def extract_summary(context: dict[str, Any]) -> list[Message]:
         Message(
             role='user',
             content=f"""
+        Given the MESSAGES and the ENTITY, update the summary that combines relevant information about the entity
+        from the messages and relevant information from the existing summary.
+
+        {summary_instructions}
 
         <MESSAGES>
-        {to_prompt_json(context['previous_episodes'], indent=2)}
-        {to_prompt_json(context['episode_content'], indent=2)}
+        {to_prompt_json(context['previous_episodes'])}
+        {to_prompt_json(context['episode_content'])}
         </MESSAGES>
 
-        Given the above MESSAGES and the following ENTITY, update the summary that combines relevant information about the entity
-        from the messages and relevant information from the existing summary.
-        
-        Guidelines:
-        1. Do not hallucinate entity summary information if they cannot be found in the current context.
-        2. Only use the provided MESSAGES and ENTITY to set attribute values.
-        3. The summary attribute represents a summary of the ENTITY, and should be updated with new information about the Entity from the MESSAGES. 
-            Summaries must be no longer than 250 words.
-
         <ENTITY>
         {context['node']}
         </ENTITY>

graphiti_core/prompts/prompt_helpers.py

@@ -1,17 +1,33 @@
+"""
+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
 from typing import Any
 
 DO_NOT_ESCAPE_UNICODE = '\nDo not escape unicode characters.\n'
 
 
-def to_prompt_json(data: Any, ensure_ascii: bool = False, indent: int = 2) -> str:
+def to_prompt_json(data: Any, ensure_ascii: bool = False, indent: int | None = None) -> str:
     """
     Serialize data to JSON for use in prompts.
 
     Args:
         data: The data to serialize
         ensure_ascii: If True, escape non-ASCII characters. If False (default), preserve them.
-        indent: Number of spaces for indentation
+        indent: Number of spaces for indentation. Defaults to None (minified).
 
     Returns:
         JSON string representation of the data