memgraph-toolbox 0.2.0__tar.gz → 0.2.2__tar.gz

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/.gitignore

@@ -181,12 +181,9 @@ cython_debug/
 # Local agent hook configuration
 /.codex/
 /.claude/settings.local.json
-/AGENTS.md
-/docs/agents/
-/context-graph/CONTEXT-MAP.md
-/context-graph/*/CONTEXT.md
-/context-graph/docs/adr/
-/context-graph/*/docs/adr/
 
 # Project specific files
 /enterprise-context/sic-agent/sic-scrapper/output/*
+
+# deepeval writes run caches and telemetry here; build artifacts, not source.
+.deepeval/

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: memgraph-toolbox
-Version: 0.2.0
+Version: 0.2.2
 Summary: Memgraph toolbox library for Memgraph AI tools and utilities
 Project-URL: Homepage, https://github.com/memgraph/ai-toolkit
 Project-URL: Source, https://github.com/memgraph/ai-toolkit

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "memgraph-toolbox"
-version = "0.2.0"
+version = "0.2.2"
 description = "Memgraph toolbox library for Memgraph AI tools and utilities"
 readme = "README.md"
 authors = [

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/api/memgraph.py

@@ -1,10 +1,11 @@
 import os
+import time
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, cast
 
 from neo4j import AsyncGraphDatabase, GraphDatabase
 
-from ..utils.serialization import serialize_record_data
+from ..utils.serialization import serialize_records
 
 MEMGRAPH_ENV_DEFAULTS = {
     "MEMGRAPH_URL": "bolt://localhost:7687",
@@ -78,6 +79,39 @@ def memgraph_env(
     }
 
 
+def merge_with_retry(db: Any, query: str, params: dict | None = None, *, attempts: int = 3) -> list[dict[str, Any]]:
+    """Run *query*, a MERGE of a uniqueness-constrained node, through *db*, surviving a concurrent create.
+
+    Two transactions that MERGE the same new node both find it missing and
+    both create it; the later commit fails the uniqueness constraint. Run
+    again, its MERGE matches the node the other one committed. Writers that
+    race this way are typically one process per event (agent hooks).
+
+    Args:
+        db: Anything with ``query(query, params)``, normally a :class:`Memgraph`.
+        query: The Cypher to run; it must be safe to run more than once.
+        params: Its parameters.
+        attempts: Runs before giving up, the first included.
+
+    Returns:
+        The rows of the run that committed.
+
+    Raises:
+        neo4j.exceptions.ClientError: a uniqueness violation on the last
+            attempt, or any other client error at once.
+    """
+    from neo4j.exceptions import ClientError
+
+    for attempt in range(attempts):
+        try:
+            return db.query(query, params=params)
+        except ClientError as exc:
+            if "unique constraint" not in str(exc) or attempt == attempts - 1:
+                raise
+            time.sleep(0.05 * (attempt + 1))
+    raise ValueError(f"attempts must be at least 1, got {attempts}")
+
+
 class Memgraph:
     """
     Base Memgraph client for interacting with Memgraph database.
@@ -87,10 +121,10 @@ class Memgraph:
 
     def __init__(
         self,
-        url: str = None,
-        username: str = None,
-        password: str = None,
-        database: str = None,
+        url: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        database: str | None = None,
         driver_config: dict | None = None,
         user_agent: str | None = None,
     ):
@@ -140,37 +174,43 @@ class Memgraph:
                 f"Could not connect to Memgraph database. Authentication failed for user '{username}'"
             ) from e
 
-    def query(self, query: str, params: dict = None) -> list[dict[str, Any]]:
+    def query(self, query: str, params: dict | None = None) -> list[dict[str, Any]]:
         """
-        Execute a Cypher query and return results as a list of dictionaries.
+        Execute a Cypher query and return type-preserving result rows.
+
+        Each value keeps its type: nodes, relationships and paths become
+        ``_type``-tagged objects (id, labels/type, endpoints, properties);
+        primitives, lists, maps, temporals and points keep their shape.
 
         Args:
             query: The Cypher query to execute
+            params: Optional query parameters
 
         Returns:
-            List of dictionaries containing query results
+            List of dictionaries, one per result row
         """
         from neo4j.exceptions import Neo4jError
 
         if params is None:
             params = {}
         try:
-            data, _, _ = self.driver.execute_query(
-                query,
+            # neo4j's driver stubs require query_: LiteralString | Query, a hint aimed at
+            # discouraging unparameterized query construction. This is a generic passthrough
+            # executor that legitimately accepts caller-built Cypher text (injection is
+            # already guarded by `params`), which a `str` can never satisfy statically.
+            records, _, _ = self.driver.execute_query(
+                cast("Any", query),
                 parameters_=params,
                 database_=self.database,
             )
-            json_data = [serialize_record_data(r.data()) for r in data]
-            return json_data
+            return serialize_records(records)
         except Neo4jError as e:
             if not _is_implicit_transaction_fallback_error(e):
                 raise
 
         # fallback to allow implicit transactions
         with self.driver.session(database=self.database) as session:
-            data = session.run(query, params)
-            json_data = [serialize_record_data(r.data()) for r in data]
-            return json_data
+            return serialize_records(session.run(cast("Any", query), params))
 
     def close(self) -> None:
         """
@@ -191,10 +231,10 @@ class AsyncMemgraph:
 
     def __init__(
         self,
-        url: str = None,
-        username: str = None,
-        password: str = None,
-        database: str = None,
+        url: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        database: str | None = None,
         driver_config: dict | None = None,
         user_agent: str | None = None,
     ):
@@ -236,41 +276,45 @@ class AsyncMemgraph:
         """
         await self.driver.verify_connectivity()
 
-    async def query(self, query: str, params: dict = None) -> list[dict[str, Any]]:
+    async def query(self, query: str, params: dict | None = None) -> list[dict[str, Any]]:
         """
-        Execute a Cypher query and return results as a list of dictionaries.
+        Execute a Cypher query and return type-preserving result rows.
 
-        Mirrors :meth:`Memgraph.query`, including the fallback to an implicit
-        transaction for statements that cannot run inside a managed transaction.
+        Async mirror of :meth:`Memgraph.query`, including the fallback to an
+        implicit transaction for statements that cannot run inside a managed
+        transaction.
 
         Args:
             query: The Cypher query to execute
             params: Optional query parameters
 
         Returns:
-            List of dictionaries containing query results
+            List of dictionaries, one per result row
         """
         from neo4j.exceptions import Neo4jError
 
         if params is None:
             params = {}
         try:
-            data, _, _ = await self.driver.execute_query(
-                query,
+            # See the matching comment in Memgraph.query: query_ is typed LiteralString |
+            # Query in the driver stubs, but this is a generic passthrough executor that
+            # legitimately accepts caller-built Cypher text -- a `str` can never satisfy
+            # LiteralString statically, regardless of how it was validated at runtime.
+            records, _, _ = await self.driver.execute_query(
+                cast("Any", query),
                 parameters_=params,
                 database_=self.database,
             )
-            json_data = [serialize_record_data(r.data()) for r in data]
-            return json_data
+            return serialize_records(records)
         except Neo4jError as e:
             if not _is_implicit_transaction_fallback_error(e):
                 raise
 
         # fallback to allow implicit transactions
         async with self.driver.session(database=self.database) as session:
-            result = await session.run(query, params)
-            json_data = [serialize_record_data(r.data()) async for r in result]
-            return json_data
+            result = await session.run(cast("Any", query), params)
+            records = [record async for record in result]
+            return serialize_records(records)
 
     async def close(self) -> None:
         """

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/api/tool.py

@@ -17,7 +17,7 @@ class BaseTool(ABC):
         self.input_schema = input_schema
 
     @abstractmethod
-    def call(self, arguments: dict[str, Any]) -> list[Any]:
+    def call(self, arguments: dict[str, Any]) -> list[Any] | dict[str, Any]:
         """
         Execute the tool with the provided arguments.
 
@@ -25,7 +25,9 @@ class BaseTool(ABC):
             arguments (dict): A dictionary of arguments as defined by the input schema.
 
         Returns:
-            List containing one or more output values.
+            Either a list containing one or more output values, or -- for tools whose
+            "found" result is naturally a single object (e.g. schema lookups) -- a bare
+            dict.
         """
         pass
 

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/client/mcp_prompt.py

@@ -3,6 +3,7 @@ import json
 import logging
 import os
 import sys
+from typing import cast
 
 try:
     from litellm import acompletion, experimental_mcp_client
@@ -19,9 +20,9 @@ logger = logging.getLogger(__name__)
 async def prompt_with_tools(
     prompt: str,
     tool_selection_model_name: str = "openai/gpt-4o",
-    tool_selection_system_messages: list[dict[str, str]] = None,
+    tool_selection_system_messages: list[dict[str, str]] | None = None,
     response_model_name: str = "openai/gpt-4o",
-    response_system_messages: list[dict[str, str]] = None,
+    response_system_messages: list[dict[str, str]] | None = None,
     mcp_server_params: StdioServerParameters | None = None,
 ):
     """
@@ -29,7 +30,6 @@ async def prompt_with_tools(
     First there is a tool selection phase where the LLM is asked to select the tools to use.
     Second there is a response phase where the LLM is asked to generate a response.
     """
-    # Configure MCP server connection.
     if response_system_messages is None:
         response_system_messages = []
     if tool_selection_system_messages is None:
@@ -48,18 +48,17 @@ async def prompt_with_tools(
             ],
             env={
                 "MEMGRAPH_URL": os.environ.get("MEMGRAPH_URL", "bolt://localhost:7687"),
-                # NOTE: If you want to control where uv pulls the dependencies from inject UV_PYTHON or VIRTUAL_ENV.
+                # Set UV_PYTHON or VIRTUAL_ENV when the MCP subprocess must use a specific environment.
             },
         )
     else:
         server_params = mcp_server_params
 
-    # Establish MCP server connection.
     async with stdio_client(server_params) as (read, write), ClientSession(read, write) as session:
         await session.initialize()
         # https://docs.litellm.ai/docs/mcp#litellm-python-sdk-mcp-bridge
         tools = await experimental_mcp_client.load_mcp_tools(
-            # NOTE: format="openai" works with everything tested, "mcp" does NOT work.
+            # LiteLLM's OpenAI format is required for the tool-call shape consumed below.
             session=session,
             format="openai",
         )
@@ -69,7 +68,6 @@ async def prompt_with_tools(
         )
         logger.info(f"{len(tools)} MCP tools loaded")
 
-        # Tool selection by an LLM.
         messages = [*tool_selection_system_messages, {"role": "user", "content": prompt}]
         resp = await acompletion(
             model=tool_selection_model_name,
@@ -77,7 +75,6 @@ async def prompt_with_tools(
             tools=tools,
         )
         msg = resp["choices"][0]["message"]
-        # Tool calls by the MCP server.
         tool_calls = msg.get("tool_calls", [])
         if not tool_calls:
             logger.info("LLM Response: %s", msg.get("content", ""))
@@ -89,7 +86,7 @@ async def prompt_with_tools(
                     tc["function"]["name"],
                     tc["function"].get("description", "No description"),
                 )
-            # Ensure all tool call IDs are within OpenAI's 40-character limit
+            # OpenAI limits tool-call IDs to 40 characters.
             for tc in tool_calls:
                 if len(tc["id"]) > 40:
                     logger.warning(
@@ -113,10 +110,11 @@ async def prompt_with_tools(
                     result = await session.call_tool(name=tc["function"]["name"], arguments=arguments)
                     if hasattr(result, "content") and result.content:
                         if isinstance(result.content, list):
-                            content_text = []
+                            content_text: list[str] = []
                             for content_item in result.content:
                                 if hasattr(content_item, "text"):
-                                    content_text.append(content_item.text)
+                                    # Only TextContent exposes `.text`; other content blocks are serialized below.
+                                    content_text.append(cast("str", content_item.text))
                                 elif isinstance(content_item, str):
                                     content_text.append(content_item)
                                 else:
@@ -148,7 +146,6 @@ async def prompt_with_tools(
                         tool_error,
                     )
 
-            # Generating the final response using LLM.
             messages = messages + response_system_messages
             final_resp = await acompletion(
                 model=response_model_name,
@@ -177,7 +174,6 @@ async def __async_context():
 
 
 if __name__ == "__main__":
-    # NOTE: Chaning how mcp-memgraph logs from here is not easy because the
-    # server is started in another process.
+    # The MCP server runs in a separate process, so configure its client-side logger here.
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
     asyncio.run(__async_context())

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/evals/coherence.py

@@ -110,8 +110,11 @@ class CoherenceEmbeddingsBasedMetric(BaseMetric):
         Returns:
             Array of embeddings with shape (n_sentences, embedding_dim)
         """
-        # TODO(gitbuda): This should use the utils.embeddings because GPU could be used.
-        return self.model.encode(sentences)
+        # Keep this local model call until the shared embedding utility supports this metric's API.
+        # convert_to_numpy is already the model's default; pinning it explicitly picks the
+        # np.ndarray-returning overload instead of leaving overload resolution to fall
+        # through to a Tensor-returning one.
+        return self.model.encode(sentences, convert_to_numpy=True)
 
     def _compute_cosine_similarity(self, embeddings: np.ndarray) -> list[float]:
         """Compute cosine similarity between consecutive embeddings.
@@ -221,7 +224,9 @@ class CoherenceEmbeddingsBasedMetric(BaseMetric):
         Returns:
             True if score meets threshold, False otherwise
         """
-        return self.success
+        # self.success is bool | None (None until measure() has run); treat "never measured"
+        # as "not successful" rather than lying about the return type or leaking None.
+        return self.success or False
 
     @property
     def __name__(self) -> str:

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/memgraph_toolbox.py

@@ -1,7 +1,5 @@
 from .api.memgraph import Memgraph
 from .api.toolbox import BaseToolbox
-
-# Import all tool classes
 from .tools.cypher import CypherTool
 from .tools.schema import EnumSchemaTool, NodeSchemaTool, RelationshipSchemaTool, SearchSchemaTool
 

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/tests/test_async_memgraph.py

@@ -42,12 +42,12 @@ async def test_async_memgraph_query_live():
     """Live round-trip; skips cleanly when no server is reachable."""
     import neo4j
 
-    client = AsyncMemgraph(url="bolt://localhost:7687", username="", password="")
+    client = AsyncMemgraph()
     try:
         await client.verify_connectivity()
     except (neo4j.exceptions.ServiceUnavailable, neo4j.exceptions.SessionExpired, OSError):
         await client.close()
-        pytest.skip("No Memgraph server reachable at bolt://localhost:7687")
+        pytest.skip("No Memgraph server reachable at the configured MEMGRAPH_URL")
 
     try:
         result = await client.query("RETURN 1 AS value")

memgraph_toolbox-0.2.2/src/memgraph_toolbox/tests/test_merge_with_retry.py

@@ -0,0 +1,69 @@
+"""Tests for merge_with_retry against a live Memgraph; they skip when none is reachable."""
+
+from concurrent.futures import ThreadPoolExecutor
+
+import pytest
+from neo4j.exceptions import ClientError
+
+from ..api.memgraph import Memgraph, merge_with_retry
+
+_LABEL = "MergeWithRetryProbe"
+
+
+@pytest.fixture
+def db():
+    try:
+        client = Memgraph()
+    except Exception as exc:
+        pytest.skip(f"No Memgraph server reachable at the configured MEMGRAPH_URL: {exc}")
+    client.query(f"MATCH (n:{_LABEL}) DETACH DELETE n")
+    client.query(f"CREATE CONSTRAINT ON (n:{_LABEL}) ASSERT n.key IS UNIQUE;")
+    yield client
+    client.query(f"MATCH (n:{_LABEL}) DETACH DELETE n")
+    client.query(f"DROP CONSTRAINT ON (n:{_LABEL}) ASSERT n.key IS UNIQUE;")
+    client.close()
+
+
+def test_concurrent_merges_of_new_nodes_converge_on_one_node_each(db):
+    keys = [f"k{i}" for i in range(10)]
+
+    def merge(key: str) -> None:
+        merge_with_retry(db, f"MERGE (n:{_LABEL} {{key: $key}})", {"key": key})
+
+    with ThreadPoolExecutor(max_workers=16) as pool:
+        list(pool.map(merge, keys * 16))
+
+    rows = db.query(f"MATCH (n:{_LABEL}) RETURN n.key AS key, count(*) AS c ORDER BY key")
+    assert rows == [{"key": key, "c": 1} for key in sorted(keys)]
+
+
+def test_returns_the_rows_of_the_run_that_committed(db):
+    rows = merge_with_retry(db, f"MERGE (n:{_LABEL} {{key: 'a'}}) RETURN n.key AS key")
+    assert rows == [{"key": "a"}]
+
+
+def test_any_other_client_error_is_raised_without_retrying(db):
+    calls = []
+
+    class Counting:
+        def query(self, query, params=None):
+            calls.append(query)
+            return db.query(query, params)
+
+    with pytest.raises(ClientError):
+        merge_with_retry(Counting(), "MERGE (n:")
+    assert len(calls) == 1
+
+
+def test_a_conflict_on_every_attempt_is_raised_after_the_last():
+    """A conflict that never clears can't be produced on demand by a real server."""
+    calls = []
+
+    class AlwaysConflicting:
+        def query(self, query, params=None):
+            calls.append(query)
+            raise ClientError("Unable to commit due to unique constraint violation on :X(key)")
+
+    with pytest.raises(ClientError, match="unique constraint"):
+        merge_with_retry(AlwaysConflicting(), "MERGE (n:X {key: 1})", attempts=2)
+    assert len(calls) == 2

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/tests/test_toolbox.py

@@ -40,11 +40,7 @@ def test_base_toolbox():
 def test_memgraph_toolbox():
     """Test the Memgraph Toolbox."""
 
-    url = "bolt://localhost:7687"
-    user = ""
-    password = ""
-
-    memgraph_client = Memgraph(url=url, username=user, password=password)
+    memgraph_client = Memgraph()
 
     toolkit = MemgraphToolbox(db=memgraph_client)
 

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/tests/test_tools.py

@@ -14,7 +14,7 @@ from ..tools.schema import (
 
 @pytest.fixture()
 def db():
-    client = Memgraph(url="bolt://localhost:7687", username="", password="")
+    client = Memgraph()
     yield client
     client.close()
 
@@ -72,6 +72,71 @@ def test_cypher_date_time_serialization(db):
     assert "PT2M2.33S" in record["test_duration"]
 
 
+def test_query_tags_graph_entities(db, schema_graph):
+    """Nodes, relationships and paths are returned as _type-tagged objects."""
+    rows = db.query("MATCH p = (a:Person {name: 'Alice'})-[r:KNOWS]->(b) RETURN a, r, b, p")
+
+    assert len(rows) == 1
+    row = rows[0]
+
+    assert row["a"]["_type"] == "node"
+    assert "Person" in row["a"]["labels"]
+    assert row["a"]["properties"]["name"] == "Alice"
+    assert "id" in row["a"]
+
+    assert row["r"]["_type"] == "relationship"
+    assert row["r"]["type"] == "KNOWS"
+    assert row["r"]["start"] == row["a"]["id"]
+    assert row["r"]["end"] == row["b"]["id"]
+
+    assert row["p"]["_type"] == "path"
+    assert len(row["p"]["nodes"]) == 2
+    assert len(row["p"]["relationships"]) == 1
+
+
+def test_query_mixed_row(db, schema_graph):
+    """A row can mix a typed node with a scalar aggregation in the same result."""
+    rows = db.query("MATCH (n:Person)-[e]->() RETURN n, count(e) AS out_deg")
+
+    assert len(rows) > 0
+    row = rows[0]
+    assert row["n"]["_type"] == "node"
+    assert isinstance(row["out_deg"], int)
+
+
+def test_query_primitives_and_containers(db):
+    """Primitives, lists, maps and temporals pass through with their shape."""
+    rows = db.query("RETURN 1 AS i, 'x' AS s, [1, 2] AS lst, {k: 3} AS m, date('2024-01-15') AS d")
+
+    row = rows[0]
+    assert row["i"] == 1
+    assert row["s"] == "x"
+    assert row["lst"] == [1, 2]
+    assert row["m"] == {"k": 3}
+    assert row["d"] == "2024-01-15"
+
+
+def test_query_point(db):
+    """A geo/point value is serialized to a tagged srid + coordinates record."""
+    rows = db.query("RETURN point({x: 1.5, y: 2.5}) AS p")
+
+    point = rows[0]["p"]
+    assert point["_type"] == "point"
+    assert "srid" in point
+    assert point["coordinates"] == [1.5, 2.5]
+
+
+def test_query_temporals(db):
+    """Temporal types serialize to ISO strings, zoned datetimes keeping the offset."""
+    rows = db.query("RETURN datetime('2024-01-15T10:30:45+01:00') AS dt, duration('PT2M') AS dur")
+
+    row = rows[0]
+    assert "2024-01-15T10:30:45" in row["dt"]
+    assert "+01:00" in row["dt"]
+    assert isinstance(row["dur"], str)
+    assert "PT2M" in row["dur"]
+
+
 def test_get_node_schema_found(db, schema_graph):
     tool = NodeSchemaTool(db=db)
     result = tool.call({"node_labels": ["Person"]})

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/tools/schema.py

@@ -33,7 +33,7 @@ class NodeSchemaTool(BaseTool):
         )
         self.db = db
 
-    def call(self, arguments: dict[str, Any]) -> list[dict[str, Any]]:
+    def call(self, arguments: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
         try:
             node_labels = arguments["node_labels"]
             schema = _get_schema(self.db)
@@ -96,7 +96,7 @@ class RelationshipSchemaTool(BaseTool):
         )
         self.db = db
 
-    def call(self, arguments: dict[str, Any]) -> list[dict[str, Any]]:
+    def call(self, arguments: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
         try:
             edge_type = arguments["relationship_type"]
             start_node_labels = arguments["start_node_labels"]
@@ -138,7 +138,7 @@ class EnumSchemaTool(BaseTool):
         )
         self.db = db
 
-    def call(self, arguments: dict[str, Any]) -> list[dict[str, Any]]:
+    def call(self, arguments: dict[str, Any]) -> dict[str, Any] | list[dict[str, Any]]:
         try:
             enum_name = arguments["enum_name"]
             schema = _get_schema(self.db)

{memgraph_toolbox-0.2.0 → memgraph_toolbox-0.2.2}/src/memgraph_toolbox/utils/embedding.py

@@ -7,7 +7,6 @@ except ImportError as e:
     ) from e
 
 
-# NOTE: HF_TOKEN has to be set in the environment variables.
 def get_sentence_transformer_model(
     name: str = "all-MiniLM-L6-v2",
     device: str | None = None,
@@ -26,7 +25,6 @@ def get_sentence_transformer_model(
         SentenceTransformer model. Always returns the underlying model, not DataParallel
         wrapper because then it's possible to use the whole SentenceTransformer API.
     """
-    # Auto-detect device if not specified
     if device is None:
         if torch.cuda.is_available():
             device = "cuda"
@@ -35,13 +33,11 @@ def get_sentence_transformer_model(
         else:
             device = "cpu"
     model = SentenceTransformer(name, device=device)
-    # Wrap with DataParallel if requested and CUDA is available
     if use_data_parallel and torch.cuda.is_available():
         if gpu_ids is None:
             gpu_ids = list(range(torch.cuda.device_count()))
-        # Create DataParallel wrapper but return the underlying model
         parallel_model = torch.nn.DataParallel(model, device_ids=gpu_ids)
-        # Return the underlying model to avoid DataParallel wrapper issues
+        # Keep the public API as SentenceTransformer rather than exposing the wrapper.
         return parallel_model.module
     return model
 
@@ -57,13 +53,11 @@ if __name__ == "__main__":
     print(f"  MPS: {hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()}")
     print("  CPU: True (always available)")
 
-    # Example usage
     print("\nExample usage:")
     model = get_sentence_transformer_model()
     print(f"Model device: {get_model_device(model)}")
     print(f"Model type: {type(model).__name__}")
 
-    # Example of using the model directly (no wrapper needed)
     print("\nDirect usage example:")
     sample_texts = ["Hello world", "Graph databases are fast"]
     print(f"Sample texts: {sample_texts}")
@@ -73,7 +67,6 @@ if __name__ == "__main__":
     print("  - model.encode(texts, normalize_embeddings=True)")
     print("  - model.encode(texts, show_progress_bar=True)")
 
-    # Test with DataParallel (if multiple GPUs available)
     if torch.cuda.is_available() and torch.cuda.device_count() > 1:
         print("\nTesting DataParallel:")
         parallel_model = get_sentence_transformer_model(use_data_parallel=True)

memgraph_toolbox-0.2.2/src/memgraph_toolbox/utils/serialization.py

@@ -0,0 +1,160 @@
+"""
+Serialization utilities for Neo4j date/time types.
+"""
+
+from collections.abc import Iterable
+from typing import Any
+
+from neo4j.graph import Node, Path, Relationship
+from neo4j.spatial import Point
+
+
+def serialize_neo4j_types(value: Any) -> Any:
+    """
+    Convert Neo4j date/time types to JSON-serializable strings.
+
+    Args:
+        value: The value to serialize
+
+    Returns:
+        The serialized value with Neo4j types converted to strings
+    """
+    try:
+        class_name = value.__class__.__name__
+        module_name = getattr(value.__class__, "__module__", "")
+
+        # Check if it's a Neo4j temporal type
+        if "neo4j" in module_name and hasattr(value, "iso_format"):
+            neo4j_temporal_types = [
+                "Date",
+                "Time",
+                "DateTime",
+                "LocalTime",
+                "LocalDateTime",
+            ]
+            if class_name in neo4j_temporal_types:
+                return value.iso_format()
+
+        # Handle Neo4j Duration type
+        if "neo4j" in module_name and class_name == "Duration":
+            return str(value)
+
+    except (AttributeError, TypeError):
+        # If we can't access the class name or module, return as-is
+        pass
+
+    return value
+
+
+def serialize_record_data(record_data: dict) -> dict:
+    """
+    Serialize a single record's data, handling Neo4j types recursively.
+
+    Args:
+        record_data: Dictionary from a Neo4j record.data()
+
+    Returns:
+        Dictionary with Neo4j types serialized to JSON-safe values
+    """
+    serialized = {}
+    for key, value in record_data.items():
+        if isinstance(value, dict):
+            serialized[key] = serialize_record_data(value)
+        elif isinstance(value, list):
+            serialized[key] = [serialize_neo4j_types(item) for item in value]
+        else:
+            serialized[key] = serialize_neo4j_types(value)
+
+    return serialized
+
+
+def serialize_node(node: Node) -> dict:
+    """
+    Project a neo4j Node into a JSON-safe node record.
+
+    Preserves the identity that ``Record.data()`` discards: the stable
+    ``element_id`` and ``labels``. Property values reuse
+    :func:`serialize_record_data` so temporal/duration types stay JSON-safe.
+    """
+    return {
+        "_type": "node",
+        "id": node.element_id,
+        "labels": sorted(node.labels),
+        "properties": serialize_record_data(dict(node)),
+    }
+
+
+def serialize_relationship(rel: Relationship) -> dict:
+    """
+    Project a neo4j Relationship into a JSON-safe edge record.
+
+    Keeps the type and the start/end endpoints, which ``Record.data()`` drops.
+    ``start``/``end`` are the endpoint nodes' ``element_id``s, or ``None`` when
+    the driver did not hydrate an endpoint.
+    """
+    start = rel.start_node
+    end = rel.end_node
+    return {
+        "_type": "relationship",
+        "id": rel.element_id,
+        "type": rel.type,
+        "start": start.element_id if start is not None else None,
+        "end": end.element_id if end is not None else None,
+        "properties": serialize_record_data(dict(rel)),
+    }
+
+
+def serialize_path(path: Path) -> dict:
+    """Project a neo4j Path into a JSON-safe record of its nodes and relationships."""
+    return {
+        "_type": "path",
+        "nodes": [serialize_node(node) for node in path.nodes],
+        "relationships": [serialize_relationship(rel) for rel in path.relationships],
+    }
+
+
+def serialize_point(point: Point) -> dict:
+    """Project a neo4j spatial Point into a JSON-safe record."""
+    return {
+        "_type": "point",
+        "srid": point.srid,
+        "coordinates": [float(c) for c in point],
+    }
+
+
+def serialize_value(value: Any) -> Any:
+    """Serialize a single Cypher result value, preserving its type.
+
+    The type-preserving counterpart to ``Record.data()`` (which flattens graph
+    entities into bare property maps): nodes, relationships and paths become
+    ``_type``-tagged objects with their identity and topology intact; points and
+    temporals become JSON-safe structures; lists and maps recurse; primitives
+    pass through.
+    """
+    if isinstance(value, Node):
+        return serialize_node(value)
+    if isinstance(value, Relationship):
+        return serialize_relationship(value)
+    if isinstance(value, Path):
+        return serialize_path(value)
+    if isinstance(value, Point):
+        return serialize_point(value)
+    if isinstance(value, list):
+        return [serialize_value(item) for item in value]
+    if isinstance(value, dict):
+        return {key: serialize_value(item) for key, item in value.items()}
+    return serialize_neo4j_types(value)
+
+
+def serialize_records(records: Iterable[Any]) -> list[dict]:
+    """Serialize raw neo4j records into type-preserving rows (one dict per row).
+
+    Each row keeps its ``RETURN`` column names; each value is serialized with
+    :func:`serialize_value`, so nodes/edges/paths retain their identity instead
+    of collapsing to bare property maps.
+
+    Accepts anything iterable of records (a plain ``list``, or a neo4j
+    ``Result``/``AsyncResult`` consumed lazily) -- the body only ever iterates
+    once and calls ``.items()`` per record, it never needs list operations.
+    """
+    return [{key: serialize_value(value) for key, value in record.items()} for record in records]

memgraph_toolbox-0.2.0/src/memgraph_toolbox/utils/serialization.py

@@ -1,64 +0,0 @@
-"""
-Serialization utilities for Neo4j date/time types.
-"""
-
-from typing import Any
-
-
-def serialize_neo4j_types(value: Any) -> Any:
-    """
-    Convert Neo4j date/time types to JSON-serializable strings.
-
-    Args:
-        value: The value to serialize
-
-    Returns:
-        The serialized value with Neo4j types converted to strings
-    """
-    try:
-        class_name = value.__class__.__name__
-        module_name = getattr(value.__class__, "__module__", "")
-
-        # Check if it's a Neo4j temporal type
-        if "neo4j" in module_name and hasattr(value, "iso_format"):
-            neo4j_temporal_types = [
-                "Date",
-                "Time",
-                "DateTime",
-                "LocalTime",
-                "LocalDateTime",
-            ]
-            if class_name in neo4j_temporal_types:
-                return value.iso_format()
-
-        # Handle Neo4j Duration type
-        if "neo4j" in module_name and class_name == "Duration":
-            return str(value)
-
-    except (AttributeError, TypeError):
-        # If we can't access the class name or module, return as-is
-        pass
-
-    return value
-
-
-def serialize_record_data(record_data: dict) -> dict:
-    """
-    Serialize a single record's data, handling Neo4j types recursively.
-
-    Args:
-        record_data: Dictionary from a Neo4j record.data()
-
-    Returns:
-        Dictionary with Neo4j types serialized to JSON-safe values
-    """
-    serialized = {}
-    for key, value in record_data.items():
-        if isinstance(value, dict):
-            serialized[key] = serialize_record_data(value)
-        elif isinstance(value, list):
-            serialized[key] = [serialize_neo4j_types(item) for item in value]
-        else:
-            serialized[key] = serialize_neo4j_types(value)
-
-    return serialized