memgraph-toolbox 0.2.1__tar.gz → 0.2.3__tar.gz

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.3}/.gitignore

@@ -187,3 +187,6 @@ cython_debug/
 
 # deepeval writes run caches and telemetry here; build artifacts, not source.
 .deepeval/
+
+# Eval runs hold benchmark text (questions, answers, retrieval); only published baselines are committed.
+context-graph/eval/runs/

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: memgraph-toolbox
-Version: 0.2.1
+Version: 0.2.3
 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.1 → memgraph_toolbox-0.2.3}/pyproject.toml

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

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.3}/src/memgraph_toolbox/api/memgraph.py

@@ -1,4 +1,5 @@
 import os
+import time
 from collections.abc import Mapping
 from typing import Any, cast
 
@@ -78,6 +79,49 @@ 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}")
+
+
+def _driver_database(name: str) -> str | None:
+    """The database to hand the driver: *name*, or None for an empty one (the login's main database).
+
+    Over a direct ``bolt://`` connection ``""`` and None both leave the database to the
+    server. A routing ``neo4j://`` driver instead asks the cluster for the routing table
+    of a database named ``""``, which has no writer, so every write fails.
+    """
+    return name or None
+
+
 class Memgraph:
     """
     Base Memgraph client for interacting with Memgraph database.
@@ -107,7 +151,7 @@ class Memgraph:
             url: The Memgraph connection URL
             username: Username for authentication
             password: Password for authentication
-            database: The database name to connect to (default: "memgraph")
+            database: The database name to connect to (default: "memgraph"); "" for the login's main database
             driver_config: Additional Neo4j driver configuration
             user_agent: Client name sent to the server (e.g. "mcp-memgraph", "langchain-memgraph", "sql2graph", etc.)
         """
@@ -123,7 +167,7 @@ class Memgraph:
 
         self.driver = GraphDatabase.driver(url, auth=(username, password), **config)
 
-        self.database = env["MEMGRAPH_DATABASE"]
+        self.database = _driver_database(env["MEMGRAPH_DATABASE"])
 
         try:
             import neo4j
@@ -217,7 +261,7 @@ class AsyncMemgraph:
             url: The Memgraph connection URL
             username: Username for authentication
             password: Password for authentication
-            database: The database name to connect to (default: "memgraph")
+            database: The database name to connect to (default: "memgraph"); "" for the login's main database
             driver_config: Additional Neo4j driver configuration
             user_agent: Client name sent to the server (e.g. "mcp-memgraph", "langchain-memgraph", "sql2graph", etc.)
         """
@@ -232,7 +276,7 @@ class AsyncMemgraph:
         config.setdefault("user_agent", user_agent or self.DEFAULT_USER_AGENT)
 
         self.driver = AsyncGraphDatabase.driver(url, auth=(username, password), **config)
-        self.database = env["MEMGRAPH_DATABASE"]
+        self.database = _driver_database(env["MEMGRAPH_DATABASE"])
 
     async def verify_connectivity(self) -> None:
         """

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.3}/src/memgraph_toolbox/client/mcp_prompt.py

@@ -6,7 +6,7 @@ import sys
 from typing import cast
 
 try:
-    from litellm import acompletion, experimental_mcp_client
+    from litellm import ModelResponse, acompletion, experimental_mcp_client
     from mcp import ClientSession, StdioServerParameters
     from mcp.client.stdio import stdio_client
 except ImportError as e:
@@ -69,10 +69,13 @@ async def prompt_with_tools(
         logger.info(f"{len(tools)} MCP tools loaded")
 
         messages = [*tool_selection_system_messages, {"role": "user", "content": prompt}]
-        resp = await acompletion(
-            model=tool_selection_model_name,
-            messages=messages,
-            tools=tools,
+        resp = cast(
+            "ModelResponse",
+            await acompletion(
+                model=tool_selection_model_name,
+                messages=messages,
+                tools=tools,
+            ),
         )
         msg = resp["choices"][0]["message"]
         tool_calls = msg.get("tool_calls", [])
@@ -147,9 +150,12 @@ async def prompt_with_tools(
                     )
 
             messages = messages + response_system_messages
-            final_resp = await acompletion(
-                model=response_model_name,
-                messages=messages,
+            final_resp = cast(
+                "ModelResponse",
+                await acompletion(
+                    model=response_model_name,
+                    messages=messages,
+                ),
             )
             return final_resp["choices"][0]["message"].content
 

{memgraph_toolbox-0.2.1 → memgraph_toolbox-0.2.3}/src/memgraph_toolbox/tests/test_async_memgraph.py

@@ -32,6 +32,11 @@ def test_async_memgraph_resolves_config_from_env(monkeypatch):
     assert explicit.database == "arg-db"
 
 
+def test_an_empty_database_reaches_the_driver_as_none():
+    """An empty name means the login's main database; a routing driver would look up a database named ""."""
+    assert AsyncMemgraph(url="neo4j://localhost:7687", database="").database is None
+
+
 def test_async_memgraph_mirrors_sync_defaults():
     """AsyncMemgraph shares the sync client's user-agent default."""
     assert AsyncMemgraph.DEFAULT_USER_AGENT == Memgraph.DEFAULT_USER_AGENT

memgraph_toolbox-0.2.3/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