feat(core): cleaner code separation + usage of nodeLabel

- big refactor to properly use GraphService, and enforce dependency injection

BREAKING CHANGE: The usage of node.label is no longer supported in the app, a specific node format
is used throughout the app and for inserting in the neo4j.

flowsint-core/src/flowsint_core/__init__.py

@@ -7,7 +7,7 @@ from .core.celery import *
 from .core.config import *
 from .core.enums import *
 from .core.events import *
-from .core.graph_db import *
+from .core.graph import *
 from .core.logger import *
 from .core.models import *
 from .core.postgre_db import *

flowsint-core/src/flowsint_core/core/enricher_base.py

@@ -1,12 +1,13 @@
 from abc import ABC, abstractmethod
-from typing import List, Dict, Any, Optional
-from pydantic import ValidationError, BaseModel, Field, create_model, TypeAdapter
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model
 from pydantic.config import ConfigDict
-from .graph_db import Neo4jConnection
+
+from ..utils import resolve_type
+from .graph import GraphService, create_graph_service
 from .logger import Logger
 from .vault import VaultProtocol
-from .graph_service import GraphService, create_graph_service
-from ..utils import resolve_type
 
 
 class InvalidEnricherParams(Exception):
@@ -111,7 +112,6 @@ class Enricher(ABC):
         self,
         sketch_id: Optional[str] = None,
         scan_id: Optional[str] = None,
-        neo4j_conn: Optional[Neo4jConnection] = None,
         params_schema: Optional[List[Dict[str, Any]]] = None,
         vault: Optional[VaultProtocol] = None,
         params: Optional[Dict[str, Any]] = None,
@@ -119,20 +119,17 @@ class Enricher(ABC):
     ):
         self.scan_id = scan_id or "default"
         self.sketch_id = sketch_id or "system"
-        self.neo4j_conn = neo4j_conn  # Kept for backward compatibility
         self.vault = vault
         self.params_schema = params_schema or []
         self.ParamsModel = build_params_model(self.params_schema)
         self.params: Dict[str, Any] = params or {}
 
-        # Initialize graph service (new architecture)
+        # Initialize graph service (uses singleton connection by default)
         if graph_service:
             self._graph_service = graph_service
         else:
-            # Create graph service with the provided or singleton connection
             self._graph_service = create_graph_service(
                 sketch_id=self.sketch_id,
-                neo4j_connection=neo4j_conn,
                 enable_batching=True,
             )
 
@@ -439,7 +436,7 @@ class Enricher(ABC):
                 )
             return []
 
-    def create_node(self, node_obj, **properties) -> None:
+    def create_node(self, node_obj) -> None:
         """
         Create a single Neo4j node.
 
@@ -458,7 +455,7 @@ class Enricher(ABC):
             node_obj: Either a Pydantic object or node label string
             **properties: Additional node properties or overrides
         """
-        self._graph_service.create_node(node_obj=node_obj, **properties)
+        self._graph_service.create_node_from_flowsint_type(node_obj=node_obj)
 
     def create_relationship(
         self,

flowsint-core/src/flowsint_core/core/graph/__init__.py

@@ -0,0 +1,43 @@
+"""
+Graph module for Neo4j operations.
+
+This module provides all graph-related functionality including:
+- Neo4j connection management
+- Repository pattern for graph operations
+- Serialization utilities
+- High-level graph service
+"""
+
+from .connection import Neo4jConnection, neo4j_connection
+from .repository import Neo4jGraphRepository
+from .repository_protocol import GraphRepositoryProtocol
+from .serializer import GraphSerializer
+from .service import GraphService, create_graph_service, LoggerProtocol
+from .types import (
+    GraphData,
+    GraphEdge,
+    GraphNode,
+    Neo4jDict,
+    NodeMetadata,
+)
+
+__all__ = [
+    # Connection
+    "Neo4jConnection",
+    "neo4j_connection",
+    # Repository
+    "Neo4jGraphRepository",
+    "GraphRepositoryProtocol",
+    # Serializer
+    "GraphSerializer",
+    # Service
+    "GraphService",
+    "create_graph_service",
+    "LoggerProtocol",
+    # Types
+    "GraphData",
+    "GraphEdge",
+    "GraphNode",
+    "Neo4jDict",
+    "NodeMetadata",
+]

flowsint-core/src/flowsint_core/core/graph/connection.py

@@ -1,8 +1,16 @@
+"""
+Neo4j connection management.
+
+This module provides a singleton connection manager for Neo4j with proper
+connection pooling and transaction management.
+"""
+
 import os
 from threading import Lock
-from typing import Optional, Dict, Any, List
-from neo4j import GraphDatabase, Driver, Session
+from typing import Any, Dict, List, Optional
+
 from dotenv import load_dotenv
+from neo4j import Driver, GraphDatabase
 
 load_dotenv()
 
@@ -16,7 +24,7 @@ class Neo4jConnection:
     connection pooling and resource management.
     """
 
-    _instance: Optional['Neo4jConnection'] = None
+    _instance: Optional["Neo4jConnection"] = None
     _lock: Lock = Lock()
     _driver: Optional[Driver] = None
 
@@ -44,22 +52,26 @@ class Neo4jConnection:
         """
         # Only initialize once
         if self._driver is None:
-            self._uri = uri or os.getenv("NEO4J_URI_BOLT")
-            self._user = user or os.getenv("NEO4J_USERNAME")
-            self._password = password or os.getenv("NEO4J_PASSWORD")
+            resolved_uri = uri or os.getenv("NEO4J_URI_BOLT")
+            resolved_user = user or os.getenv("NEO4J_USERNAME")
+            resolved_password = password or os.getenv("NEO4J_PASSWORD")
 
-            if not all([self._uri, self._user, self._password]):
+            if not resolved_uri or not resolved_user or not resolved_password:
                 raise ValueError("Neo4j connection credentials are required")
 
+            self._uri = resolved_uri
+            self._user = resolved_user
+            self._password = resolved_password
+
             self._driver = GraphDatabase.driver(
                 self._uri,
                 auth=(self._user, self._password),
                 max_connection_pool_size=50,
-                connection_acquisition_timeout=60.0
+                connection_acquisition_timeout=60.0,
             )
 
     @classmethod
-    def get_instance(cls) -> 'Neo4jConnection':
+    def get_instance(cls) -> "Neo4jConnection":
         """
         Get the singleton instance.
 
@@ -79,7 +91,9 @@ class Neo4jConnection:
         """
         return self._driver
 
-    def query(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
+    def query(
+        self, query: str, parameters: Dict[str, Any] = None
+    ) -> List[Dict[str, Any]]:
         """
         Execute a single query.
 
@@ -102,7 +116,9 @@ class Neo4jConnection:
             return {}
         return {k: v for k, v in parameters.items() if k is not None}
 
-    def execute_write(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
+    def execute_write(
+        self, query: str, parameters: Dict[str, Any] = None
+    ) -> List[Dict[str, Any]]:
         """
         Execute a write query within a write transaction.
 
@@ -113,6 +129,7 @@ class Neo4jConnection:
         Returns:
             List of result records as dictionaries
         """
+
         def _execute(tx):
             cleaned_params = self._clean_parameters(parameters)
             result = tx.run(query, cleaned_params)
@@ -121,7 +138,9 @@ class Neo4jConnection:
         with self._driver.session() as session:
             return session.execute_write(_execute)
 
-    def execute_batch(self, queries: List[tuple[str, Dict[str, Any]]]) -> List[List[Dict[str, Any]]]:
+    def execute_batch(
+        self, queries: List[tuple[str, Dict[str, Any]]]
+    ) -> List[List[Dict[str, Any]]]:
         """
         Execute multiple queries in a single transaction.
 
@@ -131,6 +150,7 @@ class Neo4jConnection:
         Returns:
             List of results, one for each query
         """
+
         def _execute_batch(tx):
             results = []
             for query, params in queries:
@@ -189,5 +209,6 @@ try:
         neo4j_connection = None
 except Exception as e:
     import logging
+
     logging.getLogger(__name__).warning(f"Failed to initialize Neo4j connection: {e}")
     neo4j_connection = None

flowsint-core/src/flowsint_core/core/graph/repository.py

@@ -2,24 +2,24 @@
 Graph database repository for Neo4j operations.
 
 This module provides a repository pattern implementation for Neo4j,
-handling node and relationship operations with batching support.
+handling raw Neo4jDict object and operations with batching support.
 """
 
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
-from pydantic import BaseModel
-
-from .graph_db import Neo4jConnection
-from .graph_serializer import GraphSerializer
+from .connection import Neo4jConnection
+from .types import Neo4jDict
 
 
-class GraphRepository:
+class Neo4jGraphRepository:
     """
-    Repository for Neo4j graph database operations.
+    Neo4j main implementation of the graph repository.
 
     This class follows the Repository pattern, providing a clean abstraction
     over Neo4j operations and handling batching for improved performance.
+
+    Implements GraphRepositoryProtocol for dependency injection.
     """
 
     def __init__(self, neo4j_connection: Optional[Neo4jConnection] = None):
@@ -34,64 +34,20 @@ class GraphRepository:
         self._batch_operations: List[Tuple[str, Dict[str, Any]]] = []
         self._batch_size = 100
 
-    @staticmethod
-    def _get_primary_field(obj: BaseModel) -> str:
-        """Get the primary field of a Pydantic object."""
-        model_fields = obj.__class__.model_fields
-
-        primary_field = None
-        for field_name, field_info in model_fields.items():
-            if field_info.json_schema_extra and field_info.json_schema_extra.get(
-                "primary"
-            ):
-                primary_field = field_name
-                break
-
-        if primary_field is None:
-            for field_name, field_info in model_fields.items():
-                if field_info.is_required():
-                    primary_field = field_name
-                    break
-            if primary_field is None:
-                primary_field = next(iter(model_fields.keys()))
-
-        return primary_field
-
-    @staticmethod
-    def _extract_primary_value(obj: BaseModel) -> Any:
-        """
-        Extract the primitive value from a Pydantic object recursively.
-        Uses model_dump to properly serialize Pydantic types like HttpUrl.
-        """
-        primary_field = GraphRepository._get_primary_field(obj)
-        obj_dict = (
-            obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict()
-        )
-        value = obj_dict.get(primary_field)
-
-        if isinstance(value, dict):
-            nested_obj = getattr(obj, primary_field)
-            if isinstance(nested_obj, BaseModel):
-                return GraphRepository._extract_primary_value(nested_obj)
-
-        return value
-
     def create_node(
         self,
-        node_obj: BaseModel,
-        sketch_id: Optional[str] = None,
-        **properties: Any,
+        node_obj: Neo4jDict,
+        sketch_id: str,
     ) -> Optional[str]:
         """
         Create or update a single node in Neo4j.
 
-        Supports 1 signature:
-        1. Pydantic object: create_node(obj, sketch_id="...", **overrides)
+        Supported signature:
+        1. Neo4jDict object: create_node(obj)
 
         Args:
-            node_obj: a Pydantic object
-            sketch_id: Investigation sketch ID (required)
-            **properties: Additional node properties or overrides
+            node_obj: a Neo4jDict object
+            sketch_id: str
 
         Returns:
             Element ID of created/updated node
@@ -99,37 +55,29 @@ class GraphRepository:
         if not self._connection:
             return None
 
-        query, params = self._build_node_query(node_obj, sketch_id, **properties)
+        query, params = self._build_node_query(node_obj, sketch_id)
         result = self._connection.query(query, params)
         return result[0]["id"] if result else None
 
     def create_relationship(
         self,
-        from_obj: BaseModel,
-        to_obj: BaseModel,
-        rel_label: Optional[str] = None,
-        sketch_id: Optional[str] = None,
-        **properties: Any,
+        rel_obj: Neo4jDict,
+        sketch_id: str,
     ) -> None:
         """
         Create a relationship between two nodes.
 
         Supports one signature:
-         - Pydantic objects: create_relationship(obj1, obj2, "REL_TYPE", sketch_id="...")
+         - Pydantic objects: create_relationship(relation, sketch_id="...")
 
         Args:
-            from_obj: a Pydantic object (source)
-            to_obj: a Pydantic object (target)
-            rel_label: Relationship label (ex: "IS_CONNECTED_TO")
+            rel_obj: the Neo4jDict object
             sketch_id: Investigation sketch ID (required)
-            **properties: Additional relationship properties
         """
         if not self._connection:
             return
 
-        query, params = self._build_relationship_query(
-            from_obj, to_obj, rel_label, sketch_id, **properties
-        )
+        query, params = self._build_relationship_query(rel_obj, sketch_id)
 
         self._connection.execute_write(query, params)
 
@@ -155,155 +103,53 @@ class GraphRepository:
             self.flush_batch()
 
     def _build_node_query(
-        self,
-        node_obj: BaseModel,
-        sketch_id: str,
-        **properties: Any,
+        self, node_obj: Neo4jDict, sketch_id: str
     ) -> Tuple[str, Dict[str, Any]]:
-        """Build a node creation query."""
+        node_label = node_obj.get("nodeLabel")
+        node_type = node_obj.get("nodeType")
 
-        node_type = node_obj.__class__.__name__.lower()
-        primary_field = self._get_primary_field(node_obj)
-        key_prop = primary_field
-        key_value = getattr(node_obj, primary_field, None)
-
-        # If primary key value is None or empty, fallback to label
-        if key_value is None or (isinstance(key_value, str) and not key_value):
-            key_prop = "label"
-            key_value = properties.get("label") or getattr(node_obj, "label", "Node")
-
-        if isinstance(key_value, BaseModel):
-            key_value = self._extract_primary_value(key_value)
-
-        obj_dict = (
-            node_obj.model_dump(mode="json")
-            if hasattr(node_obj, "model_dump")
-            else node_obj.dict()
-        )
-        # Extract only non-dict values, skip None keys
-        obj_properties = {
-            k: (v if v is not None else "")
-            for k, v in obj_dict.items()
-            if k is not None and not isinstance(v, dict)
-        }
-        obj_properties.update(properties)
-        properties = obj_properties
-
-        if not sketch_id:
-            raise ValueError("sketch_id is required")
-        if not key_prop or key_prop is None:
-            raise ValueError(f"key_prop cannot be None for node type {node_type}")
-
-        serialized_props = GraphSerializer.serialize_properties(properties)
-        serialized_props["type"] = node_type.lower()
-        serialized_props["sketch_id"] = sketch_id
-        label = serialized_props.get("label", key_value)
-        serialized_props["label"] = label
-
-        # Clean None keys from serialized_props
-        serialized_props = {k: v for k, v in serialized_props.items() if k is not None}
-
-        set_clauses = [
-            f"n.{prop} = ${prop}"
-            for prop in serialized_props.keys()
-            if prop != "sketch_id"
-        ]
-
-        # Build params, ensuring no None keys
+        # paramètres Neo4j
         params = {
-            key_prop: key_value,
+            "props": node_obj,  # flat with keys containing "."
+            "node_label": node_label,
+            "sketch_id": sketch_id,
             "created_at": datetime.now(timezone.utc).isoformat(),
-            **serialized_props,
         }
 
-        # MERGE on both key_prop AND sketch_id for uniqueness per sketch
-        # Use ON CREATE SET to only set created_at when creating, not updating
         query = f"""
-        MERGE (n:{node_type} {{{key_prop}: ${key_prop}, sketch_id: $sketch_id}})
+        MERGE (n:{node_type} {{ nodeLabel: $node_label, sketch_id: $sketch_id }})
         ON CREATE SET n.created_at = $created_at
-        SET {", ".join(set_clauses)}
-        RETURN elementId(n) as id
+        SET n += $props
+        RETURN elementId(n) AS id
         """
 
         return query, params
 
     def _build_relationship_query(
         self,
-        from_obj: BaseModel,
-        to_obj: BaseModel,
-        rel_label: Optional[str] = None,
-        sketch_id: Optional[str] = None,
-        **properties: Any,
+        rel_obj: Neo4jDict,
+        sketch_id: str,
     ) -> Tuple[str, Dict[str, Any]]:
-        """Build a relationship creation query."""
-        # From object
-        from_node_type = from_obj.__class__.__name__.lower()
-        from_primary_field = self._get_primary_field(from_obj)
-        from_obj_dict = (
-            from_obj.model_dump(mode="json")
-            if hasattr(from_obj, "model_dump")
-            else from_obj.dict()
-        )
-        from_key_value = from_obj_dict.get(from_primary_field)
+        from_type = rel_obj["from_type"]
+        from_label = rel_obj["from_label"]
+        to_type = rel_obj["to_type"]
+        to_label = rel_obj["to_label"]
+        rel_label = rel_obj["rel_label"]
 
-        if isinstance(from_key_value, dict):
-            nested_obj = getattr(from_obj, from_primary_field)
-            if isinstance(nested_obj, BaseModel):
-                from_key_value = self._extract_primary_value(nested_obj)
-        # To object
-        to_node_type = to_obj.__class__.__name__.lower()
-        to_primary_field = self._get_primary_field(to_obj)
-        to_obj_dict = (
-            to_obj.model_dump(mode="json")
-            if hasattr(to_obj, "model_dump")
-            else to_obj.dict()
-        )
-        to_key_value = to_obj_dict.get(to_primary_field)
-
-        if isinstance(to_key_value, dict):
-            nested_obj = getattr(to_obj, to_primary_field)
-            if isinstance(nested_obj, BaseModel):
-                to_key_value = self._extract_primary_value(nested_obj)
-
-        from_type = from_node_type
-        from_key = from_primary_field
-        from_value = from_key_value
-        to_type = to_node_type
-        to_key = to_primary_field
-        to_value = to_key_value
-        rel_type = rel_label
-
-        if not sketch_id:
-            raise ValueError("sketch_id is required")
-        if not from_key or from_key is None:
-            raise ValueError(
-                f"from_key cannot be None for relationship type {rel_type}"
-            )
-        if not to_key or to_key is None:
-            raise ValueError(f"to_key cannot be None for relationship type {rel_type}")
-
-        serialized_props = GraphSerializer.serialize_properties(properties)
-        serialized_props["sketch_id"] = sketch_id
-
-        # Clean None keys
-        serialized_props = {k: v for k, v in serialized_props.items() if k is not None}
-
-        if serialized_props:
-            props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()])
-            rel_props = f"{{{props_str}}}"
-        else:
-            rel_props = "{sketch_id: $sketch_id}"
+        params = {
+            "from_label": from_label,
+            "to_label": to_label,
+            "sketch_id": sketch_id,
+            "props": rel_obj,
+        }
 
         query = f"""
-        MATCH (from:{from_type} {{{from_key}: $from_value, sketch_id: $sketch_id}})
-        MATCH (to:{to_type} {{{to_key}: $to_value, sketch_id: $sketch_id}})
-        MERGE (from)-[:{rel_type} {rel_props}]->(to)
+        MATCH (from:{from_type} {{nodeLabel: $from_label, sketch_id: $sketch_id}})
+        MATCH (to:{to_type} {{nodeLabel: $to_label, sketch_id: $sketch_id}})
+        MERGE (from)-[r:{rel_label}]->(to)
+        SET r += $props
         """
 
-        params = {"from_value": from_value, "to_value": to_value, **serialized_props}
-        # Clean None keys from params
-        params = {k: v for k, v in params.items() if k is not None}
-
         return query, params
 
     def flush_batch(self) -> None:
@@ -337,15 +183,15 @@ class GraphRepository:
 
     def batch_create_nodes(
         self,
-        nodes: List[BaseModel],
+        nodes: List[Neo4jDict],
         sketch_id: str,
     ) -> Dict[str, Any]:
         """
         Create multiple nodes in a single batch transaction.
 
         Args:
-            nodes: List of Pydantic model instances to create
-            sketch_id: Investigation sketch ID
+            nodes: List of Neo4jDict model instances to insert
+            sketch_id: investigation sketch id
 
         Returns:
             Dictionary with:
@@ -370,8 +216,7 @@ class GraphRepository:
         for idx, node_obj in enumerate(nodes):
             try:
                 query, params = self._build_node_query(
-                    node_obj=node_obj,
-                    sketch_id=sketch_id,
+                    node_obj=node_obj, sketch_id=sketch_id
                 )
                 batch_operations.append((query, params))
             except Exception as e:
@@ -411,11 +256,7 @@ class GraphRepository:
 
         Args:
             edges: List of edge dictionaries, each with:
-                - from_obj: Source node Pydantic model
-                - to_obj: Target node Pydantic model
-                - rel_label: Relationship type/label
-                - properties: Optional dict of relationship properties
-            sketch_id: Investigation sketch ID
+                - rel_obj: Source node Neo4jDict model
 
         Returns:
             Dictionary with:
@@ -434,23 +275,8 @@ class GraphRepository:
 
         for idx, edge in enumerate(edges):
             try:
-                from_obj = edge.get("from_obj")
-                to_obj = edge.get("to_obj")
-                rel_label = edge.get("rel_label")
-                properties = edge.get("properties", {})
-
-                if not from_obj or not to_obj or not rel_label:
-                    errors.append(
-                        f"Edge {idx}: Missing required fields (from_obj, to_obj, or rel_label)"
-                    )
-                    continue
-
                 query, params = self._build_relationship_query(
-                    from_obj=from_obj,
-                    to_obj=to_obj,
-                    rel_label=rel_label,
-                    sketch_id=sketch_id,
-                    **properties,
+                    rel_obj=edge, sketch_id=sketch_id
                 )
                 batch_operations.append((query, params))
             except Exception as e:
@@ -474,9 +300,9 @@ class GraphRepository:
 
     def batch_create_edges_by_element_id(
         self,
-        edges: List[Dict[str, Any]],
+        edges: List[Neo4jDict],
         sketch_id: str,
-    ) -> Dict[str, Any]:
+    ) -> Neo4jDict:
         """
         Create multiple edges/relationships using element IDs in a single batch transaction.
 
@@ -487,7 +313,7 @@ class GraphRepository:
             edges: List of edge dictionaries, each with:
                 - from_element_id: Source node element ID
                 - to_element_id: Target node element ID
-                - rel_type: Relationship type/label
+                - rel_label: Relationship type/label
                 - properties: Optional dict of relationship properties
             sketch_id: Investigation sketch ID
 
@@ -510,8 +336,7 @@ class GraphRepository:
             try:
                 from_element_id = edge.get("from_element_id")
                 to_element_id = edge.get("to_element_id")
-                rel_type = edge.get("rel_type", "RELATED_TO")
-                properties = edge.get("properties", {})
+                rel_label = edge.get("rel_label", "RELATED_TO")
 
                 if not from_element_id or not to_element_id:
                     errors.append(
@@ -519,14 +344,10 @@ class GraphRepository:
                     )
                     continue
 
-                # Serialize properties
-                serialized_props = GraphSerializer.serialize_properties(properties)
-                serialized_props["sketch_id"] = sketch_id
+                edge["sketch_id"] = sketch_id
 
                 # Build relationship properties string
-                props_str = ", ".join(
-                    [f"{k}: ${k}_{idx}" for k in serialized_props.keys()]
-                )
+                props_str = ", ".join([f"{k}: ${k}_{idx}" for k in edge.keys()])
                 rel_props = (
                     f"{{{props_str}}}" if props_str else "{sketch_id: $sketch_id}"
                 )
@@ -534,7 +355,7 @@ class GraphRepository:
                 query = f"""
                 MATCH (from) WHERE elementId(from) = $from_id_{idx}
                 MATCH (to) WHERE elementId(to) = $to_id_{idx}
-                MERGE (from)-[r:`{rel_type}` {rel_props}]->(to)
+                MERGE (from)-[r:`{rel_label}` {rel_props}]->(to)
                 """
 
                 # Build params with unique keys for batch execution
@@ -543,7 +364,7 @@ class GraphRepository:
                     f"to_id_{idx}": to_element_id,
                 }
                 # Add serialized properties with index suffix
-                for k, v in serialized_props.items():
+                for k, v in edge.items():
                     params[f"{k}_{idx}"] = v
 
                 batch_operations.append((query, params))
@@ -567,54 +388,24 @@ class GraphRepository:
             return {"edges_created": 0, "errors": errors}
 
     def update_node(
-        self,
-        element_id: str,
-        node_obj: BaseModel,
-        sketch_id: str,
+        self, element_id: str, updates: Neo4jDict, sketch_id: str
     ) -> Optional[str]:
-        """
-        Update an existing node's properties using a Pydantic model.
-
-        Args:
-            element_id: Neo4j element ID of the node to update
-            node_obj: Pydantic model instance with updated data
-            sketch_id: Investigation sketch ID (for safety)
-
-        Returns:
-            Element ID of the updated node or None if not found
-        """
         if not self._connection:
             return None
 
-        # Extract properties from Pydantic object
-        obj_dict = (
-            node_obj.model_dump(mode="json")
-            if hasattr(node_obj, "model_dump")
-            else node_obj.dict()
-        )
-
-        # Extract only non-dict values, skip None keys
-        properties = {
-            k: (v if v is not None else "")
-            for k, v in obj_dict.items()
-            if k is not None and not isinstance(v, dict)
-        }
-
-        # Serialize properties
-        serialized_props = GraphSerializer.serialize_properties(properties)
-        serialized_props["type"] = node_obj.__class__.__name__.lower()
-
-        # Build SET clauses
-        set_clauses = [f"n.{prop} = ${prop}" for prop in serialized_props.keys()]
-
-        query = f"""
+        query = """
         MATCH (n)
         WHERE elementId(n) = $element_id AND n.sketch_id = $sketch_id
-        SET {", ".join(set_clauses)}
-        RETURN elementId(n) as id
+        SET n += $props
+        RETURN elementId(n) AS id
         """
 
-        params = {"element_id": element_id, "sketch_id": sketch_id, **serialized_props}
+        params = {
+            "element_id": element_id,
+            "sketch_id": sketch_id,
+            "props": updates,
+        }
+
         result = self._connection.query(query, params)
         return result[0]["id"] if result else None
 
@@ -699,17 +490,17 @@ class GraphRepository:
         self, sketch_id: str, limit: int = 100000
     ) -> Dict[str, List[Dict[str, Any]]]:
         """
-        Get all nodes and relationships for a sketch.
+        Get all nodes and edges for a sketch.
 
         Args:
             sketch_id: Investigation sketch ID
             limit: Maximum number of nodes to return
 
         Returns:
-            Dictionary with 'nodes' and 'relationships' lists
+            Dictionary with 'nodes' and 'edges' lists
         """
         if not self._connection:
-            return {"nodes": [], "relationships": []}
+            return {"nodes": [], "edges": []}
 
         # Get all nodes for the sketch
         # Use OPTIONAL MATCH to avoid Neo4j warning when sketch_id property doesn't exist yet
@@ -726,11 +517,10 @@ class GraphRepository:
         )
 
         if not nodes_result:
-            return {"nodes": [], "relationships": []}
+            return {"nodes": [], "edges": []}
 
         node_ids = [record["id"] for record in nodes_result]
-
-        # Get all relationships between these nodes
+        # Get all edges between these nodes
         rels_query = """
         UNWIND $node_ids AS nid
         MATCH (a)-[r]->(b)
@@ -739,50 +529,30 @@ class GraphRepository:
                elementId(b) as target, properties(r) as data
         """
         rels_result = self._connection.query(rels_query, {"node_ids": node_ids})
-
-        return {"nodes": nodes_result, "relationships": rels_result or []}
+        return {"nodes": nodes_result, "edges": rels_result or []}
 
     def update_relationship(
-        self,
-        element_id: str,
-        properties: Dict[str, Any],
-        sketch_id: str,
+        self, element_id: str, rel_obj: Neo4jDict, sketch_id: str
     ) -> Optional[Dict[str, Any]]:
-        """
-        Update an existing relationship's properties.
-
-        Args:
-            element_id: Neo4j element ID of the edge to update
-            properties: Dictionary of properties to update
-            sketch_id: Investigation sketch ID (for safety)
-
-        Returns:
-            Updated edge data or None if not found
-        """
         if not self._connection:
             return None
 
-        # Filter out None values
-        filtered_props = {
-            k: (v if v is not None else "")
-            for k, v in properties.items()
-            if k is not None
+        query = """
+        MATCH ()-[r]->()
+        WHERE elementId(r) = $element_id AND r.sketch_id = $sketch_id
+        SET r += $props
+        RETURN
+            elementId(r) AS id,
+            type(r) AS type,
+            properties(r) AS data
+        """
+
+        params = {
+            "element_id": element_id,
+            "sketch_id": sketch_id,
+            "props": rel_obj,
         }
 
-        # Serialize properties
-        serialized_props = GraphSerializer.serialize_properties(filtered_props)
-
-        # Build SET clauses
-        set_clauses = [f"r.{prop} = ${prop}" for prop in serialized_props.keys()]
-
-        query = f"""
-            MATCH ()-[r]->()
-            WHERE elementId(r) = $element_id AND r.sketch_id = $sketch_id
-            SET {", ".join(set_clauses)}
-            RETURN elementId(r) as id, COALESCE(r.label, type(r)) as type, properties(r) as data
-            """
-
-        params = {"element_id": element_id, "sketch_id": sketch_id, **serialized_props}
         result = self._connection.query(query, params)
         return result[0] if result else None
 
@@ -790,9 +560,8 @@ class GraphRepository:
         self,
         from_element_id: str,
         to_element_id: str,
-        rel_type: str,
+        rel_label: str,
         sketch_id: str,
-        **properties: Any,
     ) -> Optional[Dict[str, Any]]:
         """
         Create a relationship between two nodes using their element IDs.
@@ -800,17 +569,15 @@ class GraphRepository:
         Args:
             from_element_id: Source node element ID
             to_element_id: Target node element ID
-            rel_type: Relationship type
+            rel_label: Relationship type
             sketch_id: Investigation sketch ID
-            **properties: Additional relationship properties
 
         Returns:
             Created relationship properties or None
         """
         if not self._connection:
             return None
-
-        serialized_props = GraphSerializer.serialize_properties(properties)
+        serialized_props = {}
         serialized_props["sketch_id"] = sketch_id
 
         props_str = ", ".join([f"{k}: ${k}" for k in serialized_props.keys()])
@@ -819,14 +586,14 @@ class GraphRepository:
         query = f"""
         MATCH (a) WHERE elementId(a) = $from_id
         MATCH (b) WHERE elementId(b) = $to_id
-        MERGE (a)-[r:`{rel_type}` {rel_props}]->(b)
+        MERGE (a)-[r:`{rel_label}` {rel_props}]->(b)
         RETURN properties(r) as rel
         """
 
         params = {
             "from_id": from_element_id,
             "to_id": to_element_id,
-            **serialized_props,
+            "sketch_id": sketch_id,
         }
 
         result = self._connection.query(query, params)
@@ -906,7 +673,7 @@ class GraphRepository:
             query, {"node_ids": node_ids, "sketch_id": sketch_id}
         )
 
-        return [record["data"] for record in result] if result else []
+        return result
 
     def merge_nodes(
         self,
@@ -931,7 +698,7 @@ class GraphRepository:
             return None
 
         node_type = new_node_data.get("type", "Node")
-        properties = GraphSerializer.serialize_properties(new_node_data)
+        properties = {}
         properties["sketch_id"] = sketch_id
 
         is_reusing_node = new_node_id and new_node_id in old_node_ids
@@ -1008,131 +775,93 @@ class GraphRepository:
 
         return new_node_element_id
 
-    def get_related_nodes(self, node_id: str, sketch_id: str) -> Dict[str, Any]:
+    def get_neighbors(self, node_id: str, sketch_id: str) -> Dict[str, Any]:
         """
-        Get a node and all its direct relationships and connected nodes.
-
-        Args:
-            node_id: Element ID of the center node
-            sketch_id: Investigation sketch ID
+        Get a node and all its direct relationships and connected nodes
+        within the same sketch.
 
         Returns:
-            Dictionary with 'nds' (nodes) and 'rls' (relationships) lists
+            {
+                "nodes": [ { "id", "data" } ],
+                "edges": [ { "id", "source", "target", "label" } ]
+            }
         """
+
         if not self._connection:
-            return {"nds": [], "rls": []}
+            return {"nodes": [], "edges": []}
 
-        center_query = """
+        query = """
         MATCH (n)
         WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id
-        RETURN elementId(n) as id, labels(n) as labels, properties(n) as data
-        """
 
-        center_result = self._connection.query(
-            center_query, {"sketch_id": sketch_id, "node_id": node_id}
-        )
+        OPTIONAL MATCH (n)-[r]-(other)
+        WHERE other.sketch_id = $sketch_id AND other <> n
 
-        if not center_result:
-            return {"nds": [], "rls": []}
-
-        relationships_query = """
-        MATCH (n)
-        WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id
-        OPTIONAL MATCH (n)-[r]->(other)
-        WHERE other.sketch_id = $sketch_id
-        OPTIONAL MATCH (other)-[r2]->(n)
-        WHERE other.sketch_id = $sketch_id
         RETURN
-            elementId(r) as rel_id,
-            type(r) as rel_type,
-            properties(r) as rel_data,
-            elementId(other) as other_node_id,
-            labels(other) as other_node_labels,
-            properties(other) as other_node_data,
-            'outgoing' as direction
-        UNION
-        MATCH (n)
-        WHERE elementId(n) = $node_id AND n.sketch_id = $sketch_id
-        OPTIONAL MATCH (other)-[r]->(n)
-        WHERE other.sketch_id = $sketch_id
-        RETURN
-            elementId(r) as rel_id,
-            type(r) as rel_type,
-            properties(r) as rel_data,
-            elementId(other) as other_node_id,
-            labels(other) as other_node_labels,
-            properties(other) as other_node_data,
-            'incoming' as direction
+            elementId(n)     AS center_id,
+            properties(n)    AS center_data,
+
+            elementId(r)     AS rel_id,
+            type(r)          AS rel_label,
+
+            elementId(other) AS other_id,
+            properties(other) AS other_data,
+
+            CASE
+                WHEN r IS NULL THEN NULL
+                WHEN startNode(r) = n THEN 'outgoing'
+                ELSE 'incoming'
+            END AS direction
         """
 
         result = self._connection.query(
-            relationships_query, {"sketch_id": sketch_id, "node_id": node_id}
+            query,
+            {"node_id": node_id, "sketch_id": sketch_id},
         )
 
-        center_record = center_result[0]
+        if not result:
+            return {"nodes": [], "edges": []}
+
+        first = result[0]
         center_node = {
-            "id": center_record["id"],
-            "labels": center_record["labels"],
-            "data": center_record["data"],
-            "label": center_record["data"].get("label", "Node"),
-            "type": "custom",
-            "caption": center_record["data"].get("label", "Node"),
+            "id": first["center_id"],
+            "data": first["center_data"],
         }
 
-        related_nodes = []
-        relationships = []
-        seen_nodes = set()
-        seen_relationships = set()
+        nodes = {center_node["id"]: center_node}
+        edges = {}
 
         for record in result:
             if not record["rel_id"]:
                 continue
 
-            if record["rel_id"] not in seen_relationships:
+            other_id = record["other_id"]
+
+            # nodes
+            if other_id not in nodes:
+                nodes[other_id] = {
+                    "id": other_id,
+                    "data": record["other_data"],
+                }
+
+            # edges
+            if record["rel_id"] not in edges:
                 if record["direction"] == "outgoing":
-                    relationships.append(
-                        {
-                            "id": record["rel_id"],
-                            "type": "straight",
-                            "source": center_node["id"],
-                            "target": record["other_node_id"],
-                            "data": record["rel_data"],
-                            "caption": record["rel_type"],
-                        }
-                    )
+                    source, target = center_node["id"], other_id
                 else:
-                    relationships.append(
-                        {
-                            "id": record["rel_id"],
-                            "type": "straight",
-                            "source": record["other_node_id"],
-                            "target": center_node["id"],
-                            "data": record["rel_data"],
-                            "caption": record["rel_type"],
-                        }
-                    )
-                seen_relationships.add(record["rel_id"])
+                    source, target = other_id, center_node["id"]
 
-            if (
-                record["other_node_id"]
-                and record["other_node_id"] not in seen_nodes
-                and record["other_node_id"] != center_node["id"]
-            ):
-                related_nodes.append(
-                    {
-                        "id": record["other_node_id"],
-                        "labels": record["other_node_labels"],
-                        "data": record["other_node_data"],
-                        "label": record["other_node_data"].get("label", "Node"),
-                        "type": "custom",
-                        "caption": record["other_node_data"].get("label", "Node"),
-                    }
-                )
-                seen_nodes.add(record["other_node_id"])
+                edges[record["rel_id"]] = {
+                    "id": record["rel_id"],
+                    "source": source,
+                    "target": target,
+                    "label": record["rel_label"],
+                }
 
-        all_nodes = [center_node] + related_nodes
-
-        return {"nds": all_nodes, "rls": relationships}
+        return {
+            "nodes": list(nodes.values()),
+            "edges": list(edges.values()),
+        }
 
     def count_nodes_by_sketch(self, sketch_id: str) -> int:
         """

flowsint-core/src/flowsint_core/core/graph/repository_protocol.py

@@ -0,0 +1,131 @@
+"""
+Protocol for graph repository implementations.
+
+This module defines the interface contract that all graph repository
+implementations must follow, enabling dependency injection and easier testing.
+"""
+
+from typing import Any, Dict, List, Optional, Protocol, Tuple
+
+from .types import Neo4jDict
+
+
+class GraphRepositoryProtocol(Protocol):
+    """
+    Protocol defining the interface for graph repository implementations.
+
+    This protocol enables:
+    - Dependency injection in GraphService
+    - Easy mocking and testing without patching
+    - Alternative implementations (e.g., InMemoryGraphRepository)
+    """
+
+    # Core node operations
+    def create_node(self, node_obj: Neo4jDict, sketch_id: str) -> Optional[str]:
+        """Create or update a single node. Returns element ID."""
+        ...
+
+    def update_node(
+        self, element_id: str, updates: Neo4jDict, sketch_id: str
+    ) -> Optional[str]:
+        """Update a node by its element ID. Returns element ID."""
+        ...
+
+    def delete_nodes(self, node_ids: List[str], sketch_id: str) -> int:
+        """Delete nodes by their element IDs. Returns count deleted."""
+        ...
+
+    def delete_all_sketch_nodes(self, sketch_id: str) -> int:
+        """Delete all nodes for a sketch. Returns count deleted."""
+        ...
+
+    def get_nodes_by_ids(
+        self, node_ids: List[str], sketch_id: str
+    ) -> List[Dict[str, Any]]:
+        """Get nodes by their element IDs."""
+        ...
+
+    def update_nodes_positions(
+        self, positions: List[Dict[str, Any]], sketch_id: str
+    ) -> int:
+        """Update positions for multiple nodes. Returns count updated."""
+        ...
+
+    # Core relationship operations
+    def create_relationship(self, rel_obj: Neo4jDict, sketch_id: str) -> None:
+        """Create a relationship between two nodes."""
+        ...
+
+    def create_relationship_by_element_id(
+        self,
+        from_element_id: str,
+        to_element_id: str,
+        rel_label: str,
+        sketch_id: str,
+    ) -> Optional[Dict[str, Any]]:
+        """Create a relationship using element IDs."""
+        ...
+
+    def update_relationship(
+        self, element_id: str, rel_obj: Neo4jDict, sketch_id: str
+    ) -> Optional[Dict[str, Any]]:
+        """Update a relationship by its element ID."""
+        ...
+
+    def delete_relationships(self, relationship_ids: List[str], sketch_id: str) -> int:
+        """Delete relationships by their element IDs. Returns count deleted."""
+        ...
+
+    # Graph queries
+    def get_sketch_graph(
+        self, sketch_id: str, limit: int = 100000
+    ) -> Dict[str, List[Dict[str, Any]]]:
+        """Get all nodes and edges for a sketch."""
+        ...
+
+    def get_neighbors(self, node_id: str, sketch_id: str) -> Dict[str, Any]:
+        """Get a node and all its direct relationships."""
+        ...
+
+    # Merge operations
+    def merge_nodes(
+        self,
+        old_node_ids: List[str],
+        new_node_data: Dict[str, Any],
+        new_node_id: Optional[str],
+        sketch_id: str,
+    ) -> Optional[str]:
+        """Merge multiple nodes into one. Returns new element ID."""
+        ...
+
+    # Batch operations
+    def batch_create_nodes(
+        self, nodes: List[Neo4jDict], sketch_id: str
+    ) -> Dict[str, Any]:
+        """Create multiple nodes in a single batch."""
+        ...
+
+    def batch_create_edges_by_element_id(
+        self, edges: List[Neo4jDict], sketch_id: str
+    ) -> Dict[str, Any]:
+        """Create multiple edges using element IDs in a single batch."""
+        ...
+
+    def add_to_batch(self, operation_type: str, **kwargs: Any) -> None:
+        """Add an operation to the batch queue."""
+        ...
+
+    def flush_batch(self) -> None:
+        """Execute all batched operations."""
+        ...
+
+    def set_batch_size(self, size: int) -> None:
+        """Set the batch size for auto-flushing."""
+        ...
+
+    # Custom queries
+    def query(
+        self, cypher: str, parameters: Dict[str, Any] = {}
+    ) -> List[Dict[str, Any]]:
+        """Execute a custom Cypher query."""
+        ...

flowsint-core/src/flowsint_core/core/graph/serializer.py

@@ -0,0 +1,199 @@
+"""
+Graph property serialization utilities.
+
+This module provides utilities for serializing complex Python objects
+into Neo4j-compatible primitive types, following the Single Responsibility Principle.
+"""
+
+from typing import Any, Dict, List
+
+from flowsint_types import TYPE_REGISTRY, FlowsintType
+from pydantic import BaseModel
+
+from flowsint_core.utils import flatten, unflatten
+
+from .types import GraphEdge, GraphNode, NodeMetadata
+
+
+class GraphSerializer:
+    """
+    Handles serialization of complex objects to Neo4j-compatible types.
+
+    This class is responsible for converting Pydantic models, nested objects,
+    and other complex types into primitive types that can be stored in Neo4j.
+    """
+
+    @staticmethod
+    def _clean_empty_values(data: Dict[str, Any]) -> dict:
+        """Remove empty string values from dict to avoid Pydantic validation errors."""
+        cleaned = {}
+        for key, value in data.items():
+            if value == "" or value is None:
+                continue
+            if isinstance(value, dict):
+                cleaned_nested = GraphSerializer._clean_empty_values(value)
+                if cleaned_nested:
+                    cleaned[key] = cleaned_nested
+            elif isinstance(value, list):
+                cleaned_list = [
+                    GraphSerializer._clean_empty_values(item)
+                    if isinstance(item, dict)
+                    else item
+                    for item in value
+                    if item != "" and item is not None
+                ]
+                if cleaned_list:
+                    cleaned[key] = cleaned_list
+            else:
+                cleaned[key] = value
+        return cleaned
+
+    @staticmethod
+    def flatten(dict: Dict[str, Any]):
+        return flatten(dict, remove_empty=False)
+
+    @staticmethod
+    def parse_flowsint_type(entity: Dict, nodeType: str) -> FlowsintType:
+        DetectedType = TYPE_REGISTRY.get_lowercase(nodeType)
+        if not DetectedType:
+            raise ValueError(f"Unknown type: {nodeType}")
+        properties = GraphSerializer._clean_empty_values(entity)
+        return DetectedType(**properties)
+
+    @staticmethod
+    def graph_node_to_flowsint_type(node: GraphNode) -> FlowsintType:
+        return node.nodeProperties
+
+    @staticmethod
+    def neo4j_dict_to_graph_node(node_dict: Dict[str, Any]) -> GraphNode:
+        """Convert a flattened Neo4j node record to a GraphNode instance.
+
+        Unflattens the data, parses the nodeProperties into the appropriate
+        FlowsintType subclass, and constructs a complete GraphNode.
+        """
+        data = node_dict.get("data")
+        node_id = str(node_dict.get("id"))
+        if not data:
+            raise Exception("Could not find node data to extract.")
+        # unflatten object from neo4j
+        node_dict = unflatten(data)
+        node_type = node_dict.get(
+            "nodeType", node_dict.get("type", "")
+        )  # legacy support type
+        nodeLabel = str(node_dict.get("nodeLabel", node_dict.get("label", "")))
+        node_properties = node_dict.get("nodeProperties", {})
+        node_metadata = node_dict.get("nodeMetadata", {})
+
+        node_properties.pop(
+            "nodeLabel", None
+        )  # remove nodeLabel from original pydantic
+
+        entity = GraphSerializer.parse_flowsint_type(node_properties, node_type)
+        return GraphNode(
+            id=node_id,
+            nodeLabel=nodeLabel,
+            nodeType=node_type,
+            nodeColor=data.get("nodeColor"),
+            nodeSize=data.get("nodeSize"),
+            nodeImage=data.get("nodeImage"),
+            nodeIcon=data.get("nodeIcon"),
+            nodeFlag=data.get("nodeFlag"),
+            x=data.get("x"),
+            y=data.get("y"),
+            nodeProperties=entity,
+            nodeMetadata=node_metadata,
+        )
+
+    @staticmethod
+    def flowsint_type_to_neo4j_dict(entity: FlowsintType) -> Dict[str, Any]:
+        node_type = entity.__class__.__name__.lower()
+        node_label = entity.nodeLabel
+        graph_node = GraphNode(
+            id="",
+            nodeColor=None,
+            nodeIcon=None,
+            nodeImage=None,
+            nodeFlag=None,
+            nodeLabel=node_label or "",
+            nodeType=node_type,
+            nodeProperties=entity,
+            nodeMetadata=NodeMetadata(),
+        )
+        return GraphSerializer.graph_node_to_neo4j_dict(graph_node)
+
+    @staticmethod
+    def graph_node_to_neo4j_dict(node: GraphNode) -> Dict[str, Any]:
+        """Convert a GraphNode to a flattened Neo4j-compatible dict.
+
+        Serializes the model to JSON-compatible types and flattens nested
+        structures into dot-notation keys for Neo4j property storage.
+        """
+        neo4j_dict = node.model_dump(mode="json", serialize_as_any=True)
+        neo4j_dict_flatten = flatten(neo4j_dict, remove_empty=False)
+        neo4j_dict_flatten.pop(
+            "nodeProperties.nodeLabel", None
+        )  # remove nodeLabel from original pydantic
+        return neo4j_dict_flatten
+
+    @staticmethod
+    def neo4j_dict_to_graph_edge(edge_dict: Dict[str, Any]) -> GraphEdge:
+        """Convert a Neo4j relationship record to a GraphEdge instance."""
+        return GraphEdge(
+            id=str(edge_dict.get("id")),
+            source=str(edge_dict.get("source")),
+            target=str(edge_dict.get("target")),
+            label=str(edge_dict.get("type")),
+        )
+
+    @staticmethod
+    def graph_edge_to_neo4j_dict(
+        from_obj: BaseModel, to_obj: BaseModel, label: str
+    ) -> Dict[str, Any]:
+        """Build a Neo4j relationship dict for matching nodes by type and label.
+
+        Creates a relationship descriptor that identifies source and target nodes
+        by their nodeType and nodeLabel, allowing relationship creation without IDs.
+        """
+        from_type = (
+            from_obj.__class__.__name__.lower()
+            if isinstance(from_obj, FlowsintType)
+            else from_obj.nodeType
+        )
+        to_type = (
+            to_obj.__class__.__name__.lower()
+            if isinstance(to_obj, FlowsintType)
+            else to_obj.nodeType
+        )
+        return {
+            "from_type": from_type,
+            "from_label": from_obj.nodeLabel,
+            "to_type": to_type,
+            "to_label": to_obj.nodeLabel,
+            "rel_label": label,
+        }
+
+    @staticmethod
+    def deserialize_nodes(node_dicts: List[Dict[str, Any]]) -> List[GraphNode]:
+        """Convert a list of Neo4j node records to GraphNode instances."""
+        return [
+            GraphSerializer.neo4j_dict_to_graph_node(node_dict)
+            for node_dict in node_dicts
+        ]
+
+    @staticmethod
+    def serialize_nodes(nodes: List[GraphNode]) -> List[Dict[str, Any]]:
+        """Convert a list of Neo4j node records to GraphNode instances."""
+        return [GraphSerializer.graph_node_to_neo4j_dict(node) for node in nodes]
+
+    @staticmethod
+    def serialize_flowsint_types(nodes: List[FlowsintType]) -> List[Dict[str, Any]]:
+        """Convert a list of Neo4j node records to GraphNode instances."""
+        return [GraphSerializer.flowsint_type_to_neo4j_dict(node) for node in nodes]
+
+    @staticmethod
+    def deserialize_edges(edge_dicts: List[Dict[str, Any]]) -> List[GraphEdge]:
+        """Convert a list of Neo4j relationship records to GraphEdge instances."""
+        return [
+            GraphSerializer.neo4j_dict_to_graph_edge(edge_dict)
+            for edge_dict in edge_dicts
+        ]

flowsint-core/src/flowsint_core/core/graph/service.py

@@ -0,0 +1,361 @@
+"""
+Graph service for high-level graph operations.
+
+This module provides a service layer for graph operations,
+integrating repository and logging functionality.
+"""
+
+from typing import Any, Dict, List, Optional, Protocol
+
+from flowsint_types import FlowsintType
+from pydantic import BaseModel
+
+from .repository import Neo4jGraphRepository
+from .repository_protocol import GraphRepositoryProtocol
+from .serializer import GraphSerializer
+from .types import GraphData, GraphNode, Neo4jDict
+
+
+class LoggerProtocol(Protocol):
+    """Protocol for logger implementations."""
+
+    @staticmethod
+    def graph_append(sketch_id: str, message: Dict[str, Any]) -> None:
+        """Log a graph append message."""
+        ...
+
+
+class GraphService:
+    """
+    High-level service for graph operations.
+
+    This service provides a clean interface for enricher operations,
+    handling both graph persistence and logging with proper separation of concerns.
+
+    This service can support FlowsintType input or GraphNode input.
+    """
+
+    def __init__(
+        self,
+        sketch_id: str,
+        repository: GraphRepositoryProtocol,
+        logger: Optional[LoggerProtocol] = None,
+        enable_batching: bool = False,
+    ):
+        """
+        Initialize the graph service.
+
+        Args:
+            sketch_id: Investigation sketch ID
+            repository: Repository instance (required - dependency injection)
+            logger: Optional logger instance
+            enable_batching: Enable batch operations
+
+        Raises:
+            ValueError: If repository is not provided
+        """
+        if repository is None:
+            raise ValueError(
+                "repository is required. Use create_graph_service() factory "
+                "or provide a GraphRepositoryProtocol implementation."
+            )
+        self._sketch_id = sketch_id
+        self._repository = repository
+        self._logger = logger
+        self._enable_batching = enable_batching
+
+    @property
+    def sketch_id(self) -> str:
+        """Get the sketch ID."""
+        return self._sketch_id
+
+    @property
+    def repository(self) -> GraphRepositoryProtocol:
+        """Get the underlying repository."""
+        return self._repository
+
+    def create_node(self, node_obj: GraphNode) -> str | None:
+        """
+        Create or update a node in the graph.
+
+        Supports one signatures:
+         - GraphNode object: create_node(obj)
+
+        Args:
+            node_obj: a GraphNode object
+        """
+
+        if isinstance(node_obj, FlowsintType):
+            raise Exception(
+                "create_node method takes a GraphNode as input. If you want to insert a node from a FlowsintType, please use create_node_from_flowsint_type method."
+            )
+
+        neo4j_node_dict: Neo4jDict = GraphSerializer.graph_node_to_neo4j_dict(node_obj)
+
+        if self._enable_batching:
+            self._repository.add_to_batch(
+                "node",
+                node_obj=neo4j_node_dict,
+                sketch_id=self._sketch_id,
+            )
+        else:
+            return self._repository.create_node(
+                node_obj=neo4j_node_dict,
+                sketch_id=self._sketch_id,
+            )
+
+    def create_node_from_flowsint_type(self, node_obj: FlowsintType) -> str | None:
+        """
+        Create or update a node in the graph.
+
+        Supports one signatures:
+         - FlowsintType object: create_node(obj)
+
+        Args:
+            node_obj: a FlowsintType object
+        """
+
+        if isinstance(node_obj, GraphNode):
+            raise Exception(
+                "create_node_from_flowsint_type method takes a FlowsintType as input. If you want to insert a node from a GraphNode, please use create_node method."
+            )
+
+        neo4j_node_dict: Neo4jDict = GraphSerializer.flowsint_type_to_neo4j_dict(
+            node_obj
+        )
+
+        if self._enable_batching:
+            self._repository.add_to_batch(
+                "node",
+                node_obj=neo4j_node_dict,
+                sketch_id=self._sketch_id,
+            )
+        else:
+            return self._repository.create_node(
+                node_obj=neo4j_node_dict,
+                sketch_id=self._sketch_id,
+            )
+
+    def get_sketch_graph(self) -> GraphData:
+        graph_data = self.repository.get_sketch_graph(self.sketch_id)
+        nodes = GraphSerializer.deserialize_nodes(graph_data.get("nodes", []))
+        edges = GraphSerializer.deserialize_edges(graph_data.get("edges", []))
+        return GraphData(nodes=nodes, edges=edges)
+
+    def get_nodes_by_ids(self, node_ids: List[str]) -> List[GraphNode]:
+        nodes = self.repository.get_nodes_by_ids(node_ids, self.sketch_id)
+        return GraphSerializer.deserialize_nodes(nodes)
+
+    def get_nodes_by_ids_for_task(self, node_ids: List[str]) -> List[BaseModel]:
+        nodes = self.get_nodes_by_ids(node_ids)
+        return [GraphSerializer.graph_node_to_flowsint_type(node) for node in nodes]
+
+    def create_relationship(
+        self,
+        from_obj: BaseModel,
+        to_obj: BaseModel,
+        rel_label: str = "IS_RELATED_TO",
+    ) -> None:
+        """
+        Create a relationship between two nodes.
+
+        Supports 1 signature:
+         - Pydantic objects: create_relationship(obj1, obj2, "rel_label")
+
+        Args:
+            from_obj: A GraphNode object (source)
+            to_obj: A GraphNode object (target)
+            rel_label: Relationship label (ex: "IS_CONNECTED_TO")
+            **properties: Additional relationship properties
+        """
+
+        neo4j_rel_dict: Neo4jDict = GraphSerializer.graph_edge_to_neo4j_dict(
+            from_obj, to_obj, rel_label
+        )
+
+        if self._enable_batching:
+            self._repository.add_to_batch(
+                "relationship",
+                rel_obj=neo4j_rel_dict,
+                sketch_id=self._sketch_id,
+            )
+        else:
+            self._repository.create_relationship(
+                rel_obj=neo4j_rel_dict,
+                sketch_id=self._sketch_id,
+            )
+
+    def create_relationship_by_element_id(
+        self,
+        from_element_id: str,
+        to_element_id: str,
+        rel_label: str = "IS_RELATED_TO",
+    ):
+        return self._repository.create_relationship_by_element_id(
+            from_element_id=from_element_id,
+            to_element_id=to_element_id,
+            rel_label=rel_label,
+            sketch_id=self._sketch_id,
+        )
+
+    def get_neighbors(self, node_id: str) -> GraphData:
+        graph_data = self._repository.get_neighbors(
+            node_id=node_id,
+            sketch_id=self._sketch_id,
+        )
+        nodes = GraphSerializer.deserialize_nodes(graph_data.get("nodes", []))
+        edges = GraphSerializer.deserialize_edges(graph_data.get("edges", []))
+        return GraphData(nodes=nodes, edges=edges)
+
+    def update_node(self, element_id: str, updates: Dict[str, Any]) -> str | None:
+        flatten_updates = GraphSerializer.flatten(updates)
+        """Update a node by its element ID."""
+        return self._repository.update_node(
+            element_id=element_id,
+            updates=flatten_updates,
+            sketch_id=self._sketch_id,
+        )
+
+    def update_nodes_positions(self, positions: List[Dict[str, Any]]) -> int:
+        """Update positions (x, y) for multiple nodes in batch."""
+        return self._repository.update_nodes_positions(
+            positions=positions,
+            sketch_id=self._sketch_id,
+        )
+
+    def delete_nodes(self, node_ids: List[str]) -> int:
+        """Delete nodes by their element IDs."""
+        return self._repository.delete_nodes(
+            node_ids=node_ids,
+            sketch_id=self._sketch_id,
+        )
+
+    def delete_relationships(self, relationship_ids: List[str]) -> int:
+        """Delete relationships by their element IDs."""
+        return self._repository.delete_relationships(
+            relationship_ids=relationship_ids,
+            sketch_id=self._sketch_id,
+        )
+
+    def delete_all_sketch_nodes(self) -> int:
+        """Delete all nodes and relationships for the sketch."""
+        return self._repository.delete_all_sketch_nodes(
+            sketch_id=self._sketch_id,
+        )
+
+    def update_relationship(
+        self, element_id: str, properties: Dict[str, Any]
+    ) -> Dict[str, Any] | None:
+        """Update a relationship by its element ID."""
+        return self._repository.update_relationship(
+            element_id=element_id,
+            rel_obj=properties,
+            sketch_id=self._sketch_id,
+        )
+
+    def merge_nodes(
+        self,
+        old_node_ids: List[str],
+        new_node_data: Dict[str, Any],
+        new_node_id: str | None = None,
+    ) -> str | None:
+        """Merge multiple nodes into one, transferring all relationships."""
+        return self._repository.merge_nodes(
+            old_node_ids=old_node_ids,
+            new_node_data=new_node_data,
+            new_node_id=new_node_id,
+            sketch_id=self._sketch_id,
+        )
+
+    def batch_create_nodes(self, nodes: List[Neo4jDict]) -> Dict[str, Any]:
+        """Create multiple nodes in a single batch transaction."""
+        return self._repository.batch_create_nodes(
+            nodes=nodes,
+            sketch_id=self._sketch_id,
+        )
+
+    def batch_create_edges_by_element_id(
+        self, edges: List[Neo4jDict]
+    ) -> Dict[str, Any]:
+        """Create multiple edges using element IDs in a single batch transaction."""
+        return self._repository.batch_create_edges_by_element_id(
+            edges=edges,
+            sketch_id=self._sketch_id,
+        )
+
+    def log_graph_message(self, message: str) -> None:
+        """
+        Log a graph operation message.
+
+        Args:
+            message: Message to log
+        """
+        if self._logger:
+            self._logger.graph_append(self._sketch_id, {"message": message})
+
+    def flush(self) -> None:
+        """Flush any pending batch operations."""
+        if self._enable_batching:
+            self._repository.flush_batch()
+
+    def query(self, cypher: str, parameters: Dict[str, Any] = None) -> list:
+        """
+        Execute a custom Cypher query.
+
+        Args:
+            cypher: Cypher query string
+            parameters: Query parameters
+
+        Returns:
+            List of result records
+        """
+        return self._repository.query(cypher, parameters)
+
+    def set_batch_size(self, size: int) -> None:
+        """
+        Set the batch size for operations.
+
+        Args:
+            size: Number of operations to batch
+        """
+        self._repository.set_batch_size(size)
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - auto-flush batch."""
+        if exc_type is None:
+            self.flush()
+
+
+def create_graph_service(
+    sketch_id: str,
+    enable_batching: bool = True,
+) -> GraphService:
+    """
+    Factory function to create a GraphService instance with Neo4j repository.
+
+    This is the recommended way to create a GraphService for production use.
+    For testing, inject an InMemoryGraphRepository or mock directly into GraphService.
+
+    Args:
+        sketch_id: Investigation sketch ID
+        enable_batching: Enable batch operations
+
+    Returns:
+        Configured GraphService instance
+    """
+    # Import Logger here to avoid circular imports
+    from flowsint_core.core.logger import Logger
+
+    # Neo4jGraphRepository uses Neo4jConnection.get_instance() singleton
+    repository = Neo4jGraphRepository()
+
+    return GraphService(
+        sketch_id=sketch_id,
+        repository=repository,
+        logger=Logger,
+        enable_batching=enable_batching,
+    )

flowsint-core/src/flowsint_core/core/graph/types.py

@@ -0,0 +1,59 @@
+"""
+Graph-related type definitions.
+
+This module contains Pydantic models for graph nodes, edges, and related data structures.
+"""
+
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Union
+
+from flowsint_types import FlowsintType
+from pydantic import BaseModel, Field
+
+Neo4jDict = Dict[str, Any]
+
+
+class NodeMetadata(BaseModel):
+    created_at: datetime = Field(default_factory=datetime.now)
+
+    class Config:
+        extra = "allow"
+
+
+class GraphNode(BaseModel):
+    """
+    Represents object that's being manipulated in the frontend, and throughout the app.
+    It represents a complete node object, close to what is stored in the neo4j db.
+    """
+
+    id: Optional[str]
+    nodeLabel: str
+    nodeType: str
+    nodeSize: Optional[int] = None
+    nodeColor: Optional[str] = None
+    nodeIcon: Optional[str] = None
+    nodeImage: Optional[str] = None
+    nodeFlag: Optional[str] = None
+
+    nodeMetadata: NodeMetadata
+    nodeProperties: FlowsintType
+
+    x: Optional[float] = 100.0
+    y: Optional[float] = 100.0
+
+
+class GraphEdge(BaseModel):
+    id: str
+    source: str
+    target: str
+    label: str
+    date: Optional[str] = None
+    caption: Optional[str] = None
+    type: Optional[str] = None
+    weight: Optional[float] = None
+    confidence_level: Optional[Union[float, str]] = None
+
+
+class GraphData(BaseModel):
+    nodes: List[GraphNode]
+    edges: List[GraphEdge]

flowsint-core/src/flowsint_core/core/graph_serializer.py

@@ -1,177 +0,0 @@
-"""
-Graph property serialization utilities.
-
-This module provides utilities for serializing complex Python objects
-into Neo4j-compatible primitive types, following the Single Responsibility Principle.
-"""
-
-from typing import Any, Dict, List
-
-
-class GraphSerializer:
-    """
-    Handles serialization of complex objects to Neo4j-compatible types.
-
-    This class is responsible for converting Pydantic models, nested objects,
-    and other complex types into primitive types that can be stored in Neo4j.
-    """
-
-    @staticmethod
-    def serialize_properties(properties: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Convert properties to Neo4j-compatible values.
-
-        Handles:
-        - Pydantic models (flattened into individual properties)
-        - Nested objects (flattened with prefixed keys)
-        - Lists (converted to primitive types)
-        - Dictionaries (flattened with prefixed keys)
-        - None values (skipped)
-
-        Args:
-            properties: Dictionary of properties to serialize
-
-        Returns:
-            Dictionary of Neo4j-compatible properties
-        """
-        serialized = {}
-
-        for key, value in properties.items():
-            if key is None:
-                continue
-            if value is None:
-                serialized[key] = ""
-            elif GraphSerializer._is_pydantic_model(value):
-                # Flatten Pydantic models
-                flattened = GraphSerializer._flatten_pydantic(key, value)
-                serialized.update(flattened)
-            elif isinstance(value, dict):
-                # Flatten dictionaries
-                flattened = GraphSerializer._flatten_dict(key, value)
-                serialized.update(flattened)
-            elif isinstance(value, list):
-                # Handle lists
-                serialized[key] = GraphSerializer._serialize_list(value)
-            elif isinstance(value, (str, int, float, bool)):
-                # Keep primitive types as-is
-                serialized[key] = "" if value == "None" else value
-            else:
-                # Convert other complex types to strings
-                serialized[key] = str(value) if value != "None" else ""
-
-        return serialized
-
-    @staticmethod
-    def _is_pydantic_model(obj: Any) -> bool:
-        """Check if an object is a Pydantic model."""
-        return (
-            hasattr(obj, "__dict__")
-            and not isinstance(obj, (str, int, float, bool))
-            and (hasattr(obj, "model_dump") or hasattr(obj, "dict"))
-        )
-
-    @staticmethod
-    def _flatten_pydantic(prefix: str, model: Any) -> Dict[str, Any]:
-        """
-        Flatten a Pydantic model into individual properties.
-
-        Args:
-            prefix: Key prefix for nested properties
-            model: Pydantic model instance
-
-        Returns:
-            Flattened dictionary
-        """
-        flattened = {}
-
-        # Try Pydantic v2 first, then v1
-        if hasattr(model, "model_dump"):
-            data = model.model_dump(mode="json")
-        elif hasattr(model, "dict"):
-            data = model.dict()
-        else:
-            # Fallback to __dict__
-            data = {k: v for k, v in model.__dict__.items() if k is not None}
-
-        for nested_key, nested_value in data.items():
-            if nested_key is None:
-                continue
-            new_key = f"{prefix}_{nested_key}"
-            if nested_value is None:
-                flattened[new_key] = ""
-            elif isinstance(nested_value, (str, int, float, bool)):
-                flattened[new_key] = nested_value
-            else:
-                # Recursively handle nested complex types
-                flattened[new_key] = str(nested_value)
-
-        return flattened
-
-    @staticmethod
-    def _flatten_dict(prefix: str, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Flatten a dictionary into individual properties.
-
-        Args:
-            prefix: Key prefix for nested properties
-            data: Dictionary to flatten
-
-        Returns:
-            Flattened dictionary
-        """
-        flattened = {}
-
-        for dict_key, dict_value in data.items():
-            if dict_key is None:
-                continue
-            new_key = f"{prefix}_{dict_key}"
-            if dict_value is None:
-                flattened[new_key] = ""
-            elif isinstance(dict_value, (str, int, float, bool)):
-                flattened[new_key] = dict_value
-            else:
-                flattened[new_key] = str(dict_value)
-
-        return flattened
-
-    @staticmethod
-    def _serialize_list(items: List[Any]) -> List[Any]:
-        """
-        Serialize a list to Neo4j-compatible types.
-
-        Args:
-            items: List of items to serialize
-
-        Returns:
-            List of serialized items
-        """
-        serialized = []
-
-        for item in items:
-            if isinstance(item, (str, int, float, bool)):
-                serialized.append(item)
-            elif GraphSerializer._is_pydantic_model(item):
-                # Convert complex objects to strings
-                serialized.append(str(item))
-            else:
-                serialized.append(str(item))
-
-        return serialized
-
-    @staticmethod
-    def normalize_property_value(value: Any) -> Any:
-        """
-        Normalize a single property value for Neo4j.
-
-        Args:
-            value: Value to normalize
-
-        Returns:
-            Normalized value
-        """
-        if value is None:
-            return ""
-        elif isinstance(value, (str, int, float, bool)):
-            return "" if value == "None" else value
-        else:
-            return str(value)

flowsint-core/src/flowsint_core/core/graph_service.py

@@ -1,198 +0,0 @@
-"""
-Graph service for high-level graph operations.
-
-This module provides a service layer for graph operations,
-integrating repository and logging functionality.
-"""
-
-from typing import Any, Dict, Optional, Protocol, Union
-from uuid import UUID
-
-from pydantic import BaseModel
-
-from .graph_db import Neo4jConnection
-from .graph_repository import GraphRepository
-
-
-class LoggerProtocol(Protocol):
-    """Protocol for logger implementations."""
-
-    @staticmethod
-    def graph_append(sketch_id: str, message: Dict[str, Any]) -> None:
-        """Log a graph append message."""
-        ...
-
-
-class GraphService:
-    """
-    High-level service for graph operations.
-
-    This service provides a clean interface for enricher operations,
-    handling both graph persistence and logging with proper separation of concerns.
-    """
-
-    def __init__(
-        self,
-        sketch_id: str,
-        neo4j_connection: Optional[Neo4jConnection] = None,
-        logger: Optional[LoggerProtocol] = None,
-        enable_batching: bool = True,
-    ):
-        """
-        Initialize the graph service.
-
-        Args:
-            sketch_id: Investigation sketch ID
-            neo4j_connection: Optional Neo4j connection
-            logger: Optional logger instance
-            enable_batching: Enable batch operations
-        """
-        self._sketch_id = sketch_id
-        self._repository = GraphRepository(neo4j_connection)
-        self._logger = logger
-        self._enable_batching = enable_batching
-
-    @property
-    def sketch_id(self) -> str:
-        """Get the sketch ID."""
-        return self._sketch_id
-
-    @property
-    def repository(self) -> GraphRepository:
-        """Get the underlying repository."""
-        return self._repository
-
-    def create_node(self, node_obj: BaseModel, **properties: Any) -> None:
-        """
-        Create or update a node in the graph.
-
-        Supports one signatures:
-         - Pydantic object: create_node(obj, **overrides)
-
-        Args:
-            from_obj: a Pydantic object
-            **properties: Additional node properties or overrides
-        """
-        if self._enable_batching:
-            self._repository.add_to_batch(
-                "node",
-                node_obj=node_obj,
-                sketch_id=self._sketch_id,
-                **properties,
-            )
-        else:
-            self._repository.create_node(
-                node_obj=node_obj,
-                sketch_id=self._sketch_id,
-                **properties,
-            )
-
-    def create_relationship(
-        self,
-        from_obj: BaseModel,
-        to_obj: BaseModel,
-        rel_label: Optional[str] = None,
-        **properties: Any,
-    ) -> None:
-        """
-        Create a relationship between two nodes.
-
-        Supports 1 signature:
-         - Pydantic objects: create_relationship(obj1, obj2, "REL_TYPE")
-
-        Args:
-            from_obj: Either a Pydantic object (source)
-            to_obj: Either a Pydantic object (target)
-            rel_label: Relationship label (ex: "IS_CONNECTED_TO")
-            **properties: Additional relationship properties
-        """
-        if self._enable_batching:
-            self._repository.add_to_batch(
-                "relationship",
-                from_obj=from_obj,
-                to_obj=to_obj,
-                rel_label=rel_label,
-                sketch_id=self._sketch_id,
-                **properties,
-            )
-        else:
-            self._repository.create_relationship(
-                from_obj=from_obj,
-                to_obj=to_obj,
-                rel_label=rel_label,
-                sketch_id=self._sketch_id,
-                **properties,
-            )
-
-    def log_graph_message(self, message: str) -> None:
-        """
-        Log a graph operation message.
-
-        Args:
-            message: Message to log
-        """
-        if self._logger:
-            self._logger.graph_append(self._sketch_id, {"message": message})
-
-    def flush(self) -> None:
-        """Flush any pending batch operations."""
-        if self._enable_batching:
-            self._repository.flush_batch()
-
-    def query(self, cypher: str, parameters: Dict[str, Any] = None) -> list:
-        """
-        Execute a custom Cypher query.
-
-        Args:
-            cypher: Cypher query string
-            parameters: Query parameters
-
-        Returns:
-            List of result records
-        """
-        return self._repository.query(cypher, parameters)
-
-    def set_batch_size(self, size: int) -> None:
-        """
-        Set the batch size for operations.
-
-        Args:
-            size: Number of operations to batch
-        """
-        self._repository.set_batch_size(size)
-
-    def __enter__(self):
-        """Context manager entry."""
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        """Context manager exit - auto-flush batch."""
-        if exc_type is None:
-            self.flush()
-
-
-def create_graph_service(
-    sketch_id: str,
-    neo4j_connection: Optional[Neo4jConnection] = None,
-    enable_batching: bool = True,
-) -> GraphService:
-    """
-    Factory function to create a GraphService instance.
-
-    Args:
-        sketch_id: Investigation sketch ID
-        neo4j_connection: Optional Neo4j connection
-        enable_batching: Enable batch operations
-
-    Returns:
-        Configured GraphService instance
-    """
-    # Import Logger here to avoid circular imports
-    from .logger import Logger
-
-    return GraphService(
-        sketch_id=sketch_id,
-        neo4j_connection=neo4j_connection,
-        logger=Logger,
-        enable_batching=enable_batching,
-    )

flowsint-core/src/flowsint_core/core/logger.py

@@ -8,17 +8,19 @@ Architecture:
 - SOLID principles: Dependency injection via protocols
 - Ordering: Monotonic sequence number + application timestamp
 """
-from typing import Dict, Union, Optional, Tuple
-from uuid import UUID
+
+import atexit
 import threading
 import time
-import atexit
-from queue import Queue
 from datetime import datetime, timezone
+from queue import Queue
+from typing import Dict, Optional, Union
+from uuid import UUID
+
+from ..tasks.event import emit_event_task
 from .enums import EventLevel
 from .models import Log
 from .postgre_db import get_db
-from ..tasks.event import emit_event_task
 
 
 class LoggerSingleton:
@@ -32,14 +34,11 @@ class LoggerSingleton:
     - Automatic flush on shutdown
     """
 
-    _instance: Optional['LoggerSingleton'] = None
+    _instance: Optional["LoggerSingleton"] = None
     _lock = threading.Lock()
 
     def __new__(
-        cls,
-        batch_size: int = 50,
-        flush_interval: float = 2.0,
-        auto_start: bool = True
+        cls, batch_size: int = 50, flush_interval: float = 2.0, auto_start: bool = True
     ):
         """Thread-safe singleton implementation using double-checked locking."""
         if cls._instance is None:
@@ -51,10 +50,7 @@ class LoggerSingleton:
         return cls._instance
 
     def __init__(
-        self,
-        batch_size: int = 50,
-        flush_interval: float = 2.0,
-        auto_start: bool = True
+        self, batch_size: int = 50, flush_interval: float = 2.0, auto_start: bool = True
     ):
         """
         Initialize the Logger singleton.
@@ -100,9 +96,7 @@ class LoggerSingleton:
         if self._worker_thread is None or not self._worker_thread.is_alive():
             self._shutdown_event.clear()
             self._worker_thread = threading.Thread(
-                target=self._batch_worker,
-                daemon=True,
-                name="LoggerBatchWorker"
+                target=self._batch_worker, daemon=True, name="LoggerBatchWorker"
             )
             self._worker_thread.start()
 
@@ -133,7 +127,9 @@ class LoggerSingleton:
         logs_to_insert = []
 
         # Collect logs from queue
-        while not self._log_queue.empty() and (force or len(logs_to_insert) < self._batch_size):
+        while not self._log_queue.empty() and (
+            force or len(logs_to_insert) < self._batch_size
+        ):
             try:
                 logs_to_insert.append(self._log_queue.get_nowait())
             except:
@@ -151,6 +147,7 @@ class LoggerSingleton:
         except (StopIteration, RuntimeError) as e:
             # Generator exhausted or error getting DB - log to standard logging
             import logging
+
             logging.error(f"Failed to get database session: {e}")
             return
 
@@ -161,7 +158,7 @@ class LoggerSingleton:
                     sketch_id=str(sketch_id),
                     type=level.value,
                     content=content,
-                    created_at=timestamp  # Use application-side timestamp
+                    created_at=timestamp,  # Use application-side timestamp
                 )
                 log_objects.append(log)
 
@@ -176,16 +173,13 @@ class LoggerSingleton:
             db.rollback()
             # Log to standard logging as fallback
             import logging
+
             logging.error(f"Failed to batch insert logs: {e}")
         finally:
             db.close()
 
     def _emit_event(
-        self,
-        log_id: str,
-        sketch_id: str,
-        level: EventLevel,
-        content: Dict
+        self, log_id: str, sketch_id: str, level: EventLevel, content: Dict
     ) -> None:
         """
         Emit event immediately for real-time display.
@@ -197,19 +191,15 @@ class LoggerSingleton:
             content: Log content
         """
         try:
-            emit_event_task.apply(
-                args=[log_id, str(sketch_id), level, content]
-            )
+            emit_event_task.apply(args=[log_id, str(sketch_id), level, content])
         except Exception as e:
             # Don't let event emission errors break logging
             import logging
+
             logging.error(f"Failed to emit event: {e}")
 
     def _log(
-        self,
-        sketch_id: Union[str, UUID],
-        level: EventLevel,
-        content: Dict
+        self, sketch_id: Union[str, UUID], level: EventLevel, content: Dict
     ) -> None:
         """
         Internal logging method.
@@ -230,6 +220,7 @@ class LoggerSingleton:
 
         # Generate a temporary ID for immediate event emission
         import uuid
+
         temp_log_id = str(uuid.uuid4())
 
         # 1. IMMEDIATE: Emit event for real-time display
@@ -272,13 +263,16 @@ class LoggerSingleton:
         # Also publish to status channel for graph refresh
         try:
             import uuid
+
             temp_log_id = str(uuid.uuid4())
             from ..tasks.event import emit_status_event_task
+
             emit_status_event_task.apply(
                 args=[temp_log_id, str(sketch_id), EventLevel.COMPLETED, message]
             )
         except Exception as e:
             import logging
+
             logging.error(f"Failed to emit status event: {e}")
 
     def pending(self, sketch_id: Union[str, UUID], message: Dict) -> None:

flowsint-core/src/flowsint_core/core/orchestrator.py

@@ -22,10 +22,9 @@ class FlowOrchestrator(Enricher):
         sketch_id: str,
         scan_id: str,
         enricher_branches: List[FlowBranch],
-        neo4j_conn=None,
         vault=None,
     ):
-        super().__init__(sketch_id, scan_id, neo4j_conn=neo4j_conn, vault=vault)
+        super().__init__(sketch_id, scan_id, vault=vault)
         self.enricher_branches = enricher_branches
         self.enrichers = {}  # Map of nodeId -> enricher instance
         self.execution_log_file = None  # Path to the execution log file
@@ -243,7 +242,6 @@ class FlowOrchestrator(Enricher):
                 enricher_name,
                 self.sketch_id,
                 self.scan_id,
-                neo4j_conn=self.neo4j_conn,
                 vault=self.vault,
                 params=enricher_params,
             )

flowsint-core/src/flowsint_core/core/types.py

@@ -1,8 +1,16 @@
+"""
+Core type definitions.
+
+This module contains Pydantic models for core types like events, flows, and roles.
+Graph-related types are in flowsint_core.core.graph.types.
+"""
+
 import enum
-from pydantic import BaseModel, Field
-from typing import Dict, Any
 from datetime import datetime
-from typing import Dict, List, Any, Optional, Literal
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
 from .enums import EventLevel
 
 
@@ -16,7 +24,7 @@ class Event(BaseModel):
     )
 
 
-class Node(BaseModel):
+class FlowNode(BaseModel):
     """Represents a node in a transformation flow with position and data."""
 
     id: str = Field(..., description="Unique identifier for the node", title="Node ID")
@@ -33,7 +41,7 @@ class Node(BaseModel):
     )
 
 
-class Edge(BaseModel):
+class FlowEdge(BaseModel):
     """Represents an edge connecting two nodes in a transformation flow."""
 
     id: str = Field(..., description="Unique identifier for the edge", title="Edge ID")

flowsint-core/src/flowsint_core/imports/__init__.py

@@ -4,6 +4,12 @@ Import utilities for entity parsing and type detection.
 
 from .entity_detection import detect_type
 from .file_parser import FileParseResult, parse_import_file
+from .import_service import (
+    EntityMapping,
+    ImportResult,
+    ImportService,
+    create_import_service,
+)
 from .types import EntityPreview
 
 __all__ = [
@@ -11,4 +17,9 @@ __all__ = [
     "parse_import_file",
     "FileParseResult",
     "EntityPreview",
+    # Import service
+    "EntityMapping",
+    "ImportResult",
+    "ImportService",
+    "create_import_service",
 ]

flowsint-core/src/flowsint_core/imports/import_service.py

@@ -0,0 +1,248 @@
+"""
+Import service for handling file imports into sketches.
+
+This module provides a service layer for import operations,
+handling file parsing, entity conversion, and batch creation.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from flowsint_types import FlowsintType
+
+from flowsint_core.core.graph import GraphSerializer, GraphService
+
+
+@dataclass
+class EntityMapping:
+    """Mapping configuration for an entity."""
+
+    id: str
+    entity_type: str
+    nodeLabel: str
+    data: Dict[str, Any]
+    include: bool = True
+    node_id: Optional[str] = None
+
+
+@dataclass
+class ImportResult:
+    """Result of an import execution."""
+
+    status: str
+    nodes_created: int
+    nodes_skipped: int
+    errors: List[str]
+
+
+class ImportService:
+    """
+    Service for handling file imports into sketches.
+
+    This service handles:
+    - File analysis and parsing
+    - Entity conversion to FlowsintTypes
+    - Batch node and edge creation
+    """
+
+    def __init__(self, graph_service: GraphService):
+        """
+        Initialize the import service.
+
+        Args:
+            graph_service: GraphService instance for database operations
+        """
+        self._graph_service = graph_service
+
+    @staticmethod
+    def analyze_file(
+        file_content: bytes,
+        filename: str,
+        max_preview_rows: int = 10000000,
+    ):
+        """
+        Analyze an uploaded file for import.
+
+        Args:
+            file_content: Raw file content as bytes
+            filename: Name of the file (used for extension detection)
+            max_preview_rows: Maximum number of rows to preview
+
+        Returns:
+            FileParseResult with detected entities and edges
+
+        Raises:
+            ValueError: If file format is unsupported or parsing fails
+        """
+        from flowsint_core.imports import parse_import_file
+
+        return parse_import_file(
+            file_content=file_content,
+            filename=filename,
+            max_preview_rows=max_preview_rows,
+        )
+
+    def execute_import(
+        self,
+        entity_mappings: List[EntityMapping],
+        edges: Optional[List[Dict[str, Any]]] = None,
+    ) -> ImportResult:
+        """
+        Execute the import of entities into the sketch.
+
+        Args:
+            entity_mappings: List of EntityMapping objects to import
+            edges: Optional list of edge definitions with from_id, to_id, label
+
+        Returns:
+            ImportResult with status, counts, and any errors
+        """
+        # Filter only entities marked for inclusion
+        entities_to_import = [m for m in entity_mappings if m.include]
+
+        # Convert entity mappings to FlowsintType objects
+        conversion_result = self._convert_entities(entities_to_import)
+        pydantic_nodes = conversion_result["nodes"]
+        nodes_mapping_indices = conversion_result["mapping_indices"]
+        conversion_errors = conversion_result["errors"]
+
+        if not pydantic_nodes:
+            return ImportResult(
+                status="completed_with_errors" if conversion_errors else "completed",
+                nodes_created=0,
+                nodes_skipped=len(entities_to_import),
+                errors=conversion_errors[:50],
+            )
+
+        # Batch create nodes
+        try:
+            nodes = GraphSerializer.serialize_flowsint_types(pydantic_nodes)
+            nodes_result = self._graph_service.batch_create_nodes(nodes=nodes)
+            nodes_created = nodes_result["nodes_created"]
+            node_element_ids = nodes_result.get("node_ids", [])
+            batch_errors = nodes_result.get("errors", [])
+        except Exception as e:
+            return ImportResult(
+                status="failed",
+                nodes_created=0,
+                nodes_skipped=len(entities_to_import),
+                errors=[f"Batch node creation failed: {str(e)}"],
+            )
+
+        # Create edges if provided
+        edge_errors = []
+        if edges and nodes_mapping_indices and node_element_ids:
+            edge_errors = self._create_edges(
+                edges=edges,
+                nodes_mapping_indices=nodes_mapping_indices,
+                node_element_ids=node_element_ids,
+            )
+
+        all_errors = conversion_errors + batch_errors + edge_errors
+        nodes_skipped = len(entities_to_import) - nodes_created
+
+        return ImportResult(
+            status="completed" if not all_errors else "completed_with_errors",
+            nodes_created=nodes_created,
+            nodes_skipped=nodes_skipped,
+            errors=all_errors[:50],
+        )
+
+    def _convert_entities(self, entities: List[EntityMapping]) -> Dict[str, Any]:
+        """
+        Convert entity mappings to FlowsintType objects.
+
+        Returns:
+            Dict with nodes, mapping_indices, and errors
+        """
+        pydantic_nodes: List[FlowsintType] = []
+        nodes_mapping_indices: Dict[str, int] = {}
+        errors: List[str] = []
+
+        for idx, mapping in enumerate(entities):
+            entity_data = mapping.data.copy()
+
+            try:
+                pydantic_obj = GraphSerializer.parse_flowsint_type(
+                    entity=entity_data,
+                    nodeType=mapping.entity_type,
+                )
+                pydantic_nodes.append(pydantic_obj)
+
+                if mapping.node_id:
+                    nodes_mapping_indices[mapping.node_id] = len(pydantic_nodes) - 1
+
+            except ValueError as e:
+                errors.append(f"Entity {idx + 1} ({mapping.nodeLabel}): {str(e)}")
+            except Exception as e:
+                errors.append(f"Entity {idx + 1} ({mapping.nodeLabel}): {str(e)}")
+
+        return {
+            "nodes": pydantic_nodes,
+            "mapping_indices": nodes_mapping_indices,
+            "errors": errors,
+        }
+
+    def _create_edges(
+        self,
+        edges: List[Dict[str, Any]],
+        nodes_mapping_indices: Dict[str, int],
+        node_element_ids: List[str],
+    ) -> List[str]:
+        """
+        Create edges between imported nodes.
+
+        Returns:
+            List of error messages
+        """
+        errors: List[str] = []
+        edges_to_insert: List[Dict[str, Any]] = []
+
+        for idx, edge in enumerate(edges):
+            from_id = edge.get("from_id")
+            to_id = edge.get("to_id")
+
+            from_idx = nodes_mapping_indices.get(from_id)
+            to_idx = nodes_mapping_indices.get(to_id)
+
+            if from_idx is None or to_idx is None:
+                errors.append(
+                    f"Edge {idx}: Missing source or target node (from: {from_id}, to: {to_id})"
+                )
+                continue
+
+            if from_idx >= len(node_element_ids) or to_idx >= len(node_element_ids):
+                errors.append(f"Edge {idx}: Node index out of range")
+                continue
+
+            edges_to_insert.append(
+                {
+                    "from_element_id": node_element_ids[from_idx],
+                    "to_element_id": node_element_ids[to_idx],
+                    "rel_label": edge.get("label", "RELATED_TO"),
+                }
+            )
+
+        if edges_to_insert:
+            try:
+                edges_result = self._graph_service.batch_create_edges_by_element_id(
+                    edges=edges_to_insert
+                )
+                errors.extend(edges_result.get("errors", []))
+            except Exception as e:
+                errors.append(f"Batch edge creation failed: {str(e)}")
+
+        return errors
+
+
+def create_import_service(graph_service: GraphService) -> ImportService:
+    """
+    Factory function to create an ImportService instance.
+
+    Args:
+        graph_service: GraphService instance for database operations
+
+    Returns:
+        Configured ImportService instance
+    """
+    return ImportService(graph_service=graph_service)

flowsint-core/src/flowsint_core/tasks/enricher.py

@@ -1,13 +1,10 @@
-import os
 import uuid
 import asyncio
-from dotenv import load_dotenv
 from typing import List, Optional
 from celery import states
 from flowsint_enrichers import ENRICHER_REGISTRY, load_all_enrichers
 from ..core.celery import celery
 from ..core.postgre_db import SessionLocal, get_db
-from ..core.graph_db import Neo4jConnection
 from ..core.vault import Vault
 from ..core.models import Scan
 from sqlalchemy.orm import Session
@@ -15,16 +12,9 @@ from ..core.logger import Logger
 from ..core.enums import EventLevel
 from flowsint_core.utils import to_json_serializable
 
-load_dotenv()
-
 # Auto-discover and register all enrichers
 load_all_enrichers()
 
-URI = os.getenv("NEO4J_URI_BOLT")
-USERNAME = os.getenv("NEO4J_USERNAME")
-PASSWORD = os.getenv("NEO4J_PASSWORD")
-
-neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD)
 db: Session = next(get_db())
 
 
@@ -67,7 +57,6 @@ def run_enricher(
             name=enricher_name,
             sketch_id=sketch_id,
             scan_id=scan_id,
-            neo4j_conn=neo4j_connection,
             vault=vault,
         )
 

flowsint-core/src/flowsint_core/tasks/flow.py

@@ -1,12 +1,9 @@
-import os
 import uuid
-from dotenv import load_dotenv
 from typing import List, Optional
 from celery import states
 from ..core.celery import celery
 from ..core.orchestrator import FlowOrchestrator
 from ..core.postgre_db import SessionLocal, get_db
-from ..core.graph_db import Neo4jConnection
 from ..core.vault import Vault
 from ..core.types import FlowBranch
 from ..core.models import Scan
@@ -15,13 +12,6 @@ from ..core.logger import Logger
 from ..core.enums import EventLevel
 from flowsint_core.utils import to_json_serializable
 
-load_dotenv()
-
-URI = os.getenv("NEO4J_URI_BOLT")
-USERNAME = os.getenv("NEO4J_USERNAME")
-PASSWORD = os.getenv("NEO4J_PASSWORD")
-
-neo4j_connection = Neo4jConnection(URI, USERNAME, PASSWORD)
 db: Session = next(get_db())
 
 
@@ -64,7 +54,6 @@ def run_flow(
             sketch_id=sketch_id,
             scan_id=str(scan_id),
             enricher_branches=enricher_branches,
-            neo4j_conn=neo4j_connection,
             vault=vault,
         )
 

flowsint-core/src/flowsint_core/utils.py

@@ -1,18 +1,17 @@
-from urllib.parse import urlparse
-import phonenumbers
-import ipaddress
-from phonenumbers import NumberParseException
-from pydantic import TypeAdapter, BaseModel
-from urllib.parse import urlparse
-import re
-import ssl
-import socket
-from typing import Dict, Any, List, Type
-from pydantic import BaseModel
 import inspect
-from typing import Any, Dict, Type
+import ipaddress
+import re
+import socket
+import ssl
+from typing import Any, Dict, List, Type
+from urllib.parse import urlparse
+
+import phonenumbers
+from phonenumbers import NumberParseException
 from pydantic import BaseModel, TypeAdapter
 
+from .core.graph.types import GraphEdge, GraphNode
+
 
 def is_valid_ip(address: str) -> bool:
     try:
@@ -35,7 +34,6 @@ def is_valid_email(email: str) -> bool:
 
 
 def is_valid_domain(url_or_domain: str) -> str:
-
     try:
         parsed = urlparse(
             url_or_domain if "://" in url_or_domain else "http://" + url_or_domain
@@ -248,38 +246,73 @@ def get_label_color(label: str) -> str:
     return color_map.get(label, color_map["default"])
 
 
-def flatten(data_dict, prefix=""):
-    """
-    Flattens a dictionary to contain only Neo4j-compatible property values.
-    Neo4j supports primitive types (string, number, boolean) and arrays of those types.
-    Args:
-        data_dict (dict): Dictionary to flatten
-    Returns:
-        dict: Flattened dictionary with only Neo4j-compatible values
-    """
-    flattened = {}
-    if not isinstance(data_dict, dict):
-        return flattened
-    for key, value in data_dict.items():
-        if value is None:
-            continue
-        if isinstance(value, (str, int, float, bool)) or (
-            isinstance(value, list)
-            and all(isinstance(item, (str, int, float, bool)) for item in value)
-        ):
-            key = f"{prefix}{key}"
-            flattened[key] = value
+Primitive = (str, int, float, bool)
+
+
+def flatten(
+    data: Any, prefix: str = "", *, remove_empty: bool = False, separator: str = "."
+) -> Dict[str, Any]:
+    flattened: Dict[str, Any] = {}
+
+    if isinstance(data, dict):
+        for key, value in data.items():
+            new_key = f"{prefix}{key}" if prefix == "" else f"{prefix}{separator}{key}"
+
+            # Handle None values
+            if value is None:
+                if not remove_empty:
+                    flattened[new_key] = None
+                continue
+
+            # Ignore empty strings if option enabled
+            if remove_empty and value == "":
+                continue
+
+            if isinstance(value, Primitive):
+                flattened[new_key] = value
+
+            elif isinstance(value, list):
+                if all(isinstance(item, Primitive) for item in value):
+                    if remove_empty:
+                        value = [item for item in value if item != ""]
+                        if not value:
+                            continue
+                    flattened[new_key] = value
+                # else ignore (Neo4j incompatible)
+
+            elif isinstance(value, dict):
+                flattened.update(flatten(value, new_key, remove_empty=remove_empty, separator=separator))
+
     return flattened
 
 
-def get_inline_relationships(nodes: List[Any], edges: List[Any]) -> List[str]:
+def unflatten(data: Dict[str, Any], *, separator: str = ".") -> Dict[str, Any]:
+    result: Dict[str, Any] = {}
+
+    for flat_key, value in data.items():
+        parts = flat_key.split(separator)
+        current = result
+
+        for part in parts[:-1]:
+            if part not in current or not isinstance(current[part], dict):
+                current[part] = {}
+            current = current[part]
+
+        current[parts[-1]] = value
+
+    return result
+
+
+def get_inline_relationships(
+    nodes: List[GraphNode], edges: List[GraphEdge]
+) -> List[str]:
     """
     Get the inline relationships for a list of nodes and edges.
     """
     relationships = []
     for edge in edges:
-        source = next((node for node in nodes if node["id"] == edge["source"]), None)
-        target = next((node for node in nodes if node["id"] == edge["target"]), None)
+        source = next((node for node in nodes if node.id == edge.source), None)
+        target = next((node for node in nodes if node.id == edge.target), None)
         if source and target:
             relationships.append({"source": source, "edge": edge, "target": target})
     return relationships
@@ -288,6 +321,7 @@ def get_inline_relationships(nodes: List[Any], edges: List[Any]) -> List[str]:
 def to_json_serializable(obj):
     """Convert any object to a JSON-serializable format."""
     import json
+
     from pydantic import BaseModel
 
     try:

flowsint-core/tests/core/__init__.py

@@ -0,0 +1 @@
+# Tests for core module

flowsint-core/tests/core/graph/__init__.py

@@ -0,0 +1 @@
+# Tests for graph module

flowsint-core/tests/core/graph/in_memory_graph_repository.py

@@ -0,0 +1,477 @@
+"""
+In-memory implementation of GraphRepositoryProtocol for testing.
+
+This module provides a lightweight graph repository that stores data in memory,
+enabling fast unit tests without requiring a Neo4j database connection.
+"""
+
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+
+class InMemoryGraphRepository:
+    """
+    In-memory implementation of GraphRepositoryProtocol for testing.
+
+    Stores nodes and edges in dictionaries, no Neo4j required.
+    All operations are synchronous and data is lost when the instance is destroyed.
+    """
+
+    def __init__(self):
+        self._nodes: Dict[str, Dict[str, Any]] = {}  # element_id -> node_data
+        self._edges: Dict[str, Dict[str, Any]] = {}  # element_id -> edge_data
+        self._batch_operations: List[tuple] = []
+        self._batch_size = 100
+
+    def _generate_element_id(self, prefix: str = "mem") -> str:
+        """Generate a unique element ID."""
+        return f"{prefix}:{uuid4()}"
+
+    # -------------------------------------------------------------------------
+    # Core node operations
+    # -------------------------------------------------------------------------
+
+    def create_node(self, node_obj: Dict[str, Any], sketch_id: str) -> Optional[str]:
+        """Create or update a single node. Returns element ID."""
+        node_label = node_obj.get("nodeLabel")
+        node_type = node_obj.get("nodeType")
+
+        # Check if node already exists (MERGE behavior)
+        for element_id, data in self._nodes.items():
+            if (
+                data.get("nodeLabel") == node_label
+                and data.get("sketch_id") == sketch_id
+            ):
+                # Update existing node
+                self._nodes[element_id].update(node_obj)
+                return element_id
+
+        # Create new node
+        element_id = self._generate_element_id("node")
+        self._nodes[element_id] = {
+            **node_obj,
+            "sketch_id": sketch_id,
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "_labels": [node_type] if node_type else ["Node"],
+        }
+        return element_id
+
+    def update_node(
+        self, element_id: str, updates: Dict[str, Any], sketch_id: str
+    ) -> Optional[str]:
+        """Update a node by its element ID. Returns element ID."""
+        if element_id not in self._nodes:
+            return None
+        if self._nodes[element_id].get("sketch_id") != sketch_id:
+            return None
+        self._nodes[element_id].update(updates)
+        return element_id
+
+    def delete_nodes(self, node_ids: List[str], sketch_id: str) -> int:
+        """Delete nodes by their element IDs. Returns count deleted."""
+        deleted = 0
+        for node_id in node_ids:
+            if node_id in self._nodes:
+                if self._nodes[node_id].get("sketch_id") == sketch_id:
+                    # Also delete related edges
+                    edges_to_delete = [
+                        eid
+                        for eid, edge in self._edges.items()
+                        if edge.get("source") == node_id
+                        or edge.get("target") == node_id
+                    ]
+                    for eid in edges_to_delete:
+                        del self._edges[eid]
+                    del self._nodes[node_id]
+                    deleted += 1
+        return deleted
+
+    def delete_all_sketch_nodes(self, sketch_id: str) -> int:
+        """Delete all nodes for a sketch. Returns count deleted."""
+        to_delete = [
+            eid
+            for eid, data in self._nodes.items()
+            if data.get("sketch_id") == sketch_id
+        ]
+        for eid in to_delete:
+            del self._nodes[eid]
+
+        # Also delete related edges
+        edges_to_delete = [
+            eid
+            for eid, data in self._edges.items()
+            if data.get("sketch_id") == sketch_id
+        ]
+        for eid in edges_to_delete:
+            del self._edges[eid]
+
+        return len(to_delete)
+
+    def get_nodes_by_ids(
+        self, node_ids: List[str], sketch_id: str
+    ) -> List[Dict[str, Any]]:
+        """Get nodes by their element IDs."""
+        result = []
+        for node_id in node_ids:
+            if node_id in self._nodes:
+                node = self._nodes[node_id]
+                if node.get("sketch_id") == sketch_id:
+                    result.append({"data": node})
+        return result
+
+    def update_nodes_positions(
+        self, positions: List[Dict[str, Any]], sketch_id: str
+    ) -> int:
+        """Update positions for multiple nodes. Returns count updated."""
+        updated = 0
+        for pos in positions:
+            node_id = pos.get("nodeId")
+            if node_id in self._nodes:
+                if self._nodes[node_id].get("sketch_id") == sketch_id:
+                    self._nodes[node_id]["x"] = pos.get("x")
+                    self._nodes[node_id]["y"] = pos.get("y")
+                    updated += 1
+        return updated
+
+    # -------------------------------------------------------------------------
+    # Core relationship operations
+    # -------------------------------------------------------------------------
+
+    def create_relationship(self, rel_obj: Dict[str, Any], sketch_id: str) -> None:
+        """Create a relationship between two nodes."""
+        from_type = rel_obj.get("from_type")
+        from_label = rel_obj.get("from_label")
+        to_type = rel_obj.get("to_type")
+        to_label = rel_obj.get("to_label")
+        rel_label = rel_obj.get("rel_label", "RELATED_TO")
+
+        # Find source and target nodes
+        source_id = None
+        target_id = None
+
+        for eid, node in self._nodes.items():
+            if node.get("sketch_id") != sketch_id:
+                continue
+            if node.get("nodeLabel") == from_label:
+                source_id = eid
+            if node.get("nodeLabel") == to_label:
+                target_id = eid
+
+        if source_id and target_id:
+            element_id = self._generate_element_id("rel")
+            self._edges[element_id] = {
+                **rel_obj,
+                "source": source_id,
+                "target": target_id,
+                "type": rel_label,
+                "sketch_id": sketch_id,
+            }
+
+    def create_relationship_by_element_id(
+        self,
+        from_element_id: str,
+        to_element_id: str,
+        rel_label: str,
+        sketch_id: str,
+    ) -> Optional[Dict[str, Any]]:
+        """Create a relationship using element IDs."""
+        if from_element_id not in self._nodes or to_element_id not in self._nodes:
+            return None
+
+        element_id = self._generate_element_id("rel")
+        edge_data = {
+            "source": from_element_id,
+            "target": to_element_id,
+            "type": rel_label,
+            "sketch_id": sketch_id,
+        }
+        self._edges[element_id] = edge_data
+        return {"sketch_id": sketch_id}
+
+    def update_relationship(
+        self, element_id: str, rel_obj: Dict[str, Any], sketch_id: str
+    ) -> Optional[Dict[str, Any]]:
+        """Update a relationship by its element ID."""
+        if element_id not in self._edges:
+            return None
+        if self._edges[element_id].get("sketch_id") != sketch_id:
+            return None
+        self._edges[element_id].update(rel_obj)
+        return {
+            "id": element_id,
+            "type": self._edges[element_id].get("type"),
+            "data": self._edges[element_id],
+        }
+
+    def delete_relationships(self, relationship_ids: List[str], sketch_id: str) -> int:
+        """Delete relationships by their element IDs. Returns count deleted."""
+        deleted = 0
+        for rel_id in relationship_ids:
+            if rel_id in self._edges:
+                if self._edges[rel_id].get("sketch_id") == sketch_id:
+                    del self._edges[rel_id]
+                    deleted += 1
+        return deleted
+
+    # -------------------------------------------------------------------------
+    # Graph queries
+    # -------------------------------------------------------------------------
+
+    def get_sketch_graph(
+        self, sketch_id: str, limit: int = 100000
+    ) -> Dict[str, List[Dict[str, Any]]]:
+        """Get all nodes and edges for a sketch."""
+        nodes = []
+        node_ids = set()
+
+        for eid, data in self._nodes.items():
+            if data.get("sketch_id") == sketch_id:
+                nodes.append(
+                    {
+                        "id": eid,
+                        "labels": data.get("_labels", ["Node"]),
+                        "data": data,
+                    }
+                )
+                node_ids.add(eid)
+                if len(nodes) >= limit:
+                    break
+
+        edges = []
+        for eid, data in self._edges.items():
+            if data.get("sketch_id") == sketch_id:
+                if data.get("source") in node_ids and data.get("target") in node_ids:
+                    edges.append(
+                        {
+                            "id": eid,
+                            "type": data.get("type", "RELATED_TO"),
+                            "source": data.get("source"),
+                            "target": data.get("target"),
+                            "data": data,
+                        }
+                    )
+
+        return {"nodes": nodes, "edges": edges}
+
+    def get_neighbors(self, node_id: str, sketch_id: str) -> Dict[str, Any]:
+        """Get a node and all its direct relationships."""
+        if node_id not in self._nodes:
+            return {"nodes": [], "edges": []}
+
+        center = self._nodes[node_id]
+        if center.get("sketch_id") != sketch_id:
+            return {"nodes": [], "edges": []}
+
+        nodes = {node_id: {"id": node_id, "data": center}}
+        edges = {}
+
+        for eid, edge in self._edges.items():
+            if edge.get("sketch_id") != sketch_id:
+                continue
+
+            source = edge.get("source")
+            target = edge.get("target")
+
+            if source == node_id:
+                # Outgoing edge
+                if target in self._nodes:
+                    nodes[target] = {"id": target, "data": self._nodes[target]}
+                    edges[eid] = {
+                        "id": eid,
+                        "source": source,
+                        "target": target,
+                        "label": edge.get("type"),
+                    }
+            elif target == node_id:
+                # Incoming edge
+                if source in self._nodes:
+                    nodes[source] = {"id": source, "data": self._nodes[source]}
+                    edges[eid] = {
+                        "id": eid,
+                        "source": source,
+                        "target": target,
+                        "label": edge.get("type"),
+                    }
+
+        return {"nodes": list(nodes.values()), "edges": list(edges.values())}
+
+    # -------------------------------------------------------------------------
+    # Merge operations
+    # -------------------------------------------------------------------------
+
+    def merge_nodes(
+        self,
+        old_node_ids: List[str],
+        new_node_data: Dict[str, Any],
+        new_node_id: Optional[str],
+        sketch_id: str,
+    ) -> Optional[str]:
+        """Merge multiple nodes into one. Returns new element ID."""
+        if not old_node_ids:
+            return None
+
+        # Determine target node ID
+        if new_node_id and new_node_id in old_node_ids:
+            target_id = new_node_id
+            self._nodes[target_id].update(new_node_data)
+        else:
+            target_id = self._generate_element_id("node")
+            self._nodes[target_id] = {
+                **new_node_data,
+                "sketch_id": sketch_id,
+                "created_at": datetime.now(timezone.utc).isoformat(),
+            }
+
+        # Transfer relationships
+        for eid, edge in list(self._edges.items()):
+            if edge.get("source") in old_node_ids and edge.get("source") != target_id:
+                edge["source"] = target_id
+            if edge.get("target") in old_node_ids and edge.get("target") != target_id:
+                edge["target"] = target_id
+
+        # Delete old nodes (except target)
+        for node_id in old_node_ids:
+            if node_id != target_id and node_id in self._nodes:
+                del self._nodes[node_id]
+
+        return target_id
+
+    # -------------------------------------------------------------------------
+    # Batch operations
+    # -------------------------------------------------------------------------
+
+    def batch_create_nodes(
+        self, nodes: List[Dict[str, Any]], sketch_id: str
+    ) -> Dict[str, Any]:
+        """Create multiple nodes in a single batch."""
+        node_ids = []
+        errors = []
+
+        for idx, node_obj in enumerate(nodes):
+            try:
+                element_id = self.create_node(node_obj, sketch_id)
+                if element_id:
+                    node_ids.append(element_id)
+            except Exception as e:
+                errors.append(f"Node {idx}: {str(e)}")
+
+        return {
+            "nodes_created": len(node_ids),
+            "node_ids": node_ids,
+            "errors": errors,
+        }
+
+    def batch_create_edges_by_element_id(
+        self, edges: List[Dict[str, Any]], sketch_id: str
+    ) -> Dict[str, Any]:
+        """Create multiple edges using element IDs in a single batch."""
+        created = 0
+        errors = []
+
+        for idx, edge in enumerate(edges):
+            try:
+                from_id = edge.get("from_element_id")
+                to_id = edge.get("to_element_id")
+                rel_label = edge.get("rel_label", "RELATED_TO")
+
+                if not from_id or not to_id:
+                    errors.append(
+                        f"Edge {idx}: Missing required fields (from_element_id or to_element_id)"
+                    )
+                    continue
+
+                result = self.create_relationship_by_element_id(
+                    from_id, to_id, rel_label, sketch_id
+                )
+                if result:
+                    created += 1
+            except Exception as e:
+                errors.append(f"Edge {idx}: {str(e)}")
+
+        return {"edges_created": created, "errors": errors}
+
+    def add_to_batch(self, operation_type: str, **kwargs: Any) -> None:
+        """Add an operation to the batch queue."""
+        if operation_type not in ("node", "relationship"):
+            raise ValueError(f"Unknown operation type: {operation_type}")
+
+        self._batch_operations.append((operation_type, kwargs))
+
+        if len(self._batch_operations) >= self._batch_size:
+            self.flush_batch()
+
+    def flush_batch(self) -> None:
+        """Execute all batched operations."""
+        for op_type, kwargs in self._batch_operations:
+            if op_type == "node":
+                self.create_node(kwargs["node_obj"], kwargs["sketch_id"])
+            elif op_type == "relationship":
+                self.create_relationship(kwargs["rel_obj"], kwargs["sketch_id"])
+        self._batch_operations.clear()
+
+    def clear_batch(self) -> None:
+        """Clear the batch without executing."""
+        self._batch_operations.clear()
+
+    def set_batch_size(self, size: int) -> None:
+        """Set the batch size for auto-flushing."""
+        if size < 1:
+            raise ValueError("Batch size must be at least 1")
+        self._batch_size = size
+
+    # -------------------------------------------------------------------------
+    # Custom queries
+    # -------------------------------------------------------------------------
+
+    def query(
+        self, cypher: str, parameters: Dict[str, Any] = {}
+    ) -> List[Dict[str, Any]]:
+        """
+        Execute a custom Cypher query.
+
+        Note: In-memory implementation doesn't support Cypher.
+        Override this method or mock it for specific test cases.
+        """
+        return []
+
+    # -------------------------------------------------------------------------
+    # Context manager
+    # -------------------------------------------------------------------------
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - auto-flush batch on success."""
+        if exc_type is None:
+            self.flush_batch()
+        else:
+            self.clear_batch()
+
+    # -------------------------------------------------------------------------
+    # Test helpers
+    # -------------------------------------------------------------------------
+
+    def get_node_count(self, sketch_id: Optional[str] = None) -> int:
+        """Get total node count, optionally filtered by sketch_id."""
+        if sketch_id:
+            return sum(
+                1 for n in self._nodes.values() if n.get("sketch_id") == sketch_id
+            )
+        return len(self._nodes)
+
+    def get_edge_count(self, sketch_id: Optional[str] = None) -> int:
+        """Get total edge count, optionally filtered by sketch_id."""
+        if sketch_id:
+            return sum(
+                1 for e in self._edges.values() if e.get("sketch_id") == sketch_id
+            )
+        return len(self._edges)
+
+    def clear(self) -> None:
+        """Clear all data (useful between tests)."""
+        self._nodes.clear()
+        self._edges.clear()
+        self._batch_operations.clear()

flowsint-core/tests/core/graph/test_graph_repository.py

@@ -0,0 +1,874 @@
+"""Tests for Neo4jGraphRepository."""
+
+import pytest
+from unittest.mock import MagicMock, patch, call
+from datetime import datetime, timezone
+
+from flowsint_core.core.graph import Neo4jGraphRepository
+
+
+class TestNeo4jGraphRepositoryInit:
+    def test_init_with_connection(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+        assert repo._connection == mock_connection
+        assert repo._batch_operations == []
+        assert repo._batch_size == 100
+
+    def test_init_without_connection_uses_singleton(self):
+        with patch(
+            "flowsint_core.core.graph.repository.Neo4jConnection.get_instance"
+        ) as mock_get_instance:
+            mock_connection = MagicMock()
+            mock_get_instance.return_value = mock_connection
+            repo = Neo4jGraphRepository()
+            assert repo._connection == mock_connection
+
+
+class TestCreateNode:
+    def test_create_node_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"id": "element-123"}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        node_obj = {
+            "nodeLabel": "example.com",
+            "nodeType": "domain",
+            "nodeProperties.domain": "example.com",
+        }
+
+        result = repo.create_node(node_obj, sketch_id="sketch-1")
+
+        assert result == "element-123"
+        mock_connection.query.assert_called_once()
+
+    def test_create_node_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.create_node({"nodeLabel": "test", "nodeType": "domain"}, "sketch-1")
+
+        assert result is None
+
+    def test_create_node_empty_result(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = []
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.create_node(
+            {"nodeLabel": "test", "nodeType": "domain"}, sketch_id="sketch-1"
+        )
+
+        assert result is None
+
+
+class TestCreateRelationship:
+    def test_create_relationship_success(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        rel_obj = {
+            "from_type": "domain",
+            "from_label": "source.com",
+            "to_type": "ip",
+            "to_label": "1.1.1.1",
+            "rel_label": "RESOLVES_TO",
+        }
+
+        repo.create_relationship(rel_obj, sketch_id="sketch-1")
+
+        mock_connection.execute_write.assert_called_once()
+
+    def test_create_relationship_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        rel_obj = {
+            "from_type": "domain",
+            "from_label": "source.com",
+            "to_type": "ip",
+            "to_label": "1.1.1.1",
+            "rel_label": "RESOLVES_TO",
+        }
+
+        # Should not raise, just return early
+        repo.create_relationship(rel_obj, sketch_id="sketch-1")
+
+
+class TestBuildNodeQuery:
+    def test_build_node_query_structure(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        node_obj = {
+            "nodeLabel": "example.com",
+            "nodeType": "domain",
+        }
+
+        query, params = repo._build_node_query(node_obj, sketch_id="sketch-1")
+
+        assert "MERGE" in query
+        assert "domain" in query  # nodeType used as label
+        assert params["node_label"] == "example.com"
+        assert params["sketch_id"] == "sketch-1"
+        assert params["props"] == node_obj
+        assert "created_at" in params
+
+
+class TestBuildRelationshipQuery:
+    def test_build_relationship_query_structure(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        rel_obj = {
+            "from_type": "domain",
+            "from_label": "source.com",
+            "to_type": "ip",
+            "to_label": "1.1.1.1",
+            "rel_label": "RESOLVES_TO",
+        }
+
+        query, params = repo._build_relationship_query(rel_obj, sketch_id="sketch-1")
+
+        assert "MATCH" in query
+        assert "MERGE" in query
+        assert "domain" in query
+        assert "ip" in query
+        assert "RESOLVES_TO" in query
+        assert params["from_label"] == "source.com"
+        assert params["to_label"] == "1.1.1.1"
+        assert params["sketch_id"] == "sketch-1"
+
+
+class TestBatchOperations:
+    def test_add_to_batch_node(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        repo.add_to_batch(
+            "node",
+            node_obj={"nodeLabel": "test", "nodeType": "domain"},
+            sketch_id="sketch-1",
+        )
+
+        assert len(repo._batch_operations) == 1
+
+    def test_add_to_batch_relationship(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        repo.add_to_batch(
+            "relationship",
+            rel_obj={
+                "from_type": "domain",
+                "from_label": "a.com",
+                "to_type": "ip",
+                "to_label": "1.1.1.1",
+                "rel_label": "REL",
+            },
+            sketch_id="sketch-1",
+        )
+
+        assert len(repo._batch_operations) == 1
+
+    def test_add_to_batch_unknown_type_raises(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        with pytest.raises(ValueError, match="Unknown operation type"):
+            repo.add_to_batch("unknown", sketch_id="sketch-1")
+
+    def test_auto_flush_when_batch_full(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+        repo._batch_size = 2
+
+        repo.add_to_batch(
+            "node",
+            node_obj={"nodeLabel": "test1", "nodeType": "domain"},
+            sketch_id="sketch-1",
+        )
+        assert len(repo._batch_operations) == 1
+
+        repo.add_to_batch(
+            "node",
+            node_obj={"nodeLabel": "test2", "nodeType": "domain"},
+            sketch_id="sketch-1",
+        )
+
+        # Should have auto-flushed
+        mock_connection.execute_batch.assert_called_once()
+        assert len(repo._batch_operations) == 0
+
+    def test_flush_batch(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        # Track what was passed to execute_batch before it gets cleared
+        captured_args = []
+
+        def capture_batch(ops):
+            captured_args.extend(list(ops))
+
+        mock_connection.execute_batch.side_effect = capture_batch
+
+        repo._batch_operations = [("query1", {"p": 1}), ("query2", {"p": 2})]
+
+        repo.flush_batch()
+
+        mock_connection.execute_batch.assert_called_once()
+        assert len(captured_args) == 2
+        assert captured_args[0] == ("query1", {"p": 1})
+        assert captured_args[1] == ("query2", {"p": 2})
+        assert len(repo._batch_operations) == 0
+
+    def test_flush_batch_empty(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        repo.flush_batch()
+
+        mock_connection.execute_batch.assert_not_called()
+
+    def test_flush_batch_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+        repo._batch_operations = [("query", {})]
+
+        repo.flush_batch()
+
+        assert len(repo._batch_operations) == 0
+
+    def test_clear_batch(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+        repo._batch_operations = [("query", {})]
+
+        repo.clear_batch()
+
+        assert len(repo._batch_operations) == 0
+
+    def test_set_batch_size(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        repo.set_batch_size(50)
+
+        assert repo._batch_size == 50
+
+    def test_set_batch_size_invalid(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        with pytest.raises(ValueError, match="Batch size must be at least 1"):
+            repo.set_batch_size(0)
+
+
+class TestBatchCreateNodes:
+    def test_batch_create_nodes_success(self):
+        mock_connection = MagicMock()
+        mock_connection.execute_batch.return_value = [
+            [{"id": "id-1"}],
+            [{"id": "id-2"}],
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        nodes = [
+            {"nodeLabel": "a.com", "nodeType": "domain"},
+            {"nodeLabel": "b.com", "nodeType": "domain"},
+        ]
+
+        result = repo.batch_create_nodes(nodes, sketch_id="sketch-1")
+
+        assert result["nodes_created"] == 2
+        assert result["node_ids"] == ["id-1", "id-2"]
+        assert result["errors"] == []
+
+    def test_batch_create_nodes_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.batch_create_nodes(
+            [{"nodeLabel": "test", "nodeType": "domain"}], sketch_id="sketch-1"
+        )
+
+        assert result["nodes_created"] == 0
+        assert result["node_ids"] == []
+        assert "No database connection" in result["errors"]
+
+    def test_batch_create_nodes_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.batch_create_nodes([], sketch_id="sketch-1")
+
+        assert result == {"nodes_created": 0, "node_ids": [], "errors": []}
+
+    def test_batch_create_nodes_execution_error(self):
+        mock_connection = MagicMock()
+        mock_connection.execute_batch.side_effect = Exception("DB error")
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.batch_create_nodes(
+            [{"nodeLabel": "test", "nodeType": "domain"}], sketch_id="sketch-1"
+        )
+
+        assert result["nodes_created"] == 0
+        assert "Batch execution failed" in result["errors"][0]
+
+
+class TestBatchCreateEdges:
+    def test_batch_create_edges_success(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        edges = [
+            {
+                "from_type": "domain",
+                "from_label": "a.com",
+                "to_type": "ip",
+                "to_label": "1.1.1.1",
+                "rel_label": "RESOLVES",
+            }
+        ]
+
+        result = repo.batch_create_edges(edges, sketch_id="sketch-1")
+
+        assert result["edges_created"] == 1
+        assert result["errors"] == []
+
+    def test_batch_create_edges_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.batch_create_edges([{}], sketch_id="sketch-1")
+
+        assert result["edges_created"] == 0
+        assert "No database connection" in result["errors"]
+
+    def test_batch_create_edges_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.batch_create_edges([], sketch_id="sketch-1")
+
+        assert result == {"edges_created": 0, "errors": []}
+
+
+class TestBatchCreateEdgesByElementId:
+    def test_batch_create_edges_by_element_id_success(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        edges = [
+            {
+                "from_element_id": "elem-1",
+                "to_element_id": "elem-2",
+                "rel_label": "CONNECTS",
+            }
+        ]
+
+        result = repo.batch_create_edges_by_element_id(edges, sketch_id="sketch-1")
+
+        assert result["edges_created"] == 1
+        assert result["errors"] == []
+
+    def test_batch_create_edges_by_element_id_missing_fields(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        edges = [{"rel_label": "CONNECTS"}]  # Missing from/to element IDs
+
+        result = repo.batch_create_edges_by_element_id(edges, sketch_id="sketch-1")
+
+        assert result["edges_created"] == 0
+        assert any("Missing required fields" in e for e in result["errors"])
+
+    def test_batch_create_edges_by_element_id_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.batch_create_edges_by_element_id([{}], sketch_id="sketch-1")
+
+        assert result["edges_created"] == 0
+        assert "No database connection" in result["errors"]
+
+
+class TestUpdateNode:
+    def test_update_node_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"id": "elem-1"}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.update_node(
+            element_id="elem-1",
+            updates={"nodeLabel": "updated"},
+            sketch_id="sketch-1",
+        )
+
+        assert result == "elem-1"
+
+    def test_update_node_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.update_node("elem-1", {"nodeLabel": "x"}, "sketch-1")
+
+        assert result is None
+
+
+class TestDeleteNodes:
+    def test_delete_nodes_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"deleted_count": 3}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.delete_nodes(["id-1", "id-2", "id-3"], sketch_id="sketch-1")
+
+        assert result == 3
+
+    def test_delete_nodes_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.delete_nodes(["id-1"], sketch_id="sketch-1")
+
+        assert result == 0
+
+    def test_delete_nodes_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.delete_nodes([], sketch_id="sketch-1")
+
+        assert result == 0
+        mock_connection.query.assert_not_called()
+
+
+class TestDeleteRelationships:
+    def test_delete_relationships_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"deleted_count": 2}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.delete_relationships(["rel-1", "rel-2"], sketch_id="sketch-1")
+
+        assert result == 2
+
+    def test_delete_relationships_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.delete_relationships(["rel-1"], sketch_id="sketch-1")
+
+        assert result == 0
+
+    def test_delete_relationships_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.delete_relationships([], sketch_id="sketch-1")
+
+        assert result == 0
+
+
+class TestDeleteAllSketchNodes:
+    def test_delete_all_sketch_nodes_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"deleted_count": 10}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.delete_all_sketch_nodes(sketch_id="sketch-1")
+
+        assert result == 10
+
+    def test_delete_all_sketch_nodes_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.delete_all_sketch_nodes(sketch_id="sketch-1")
+
+        assert result == 0
+
+
+class TestGetSketchGraph:
+    def test_get_sketch_graph_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.side_effect = [
+            # First call: nodes query
+            [
+                {"id": "node-1", "labels": ["domain"], "data": {}},
+                {"id": "node-2", "labels": ["ip"], "data": {}},
+            ],
+            # Second call: edges query
+            [
+                {
+                    "id": "edge-1",
+                    "type": "RESOLVES",
+                    "source": "node-1",
+                    "target": "node-2",
+                    "data": {},
+                }
+            ],
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_sketch_graph(sketch_id="sketch-1")
+
+        assert len(result["nodes"]) == 2
+        assert len(result["edges"]) == 1
+
+    def test_get_sketch_graph_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.get_sketch_graph(sketch_id="sketch-1")
+
+        assert result == {"nodes": [], "edges": []}
+
+    def test_get_sketch_graph_no_nodes(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = []
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_sketch_graph(sketch_id="sketch-1")
+
+        assert result == {"nodes": [], "edges": []}
+
+
+class TestUpdateRelationship:
+    def test_update_relationship_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [
+            {"id": "rel-1", "type": "CONNECTS", "data": {"weight": 5}}
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.update_relationship(
+            element_id="rel-1",
+            rel_obj={"weight": 5},
+            sketch_id="sketch-1",
+        )
+
+        assert result["id"] == "rel-1"
+
+    def test_update_relationship_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.update_relationship("rel-1", {}, "sketch-1")
+
+        assert result is None
+
+
+class TestCreateRelationshipByElementId:
+    def test_create_relationship_by_element_id_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"rel": {"sketch_id": "sketch-1"}}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.create_relationship_by_element_id(
+            from_element_id="elem-1",
+            to_element_id="elem-2",
+            rel_label="CONNECTS",
+            sketch_id="sketch-1",
+        )
+
+        assert result["sketch_id"] == "sketch-1"
+
+    def test_create_relationship_by_element_id_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.create_relationship_by_element_id(
+            "elem-1", "elem-2", "CONNECTS", "sketch-1"
+        )
+
+        assert result is None
+
+
+class TestQuery:
+    def test_query_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"count": 5}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.query("MATCH (n) RETURN count(n) as count", {})
+
+        assert result == [{"count": 5}]
+
+    def test_query_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.query("MATCH (n) RETURN n", {})
+
+        assert result == []
+
+
+class TestUpdateNodesPositions:
+    def test_update_nodes_positions_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [{"updated_count": 2}]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        positions = [
+            {"nodeId": "node-1", "x": 100, "y": 200},
+            {"nodeId": "node-2", "x": 300, "y": 400},
+        ]
+
+        result = repo.update_nodes_positions(positions, sketch_id="sketch-1")
+
+        assert result == 2
+
+    def test_update_nodes_positions_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.update_nodes_positions([{"nodeId": "x", "x": 0, "y": 0}], "s")
+
+        assert result == 0
+
+    def test_update_nodes_positions_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.update_nodes_positions([], sketch_id="sketch-1")
+
+        assert result == 0
+
+
+class TestGetNodesByIds:
+    def test_get_nodes_by_ids_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [
+            {"data": {"nodeLabel": "a.com"}},
+            {"data": {"nodeLabel": "b.com"}},
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_nodes_by_ids(["id-1", "id-2"], sketch_id="sketch-1")
+
+        assert len(result) == 2
+
+    def test_get_nodes_by_ids_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.get_nodes_by_ids(["id-1"], sketch_id="sketch-1")
+
+        assert result == []
+
+    def test_get_nodes_by_ids_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_nodes_by_ids([], sketch_id="sketch-1")
+
+        assert result == []
+
+
+class TestMergeNodes:
+    def test_merge_nodes_create_new_node(self):
+        mock_connection = MagicMock()
+        mock_connection.query.side_effect = [
+            [{"newElementId": "new-elem-1"}],  # Create query
+            None,  # Copy relationships
+            None,  # Delete old nodes
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.merge_nodes(
+            old_node_ids=["old-1", "old-2"],
+            new_node_data={"type": "domain"},
+            new_node_id=None,
+            sketch_id="sketch-1",
+        )
+
+        assert result == "new-elem-1"
+
+    def test_merge_nodes_reuse_existing_node(self):
+        mock_connection = MagicMock()
+        mock_connection.query.side_effect = [
+            [{"newElementId": "old-1"}],  # Update existing
+            None,  # Copy relationships
+            None,  # Delete old nodes
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.merge_nodes(
+            old_node_ids=["old-1", "old-2"],
+            new_node_data={"type": "domain"},
+            new_node_id="old-1",  # Reusing old-1
+            sketch_id="sketch-1",
+        )
+
+        assert result == "old-1"
+
+    def test_merge_nodes_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.merge_nodes(["old-1"], {}, None, "sketch-1")
+
+        assert result is None
+
+    def test_merge_nodes_empty_list(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.merge_nodes([], {}, None, "sketch-1")
+
+        assert result is None
+
+
+class TestGetNeighbors:
+    def test_get_neighbors_success(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [
+            {
+                "center_id": "node-1",
+                "center_data": {"nodeLabel": "center"},
+                "rel_id": "rel-1",
+                "rel_label": "CONNECTS",
+                "other_id": "node-2",
+                "other_data": {"nodeLabel": "neighbor"},
+                "direction": "outgoing",
+            }
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_neighbors(node_id="node-1", sketch_id="sketch-1")
+
+        assert len(result["nodes"]) == 2
+        assert len(result["edges"]) == 1
+
+    def test_get_neighbors_no_relationships(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [
+            {
+                "center_id": "node-1",
+                "center_data": {"nodeLabel": "center"},
+                "rel_id": None,
+                "rel_label": None,
+                "other_id": None,
+                "other_data": None,
+                "direction": None,
+            }
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_neighbors(node_id="node-1", sketch_id="sketch-1")
+
+        assert len(result["nodes"]) == 1
+        assert len(result["edges"]) == 0
+
+    def test_get_neighbors_no_connection(self):
+        repo = Neo4jGraphRepository(neo4j_connection=None)
+        repo._connection = None
+
+        result = repo.get_neighbors("node-1", "sketch-1")
+
+        assert result == {"nodes": [], "edges": []}
+
+    def test_get_neighbors_not_found(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = []
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_neighbors("nonexistent", "sketch-1")
+
+        assert result == {"nodes": [], "edges": []}
+
+    def test_get_neighbors_incoming_direction(self):
+        mock_connection = MagicMock()
+        mock_connection.query.return_value = [
+            {
+                "center_id": "node-1",
+                "center_data": {"nodeLabel": "center"},
+                "rel_id": "rel-1",
+                "rel_label": "CONNECTS",
+                "other_id": "node-2",
+                "other_data": {"nodeLabel": "neighbor"},
+                "direction": "incoming",
+            }
+        ]
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.get_neighbors(node_id="node-1", sketch_id="sketch-1")
+
+        # For incoming, source should be the other node
+        edge = result["edges"][0]
+        assert edge["source"] == "node-2"
+        assert edge["target"] == "node-1"
+
+
+class TestCountNodesBySketch:
+    def test_count_nodes_by_sketch_success(self):
+        mock_connection = MagicMock()
+        mock_driver = MagicMock()
+        mock_session = MagicMock()
+        mock_result = MagicMock()
+        mock_record = MagicMock()
+
+        mock_record.__getitem__ = lambda self, key: 5 if key == "total" else None
+        mock_result.single.return_value = mock_record
+        mock_session.run.return_value = mock_result
+        mock_session.__enter__ = MagicMock(return_value=mock_session)
+        mock_session.__exit__ = MagicMock(return_value=None)
+        mock_driver.session.return_value = mock_session
+        mock_connection.get_driver.return_value = mock_driver
+
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.count_nodes_by_sketch(sketch_id="sketch-1")
+
+        assert result == 5
+
+
+class TestCountEdgesBySketch:
+    def test_count_edges_by_sketch_success(self):
+        mock_connection = MagicMock()
+        mock_driver = MagicMock()
+        mock_session = MagicMock()
+        mock_result = MagicMock()
+        mock_record = MagicMock()
+
+        mock_record.__getitem__ = lambda self, key: 3 if key == "total" else None
+        mock_result.single.return_value = mock_record
+        mock_session.run.return_value = mock_result
+        mock_session.__enter__ = MagicMock(return_value=mock_session)
+        mock_session.__exit__ = MagicMock(return_value=None)
+        mock_driver.session.return_value = mock_session
+        mock_connection.get_driver.return_value = mock_driver
+
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+
+        result = repo.count_edges_by_sketch(sketch_id="sketch-1")
+
+        assert result == 3
+
+
+class TestContextManager:
+    def test_context_manager_flushes_on_success(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+        repo._batch_operations = [("query", {})]
+
+        with repo:
+            pass
+
+        mock_connection.execute_batch.assert_called_once()
+
+    def test_context_manager_clears_on_exception(self):
+        mock_connection = MagicMock()
+        repo = Neo4jGraphRepository(neo4j_connection=mock_connection)
+        repo._batch_operations = [("query", {})]
+
+        with pytest.raises(ValueError):
+            with repo:
+                raise ValueError("Test error")
+
+        # Should have cleared, not flushed
+        mock_connection.execute_batch.assert_not_called()
+        assert len(repo._batch_operations) == 0

flowsint-core/tests/core/graph/test_graph_serializer.py

@@ -0,0 +1,320 @@
+from datetime import datetime
+
+import pytest
+from flowsint_types import Domain, Ip
+
+from flowsint_core.core.graph import (
+    GraphSerializer,
+    GraphEdge,
+    GraphNode,
+    NodeMetadata,
+)
+
+
+def test_serializer():
+    created_at = datetime.now()
+    node = GraphNode(
+        id="id",
+        nodeLabel="nodeLabel",
+        nodeFlag="blue",
+        nodeType="domain",
+        nodeColor="nodeColor",
+        nodeSize=4,
+        nodeImage="nodeImage",
+        nodeIcon="nodeIcon",
+        x=100,
+        y=100,
+        nodeProperties=Domain(domain="domain.com", nodeLabel="domain.com", root=True),
+        nodeMetadata=NodeMetadata(created_at=created_at),
+    )
+    to_neo4j = GraphSerializer.graph_node_to_neo4j_dict(node)
+
+    expected = {
+        "id": "id",
+        "nodeLabel": "nodeLabel",
+        "nodeType": "domain",
+        "nodeColor": "nodeColor",
+        "nodeFlag": "blue",
+        "nodeSize": 4,
+        "nodeImage": "nodeImage",
+        "nodeIcon": "nodeIcon",
+        "x": 100.0,
+        "y": 100.0,
+        "nodeProperties.domain": "domain.com",
+        "nodeProperties.root": True,
+        "nodeMetadata.created_at": created_at.isoformat(),
+    }
+
+    assert to_neo4j == expected
+
+
+def test_deserializer():
+    created_at = datetime.now()
+
+    node = {
+        "id": "id",
+        "data": {
+            "nodeLabel": "nodeLabel",
+            "nodeType": "domain",
+            "nodeColor": "nodeColor",
+            "nodeSize": 4,
+            "nodeImage": "nodeImage",
+            "nodeIcon": "nodeIcon",
+            "x": 100.0,
+            "y": 100.0,
+            "nodeProperties.domain": "domain.com",
+            "nodeProperties.root": True,
+            "nodeMetadata.created_at": created_at.isoformat(),
+        },
+    }
+
+    graph_node = GraphNode(
+        id="id",
+        nodeLabel="nodeLabel",
+        nodeType="domain",
+        nodeColor="nodeColor",
+        nodeSize=4,
+        nodeImage="nodeImage",
+        nodeIcon="nodeIcon",
+        x=100,
+        y=100,
+        nodeProperties=Domain(domain="domain.com", nodeLabel="domain.com", root=True),
+        nodeMetadata=NodeMetadata(created_at=created_at),
+    )
+    output = GraphSerializer.neo4j_dict_to_graph_node(node)
+
+    assert output == graph_node
+
+
+def test_serialize_from_flowsint_type():
+    domain = Domain(domain="domain.com", nodeLabel="domain.com", root=True)
+
+    neo4j_dict = GraphSerializer.flowsint_type_to_neo4j_dict(domain)
+
+    # Check static fields
+    assert neo4j_dict["id"] == ""
+    assert neo4j_dict["nodeLabel"] == "domain.com"
+    assert neo4j_dict["nodeType"] == "domain"
+    assert neo4j_dict["nodeColor"] is None
+    assert neo4j_dict["nodeSize"] is None
+    assert neo4j_dict["nodeImage"] is None
+    assert neo4j_dict["nodeIcon"] is None
+    assert neo4j_dict["nodeFlag"] is None
+    assert neo4j_dict["x"] == 100.0
+    assert neo4j_dict["y"] == 100.0
+    assert neo4j_dict["nodeProperties.domain"] == "domain.com"
+    assert neo4j_dict["nodeProperties.root"] is True
+
+    # Timestamp is generated internally, just verify it exists and is ISO format
+    assert "nodeMetadata.created_at" in neo4j_dict
+    assert "T" in neo4j_dict["nodeMetadata.created_at"]
+
+
+class TestCleanEmptyValues:
+    def test_removes_empty_strings(self):
+        data = {"key1": "value", "key2": "", "key3": "another"}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {"key1": "value", "key3": "another"}
+
+    def test_removes_none_values(self):
+        data = {"key1": "value", "key2": None, "key3": "another"}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {"key1": "value", "key3": "another"}
+
+    def test_cleans_nested_dicts(self):
+        data = {"outer": {"inner1": "value", "inner2": "", "inner3": None}}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {"outer": {"inner1": "value"}}
+
+    def test_cleans_lists(self):
+        data = {"items": ["a", "", "b", None, "c"]}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {"items": ["a", "b", "c"]}
+
+    def test_cleans_list_of_dicts(self):
+        data = {"items": [{"a": "1", "b": ""}, {"c": None, "d": "2"}]}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {"items": [{"a": "1"}, {"d": "2"}]}
+
+    def test_removes_empty_nested_dict(self):
+        data = {"outer": {"inner": ""}}
+        result = GraphSerializer._clean_empty_values(data)
+        assert result == {}
+
+
+class TestParseFlowsintType:
+    def test_parses_domain(self):
+        entity = {"domain": "example.com", "root": True}
+        result = GraphSerializer.parse_flowsint_type(entity, "domain")
+        assert isinstance(result, Domain)
+        assert result.domain == "example.com"
+        assert result.root is True
+
+    def test_parses_ip(self):
+        entity = {"address": "192.168.1.1"}
+        result = GraphSerializer.parse_flowsint_type(entity, "ip")
+        assert isinstance(result, Ip)
+        assert result.address == "192.168.1.1"
+
+    def test_cleans_empty_values_before_parsing(self):
+        entity = {"domain": "example.com", "extra": "", "other": None}
+        result = GraphSerializer.parse_flowsint_type(entity, "domain")
+        assert isinstance(result, Domain)
+        assert result.domain == "example.com"
+
+    def test_raises_on_unknown_type(self):
+        entity = {"key": "value"}
+        with pytest.raises(ValueError, match="Unknown type: unknowntype"):
+            GraphSerializer.parse_flowsint_type(entity, "unknowntype")
+
+
+class TestGraphNodeToFlowsintType:
+    def test_extracts_node_properties(self):
+        domain = Domain(domain="example.com")
+        node = GraphNode(
+            id="123",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=domain,
+            nodeMetadata=NodeMetadata(),
+        )
+        result = GraphSerializer.graph_node_to_flowsint_type(node)
+        assert result is domain
+        assert isinstance(result, Domain)
+        assert result.domain == "example.com"
+
+
+class TestNeo4jDictToGraphEdge:
+    def test_converts_edge_dict(self):
+        edge_dict = {
+            "id": "edge-123",
+            "source": "node-1",
+            "target": "node-2",
+            "type": "CONNECTED_TO",
+        }
+        result = GraphSerializer.neo4j_dict_to_graph_edge(edge_dict)
+        assert isinstance(result, GraphEdge)
+        assert result.id == "edge-123"
+        assert result.source == "node-1"
+        assert result.target == "node-2"
+        assert result.label == "CONNECTED_TO"
+
+    def test_converts_values_to_string(self):
+        edge_dict = {"id": 123, "source": 1, "target": 2, "type": "REL"}
+        result = GraphSerializer.neo4j_dict_to_graph_edge(edge_dict)
+        assert result.id == "123"
+        assert result.source == "1"
+        assert result.target == "2"
+
+
+class TestGraphEdgeToNeo4jDict:
+    def test_with_flowsint_types(self):
+        from_obj = Domain(domain="source.com")
+        to_obj = Domain(domain="target.com")
+        result = GraphSerializer.graph_edge_to_neo4j_dict(from_obj, to_obj, "LINKS_TO")
+        assert result == {
+            "from_type": "domain",
+            "from_label": "source.com",
+            "to_type": "domain",
+            "to_label": "target.com",
+            "rel_label": "LINKS_TO",
+        }
+
+    def test_with_graph_nodes(self):
+        from_obj = GraphNode(
+            id="1",
+            nodeLabel="source",
+            nodeType="ip",
+            nodeProperties=Ip(address="1.1.1.1"),
+            nodeMetadata=NodeMetadata(),
+        )
+        to_obj = GraphNode(
+            id="2",
+            nodeLabel="target",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+        result = GraphSerializer.graph_edge_to_neo4j_dict(from_obj, to_obj, "RESOLVES")
+        assert result == {
+            "from_type": "ip",
+            "from_label": "source",
+            "to_type": "domain",
+            "to_label": "target",
+            "rel_label": "RESOLVES",
+        }
+
+    def test_with_mixed_types(self):
+        from_obj = Domain(domain="source.com")
+        to_obj = GraphNode(
+            id="2",
+            nodeLabel="target",
+            nodeType="ip",
+            nodeProperties=Ip(address="1.1.1.1"),
+            nodeMetadata=NodeMetadata(),
+        )
+        result = GraphSerializer.graph_edge_to_neo4j_dict(from_obj, to_obj, "HOSTS")
+        assert result == {
+            "from_type": "domain",
+            "from_label": "source.com",
+            "to_type": "ip",
+            "to_label": "target",
+            "rel_label": "HOSTS",
+        }
+
+
+class TestDeserializeNodes:
+    def test_deserializes_multiple_nodes(self):
+        node_dicts = [
+            {
+                "id": "1",
+                "data": {
+                    "nodeLabel": "example.com",
+                    "nodeType": "domain",
+                    "nodeProperties.domain": "example.com",
+                    "nodeMetadata.created_at": "2026-01-01T00:00:00",
+                },
+            },
+            {
+                "id": "2",
+                "data": {
+                    "nodeLabel": "1.1.1.1",
+                    "nodeType": "ip",
+                    "nodeProperties.address": "1.1.1.1",
+                    "nodeMetadata.created_at": "2026-01-01T00:00:00",
+                },
+            },
+        ]
+        result = GraphSerializer.deserialize_nodes(node_dicts)
+        assert len(result) == 2
+        assert all(isinstance(node, GraphNode) for node in result)
+        assert result[0].nodeLabel == "example.com"
+        assert result[1].nodeLabel == "1.1.1.1"
+
+    def test_empty_list(self):
+        result = GraphSerializer.deserialize_nodes([])
+        assert result == []
+
+
+class TestDeserializeEdges:
+    def test_deserializes_multiple_edges(self):
+        edge_dicts = [
+            {"id": "e1", "source": "n1", "target": "n2", "type": "CONNECTS"},
+            {"id": "e2", "source": "n2", "target": "n3", "type": "LINKS"},
+        ]
+        result = GraphSerializer.deserialize_edges(edge_dicts)
+        assert len(result) == 2
+        assert all(isinstance(edge, GraphEdge) for edge in result)
+        assert result[0].label == "CONNECTS"
+        assert result[1].label == "LINKS"
+
+    def test_empty_list(self):
+        result = GraphSerializer.deserialize_edges([])
+        assert result == []
+
+
+class TestNeo4jDictToGraphNodeErrors:
+    def test_raises_on_missing_data(self):
+        node_dict = {"id": "123"}
+        with pytest.raises(Exception, match="Could not find node data"):
+            GraphSerializer.neo4j_dict_to_graph_node(node_dict)

flowsint-core/tests/core/graph/test_graph_service.py

@@ -0,0 +1,682 @@
+"""Tests for GraphService using dependency injection."""
+
+import pytest
+from unittest.mock import MagicMock, patch
+
+from flowsint_types import Domain, Ip
+from flowsint_core.core.graph import (
+    GraphService,
+    create_graph_service,
+    GraphNode,
+    GraphData,
+    NodeMetadata,
+)
+
+from .in_memory_graph_repository import InMemoryGraphRepository
+
+
+class TestGraphServiceInit:
+    def test_init_with_injected_repository(self):
+        """Test that a repository can be injected."""
+        mock_repo = MagicMock()
+        mock_logger = MagicMock()
+
+        service = GraphService(
+            sketch_id="sketch-1",
+            repository=mock_repo,
+            logger=mock_logger,
+            enable_batching=True,
+        )
+
+        assert service._sketch_id == "sketch-1"
+        assert service._repository == mock_repo
+        assert service._logger == mock_logger
+        assert service._enable_batching is True
+
+    def test_init_with_in_memory_repository(self):
+        """Test initialization with InMemoryGraphRepository."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+
+        assert service._sketch_id == "sketch-1"
+        assert service._repository is repo
+
+    def test_init_without_repository_raises(self):
+        """Test that without repository injection, ValueError is raised."""
+        with pytest.raises(ValueError, match="repository is required"):
+            GraphService(sketch_id="sketch-1", repository=None)
+
+    def test_sketch_id_property(self):
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+        assert service.sketch_id == "sketch-1"
+
+    def test_repository_property(self):
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+        assert service.repository is repo
+
+
+class TestCreateNode:
+    def test_create_node_with_graph_node(self):
+        """Test creating a node with injected mock repository."""
+        mock_repo = MagicMock()
+        mock_repo.create_node.return_value = "elem-123"
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        node = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        result = service.create_node(node)
+
+        assert result == "elem-123"
+        mock_repo.create_node.assert_called_once()
+
+    def test_create_node_with_in_memory_repository(self):
+        """Test creating a node with InMemoryGraphRepository."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+
+        node = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        result = service.create_node(node)
+
+        assert result is not None
+        assert repo.get_node_count("sketch-1") == 1
+
+    def test_create_node_with_flowsint_type_raises(self):
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+        domain = Domain(domain="example.com")
+
+        with pytest.raises(Exception, match="create_node method takes a GraphNode"):
+            service.create_node(domain)
+
+    def test_create_node_with_batching(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+
+        node = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        service.create_node(node)
+
+        mock_repo.add_to_batch.assert_called_once()
+        mock_repo.create_node.assert_not_called()
+
+
+class TestCreateNodeFromFlowsintType:
+    def test_create_node_from_flowsint_type(self):
+        mock_repo = MagicMock()
+        mock_repo.create_node.return_value = "elem-123"
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        domain = Domain(domain="example.com")
+
+        result = service.create_node_from_flowsint_type(domain)
+
+        assert result == "elem-123"
+        mock_repo.create_node.assert_called_once()
+
+    def test_create_node_from_flowsint_type_with_graph_node_raises(self):
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+        node = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        with pytest.raises(
+            Exception, match="create_node_from_flowsint_type method takes a FlowsintType"
+        ):
+            service.create_node_from_flowsint_type(node)
+
+    def test_create_node_from_flowsint_type_with_batching(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+        domain = Domain(domain="example.com")
+
+        service.create_node_from_flowsint_type(domain)
+
+        mock_repo.add_to_batch.assert_called_once()
+
+
+class TestGetSketchGraph:
+    def test_get_sketch_graph_with_mock(self):
+        mock_repo = MagicMock()
+        mock_repo.get_sketch_graph.return_value = {"nodes": [], "edges": []}
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.get_sketch_graph()
+
+        assert isinstance(result, GraphData)
+        mock_repo.get_sketch_graph.assert_called_once_with("sketch-1")
+
+    def test_get_sketch_graph_with_in_memory(self):
+        """Test getting sketch graph with actual data in InMemoryRepository."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+
+        # Add some nodes
+        node1 = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+        node2 = GraphNode(
+            id="2",
+            nodeLabel="1.1.1.1",
+            nodeType="ip",
+            nodeProperties=Ip(address="1.1.1.1"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        service.create_node(node1)
+        service.create_node(node2)
+
+        result = service.get_sketch_graph()
+
+        assert isinstance(result, GraphData)
+        assert len(result.nodes) == 2
+
+
+class TestGetNodesByIds:
+    def test_get_nodes_by_ids(self):
+        """Test that get_nodes_by_ids calls the repository correctly."""
+        mock_repo = MagicMock()
+        # Return empty list to avoid serialization issues in this unit test
+        mock_repo.get_nodes_by_ids.return_value = []
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.get_nodes_by_ids(["id-1", "id-2"])
+
+        mock_repo.get_nodes_by_ids.assert_called_once_with(["id-1", "id-2"], "sketch-1")
+        assert result == []
+
+
+class TestGetNodesByIdsForTask:
+    def test_get_nodes_by_ids_for_task(self):
+        """Test that get_nodes_by_ids_for_task calls the repository correctly."""
+        mock_repo = MagicMock()
+        # Return empty list to avoid serialization issues in this unit test
+        mock_repo.get_nodes_by_ids.return_value = []
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+
+        result = service.get_nodes_by_ids_for_task(["id-1"])
+
+        mock_repo.get_nodes_by_ids.assert_called_once_with(["id-1"], "sketch-1")
+        assert result == []
+
+
+class TestCreateRelationship:
+    def test_create_relationship(self):
+        mock_repo = MagicMock()
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+
+        from_obj = Domain(domain="example.com")
+        to_obj = Ip(address="1.1.1.1")
+
+        service.create_relationship(from_obj, to_obj, "RESOLVES")
+
+        mock_repo.create_relationship.assert_called_once()
+
+    def test_create_relationship_with_batching(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+
+        service.create_relationship(
+            Domain(domain="a.com"), Domain(domain="b.com")
+        )
+
+        mock_repo.add_to_batch.assert_called_once()
+        mock_repo.create_relationship.assert_not_called()
+
+
+class TestCreateRelationshipByElementId:
+    def test_create_relationship_by_element_id(self):
+        mock_repo = MagicMock()
+        mock_repo.create_relationship_by_element_id.return_value = {"sketch_id": "s1"}
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.create_relationship_by_element_id(
+            from_element_id="elem-1",
+            to_element_id="elem-2",
+            rel_label="CONNECTS",
+        )
+
+        assert result == {"sketch_id": "s1"}
+        mock_repo.create_relationship_by_element_id.assert_called_once_with(
+            from_element_id="elem-1",
+            to_element_id="elem-2",
+            rel_label="CONNECTS",
+            sketch_id="sketch-1",
+        )
+
+
+class TestGetNeighbors:
+    def test_get_neighbors(self):
+        mock_repo = MagicMock()
+        mock_repo.get_neighbors.return_value = {"nodes": [], "edges": []}
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.get_neighbors(node_id="node-1")
+
+        assert isinstance(result, GraphData)
+        mock_repo.get_neighbors.assert_called_once_with(
+            node_id="node-1", sketch_id="sketch-1"
+        )
+
+
+class TestUpdateNode:
+    def test_update_node(self):
+        mock_repo = MagicMock()
+        mock_repo.update_node.return_value = "elem-1"
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.update_node("elem-1", {"nodeLabel": "updated"})
+
+        assert result == "elem-1"
+        mock_repo.update_node.assert_called_once()
+
+
+class TestUpdateNodesPositions:
+    def test_update_nodes_positions(self):
+        mock_repo = MagicMock()
+        mock_repo.update_nodes_positions.return_value = 2
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        positions = [{"nodeId": "n1", "x": 100, "y": 200}]
+        result = service.update_nodes_positions(positions)
+
+        assert result == 2
+        mock_repo.update_nodes_positions.assert_called_once_with(
+            positions=positions, sketch_id="sketch-1"
+        )
+
+
+class TestDeleteNodes:
+    def test_delete_nodes(self):
+        mock_repo = MagicMock()
+        mock_repo.delete_nodes.return_value = 3
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.delete_nodes(["id-1", "id-2", "id-3"])
+
+        assert result == 3
+        mock_repo.delete_nodes.assert_called_once_with(
+            node_ids=["id-1", "id-2", "id-3"], sketch_id="sketch-1"
+        )
+
+    def test_delete_nodes_with_in_memory(self):
+        """Test delete with InMemoryRepository to verify actual behavior."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+
+        # Create nodes
+        node1 = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+        elem_id = service.create_node(node1)
+
+        assert repo.get_node_count("sketch-1") == 1
+
+        # Delete
+        deleted = service.delete_nodes([elem_id])
+
+        assert deleted == 1
+        assert repo.get_node_count("sketch-1") == 0
+
+
+class TestDeleteRelationships:
+    def test_delete_relationships(self):
+        mock_repo = MagicMock()
+        mock_repo.delete_relationships.return_value = 2
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.delete_relationships(["rel-1", "rel-2"])
+
+        assert result == 2
+        mock_repo.delete_relationships.assert_called_once_with(
+            relationship_ids=["rel-1", "rel-2"], sketch_id="sketch-1"
+        )
+
+
+class TestDeleteAllSketchNodes:
+    def test_delete_all_sketch_nodes(self):
+        mock_repo = MagicMock()
+        mock_repo.delete_all_sketch_nodes.return_value = 10
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.delete_all_sketch_nodes()
+
+        assert result == 10
+        mock_repo.delete_all_sketch_nodes.assert_called_once_with(sketch_id="sketch-1")
+
+
+class TestUpdateRelationship:
+    def test_update_relationship(self):
+        mock_repo = MagicMock()
+        mock_repo.update_relationship.return_value = {"id": "rel-1", "weight": 5}
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.update_relationship("rel-1", {"weight": 5})
+
+        assert result == {"id": "rel-1", "weight": 5}
+        mock_repo.update_relationship.assert_called_once_with(
+            element_id="rel-1", rel_obj={"weight": 5}, sketch_id="sketch-1"
+        )
+
+
+class TestMergeNodes:
+    def test_merge_nodes(self):
+        mock_repo = MagicMock()
+        mock_repo.merge_nodes.return_value = "new-elem"
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.merge_nodes(
+            old_node_ids=["old-1", "old-2"],
+            new_node_data={"type": "domain"},
+            new_node_id="old-1",
+        )
+
+        assert result == "new-elem"
+        mock_repo.merge_nodes.assert_called_once_with(
+            old_node_ids=["old-1", "old-2"],
+            new_node_data={"type": "domain"},
+            new_node_id="old-1",
+            sketch_id="sketch-1",
+        )
+
+
+class TestBatchCreateNodes:
+    def test_batch_create_nodes(self):
+        mock_repo = MagicMock()
+        mock_repo.batch_create_nodes.return_value = {
+            "nodes_created": 2,
+            "node_ids": ["id-1", "id-2"],
+            "errors": [],
+        }
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        nodes = [{"nodeLabel": "a"}, {"nodeLabel": "b"}]
+        result = service.batch_create_nodes(nodes)
+
+        assert result["nodes_created"] == 2
+        mock_repo.batch_create_nodes.assert_called_once_with(
+            nodes=nodes, sketch_id="sketch-1"
+        )
+
+
+class TestBatchCreateEdgesByElementId:
+    def test_batch_create_edges_by_element_id(self):
+        mock_repo = MagicMock()
+        mock_repo.batch_create_edges_by_element_id.return_value = {
+            "edges_created": 1,
+            "errors": [],
+        }
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        edges = [{"from_element_id": "e1", "to_element_id": "e2"}]
+        result = service.batch_create_edges_by_element_id(edges)
+
+        assert result["edges_created"] == 1
+        mock_repo.batch_create_edges_by_element_id.assert_called_once_with(
+            edges=edges, sketch_id="sketch-1"
+        )
+
+
+class TestLogGraphMessage:
+    def test_log_graph_message_with_logger(self):
+        mock_logger = MagicMock()
+        repo = InMemoryGraphRepository()
+
+        service = GraphService(
+            sketch_id="sketch-1", repository=repo, logger=mock_logger
+        )
+        service.log_graph_message("Test message")
+
+        mock_logger.graph_append.assert_called_once_with(
+            "sketch-1", {"message": "Test message"}
+        )
+
+    def test_log_graph_message_without_logger(self):
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo, logger=None)
+        # Should not raise
+        service.log_graph_message("Test message")
+
+
+class TestFlush:
+    def test_flush_with_batching_enabled(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+        service.flush()
+
+        mock_repo.flush_batch.assert_called_once()
+
+    def test_flush_with_batching_disabled(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=False
+        )
+        service.flush()
+
+        mock_repo.flush_batch.assert_not_called()
+
+
+class TestQuery:
+    def test_query(self):
+        mock_repo = MagicMock()
+        mock_repo.query.return_value = [{"count": 5}]
+
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        result = service.query("MATCH (n) RETURN count(n)", {"param": "value"})
+
+        assert result == [{"count": 5}]
+        mock_repo.query.assert_called_once_with(
+            "MATCH (n) RETURN count(n)", {"param": "value"}
+        )
+
+
+class TestSetBatchSize:
+    def test_set_batch_size(self):
+        mock_repo = MagicMock()
+        service = GraphService(sketch_id="sketch-1", repository=mock_repo)
+        service.set_batch_size(50)
+
+        mock_repo.set_batch_size.assert_called_once_with(50)
+
+
+class TestContextManager:
+    def test_context_manager_flushes_on_success(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+
+        with service:
+            pass
+
+        mock_repo.flush_batch.assert_called_once()
+
+    def test_context_manager_does_not_flush_on_exception(self):
+        mock_repo = MagicMock()
+        service = GraphService(
+            sketch_id="sketch-1", repository=mock_repo, enable_batching=True
+        )
+
+        with pytest.raises(ValueError):
+            with service:
+                raise ValueError("Test error")
+
+        mock_repo.flush_batch.assert_not_called()
+
+
+class TestCreateGraphServiceFactory:
+    def test_create_graph_service(self):
+        mock_repo = MagicMock()
+
+        with patch(
+            "flowsint_core.core.graph.service.Neo4jGraphRepository", return_value=mock_repo
+        ) as MockRepoClass:
+            with patch("flowsint_core.core.logger.Logger") as MockLogger:
+                service = create_graph_service(
+                    sketch_id="sketch-1",
+                    enable_batching=True,
+                )
+
+                # Factory should create the repository
+                MockRepoClass.assert_called_once()
+                assert service._sketch_id == "sketch-1"
+                assert service._repository == mock_repo
+                assert service._logger == MockLogger
+                assert service._enable_batching is True
+
+    def test_create_graph_service_defaults(self):
+        mock_repo = MagicMock()
+
+        with patch(
+            "flowsint_core.core.graph.service.Neo4jGraphRepository", return_value=mock_repo
+        ):
+            with patch("flowsint_core.core.logger.Logger"):
+                service = create_graph_service(sketch_id="sketch-1")
+
+                assert service._sketch_id == "sketch-1"
+                assert service._repository == mock_repo
+                assert service._enable_batching is True  # Default is True in factory
+
+
+class TestInMemoryRepositoryIntegration:
+    """Integration tests using InMemoryGraphRepository to verify full workflows."""
+
+    def test_full_workflow_create_and_query(self):
+        """Test a complete workflow: create nodes, create relationships, query."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(sketch_id="sketch-1", repository=repo)
+
+        # Create two nodes
+        node1 = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+        node2 = GraphNode(
+            id="2",
+            nodeLabel="1.1.1.1",
+            nodeType="ip",
+            nodeProperties=Ip(address="1.1.1.1"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        elem_id1 = service.create_node(node1)
+        elem_id2 = service.create_node(node2)
+
+        assert elem_id1 is not None
+        assert elem_id2 is not None
+        assert repo.get_node_count("sketch-1") == 2
+
+        # Create a relationship
+        service.create_relationship(
+            Domain(domain="example.com"),
+            Ip(address="1.1.1.1"),
+            "RESOLVES_TO",
+        )
+
+        assert repo.get_edge_count("sketch-1") == 1
+
+        # Query the graph
+        graph = service.get_sketch_graph()
+        assert len(graph.nodes) == 2
+        assert len(graph.edges) == 1
+
+    def test_batch_operations_with_in_memory(self):
+        """Test batch operations work correctly with InMemoryRepository."""
+        repo = InMemoryGraphRepository()
+        service = GraphService(
+            sketch_id="sketch-1", repository=repo, enable_batching=True
+        )
+
+        # Add nodes to batch
+        for i in range(3):
+            node = GraphNode(
+                id=str(i),
+                nodeLabel=f"node{i}.com",
+                nodeType="domain",
+                nodeProperties=Domain(domain=f"node{i}.com"),
+                nodeMetadata=NodeMetadata(),
+            )
+            service.create_node(node)
+
+        # Nodes should not be created yet (batched)
+        assert repo.get_node_count("sketch-1") == 0
+
+        # Flush the batch
+        service.flush()
+
+        # Now nodes should exist
+        assert repo.get_node_count("sketch-1") == 3
+
+    def test_sketch_isolation(self):
+        """Test that different sketches are isolated."""
+        repo = InMemoryGraphRepository()
+        service1 = GraphService(sketch_id="sketch-1", repository=repo)
+        service2 = GraphService(sketch_id="sketch-2", repository=repo)
+
+        # Create nodes in different sketches
+        node = GraphNode(
+            id="1",
+            nodeLabel="example.com",
+            nodeType="domain",
+            nodeProperties=Domain(domain="example.com"),
+            nodeMetadata=NodeMetadata(),
+        )
+
+        service1.create_node(node)
+        service2.create_node(node)
+
+        # Each sketch should have its own node
+        assert repo.get_node_count("sketch-1") == 1
+        assert repo.get_node_count("sketch-2") == 1
+        assert repo.get_node_count() == 2
+
+        # Delete all from sketch-1
+        service1.delete_all_sketch_nodes()
+
+        assert repo.get_node_count("sketch-1") == 0
+        assert repo.get_node_count("sketch-2") == 1

flowsint-core/tests/core/test_enricher_base_simplified_api.py

@@ -1,11 +1,13 @@
 """Test simplified API for create_node and create_relationship."""
 
+from typing import List
+
 import pytest
-from flowsint_core.core.enricher_base import Enricher
 from flowsint_types.domain import Domain
 from flowsint_types.email import Email
 from flowsint_types.individual import Individual
-from typing import List
+
+from flowsint_core.core.enricher_base import Enricher
 
 
 class MockEnricher(Enricher):
@@ -49,4 +51,4 @@ def test_create_node_with_property_override():
     domain = Domain(domain="example.com")
 
     # Should be able to override properties
-    enricher.create_node(domain, type="subdomain")
+    enricher.create_node(domain)

flowsint-core/tests/core/test_graph_repository.py

@@ -1,260 +0,0 @@
-"""Test GraphRepository batch operations."""
-
-import pytest
-from unittest.mock import Mock, MagicMock
-from flowsint_core.core.graph_repository import GraphRepository
-from flowsint_types.domain import Domain
-from flowsint_types.ip import Ip
-from flowsint_types.email import Email
-
-
-@pytest.fixture
-def mock_connection():
-    """Create a mock Neo4j connection."""
-    conn = Mock()
-    conn.execute_batch = Mock(return_value=[])
-    conn.query = Mock(return_value=[])
-    return conn
-
-
-@pytest.fixture
-def graph_repo(mock_connection):
-    """Create a GraphRepository with mocked connection."""
-    return GraphRepository(neo4j_connection=mock_connection)
-
-
-def test_batch_create_nodes_empty_list(graph_repo):
-    """Test batch_create_nodes with empty list."""
-    result = graph_repo.batch_create_nodes(nodes=[], sketch_id="test-sketch")
-
-    assert result["nodes_created"] == 0
-    assert result["node_ids"] == []
-    assert result["errors"] == []
-
-
-def test_batch_create_nodes_single_node(graph_repo, mock_connection):
-    """Test batch_create_nodes with a single node."""
-    domain = Domain(domain="example.com")
-
-    # Mock the batch execution to return a node ID
-    mock_connection.execute_batch.return_value = [
-        [{"id": "element-id-123"}]
-    ]
-
-    result = graph_repo.batch_create_nodes(
-        nodes=[domain],
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 1
-    assert len(result["node_ids"]) == 1
-    assert result["node_ids"][0] == "element-id-123"
-    assert result["errors"] == []
-
-    # Verify execute_batch was called once
-    assert mock_connection.execute_batch.call_count == 1
-
-
-def test_batch_create_nodes_multiple_nodes(graph_repo, mock_connection):
-    """Test batch_create_nodes with multiple nodes of different types."""
-    domain = Domain(domain="example.com")
-    ip = Ip(address="192.168.1.1")
-    email = Email(email="test@example.com")
-
-    # Mock the batch execution to return multiple node IDs
-    mock_connection.execute_batch.return_value = [
-        [{"id": "element-id-1"}],
-        [{"id": "element-id-2"}],
-        [{"id": "element-id-3"}]
-    ]
-
-    result = graph_repo.batch_create_nodes(
-        nodes=[domain, ip, email],
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 3
-    assert len(result["node_ids"]) == 3
-    assert result["errors"] == []
-
-    # Verify execute_batch was called once with 3 operations
-    assert mock_connection.execute_batch.call_count == 1
-    batch_operations = mock_connection.execute_batch.call_args[0][0]
-    assert len(batch_operations) == 3
-
-
-def test_batch_create_nodes_with_validation_errors(graph_repo, mock_connection):
-    """Test batch_create_nodes when some nodes have validation errors."""
-    valid_domain = Domain(domain="example.com")
-
-    # Mock execute_batch to succeed for valid nodes
-    mock_connection.execute_batch.return_value = [
-        [{"id": "element-id-1"}]
-    ]
-
-    result = graph_repo.batch_create_nodes(
-        nodes=[valid_domain],
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 1
-    assert len(result["errors"]) == 0
-
-
-def test_batch_create_nodes_batch_execution_failure(graph_repo, mock_connection):
-    """Test batch_create_nodes when the batch execution fails."""
-    domain = Domain(domain="example.com")
-
-    # Mock execute_batch to raise an exception
-    mock_connection.execute_batch.side_effect = Exception("Database connection error")
-
-    result = graph_repo.batch_create_nodes(
-        nodes=[domain],
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 0
-    assert result["node_ids"] == []
-    assert len(result["errors"]) > 0
-    assert "Database connection error" in result["errors"][0]
-
-
-def test_batch_create_nodes_no_connection():
-    """Test batch_create_nodes when there's no database connection."""
-    # Create a repository with no connection by manually setting it to None
-    repo = GraphRepository(neo4j_connection=Mock())
-    repo._connection = None  # Force no connection
-    domain = Domain(domain="example.com")
-
-    result = repo.batch_create_nodes(
-        nodes=[domain],
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 0
-    assert result["node_ids"] == []
-    assert len(result["errors"]) == 1
-    assert "No database connection" in result["errors"][0]
-
-
-def test_batch_create_nodes_with_label_fallback(graph_repo, mock_connection):
-    """Test batch_create_nodes with nodes that need label fallback."""
-    from pydantic import BaseModel
-
-    class CustomNode(BaseModel):
-        """Node type without a clear primary field."""
-        name: str = ""
-        label: str = "Custom Node"
-
-    node = CustomNode(name="", label="Fallback Label")
-
-    # Mock the batch execution
-    mock_connection.execute_batch.return_value = [
-        [{"id": "element-id-fallback"}]
-    ]
-
-    result = graph_repo.batch_create_nodes(
-        nodes=[node],
-        sketch_id="test-sketch"
-    )
-
-    # Should succeed with fallback to label
-    assert result["nodes_created"] == 1
-    assert len(result["node_ids"]) == 1
-
-
-def test_batch_create_nodes_large_batch(graph_repo, mock_connection):
-    """Test batch_create_nodes with a large number of nodes."""
-    # Create 1000 domains
-    domains = [Domain(domain=f"example{i}.com") for i in range(1000)]
-
-    # Mock the batch execution to return 1000 node IDs
-    mock_connection.execute_batch.return_value = [
-        [{"id": f"element-id-{i}"}] for i in range(1000)
-    ]
-
-    result = graph_repo.batch_create_nodes(
-        nodes=domains,
-        sketch_id="test-sketch"
-    )
-
-    assert result["nodes_created"] == 1000
-    assert len(result["node_ids"]) == 1000
-    assert result["errors"] == []
-
-    # Verify execute_batch was called once (single transaction)
-    assert mock_connection.execute_batch.call_count == 1
-    batch_operations = mock_connection.execute_batch.call_args[0][0]
-    assert len(batch_operations) == 1000
-
-
-# Tests for update_node
-
-
-def test_update_node_success(graph_repo, mock_connection):
-    """Test update_node with a valid Pydantic object."""
-    domain = Domain(domain="updated-example.com", label="Updated Domain")
-
-    # Mock the query to return the element ID
-    mock_connection.query.return_value = [{"id": "element-id-123"}]
-
-    result = graph_repo.update_node(
-        element_id="element-id-123",
-        node_obj=domain,
-        sketch_id="test-sketch"
-    )
-
-    assert result == "element-id-123"
-    assert mock_connection.query.call_count == 1
-
-
-def test_update_node_not_found(graph_repo, mock_connection):
-    """Test update_node when node doesn't exist."""
-    domain = Domain(domain="example.com")
-
-    # Mock the query to return empty result
-    mock_connection.query.return_value = []
-
-    result = graph_repo.update_node(
-        element_id="non-existent-id",
-        node_obj=domain,
-        sketch_id="test-sketch"
-    )
-
-    assert result is None
-
-
-def test_update_node_no_connection():
-    """Test update_node when there's no database connection."""
-    repo = GraphRepository(neo4j_connection=Mock())
-    repo._connection = None
-    domain = Domain(domain="example.com")
-
-    result = repo.update_node(
-        element_id="element-id-123",
-        node_obj=domain,
-        sketch_id="test-sketch"
-    )
-
-    assert result is None
-
-
-def test_update_node_different_types(graph_repo, mock_connection):
-    """Test update_node with different Pydantic types."""
-    ip = Ip(address="10.0.0.1", label="Updated IP")
-
-    # Mock the query
-    mock_connection.query.return_value = [{"id": "element-id-456"}]
-
-    result = graph_repo.update_node(
-        element_id="element-id-456",
-        node_obj=ip,
-        sketch_id="test-sketch"
-    )
-
-    assert result == "element-id-456"
-
-    # Verify the query contains the correct type
-    query_call = mock_connection.query.call_args
-    params = query_call[0][1]
-    assert params["type"] == "ip"

flowsint-core/tests/utils/flatten.py

@@ -0,0 +1,38 @@
+from flowsint_core.utils import flatten, unflatten
+
+my_dict = {
+    "root_key_1": "value 1",
+    "root_key_2": 2,
+    "root_key_3": "value 3",
+    "root_key_4": {
+        "child_key_1": "child 1",
+        "child_key_2": "child 2",
+        "child_key_3": {"grand_child_1": 0},
+    },
+}
+
+my_flat_dict = {
+    "root_key_1": "value 1",
+    "root_key_2": 2,
+    "root_key_3": "value 3",
+    "root_key_4.child_key_1": "child 1",
+    "root_key_4.child_key_2": "child 2",
+    "root_key_4.child_key_3.grand_child_1": 0,
+}
+
+my_flat_dict_other_separator = {
+    "root_key_1": "value 1",
+    "root_key_2": 2,
+    "root_key_3": "value 3",
+    "root_key_4_child_key_1": "child 1",
+    "root_key_4_child_key_2": "child 2",
+    "root_key_4_child_key_3_grand_child_1": 0,
+}
+
+
+def test_flatten():
+    assert flatten(my_dict) == my_flat_dict
+
+
+def test_unflatten():
+    assert unflatten(my_flat_dict) == my_dict