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

graphiti_core/llm_client/openai_generic_client.py

@@ -17,7 +17,7 @@ limitations under the License.
 import json
 import logging
 import typing
-from typing import ClassVar
+from typing import Any, ClassVar
 
 import openai
 from openai import AsyncOpenAI
@@ -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
 
@@ -59,15 +59,20 @@ class OpenAIGenericClient(LLMClient):
     MAX_RETRIES: ClassVar[int] = 2
 
     def __init__(
-        self, config: LLMConfig | None = None, cache: bool = False, client: typing.Any = None
+        self,
+        config: LLMConfig | None = None,
+        cache: bool = False,
+        client: typing.Any = None,
+        max_tokens: int = 16384,
     ):
         """
-        Initialize the OpenAIClient with the provided configuration, cache setting, and client.
+        Initialize the OpenAIGenericClient with the provided configuration, cache setting, and client.
 
         Args:
             config (LLMConfig | None): The configuration for the LLM client, including API key, model, base URL, temperature, and max tokens.
             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.
+            max_tokens (int): The maximum number of tokens to generate. Defaults to 16384 (16K) for better compatibility with local models.
 
         """
         # removed caching to simplify the `generate_response` override
@@ -79,6 +84,9 @@ class OpenAIGenericClient(LLMClient):
 
         super().__init__(config, cache)
 
+        # Override max_tokens to support higher limits for local models
+        self.max_tokens = max_tokens
+
         if client is None:
             self.client = AsyncOpenAI(api_key=config.api_key, base_url=config.base_url)
         else:
@@ -99,12 +107,25 @@ class OpenAIGenericClient(LLMClient):
             elif m.role == 'system':
                 openai_messages.append({'role': 'system', 'content': m.content})
         try:
+            # Prepare response format
+            response_format: dict[str, Any] = {'type': 'json_object'}
+            if response_model is not None:
+                schema_name = getattr(response_model, '__name__', 'structured_response')
+                json_schema = response_model.model_json_schema()
+                response_format = {
+                    'type': 'json_schema',
+                    'json_schema': {
+                        'name': schema_name,
+                        'schema': json_schema,
+                    },
+                }
+
             response = await self.client.chat.completions.create(
                 model=self.model or DEFAULT_MODEL,
                 messages=openai_messages,
                 temperature=self.temperature,
                 max_tokens=self.max_tokens,
-                response_format={'type': 'json_object'},
+                response_format=response_format,  # type: ignore[arg-type]
             )
             result = response.choices[0].message.content or ''
             return json.loads(result)
@@ -120,60 +141,74 @@ 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[
-                -1
-            ].content += (
-                f'\n\nRespond with a JSON object in the following format:\n\n{serialized_model}'
-            )
-
         # 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=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

@@ -14,43 +14,267 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from graphiti_core.driver.driver import GraphProvider
+
 EPISODIC_EDGE_SAVE = """
-        MATCH (episode:Episodic {uuid: $episode_uuid}) 
-        MATCH (node:Entity {uuid: $entity_uuid}) 
-        MERGE (episode)-[r:MENTIONS {uuid: $uuid}]->(node)
-        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-        RETURN r.uuid AS uuid"""
-
-EPISODIC_EDGE_SAVE_BULK = """
-    UNWIND $episodic_edges AS edge
-    MATCH (episode:Episodic {uuid: edge.source_node_uuid}) 
-    MATCH (node:Entity {uuid: edge.target_node_uuid}) 
-    MERGE (episode)-[r:MENTIONS {uuid: edge.uuid}]->(node)
-    SET r = {uuid: edge.uuid, group_id: edge.group_id, created_at: edge.created_at}
-    RETURN r.uuid AS uuid
+    MATCH (episode:Episodic {uuid: $episode_uuid})
+    MATCH (node:Entity {uuid: $entity_uuid})
+    MERGE (episode)-[e:MENTIONS {uuid: $uuid}]->(node)
+    SET
+        e.group_id = $group_id,
+        e.created_at = $created_at
+    RETURN e.uuid AS uuid
 """
 
-ENTITY_EDGE_SAVE = """
-        MATCH (source:Entity {uuid: $edge_data.source_uuid}) 
-        MATCH (target:Entity {uuid: $edge_data.target_uuid}) 
-        MERGE (source)-[r:RELATES_TO {uuid: $edge_data.uuid}]->(target)
-        SET r = $edge_data
-        WITH r CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", $edge_data.fact_embedding)
-        RETURN r.uuid AS uuid"""
-
-ENTITY_EDGE_SAVE_BULK = """
-    UNWIND $entity_edges AS edge
-    MATCH (source:Entity {uuid: edge.source_node_uuid}) 
-    MATCH (target:Entity {uuid: edge.target_node_uuid}) 
-    MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
-    SET r = edge
-    WITH r, edge CALL db.create.setRelationshipVectorProperty(r, "fact_embedding", edge.fact_embedding)
-    RETURN edge.uuid AS uuid
+
+def get_episodic_edge_save_bulk_query(provider: GraphProvider) -> str:
+    if provider == GraphProvider.KUZU:
+        return """
+            MATCH (episode:Episodic {uuid: $source_node_uuid})
+            MATCH (node:Entity {uuid: $target_node_uuid})
+            MERGE (episode)-[e:MENTIONS {uuid: $uuid}]->(node)
+            SET
+                e.group_id = $group_id,
+                e.created_at = $created_at
+            RETURN e.uuid AS uuid
+        """
+
+    return """
+        UNWIND $episodic_edges AS edge
+        MATCH (episode:Episodic {uuid: edge.source_node_uuid})
+        MATCH (node:Entity {uuid: edge.target_node_uuid})
+        MERGE (episode)-[e:MENTIONS {uuid: edge.uuid}]->(node)
+        SET
+            e.group_id = edge.group_id,
+            e.created_at = edge.created_at
+        RETURN e.uuid AS uuid
+    """
+
+
+EPISODIC_EDGE_RETURN = """
+    e.uuid AS uuid,
+    e.group_id AS group_id,
+    n.uuid AS source_node_uuid,
+    m.uuid AS target_node_uuid,
+    e.created_at AS created_at
 """
 
-COMMUNITY_EDGE_SAVE = """
-        MATCH (community:Community {uuid: $community_uuid}) 
-        MATCH (node:Entity | Community {uuid: $entity_uuid}) 
-        MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node)
-        SET r = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
-        RETURN r.uuid AS uuid"""
+
+def get_entity_edge_save_query(provider: GraphProvider, has_aoss: bool = False) -> str:
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                MATCH (source:Entity {uuid: $edge_data.source_uuid})
+                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:
+            return """
+                MATCH (source:Entity {uuid: $edge_data.source_uuid})
+                MATCH (target:Entity {uuid: $edge_data.target_uuid})
+                MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+                SET e = removeKeyFromMap(removeKeyFromMap($edge_data, "fact_embedding"), "episodes")
+                SET e.fact_embedding = join([x IN coalesce($edge_data.fact_embedding, []) | toString(x) ], ",")
+                SET e.episodes = join($edge_data.episodes, ",")
+                RETURN $edge_data.uuid AS uuid
+            """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (source:Entity {uuid: $source_uuid})
+                MATCH (target:Entity {uuid: $target_uuid})
+                MERGE (source)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(target)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at,
+                    e.name = $name,
+                    e.fact = $fact,
+                    e.fact_embedding = $fact_embedding,
+                    e.episodes = $episodes,
+                    e.expired_at = $expired_at,
+                    e.valid_at = $valid_at,
+                    e.invalid_at = $invalid_at,
+                    e.attributes = $attributes
+                RETURN e.uuid AS uuid
+            """
+        case _:  # Neo4j
+            save_embedding_query = (
+                """WITH e CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", $edge_data.fact_embedding)"""
+                if not has_aoss
+                else ''
+            )
+            return (
+                (
+                    """
+                        MATCH (source:Entity {uuid: $edge_data.source_uuid})
+                        MATCH (target:Entity {uuid: $edge_data.target_uuid})
+                        MERGE (source)-[e:RELATES_TO {uuid: $edge_data.uuid}]->(target)
+                        SET e = $edge_data
+                        """
+                    + save_embedding_query
+                )
+                + """
+                RETURN e.uuid AS uuid
+                """
+            )
+
+
+def get_entity_edge_save_bulk_query(provider: GraphProvider, has_aoss: bool = False) -> str:
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                UNWIND $entity_edges AS edge
+                MATCH (source:Entity {uuid: edge.source_node_uuid})
+                MATCH (target:Entity {uuid: edge.target_node_uuid})
+                MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
+                SET r = edge
+                SET r.fact_embedding = vecf32(edge.fact_embedding)
+                WITH r, edge
+                RETURN edge.uuid AS uuid
+            """
+        case GraphProvider.NEPTUNE:
+            return """
+                UNWIND $entity_edges AS edge
+                MATCH (source:Entity {uuid: edge.source_node_uuid})
+                MATCH (target:Entity {uuid: edge.target_node_uuid})
+                MERGE (source)-[r:RELATES_TO {uuid: edge.uuid}]->(target)
+                SET r = removeKeyFromMap(removeKeyFromMap(edge, "fact_embedding"), "episodes")
+                SET r.fact_embedding = join([x IN coalesce(edge.fact_embedding, []) | toString(x) ], ",")
+                SET r.episodes = join(edge.episodes, ",")
+                RETURN edge.uuid AS uuid
+            """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (source:Entity {uuid: $source_node_uuid})
+                MATCH (target:Entity {uuid: $target_node_uuid})
+                MERGE (source)-[:RELATES_TO]->(e:RelatesToNode_ {uuid: $uuid})-[:RELATES_TO]->(target)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at,
+                    e.name = $name,
+                    e.fact = $fact,
+                    e.fact_embedding = $fact_embedding,
+                    e.episodes = $episodes,
+                    e.expired_at = $expired_at,
+                    e.valid_at = $valid_at,
+                    e.invalid_at = $invalid_at,
+                    e.attributes = $attributes
+                RETURN e.uuid AS uuid
+            """
+        case _:
+            save_embedding_query = (
+                'WITH e, edge CALL db.create.setRelationshipVectorProperty(e, "fact_embedding", edge.fact_embedding)'
+                if not has_aoss
+                else ''
+            )
+            return (
+                """
+                    UNWIND $entity_edges AS edge
+                    MATCH (source:Entity {uuid: edge.source_node_uuid})
+                    MATCH (target:Entity {uuid: edge.target_node_uuid})
+                    MERGE (source)-[e:RELATES_TO {uuid: edge.uuid}]->(target)
+                    SET e = edge
+                    """
+                + save_embedding_query
+                + """
+                RETURN edge.uuid AS uuid
+            """
+            )
+
+
+def get_entity_edge_return_query(provider: GraphProvider) -> str:
+    # `fact_embedding` is not returned by default and must be manually loaded using `load_fact_embedding()`.
+
+    if provider == GraphProvider.NEPTUNE:
+        return """
+        e.uuid AS uuid,
+        n.uuid AS source_node_uuid,
+        m.uuid AS target_node_uuid,
+        e.group_id AS group_id,
+        e.name AS name,
+        e.fact AS fact,
+        split(e.episodes, ',') AS episodes,
+        e.created_at AS created_at,
+        e.expired_at AS expired_at,
+        e.valid_at AS valid_at,
+        e.invalid_at AS invalid_at,
+        properties(e) AS attributes
+    """
+
+    return """
+        e.uuid AS uuid,
+        n.uuid AS source_node_uuid,
+        m.uuid AS target_node_uuid,
+        e.group_id AS group_id,
+        e.created_at AS created_at,
+        e.name AS name,
+        e.fact AS fact,
+        e.episodes AS episodes,
+        e.expired_at AS expired_at,
+        e.valid_at AS valid_at,
+        e.invalid_at AS invalid_at,
+    """ + (
+        'e.attributes AS attributes'
+        if provider == GraphProvider.KUZU
+        else 'properties(e) AS attributes'
+    )
+
+
+def get_community_edge_save_query(provider: GraphProvider) -> str:
+    match provider:
+        case GraphProvider.FALKORDB:
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+                RETURN e.uuid AS uuid
+            """
+        case GraphProvider.NEPTUNE:
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node {uuid: $entity_uuid})
+                WHERE node:Entity OR node:Community
+                MERGE (community)-[r:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET r.uuid= $uuid
+                SET r.group_id= $group_id
+                SET r.created_at= $created_at
+                RETURN r.uuid AS uuid
+            """
+        case GraphProvider.KUZU:
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node:Entity {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at
+                RETURN e.uuid AS uuid
+                UNION
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node:Community {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET
+                    e.group_id = $group_id,
+                    e.created_at = $created_at
+                RETURN e.uuid AS uuid
+            """
+        case _:  # Neo4j
+            return """
+                MATCH (community:Community {uuid: $community_uuid})
+                MATCH (node:Entity | Community {uuid: $entity_uuid})
+                MERGE (community)-[e:HAS_MEMBER {uuid: $uuid}]->(node)
+                SET e = {uuid: $uuid, group_id: $group_id, created_at: $created_at}
+                RETURN e.uuid AS uuid
+            """
+
+
+COMMUNITY_EDGE_RETURN = """
+    e.uuid AS uuid,
+    e.group_id AS group_id,
+    n.uuid AS source_node_uuid,
+    m.uuid AS target_node_uuid,
+    e.created_at AS created_at
+"""