diff --git a/integrations/lightrag-memgraph/src/lightrag_memgraph/__init__.py b/integrations/lightrag-memgraph/src/lightrag_memgraph/__init__.py index 7df7cee..b9856db 100644 --- a/integrations/lightrag-memgraph/src/lightrag_memgraph/__init__.py +++ b/integrations/lightrag-memgraph/src/lightrag_memgraph/__init__.py @@ -6,7 +6,12 @@ from .core import MemgraphLightRAGWrapper from .docstatus_impl import MemgraphDocStatusStorage -from .embeddings import build_memgraph_sentence_embed, memgraph_sentence_embed +from .embeddings import ( + DEFAULT_EMBEDDING_DIM, + DEFAULT_MODEL_NAME, + build_memgraph_sentence_embed, + memgraph_sentence_embed, +) from .kv_impl import MemgraphKVStorage from .registry import register_memgraph_storage from .vector_impl import MemgraphVectorStorage @@ -125,6 +130,8 @@ async def _wrapped( _patch_anthropic() __all__ = [ + "DEFAULT_EMBEDDING_DIM", + "DEFAULT_MODEL_NAME", "MemgraphDocStatusStorage", "MemgraphKVStorage", "MemgraphLightRAGWrapper", diff --git a/integrations/lightrag-memgraph/tests/test_import.py b/integrations/lightrag-memgraph/tests/test_import.py index 9fa82b0..4854780 100644 --- a/integrations/lightrag-memgraph/tests/test_import.py +++ b/integrations/lightrag-memgraph/tests/test_import.py @@ -29,3 +29,13 @@ def test_wrapper_with_custom_params(): ) assert wrapper.log_level == "DEBUG" assert wrapper.full_memgraph_persistence is False + + +def test_embedding_defaults_are_exported_from_top_level(): + """DEFAULT_EMBEDDING_DIM/DEFAULT_MODEL_NAME must be reachable without + digging into the embeddings submodule, so downstream packages can treat + them as the single source of truth for the default embedding dimension.""" + from lightrag_memgraph import DEFAULT_EMBEDDING_DIM, DEFAULT_MODEL_NAME, embeddings + + assert DEFAULT_EMBEDDING_DIM == embeddings.DEFAULT_EMBEDDING_DIM + assert DEFAULT_MODEL_NAME == embeddings.DEFAULT_MODEL_NAME diff --git a/unstructured2graph/src/unstructured2graph/memgraph.py b/unstructured2graph/src/unstructured2graph/memgraph.py index ce21c9c..08738df 100644 --- a/unstructured2graph/src/unstructured2graph/memgraph.py +++ b/unstructured2graph/src/unstructured2graph/memgraph.py @@ -1,6 +1,7 @@ import logging import time +from lightrag_memgraph import DEFAULT_EMBEDDING_DIM from memgraph_toolbox.api.memgraph import Memgraph logger = logging.getLogger(__name__) @@ -127,10 +128,17 @@ def create_label_index(memgraph: Memgraph, label: str): logger.warning(f"Could not create label index on :{label}: {e}") -def create_vector_search_index(memgraph: Memgraph, label: str, property: str): +def create_vector_search_index( + memgraph: Memgraph, + label: str, + property: str, + dimension: int = DEFAULT_EMBEDDING_DIM, + index_name: str = "vs_name", +): try: memgraph.query( - f"CREATE VECTOR INDEX vs_name ON :{label}({property}) WITH CONFIG {{'dimension': 384, 'capacity': 10000}};" + f"CREATE VECTOR INDEX {index_name} ON :{label}({property}) " + f"WITH CONFIG {{'dimension': {dimension}, 'capacity': 10000}};" ) except Exception as e: logger.warning(f"Error creating vector search index: {e}") diff --git a/unstructured2graph/tests/test_memgraph.py b/unstructured2graph/tests/test_memgraph.py index 777a24f..4f80b98 100644 --- a/unstructured2graph/tests/test_memgraph.py +++ b/unstructured2graph/tests/test_memgraph.py @@ -2,7 +2,33 @@ from unittest.mock import MagicMock -from unstructured2graph.memgraph import create_nodes_from_list, create_unique_constraint +from lightrag_memgraph import DEFAULT_EMBEDDING_DIM +from unstructured2graph.memgraph import create_nodes_from_list, create_unique_constraint, create_vector_search_index + + +def test_create_vector_search_index_defaults_match_embedding_dim(): + """Default dimension must track lightrag_memgraph's own default, not a + separately-maintained literal.""" + memgraph = MagicMock() + + create_vector_search_index(memgraph, "Chunk", "embedding") + + query = memgraph.query.call_args[0][0] + assert "CREATE VECTOR INDEX vs_name ON :Chunk(embedding)" in query + assert f"'dimension': {DEFAULT_EMBEDDING_DIM}" in query + + +def test_create_vector_search_index_accepts_custom_dimension_and_name(): + """A caller using a different embedding model must be able to pass a + matching dimension and a distinct index name (e.g. for a second index + on another label).""" + memgraph = MagicMock() + + create_vector_search_index(memgraph, "Entity", "embedding", dimension=768, index_name="entity_vs") + + query = memgraph.query.call_args[0][0] + assert "CREATE VECTOR INDEX entity_vs ON :Entity(embedding)" in query + assert "'dimension': 768" in query def test_create_nodes_from_list_defaults_to_create():