diff --git a/src/memos/graph_dbs/neo4j.py b/src/memos/graph_dbs/neo4j.py index 56c3e08a0..dd8a43996 100644 --- a/src/memos/graph_dbs/neo4j.py +++ b/src/memos/graph_dbs/neo4j.py @@ -1,9 +1,12 @@ import json +import re import time from datetime import datetime from typing import Any, Literal +from neo4j.exceptions import ClientError + from memos.configs.graph_db import Neo4jGraphDBConfig from memos.dependency import require_python_package from memos.graph_dbs.base import BaseGraphDB @@ -12,6 +15,13 @@ logger = get_logger(__name__) +# Pattern for validating Cypher property names — alphanumeric + underscore, +# must start with letter or underscore (same as base._VALID_FIELD_NAME_RE) +_VALID_PROPERTY_NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + +# Lucene special characters that need escaping in fulltext queries +_LUCENE_SPECIAL_CHARS = frozenset(r'+-&|!(){}[]^"~*?:\\/') + def _compose_node(item: dict[str, Any]) -> tuple[str, str, dict[str, Any]]: node_id = item["id"] @@ -1011,10 +1021,145 @@ def search_by_fulltext( **kwargs, ) -> list[dict]: """ - TODO: Implement fulltext search for Neo4j to be compatible with TreeTextMemory's keyword/fulltext recall path. - Currently, return an empty list to avoid runtime errors due to missing methods when switching to Neo4j. + Fulltext keyword search using Neo4j's built-in Apache Lucene fulltext index. + + Creates a FULLTEXT INDEX on the ``memory`` property if it doesn't already + exist, then executes a keyword-based search via + ``db.index.fulltext.queryNodes``. + + Compatible with TreeTextMemory's keyword/fulltext recall path. + + Args: + query_words: List of search terms (keywords). Each word is combined + with OR in the Lucene query string. + top_k: Maximum number of results. + scope: Filter by memory_type (e.g. 'WorkingMemory'). + status: Filter by status (e.g. 'activated', 'archived'). + threshold: Minimum score threshold. + search_filter: Additional property filters on nodes. + user_name: User isolation filter. + filter: Advanced filter conditions. + knowledgebase_ids: KB ids for multi-tenant filtering. + tsquery_config: Unused — kept for interface compatibility with + PolarDB. + + Returns: + list[dict]: Results with ``id`` and ``score``, ordered by relevance + score descending. + """ + if not query_words: + return [] + + # Ensure fulltext index exists (lazy creation) + self._ensure_fulltext_index() + + user_name = user_name if user_name else self.config.user_name + + # Build Lucene query: escape each word and join with OR + escaped_words = [ + self._escape_lucene_query(word.strip()) + for word in query_words + if word.strip() + ] + if not escaped_words: + return [] + + lucene_query = " OR ".join(escaped_words) + logger.info( + "[search_by_fulltext] lucene_query=%s top_k=%s scope=%s status=%s user_name=%s", + lucene_query, + top_k, + scope, + status, + user_name, + ) + + # ---- WHERE clauses for post-filtering on fulltext results ---- + where_clauses: list[str] = [] + params: dict[str, Any] = { + "lucene_query": lucene_query, + "top_k": top_k, + } + + if scope: + where_clauses.append("node.memory_type = $scope") + params["scope"] = scope + if status: + where_clauses.append("node.status = $status") + params["status"] = status + + # User isolation (shared-database multi-tenant) + user_name_conditions, user_name_params = ( + self._build_user_name_and_kb_ids_conditions_cypher( + user_name=user_name, + knowledgebase_ids=knowledgebase_ids, + default_user_name=self.config.user_name, + node_alias="node", + ) + ) + if user_name_conditions: + if len(user_name_conditions) == 1: + where_clauses.append(user_name_conditions[0]) + else: + where_clauses.append(f"({' OR '.join(user_name_conditions)})") + params.update(user_name_params) + + # search_filter (property equality filters) + if search_filter: + for key in search_filter: + if not _VALID_PROPERTY_NAME_RE.match(key): + logger.warning( + "[search_by_fulltext] skipping invalid filter key: %s", key + ) + continue + param_name = f"filter_{key}" + where_clauses.append(f"node.{key} = ${param_name}") + params[param_name] = search_filter[key] + + # Advanced filter conditions + filter_conditions, filter_params = self._build_filter_conditions_cypher( + filter=filter, + param_counter_start=0, + node_alias="node", + ) + where_clauses.extend(filter_conditions) + params.update(filter_params) + + # ---- Assemble Cypher query ---- + where_clause = "" + if where_clauses or threshold is not None: + if threshold is not None: + where_clauses.append("score >= $threshold") + params["threshold"] = threshold + where_clause = "WHERE " + " AND ".join(where_clauses) + + query = f""" + CALL db.index.fulltext.queryNodes( + 'memory_fulltext_index', + $lucene_query + ) YIELD node, score + WITH node, score + {where_clause} + RETURN node.id AS id, score + ORDER BY score DESC + LIMIT $top_k """ - return [] + + logger.info("[search_by_fulltext] query=%s params_keys=%s", query, list(params.keys())) + + with self.driver.session(database=self.db_name) as session: + result = session.run(query, params) + records = [] + for record in result: + item: dict[str, Any] = {"id": record["id"], "score": record["score"]} + records.append(item) + + logger.info( + "[search_by_fulltext] returned %d results (threshold=%s)", + len(records), + threshold, + ) + return records def get_by_metadata( self, @@ -1702,6 +1847,70 @@ def _index_exists(self, index_name: str) -> bool: return True return False + def _ensure_fulltext_index(self) -> None: + """Create the fulltext index on ``memory`` if it doesn't exist. + + Uses ``CREATE FULLTEXT INDEX ... IF NOT EXISTS`` so this is safe + to call on every fulltext search invocation. + """ + index_name = "memory_fulltext_index" + if self._fulltext_index_exists(index_name): + return + try: + self._create_fulltext_index(index_name) + except Exception as e: + logger.warning("Failed to create fulltext index '%s': %s", index_name, e) + + def _fulltext_index_exists(self, index_name: str = "memory_fulltext_index") -> bool: + """Check whether a FULLTEXT index with *index_name* exists.""" + query = ( + "SHOW FULLTEXT INDEXES YIELD name WHERE name = $name RETURN name" + ) + try: + with self.driver.session(database=self.db_name) as session: + result = session.run(query, name=index_name) + return result.single() is not None + except ClientError: + # Fallback for older Neo4j versions that don't support + # SHOW FULLTEXT INDEXES — use SHOW INDEXES instead. + return self._index_exists(index_name) + + def _create_fulltext_index( + self, index_name: str = "memory_fulltext_index" + ) -> None: + """Create a FULLTEXT INDEX on ``Memory.memory`` for keyword search.""" + if not _VALID_PROPERTY_NAME_RE.match(index_name): + raise ValueError(f"Invalid fulltext index name: {index_name!r}") + if not _VALID_PROPERTY_NAME_RE.match(index_name): + raise ValueError(f"Invalid fulltext index name: {index_name!r}") + query = f""" + CREATE FULLTEXT INDEX {index_name} IF NOT EXISTS + FOR (n:Memory) ON EACH [n.memory] + """ + with self.driver.session(database=self.db_name) as session: + session.run(query) + logger.info("Fulltext index '%s' created.", index_name) + + @staticmethod + def _escape_lucene_query(term: str) -> str: + r"""Escape special characters in a single Lucene query term. + + Characters escaped: ``+ - && || ! ( ) { } [ ] ^ " ~ * ? : \ /`` + + Returns the term unmodified if it is a wildcard-only string (e.g. ``*``). + """ + if not term: + return term + _LUCENE_WILDCARDS = frozenset("*?") + if all(ch in _LUCENE_WILDCARDS for ch in term): + return term + escaped = [] + for ch in term: + if ch in _LUCENE_SPECIAL_CHARS: + escaped.append("\\") + escaped.append(ch) + return "".join(escaped) + def _build_user_name_and_kb_ids_conditions_cypher( self, user_name: str | None, diff --git a/src/memos/graph_dbs/neo4j_community.py b/src/memos/graph_dbs/neo4j_community.py index b5c92f40a..6c5e1f278 100644 --- a/src/memos/graph_dbs/neo4j_community.py +++ b/src/memos/graph_dbs/neo4j_community.py @@ -465,26 +465,6 @@ def search_by_embedding( return filtered_results - def search_by_fulltext( - self, - query_words: list[str], - top_k: int = 10, - scope: str | None = None, - status: str | None = None, - threshold: float | None = None, - search_filter: dict | None = None, - user_name: str | None = None, - filter: dict | None = None, - knowledgebase_ids: list[str] | None = None, - tsquery_config: str | None = None, - **kwargs, - ) -> list[dict]: - """ - TODO: Implement fulltext search for Neo4j to be compatible with TreeTextMemory's keyword/fulltext recall path. - Currently, return an empty list to avoid runtime errors due to missing methods when switching to Neo4j. - """ - return [] - def _normalize_date_string(self, date_str: str) -> str: """ Normalize date string to ISO 8601 format for Neo4j datetime() function. diff --git a/tests/graph_dbs/test_fulltext_search.py b/tests/graph_dbs/test_fulltext_search.py new file mode 100644 index 000000000..b7b36c2d8 --- /dev/null +++ b/tests/graph_dbs/test_fulltext_search.py @@ -0,0 +1,419 @@ +""" +Unit tests for Neo4j fulltext search implementation. + +Tests cover: +- Basic fulltext search with query words +- Filtering by scope, status, user_name +- Threshold filtering +- Empty query_words edge case +- Lucene special character escaping +- search_filter and advanced filter conditions +- Fulltext index creation (lazy) +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from memos.configs.graph_db import Neo4jGraphDBConfig + + +# ──────────────────────────────────────────────────────────────────────────── +# Fixtures +# ──────────────────────────────────────────────────────────────────────────── + + +@pytest.fixture +def shared_db_config(): + """Shared-database multi-tenant config (use_multi_db=False).""" + return Neo4jGraphDBConfig( + uri="bolt://localhost:7687", + user="neo4j", + password="test", + db_name="test_db", + auto_create=False, + use_multi_db=False, + user_name="default_user", + embedding_dimension=3, + ) + + +@pytest.fixture +def multi_db_config(): + """Multi-database config — no user_name filter in queries.""" + return Neo4jGraphDBConfig( + uri="bolt://localhost:7687", + user="neo4j", + password="test", + db_name="test_db", + auto_create=False, + use_multi_db=True, + embedding_dimension=3, + ) + + +@pytest.fixture +def shared_neo4j_db(shared_db_config): + """Shared-database Neo4j graph DB with mocked driver.""" + with patch("neo4j.GraphDatabase") as mock_gd: + mock_driver = MagicMock() + mock_gd.driver.return_value = mock_driver + from memos.graph_dbs.neo4j import Neo4jGraphDB + + db = Neo4jGraphDB(shared_db_config) + db.driver = mock_driver + yield db + + +@pytest.fixture +def multi_neo4j_db(multi_db_config): + """Multi-database Neo4j graph DB with mocked driver.""" + with patch("neo4j.GraphDatabase") as mock_gd: + mock_driver = MagicMock() + mock_gd.driver.return_value = mock_driver + from memos.graph_dbs.neo4j import Neo4jGraphDB + + db = Neo4jGraphDB(multi_db_config) + db.driver = mock_driver + yield db + + +# ──────────────────────────────────────────────────────────────────────────── +# Tests: search_by_fulltext — basic functionality +# ──────────────────────────────────────────────────────────────────────────── + + +class TestFulltextSearchBasic: + """Basic fulltext search tests.""" + + def test_search_with_single_word(self, shared_neo4j_db): + """Search with a single query word returns scored results.""" + session_mock = _mock_session_run(shared_neo4j_db, [ + {"id": "mem-1", "score": 0.95}, + {"id": "mem-2", "score": 0.80}, + ]) + + results = shared_neo4j_db.search_by_fulltext( + query_words=["strawberry"], + top_k=5, + ) + + assert len(results) == 2 + assert results[0] == {"id": "mem-1", "score": 0.95} + assert results[1] == {"id": "mem-2", "score": 0.80} + + # Verify the Cypher query structure + query = session_mock.run.call_args[0][0] + assert "db.index.fulltext.queryNodes" in query + assert "memory_fulltext_index" in query + assert "$lucene_query" in query + assert "ORDER BY score DESC" in query + assert "LIMIT $top_k" in query + + def test_search_with_multiple_words(self, shared_neo4j_db): + """Multiple query words are joined with OR in the Lucene query.""" + session_mock = _mock_session_run(shared_neo4j_db, [ + {"id": "mem-1", "score": 0.90}, + ]) + + shared_neo4j_db.search_by_fulltext( + query_words=["apple", "banana", "cherry"], + top_k=10, + ) + + params = session_mock.run.call_args[1] + lucene = params["lucene_query"] + assert " OR " in lucene + assert "apple" in lucene + assert "banana" in lucene + assert "cherry" in lucene + + def test_empty_query_words_returns_empty(self, shared_neo4j_db): + """Empty query_words should return [] without running a query.""" + session_mock = _mock_session_run(shared_neo4j_db, []) + + results = shared_neo4j_db.search_by_fulltext(query_words=[]) + + assert results == [] + # Should not have called session.run + session_mock.run.assert_not_called() + + def test_whitespace_only_words_returns_empty(self, shared_neo4j_db): + """Whitespace-only query words should be filtered out.""" + session_mock = _mock_session_run(shared_neo4j_db, []) + + results = shared_neo4j_db.search_by_fulltext(query_words=[" ", "\t"]) + + assert results == [] + session_mock.run.assert_not_called() + + def test_top_k_limits_results(self, shared_neo4j_db): + """top_k parameter is passed correctly.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "x", "score": 1.0}]) + + shared_neo4j_db.search_by_fulltext(query_words=["test"], top_k=3) + + params = session_mock.run.call_args[1] + assert params["top_k"] == 3 + + +# ──────────────────────────────────────────────────────────────────────────── +# Tests: search_by_fulltext — filtering +# ──────────────────────────────────────────────────────────────────────────── + + +class TestFulltextSearchFiltering: + """Tests for scope, status, user_name, and other filters.""" + + def test_scope_filter(self, shared_neo4j_db): + """scope filter adds 'node.memory_type = $scope' in WHERE.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "a", "score": 0.9}]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + scope="WorkingMemory", + ) + + query = session_mock.run.call_args[0][0] + params = session_mock.run.call_args[1] + assert "node.memory_type = $scope" in query + assert params["scope"] == "WorkingMemory" + + def test_status_filter(self, shared_neo4j_db): + """status filter adds 'node.status = $status' in WHERE.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "a", "score": 0.9}]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + status="activated", + ) + + query = session_mock.run.call_args[0][0] + params = session_mock.run.call_args[1] + assert "node.status = $status" in query + assert params["status"] == "activated" + + def test_user_name_filter_shared_db(self, shared_neo4j_db): + """User name filter is added in shared-database mode.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "a", "score": 0.9}]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + user_name="alice", + ) + + query = session_mock.run.call_args[0][0] + assert "user_name" in _get_where_clause(query) + + def test_no_user_name_filter_multi_db(self, multi_neo4j_db): + """No user_name filter in multi-database mode.""" + session_mock = _mock_session_run(multi_neo4j_db, [{"id": "a", "score": 0.9}]) + + multi_neo4j_db.search_by_fulltext( + query_words=["test"], + ) + + query = session_mock.run.call_args[0][0] + where = _get_where_clause(query) + assert "user_name" not in where + + def test_search_filter(self, shared_neo4j_db): + """search_filter adds equality constraints.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "a", "score": 0.9}]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + search_filter={"tags": "important"}, + ) + + query = session_mock.run.call_args[0][0] + params = session_mock.run.call_args[1] + assert "node.tags = $filter_tags" in query + assert params["filter_tags"] == "important" + + def test_search_filter_rejects_invalid_key(self, shared_neo4j_db): + """Invalid filter keys (Cypher injection attempt) are skipped.""" + session_mock = _mock_session_run(shared_neo4j_db, [{"id": "a", "score": 0.9}]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + search_filter={"x} DETACH DELETE n //": "evil"}, + ) + + query = session_mock.run.call_args[0][0] + assert "DETACH DELETE" not in query + + +# ──────────────────────────────────────────────────────────────────────────── +# Tests: search_by_fulltext — threshold +# ──────────────────────────────────────────────────────────────────────────── + + +class TestFulltextSearchThreshold: + """Threshold filtering tests (applied in Cypher, not Python).""" + + def test_threshold_added_to_cypher_query(self, shared_neo4j_db): + """When threshold is set, it's pushed into the Cypher WHERE clause.""" + session_mock = _mock_session_run(shared_neo4j_db, [ + {"id": "a", "score": 0.90}, + ]) + + shared_neo4j_db.search_by_fulltext( + query_words=["test"], + threshold=0.50, + ) + + query = session_mock.run.call_args[0][0] + params = session_mock.run.call_args[1] + assert "score >= $threshold" in query + assert params["threshold"] == 0.50 + + def test_no_threshold_returns_all(self, shared_neo4j_db): + """Without threshold, no score filter in query.""" + session_mock = _mock_session_run(shared_neo4j_db, [ + {"id": "a", "score": 0.10}, + ]) + + shared_neo4j_db.search_by_fulltext(query_words=["test"]) + + query = session_mock.run.call_args[0][0] + assert "score >= $threshold" not in query + + +# ──────────────────────────────────────────────────────────────────────────── +# Tests: Lucene escaping +# ──────────────────────────────────────────────────────────────────────────── + + +class TestLuceneEscaping: + """Tests for _escape_lucene_query static method.""" + + @staticmethod + def _escape(term): + from memos.graph_dbs.neo4j import Neo4jGraphDB + + return Neo4jGraphDB._escape_lucene_query(term) + + def test_plain_word_passes_through(self): + """Plain word is returned unchanged.""" + assert self._escape("hello") == "hello" + + def test_plus_sign_escaped(self): + """'+' is escaped.""" + assert self._escape("C++") == "C\\+\\+" + + def test_parentheses_escaped(self): + """Parentheses are escaped.""" + escaped = self._escape("func(arg)") + assert "\\(" in escaped + assert "\\)" in escaped + + def test_colon_escaped(self): + """Colon is escaped.""" + assert self._escape("key:value") == "key\\:value" + + def test_bracket_escaped(self): + """Brackets are escaped.""" + escaped = self._escape("[test]") + assert "\\[" in escaped + assert "\\]" in escaped + + def test_empty_string(self): + """Empty string is returned as-is.""" + assert self._escape("") == "" + + def test_special_chars_only_wildcard_not_escaped(self): + """A lone wildcard '*' is preserved for prefix queries.""" + escaped = self._escape("*") + assert escaped == "*" + + def test_wildcard_in_mixed_term_is_escaped(self): + """Wildcard '*' within a regular term should be escaped.""" + escaped = self._escape("foo*") + assert escaped == "foo\\*" + + +# ──────────────────────────────────────────────────────────────────────────── +# Tests: Fulltext index creation +# ──────────────────────────────────────────────────────────────────────────── + + +class TestFulltextIndexCreation: + """Tests for _ensure_fulltext_index and related methods.""" + + def test_index_creation_called_on_first_search(self, shared_neo4j_db): + """First fulltext search triggers index creation.""" + session_mock = shared_neo4j_db.driver.session.return_value + session_mock.__enter__.return_value = session_mock + + # Override run return values for the search call + def run_side_effect(query, **params): + mock_result = MagicMock() + if "SHOW FULLTEXT" in query: + mock_result.single.return_value = None + elif "CREATE FULLTEXT" in query: + mock_result.single.return_value = None + else: + mock_result.__iter__.return_value = iter([]) + return mock_result + + session_mock.run.side_effect = run_side_effect + + shared_neo4j_db.search_by_fulltext(query_words=["hello"]) + + # Verify CREATE FULLTEXT INDEX was called + create_calls = [ + c for c in session_mock.run.call_args_list + if "CREATE FULLTEXT" in str(c) + ] + assert len(create_calls) >= 1 + + +# ──────────────────────────────────────────────────────────────────────────── +# Helpers +# ──────────────────────────────────────────────────────────────────────────── + + +def _mock_session_run(db, return_rows): + """Set up a mocked session.run that dispatches by query content. + + Index-related calls (SHOW/CREATE FULLTEXT INDEX) return appropriate + stubs so they don't interfere with the search query assertions. + """ + session_mock = db.driver.session.return_value + session_mock.__enter__.return_value = session_mock + + def _make_record(row): + mock_record = MagicMock() + mock_record.__getitem__.side_effect = row.__getitem__ + mock_record.keys.return_value = row.keys() + return mock_record + + search_records = [_make_record(r) for r in return_rows] + + def _run_side_effect(query, **params): + mock_result = MagicMock() + if "SHOW FULLTEXT" in query: + mock_result.single.return_value = None # index doesn't exist + elif "CREATE FULLTEXT" in query: + mock_result.single.return_value = None + else: + # Search query — return the test data + mock_result.__iter__.return_value = iter(search_records) + return mock_result + + session_mock.run.side_effect = _run_side_effect + return session_mock + + +def _get_where_clause(query: str) -> str: + """Extract the WHERE clause from a Cypher query.""" + idx = query.find("WHERE ") + if idx == -1: + return "" + limit_idx = query.find("RETURN", idx) + if limit_idx == -1: + limit_idx = query.find("ORDER BY", idx) + if limit_idx == -1: + return query[idx:] + return query[idx:limit_idx]