graphiti-core 0.11.6rc9__tar.gz → 0.12.0__tar.gz

{graphiti_core-0.11.6rc9 → graphiti_core-0.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: graphiti-core
-Version: 0.11.6rc9
+Version: 0.12.0
 Summary: A temporal graph building library
 License: Apache-2.0
 Author: Paul Paliychuk
@@ -17,12 +17,13 @@ Provides-Extra: google-genai
 Provides-Extra: groq
 Requires-Dist: anthropic (>=0.49.0) ; extra == "anthropic"
 Requires-Dist: diskcache (>=5.6.3)
+Requires-Dist: falkordb (>=1.1.2,<2.0.0)
 Requires-Dist: google-genai (>=1.8.0) ; extra == "google-genai"
 Requires-Dist: groq (>=0.2.0) ; extra == "groq"
-Requires-Dist: neo4j (>=5.23.0)
+Requires-Dist: neo4j (>=5.26.0)
 Requires-Dist: numpy (>=1.0.0)
 Requires-Dist: openai (>=1.53.0)
-Requires-Dist: pydantic (>=2.8.2)
+Requires-Dist: pydantic (>=2.11.5)
 Requires-Dist: python-dotenv (>=1.0.1)
 Requires-Dist: tenacity (>=9.0.0)
 Project-URL: Homepage, https://help.getzep.com/graphiti/graphiti/overview
@@ -136,7 +137,7 @@ Graphiti is specifically designed to address the challenges of dynamic and frequ
 Requirements:
 
 - Python 3.10 or higher
-- Neo4j 5.26 or higher (serves as the embeddings storage backend)
+- Neo4j 5.26 / FalkorDB 1.1.2 or higher (serves as the embeddings storage backend)
 - OpenAI API key (for LLM inference and embedding)
 
 > [!IMPORTANT]
@@ -236,7 +237,7 @@ Graphiti supports Azure OpenAI for both LLM inference and embeddings. To use Azu
 ```python
 from openai import AsyncAzureOpenAI
 from graphiti_core import Graphiti
-from graphiti_core.llm_client import OpenAIClient
+from graphiti_core.llm_client import LLMConfig, OpenAIClient
 from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 
@@ -252,12 +253,19 @@ azure_openai_client = AsyncAzureOpenAI(
     azure_endpoint=azure_endpoint
 )
 
+# Create LLM Config with your Azure deployed model names
+azure_llm_config = LLMConfig(
+    small_model="gpt-4.1-nano",
+    model="gpt-4.1-mini",
+)
+
 # Initialize Graphiti with Azure OpenAI clients
 graphiti = Graphiti(
     "bolt://localhost:7687",
     "neo4j",
     "password",
     llm_client=OpenAIClient(
+        llm_config=azure_llm_config,
         client=azure_openai_client
     ),
     embedder=OpenAIEmbedder(
@@ -268,6 +276,7 @@ graphiti = Graphiti(
     ),
     # Optional: Configure the OpenAI cross encoder with Azure OpenAI
     cross_encoder=OpenAIRerankerClient(
+        llm_config=azure_llm_config,
         client=azure_openai_client
     )
 )

{graphiti_core-0.11.6rc9 → graphiti_core-0.12.0}/README.md

@@ -105,7 +105,7 @@ Graphiti is specifically designed to address the challenges of dynamic and frequ
 Requirements:
 
 - Python 3.10 or higher
-- Neo4j 5.26 or higher (serves as the embeddings storage backend)
+- Neo4j 5.26 / FalkorDB 1.1.2 or higher (serves as the embeddings storage backend)
 - OpenAI API key (for LLM inference and embedding)
 
 > [!IMPORTANT]
@@ -205,7 +205,7 @@ Graphiti supports Azure OpenAI for both LLM inference and embeddings. To use Azu
 ```python
 from openai import AsyncAzureOpenAI
 from graphiti_core import Graphiti
-from graphiti_core.llm_client import OpenAIClient
+from graphiti_core.llm_client import LLMConfig, OpenAIClient
 from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
 from graphiti_core.cross_encoder.openai_reranker_client import OpenAIRerankerClient
 
@@ -221,12 +221,19 @@ azure_openai_client = AsyncAzureOpenAI(
     azure_endpoint=azure_endpoint
 )
 
+# Create LLM Config with your Azure deployed model names
+azure_llm_config = LLMConfig(
+    small_model="gpt-4.1-nano",
+    model="gpt-4.1-mini",
+)
+
 # Initialize Graphiti with Azure OpenAI clients
 graphiti = Graphiti(
     "bolt://localhost:7687",
     "neo4j",
     "password",
     llm_client=OpenAIClient(
+        llm_config=azure_llm_config,
         client=azure_openai_client
     ),
     embedder=OpenAIEmbedder(
@@ -237,6 +244,7 @@ graphiti = Graphiti(
     ),
     # Optional: Configure the OpenAI cross encoder with Azure OpenAI
     cross_encoder=OpenAIRerankerClient(
+        llm_config=azure_llm_config,
         client=azure_openai_client
     )
 )

{graphiti_core-0.11.6rc9 → graphiti_core-0.12.0}/graphiti_core/cross_encoder/openai_reranker_client.py

@@ -106,7 +106,7 @@ class OpenAIRerankerClient(CrossEncoderClient):
                 if len(top_logprobs) == 0:
                     continue
                 norm_logprobs = np.exp(top_logprobs[0].logprob)
-                if bool(top_logprobs[0].token):
+                if top_logprobs[0].token.strip().split(' ')[0].lower() == 'true':
                     scores.append(norm_logprobs)
                 else:
                     scores.append(1 - norm_logprobs)

graphiti_core-0.12.0/graphiti_core/driver/__init__.py

@@ -0,0 +1,17 @@
+"""
+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.
+"""
+
+__all__ = ['GraphDriver', 'Neo4jDriver', 'FalkorDriver']

graphiti_core-0.12.0/graphiti_core/driver/driver.py

@@ -0,0 +1,66 @@
+"""
+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 logging
+from abc import ABC, abstractmethod
+from collections.abc import Coroutine
+from typing import Any
+
+from graphiti_core.helpers import DEFAULT_DATABASE
+
+logger = logging.getLogger(__name__)
+
+
+class GraphDriverSession(ABC):
+    async def __aenter__(self):
+        return self
+
+    @abstractmethod
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Falkor, but method must exist
+        pass
+
+    @abstractmethod
+    async def run(self, query: str, **kwargs: Any) -> Any:
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    async def execute_write(self, func, *args, **kwargs):
+        raise NotImplementedError()
+
+
+class GraphDriver(ABC):
+    provider: str
+
+    @abstractmethod
+    def execute_query(self, cypher_query_: str, **kwargs: Any) -> Coroutine:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def session(self, database: str) -> GraphDriverSession:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def delete_all_indexes(self, database_: str = DEFAULT_DATABASE) -> Coroutine:
+        raise NotImplementedError()

graphiti_core-0.12.0/graphiti_core/driver/falkordb_driver.py

@@ -0,0 +1,132 @@
+"""
+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 logging
+from collections.abc import Coroutine
+from datetime import datetime
+from typing import Any
+
+from falkordb import Graph as FalkorGraph  # type: ignore
+from falkordb.asyncio import FalkorDB  # type: ignore
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.helpers import DEFAULT_DATABASE
+
+logger = logging.getLogger(__name__)
+
+
+class FalkorDriverSession(GraphDriverSession):
+    def __init__(self, graph: FalkorGraph):
+        self.graph = graph
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        # No cleanup needed for Falkor, but method must exist
+        pass
+
+    async def close(self):
+        # No explicit close needed for FalkorDB, but method must exist
+        pass
+
+    async def execute_write(self, func, *args, **kwargs):
+        # Directly await the provided async function with `self` as the transaction/session
+        return await func(self, *args, **kwargs)
+
+    async def run(self, query: str | list, **kwargs: Any) -> Any:
+        # FalkorDB does not support argument for Label Set, so it's converted into an array of queries
+        if isinstance(query, list):
+            for cypher, params in query:
+                params = convert_datetimes_to_strings(params)
+                await self.graph.query(str(cypher), params)
+        else:
+            params = dict(kwargs)
+            params = convert_datetimes_to_strings(params)
+            await self.graph.query(str(query), params)
+        # Assuming `graph.query` is async (ideal); otherwise, wrap in executor
+        return None
+
+
+class FalkorDriver(GraphDriver):
+    provider: str = 'falkordb'
+
+    def __init__(
+        self,
+        uri: str,
+        user: str,
+        password: str,
+    ):
+        super().__init__()
+        if user and password:
+            uri_parts = uri.split('://', 1)
+            uri = f'{uri_parts[0]}://{user}:{password}@{uri_parts[1]}'
+
+        self.client = FalkorDB.from_url(
+            url=uri,
+        )
+
+    def _get_graph(self, graph_name: str | None) -> FalkorGraph:
+        # FalkorDB requires a non-None database name for multi-tenant graphs; the default is "DEFAULT_DATABASE"
+        if graph_name is None:
+            graph_name = 'DEFAULT_DATABASE'
+        return self.client.select_graph(graph_name)
+
+    async def execute_query(self, cypher_query_, **kwargs: Any):
+        graph_name = kwargs.pop('database_', DEFAULT_DATABASE)
+        graph = self._get_graph(graph_name)
+
+        # Convert datetime objects to ISO strings (FalkorDB does not support datetime objects directly)
+        params = convert_datetimes_to_strings(dict(kwargs))
+
+        try:
+            result = await graph.query(cypher_query_, params)
+        except Exception as e:
+            if 'already indexed' in str(e):
+                # check if index already exists
+                logger.info(f'Index already exists: {e}')
+                return None
+            logger.error(f'Error executing FalkorDB query: {e}')
+            raise
+
+        # Convert the result header to a list of strings
+        header = [h[1].decode('utf-8') for h in result.header]
+        return result.result_set, header, None
+
+    def session(self, database: str | None) -> GraphDriverSession:
+        return FalkorDriverSession(self._get_graph(database))
+
+    async def close(self) -> None:
+        await self.client.connection.close()
+
+    async def delete_all_indexes(self, database_: str = DEFAULT_DATABASE) -> Coroutine:
+        return self.execute_query(
+            'CALL db.indexes() YIELD name DROP INDEX name',
+            database_=database_,
+        )
+
+
+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

graphiti_core-0.12.0/graphiti_core/driver/neo4j_driver.py

@@ -0,0 +1,61 @@
+"""
+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 logging
+from collections.abc import Coroutine
+from typing import Any
+
+from neo4j import AsyncGraphDatabase
+from typing_extensions import LiteralString
+
+from graphiti_core.driver.driver import GraphDriver, GraphDriverSession
+from graphiti_core.helpers import DEFAULT_DATABASE
+
+logger = logging.getLogger(__name__)
+
+
+class Neo4jDriver(GraphDriver):
+    provider: str = 'neo4j'
+
+    def __init__(
+        self,
+        uri: str,
+        user: str | None,
+        password: str | None,
+    ):
+        super().__init__()
+        self.client = AsyncGraphDatabase.driver(
+            uri=uri,
+            auth=(user or '', password or ''),
+        )
+
+    async def execute_query(self, cypher_query_: LiteralString, **kwargs: Any) -> Coroutine:
+        params = kwargs.pop('params', None)
+        result = await self.client.execute_query(cypher_query_, parameters_=params, **kwargs)
+
+        return result
+
+    def session(self, database: str) -> GraphDriverSession:
+        return self.client.session(database=database)  # type: ignore
+
+    async def close(self) -> None:
+        return await self.client.close()
+
+    def delete_all_indexes(self, database_: str = DEFAULT_DATABASE) -> Coroutine:
+        return self.client.execute_query(
+            'CALL db.indexes() YIELD name DROP INDEX name',
+            database_=database_,
+        )

{graphiti_core-0.11.6rc9 → graphiti_core-0.12.0}/graphiti_core/edges.py

@@ -21,10 +21,10 @@ from time import time
 from typing import Any
 from uuid import uuid4
 
-from neo4j import AsyncDriver
 from pydantic import BaseModel, Field
 from typing_extensions import LiteralString
 
+from graphiti_core.driver.driver import GraphDriver
 from graphiti_core.embedder import EmbedderClient
 from graphiti_core.errors import EdgeNotFoundError, GroupsEdgesNotFoundError
 from graphiti_core.helpers import DEFAULT_DATABASE, parse_db_date
@@ -49,7 +49,9 @@ ENTITY_EDGE_RETURN: LiteralString = """
             e.episodes AS episodes,
             e.expired_at AS expired_at,
             e.valid_at AS valid_at,
-            e.invalid_at AS invalid_at"""
+            e.invalid_at AS invalid_at,
+            properties(e) AS attributes
+            """
 
 
 class Edge(BaseModel, ABC):
@@ -60,9 +62,9 @@ class Edge(BaseModel, ABC):
     created_at: datetime
 
     @abstractmethod
-    async def save(self, driver: AsyncDriver): ...
+    async def save(self, driver: GraphDriver): ...
 
-    async def delete(self, driver: AsyncDriver):
+    async def delete(self, driver: GraphDriver):
         result = await driver.execute_query(
             """
         MATCH (n)-[e:MENTIONS|RELATES_TO|HAS_MEMBER {uuid: $uuid}]->(m)
@@ -85,11 +87,11 @@ class Edge(BaseModel, ABC):
         return False
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str): ...
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str): ...
 
 
 class EpisodicEdge(Edge):
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             EPISODIC_EDGE_SAVE,
             episode_uuid=self.source_node_uuid,
@@ -100,12 +102,12 @@ class EpisodicEdge(Edge):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Episodic)-[e:MENTIONS {uuid: $uuid}]->(m:Entity)
@@ -128,7 +130,7 @@ class EpisodicEdge(Edge):
         return edges[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Episodic)-[e:MENTIONS]->(m:Entity)
@@ -154,7 +156,7 @@ class EpisodicEdge(Edge):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -209,6 +211,9 @@ class EntityEdge(Edge):
     invalid_at: datetime | None = Field(
         default=None, description='datetime of when the fact stopped being true'
     )
+    attributes: dict[str, Any] = Field(
+        default={}, description='Additional attributes of the edge. Dependent on edge name'
+    )
 
     async def generate_embedding(self, embedder: EmbedderClient):
         start = time()
@@ -221,7 +226,7 @@ class EntityEdge(Edge):
 
         return self.fact_embedding
 
-    async def load_fact_embedding(self, driver: AsyncDriver):
+    async def load_fact_embedding(self, driver: GraphDriver):
         query: LiteralString = """
             MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
             RETURN e.fact_embedding AS fact_embedding
@@ -235,30 +240,36 @@ class EntityEdge(Edge):
 
         self.fact_embedding = records[0]['fact_embedding']
 
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
+        edge_data: dict[str, Any] = {
+            'source_uuid': self.source_node_uuid,
+            'target_uuid': self.target_node_uuid,
+            'uuid': self.uuid,
+            'name': self.name,
+            'group_id': self.group_id,
+            'fact': self.fact,
+            'fact_embedding': self.fact_embedding,
+            'episodes': self.episodes,
+            'created_at': self.created_at,
+            'expired_at': self.expired_at,
+            'valid_at': self.valid_at,
+            'invalid_at': self.invalid_at,
+        }
+
+        edge_data.update(self.attributes or {})
+
         result = await driver.execute_query(
             ENTITY_EDGE_SAVE,
-            source_uuid=self.source_node_uuid,
-            target_uuid=self.target_node_uuid,
-            uuid=self.uuid,
-            name=self.name,
-            group_id=self.group_id,
-            fact=self.fact,
-            fact_embedding=self.fact_embedding,
-            episodes=self.episodes,
-            created_at=self.created_at,
-            expired_at=self.expired_at,
-            valid_at=self.valid_at,
-            invalid_at=self.invalid_at,
+            edge_data=edge_data,
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Entity)-[e:RELATES_TO {uuid: $uuid}]->(m:Entity)
@@ -276,7 +287,7 @@ class EntityEdge(Edge):
         return edges[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         if len(uuids) == 0:
             return []
 
@@ -298,7 +309,7 @@ class EntityEdge(Edge):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -331,11 +342,11 @@ class EntityEdge(Edge):
         return edges
 
     @classmethod
-    async def get_by_node_uuid(cls, driver: AsyncDriver, node_uuid: str):
+    async def get_by_node_uuid(cls, driver: GraphDriver, node_uuid: str):
         query: LiteralString = (
             """
-                                        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
-                                        """
+                                                        MATCH (n:Entity {uuid: $node_uuid})-[e:RELATES_TO]-(m:Entity)
+                                                        """
             + ENTITY_EDGE_RETURN
         )
         records, _, _ = await driver.execute_query(
@@ -348,7 +359,7 @@ class EntityEdge(Edge):
 
 
 class CommunityEdge(Edge):
-    async def save(self, driver: AsyncDriver):
+    async def save(self, driver: GraphDriver):
         result = await driver.execute_query(
             COMMUNITY_EDGE_SAVE,
             community_uuid=self.source_node_uuid,
@@ -359,12 +370,12 @@ class CommunityEdge(Edge):
             database_=DEFAULT_DATABASE,
         )
 
-        logger.debug(f'Saved edge to neo4j: {self.uuid}')
+        logger.debug(f'Saved edge to Graph: {self.uuid}')
 
         return result
 
     @classmethod
-    async def get_by_uuid(cls, driver: AsyncDriver, uuid: str):
+    async def get_by_uuid(cls, driver: GraphDriver, uuid: str):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community)-[e:HAS_MEMBER {uuid: $uuid}]->(m:Entity | Community)
@@ -385,7 +396,7 @@ class CommunityEdge(Edge):
         return edges[0]
 
     @classmethod
-    async def get_by_uuids(cls, driver: AsyncDriver, uuids: list[str]):
+    async def get_by_uuids(cls, driver: GraphDriver, uuids: list[str]):
         records, _, _ = await driver.execute_query(
             """
         MATCH (n:Community)-[e:HAS_MEMBER]->(m:Entity | Community)
@@ -409,7 +420,7 @@ class CommunityEdge(Edge):
     @classmethod
     async def get_by_group_ids(
         cls,
-        driver: AsyncDriver,
+        driver: GraphDriver,
         group_ids: list[str],
         limit: int | None = None,
         uuid_cursor: str | None = None,
@@ -452,12 +463,12 @@ def get_episodic_edge_from_record(record: Any) -> EpisodicEdge:
         group_id=record['group_id'],
         source_node_uuid=record['source_node_uuid'],
         target_node_uuid=record['target_node_uuid'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
     )
 
 
 def get_entity_edge_from_record(record: Any) -> EntityEdge:
-    return EntityEdge(
+    edge = EntityEdge(
         uuid=record['uuid'],
         source_node_uuid=record['source_node_uuid'],
         target_node_uuid=record['target_node_uuid'],
@@ -465,12 +476,27 @@ def get_entity_edge_from_record(record: Any) -> EntityEdge:
         name=record['name'],
         group_id=record['group_id'],
         episodes=record['episodes'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
         expired_at=parse_db_date(record['expired_at']),
         valid_at=parse_db_date(record['valid_at']),
         invalid_at=parse_db_date(record['invalid_at']),
+        attributes=record['attributes'],
     )
 
+    edge.attributes.pop('uuid', None)
+    edge.attributes.pop('source_node_uuid', None)
+    edge.attributes.pop('target_node_uuid', None)
+    edge.attributes.pop('fact', None)
+    edge.attributes.pop('name', None)
+    edge.attributes.pop('group_id', None)
+    edge.attributes.pop('episodes', None)
+    edge.attributes.pop('created_at', None)
+    edge.attributes.pop('expired_at', None)
+    edge.attributes.pop('valid_at', None)
+    edge.attributes.pop('invalid_at', None)
+
+    return edge
+
 
 def get_community_edge_from_record(record: Any):
     return CommunityEdge(
@@ -478,7 +504,7 @@ def get_community_edge_from_record(record: Any):
         group_id=record['group_id'],
         source_node_uuid=record['source_node_uuid'],
         target_node_uuid=record['target_node_uuid'],
-        created_at=record['created_at'].to_native(),
+        created_at=parse_db_date(record['created_at']),  # type: ignore
     )