diff --git a/src/memos/graph_dbs/polardb.py b/src/memos/graph_dbs/polardb.py index bf74fbb8b..8bd9fe76d 100644 --- a/src/memos/graph_dbs/polardb.py +++ b/src/memos/graph_dbs/polardb.py @@ -20,6 +20,19 @@ logger = get_logger(__name__) +_LIGHTWEIGHT_IDENTIFIER_FIELDS = frozenset( + { + "user_id", + "user_name", + "session_id", + "mem_cube_id", + "dialog_id", + "working_binding", + "project_id", + "manager_user_id", + } +) + def _build_lightweight_return_columns(return_fields: list[str]) -> str: columns = [] @@ -33,9 +46,20 @@ def _build_lightweight_return_columns(return_fields: list[str]) -> str: return "".join(f",\n {column}" for column in columns) -def _decode_lightweight_return_value(value: Any) -> Any: +def _decode_lightweight_return_value(value: Any, field: str | None = None) -> Any: if hasattr(value, "value"): value = value.value + # Drivers may return agtype identifiers as already-decoded scalar strings. + # Parsing a numeric ID again would change its type. + if field in _LIGHTWEIGHT_IDENTIFIER_FIELDS: + if isinstance(value, str): + try: + decoded = json.loads(value) + except (json.JSONDecodeError, TypeError): + return value + return decoded if decoded is None or isinstance(decoded, str) else value + if isinstance(value, int) and not isinstance(value, bool): + return str(value) if isinstance(value, str): try: return json.loads(value) @@ -1990,7 +2014,7 @@ def search_by_embedding( item = {"id": id_val, "score": score_val} if light_weight_mode: for offset, field in enumerate(validated_return_fields, start=5): - item[field] = _decode_lightweight_return_value(row[offset]) + item[field] = _decode_lightweight_return_value(row[offset], field=field) elif validated_return_fields: properties = row[1] # properties column item.update( diff --git a/tests/graph_dbs/test_search_return_fields.py b/tests/graph_dbs/test_search_return_fields.py index 1ed0f6173..a23374e5a 100644 --- a/tests/graph_dbs/test_search_return_fields.py +++ b/tests/graph_dbs/test_search_return_fields.py @@ -264,6 +264,73 @@ def fake_connection(): def test_polardb_declares_lightweight_vector_search_support(self, polardb_instance): assert polardb_instance.supports_lightweight_vector_search is True + @pytest.mark.parametrize( + "stored_user_id,expected_user_id", + [ + ("1234567890123456789", "1234567890123456789"), + (1234567890123456789, "1234567890123456789"), + ('"1234567890123456789"', "1234567890123456789"), + ("00123", "00123"), + (None, None), + ("null", None), + ], + ) + def test_lightweight_numeric_user_id_can_be_recalled( + self, polardb_instance, stored_user_id, expected_user_id + ): + from memos.memories.textual.tree_text_memory.retrieve.recall import GraphMemoryRetriever + + polardb_instance.db_name = "test_db" + polardb_instance.config = {"user_name": "default_user"} + polardb_instance._build_user_name_and_kb_ids_conditions_sql = MagicMock(return_value=[]) + polardb_instance._build_filter_conditions_sql = MagicMock(return_value=[]) + cursor = MagicMock() + cursor.fetchall.return_value = [ + ( + 1, + None, + None, + '"00000000-0000-4000-8000-000000000001"', + 0.2, + '"synthetic knowledge"', + '"LongTermMemory"', + stored_user_id, + '["synthetic"]', + '{"test":true}', + "0.9", + "true", + 0.8, + ) + ] + connection = MagicMock() + connection.cursor.return_value.__enter__.return_value = cursor + + @contextmanager + def fake_connection(): + yield connection + + polardb_instance._get_connection = fake_connection + hits = polardb_instance.search_by_embedding( + vector=[0.1] * 1024, + user_name="knowledge-cube", + light_weight_mode=True, + return_fields=[ + "memory", + "memory_type", + "user_id", + "tags", + "info", + "confidence", + "is_fast", + ], + ) + memory = GraphMemoryRetriever._memory_item_from_vector_hit(hits[0]) + assert memory.metadata.user_id == expected_user_id + assert memory.metadata.tags == ["synthetic"] + assert memory.metadata.info == {"test": True} + assert memory.metadata.confidence == 0.9 + assert memory.metadata.is_fast is True + class TestFieldNameValidation: """Tests for _validate_return_fields injection prevention."""